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 { dim: 512, n_layers: 12, n_heads: 4, dim_output: 10, task: Task::Classification }
23    }
24
25    /// `autogluon/mitra-regressor`: dim=512, n_layers=12, n_heads=4, dim_output=1.
26    pub fn regressor() -> Self {
27        Self { dim: 512, n_layers: 12, n_heads: 4, dim_output: 1, task: Task::Regression }
28    }
29
30    /// Parse the HF `config.json` shipped alongside the weights: `{"dim", "dim_output",
31    /// "n_layers", "n_heads", "task"}`.
32    pub fn from_json(v: &serde_json::Value) -> anyhow::Result<Self> {
33        let task = match v["task"].as_str().unwrap_or("CLASSIFICATION") {
34            "REGRESSION" => Task::Regression,
35            _ => Task::Classification,
36        };
37        Ok(Self {
38            dim: v["dim"].as_u64().unwrap_or(512) as usize,
39            n_layers: v["n_layers"].as_u64().unwrap_or(12) as usize,
40            n_heads: v["n_heads"].as_u64().unwrap_or(4) as usize,
41            dim_output: v["dim_output"].as_u64().unwrap_or(if task == Task::Regression { 1 } else { 10 }) as usize,
42            task,
43        })
44    }
45}
zsfm_mitra/config.rs

zsfm_mitra/
config.rs
