zsfm_mitra/
config.rs

1/// Mitra (Tab2D) model configuration. Both published variants (classifier, regressor) share
2/// the same architecture; only `dim_output` and `task` differ.
3#[derive(Clone, Copy, Debug, PartialEq, Eq)]
4pub enum Task {
5    Classification,
6    Regression,
7}
8
9#[derive(Clone, Debug)]
10pub struct MitraConfig {
11    pub dim: usize,
12    pub n_layers: usize,
13    pub n_heads: usize,
14    /// Max classes the classifier head was trained on (10), or 1 for the regressor.
15    pub dim_output: usize,
16    pub task: Task,
17}
18
19impl MitraConfig {
20    /// `autogluon/mitra-classifier`: dim=512, n_layers=12, n_heads=4, dim_output=10.
21    pub fn classifier() -> Self {
22        Self {
23            dim: 512,
24            n_layers: 12,
25            n_heads: 4,
26            dim_output: 10,
27            task: Task::Classification,
28        }
29    }
30
31    /// `autogluon/mitra-regressor`: dim=512, n_layers=12, n_heads=4, dim_output=1.
32    pub fn regressor() -> Self {
33        Self {
34            dim: 512,
35            n_layers: 12,
36            n_heads: 4,
37            dim_output: 1,
38            task: Task::Regression,
39        }
40    }
41
42    /// Parse the HF `config.json` shipped alongside the weights: `{"dim", "dim_output",
43    /// "n_layers", "n_heads", "task"}`.
44    pub fn from_json(v: &serde_json::Value) -> anyhow::Result<Self> {
45        let task = match v["task"].as_str().unwrap_or("CLASSIFICATION") {
46            "REGRESSION" => Task::Regression,
47            _ => Task::Classification,
48        };
49        Ok(Self {
50            dim: v["dim"].as_u64().unwrap_or(512) as usize,
51            n_layers: v["n_layers"].as_u64().unwrap_or(12) as usize,
52            n_heads: v["n_heads"].as_u64().unwrap_or(4) as usize,
53            dim_output: v["dim_output"]
54                .as_u64()
55                .unwrap_or(if task == Task::Regression { 1 } else { 10 })
56                as usize,
57            task,
58        })
59    }
60}
zsfm_mitra/config.rs

zsfm_mitra/
config.rs
