zsfm/
mitra.rs

1use std::path::PathBuf;
2
3use crate::common::DtypeArg;
4use anyhow::Context;
5use clap::{Subcommand, ValueEnum};
6use serde::Serialize;
7
8use zsfm_gguf::GGMLType;
9use zsfm_mitra::config::MitraConfig;
10use zsfm_mitra::convert::{convert, ConvertOptions};
11
12#[derive(Clone, ValueEnum)]
13pub enum TaskArg {
14    Classification,
15    Regression,
16}
17
18#[derive(Subcommand)]
19pub enum Command {
20    /// Download a Mitra variant from HuggingFace and convert to GGUF.
21    ///
22    /// The first conversion downloads the model and writes a canonical F32 GGUF
23    /// alongside it, deleting the (large) downloaded weight file afterward. Later
24    /// conversions (any --dtype) recast from that cached F32 GGUF instead of
25    /// re-downloading — pass --redownload to force a fresh download anyway (e.g.
26    /// the repo was updated).
27    Convert {
28        /// `autogluon/mitra-classifier` or `autogluon/mitra-regressor` — pick the one
29        /// matching --task (there's no single checkpoint serving both, unlike tabfm).
30        #[arg(short, long, default_value = "autogluon/mitra-classifier")]
31        model: String,
32        #[arg(long, value_enum, default_value = "classification")]
33        task: TaskArg,
34        #[arg(short, long)]
35        output: Option<PathBuf>,
36        #[arg(long, default_value = "f32")]
37        dtype: DtypeArg,
38        #[arg(long, default_value = "models")]
39        model_dir: PathBuf,
40        #[arg(long, env = "HF_TOKEN")]
41        token: Option<String>,
42        /// Force a fresh download even if a cached F32 GGUF already exists.
43        #[arg(long)]
44        redownload: bool,
45    },
46    /// Print all tensor names in a local checkpoint file (safetensors, PyTorch
47    /// pickle, GGUF, or npy/npz — format auto-detected).
48    InspectTensors { path: PathBuf },
49    /// Upload source + GGUF files to HuggingFace Hub.
50    Upload {
51        #[arg(short, long, default_value = "amaye15/mitra-gguf")]
52        repo: String,
53        #[arg(long, env = "HF_TOKEN")]
54        token: String,
55        #[arg(long, default_value = ".")]
56        root: PathBuf,
57    },
58    /// Zero-shot classification/regression from a GGUF file (no fine-tuning — see the crate's
59    /// module docs for why).
60    ///
61    /// Reads a JSON request from stdin:
62    ///   {"x_support": [[...]], "y_support": [...], "x_query": [[...]], "n_classes": N}
63    /// `n_classes` is only used for classification (ignored, may be omitted, for regression).
64    /// Outputs JSON: classification -> {"task":"classification","logits":[[...]],
65    /// "probabilities":[[...]]}; regression -> {"task":"regression","predictions":[...]}.
66    /// Mitra is the only tabular model whose `predict_classification` exposes raw
67    /// pre-softmax logits alongside probabilities — tabdpt/tabicl/tabpfn's underlying
68    /// model APIs only ever produce probabilities, so their `infer` responses have no
69    /// `logits` field.
70    Infer {
71        #[arg(short, long)]
72        gguf: PathBuf,
73        #[arg(long, value_enum, default_value = "classification")]
74        task: TaskArg,
75    },
76
77    /// Delete cached model files for this model.
78    ///
79    /// Removes the canonical F32 GGUF and config cache at
80    /// `<model_dir>/[<variant>-]<owner>__<name>/` created by `convert`.
81    /// By default only the cache is removed; pass `--output <path>` to also
82    /// delete a converted GGUF file (e.g. `gguf/...`).
83    Delete {
84        /// HuggingFace repo id (must match the `convert` --model you used).
85        #[arg(short, long, default_value = "autogluon/mitra-classifier")]
86        model: String,
87        /// Which variant to delete (must match `convert` --task).
88        #[arg(long, value_enum, default_value = "classification")]
89        task: TaskArg,
90        /// Directory where the cache was written (must match `convert` --model-dir).
91        #[arg(long, default_value = "models")]
92        model_dir: PathBuf,
93        /// Also delete this output GGUF file if it exists.
94        #[arg(short, long)]
95        output: Option<PathBuf>,
96    },
97}
98
99fn dtype_name(d: &DtypeArg) -> &'static str {
100    match d {
101        DtypeArg::F32 => "f32",
102        DtypeArg::F16 => "f16",
103        DtypeArg::Q8 => "q8",
104    }
105}
106
107pub async fn run(command: Command) -> anyhow::Result<()> {
108    match command {
109        Command::Delete {
110            model,
111            task,
112            model_dir,
113            output,
114        } => {
115            let task_str = match task {
116                TaskArg::Classification => "classification",
117                TaskArg::Regression => "regression",
118            };
119            let variant_dir = model_dir.join(format!("mitra-{task_str}"));
120            let canonical = zsfm_hub::canonical_gguf_path(&variant_dir, &model);
121            crate::common::delete_cached_model(&canonical, output.as_deref())?;
122        }
123        Command::Upload { repo, token, root } => {
124            zsfm_hub::upload_repo(&repo, &token, &root).await?;
125        }
126
127        Command::Convert {
128            model,
129            task,
130            output,
131            dtype,
132            model_dir,
133            token,
134            redownload,
135        } => {
136            let task_str = match task {
137                TaskArg::Classification => "classification",
138                TaskArg::Regression => "regression",
139            };
140            let variant_dir = model_dir.join(format!("mitra-{task_str}"));
141            let output = output.unwrap_or_else(|| {
142                PathBuf::from(format!("gguf/mitra-{task_str}-{}.gguf", dtype_name(&dtype)))
143            });
144            let canonical = zsfm_hub::canonical_gguf_path(&variant_dir, &model);
145
146            if crate::common::try_recast_from_cache(&canonical, &output, dtype.into(), redownload)?
147            {
148                return Ok(());
149            }
150
151            eprintln!("Downloading {model} into {} …", variant_dir.display());
152            let files = zsfm_hub::download_model(&model, token.as_deref(), &variant_dir)
153                .await
154                .context("download failed")?;
155            let safetensors_path = files
156                .safetensors_shards
157                .first()
158                .context("no safetensors file downloaded")?;
159
160            let config = match task {
161                TaskArg::Classification => MitraConfig::classifier(),
162                TaskArg::Regression => MitraConfig::regressor(),
163            };
164            if files.config_json.exists() {
165                let config_str =
166                    std::fs::read_to_string(&files.config_json).context("read config.json")?;
167                let v: serde_json::Value =
168                    serde_json::from_str(&config_str).context("parse config.json")?;
169                println!(
170                    "HF config.json: dim={} n_layers={} n_heads={} dim_output={}",
171                    v.get("dim").and_then(|x| x.as_u64()).unwrap_or(0),
172                    v.get("n_layers").and_then(|x| x.as_u64()).unwrap_or(0),
173                    v.get("n_heads").and_then(|x| x.as_u64()).unwrap_or(0),
174                    v.get("dim_output").and_then(|x| x.as_u64()).unwrap_or(0),
175                );
176            }
177
178            let f32_opts = ConvertOptions {
179                output_dtype: GGMLType::F32,
180            };
181            convert(
182                std::slice::from_ref(safetensors_path),
183                &config,
184                &f32_opts,
185                &canonical,
186            )?;
187            eprintln!("Wrote canonical F32 GGUF to {} …", canonical.display());
188            files.cleanup_weights();
189
190            zsfm_checkpoint::recast(&canonical, &output, dtype.into())?;
191            eprintln!("Wrote {}", output.display());
192        }
193
194        Command::InspectTensors { path } => crate::common::inspect_tensors(&path)?,
195
196        Command::Infer { gguf, task } => {
197            run_infer_burn(&gguf, &task).await?;
198        }
199    }
200
201    Ok(())
202}
203
204async fn run_infer_burn(gguf: &std::path::PathBuf, task: &TaskArg) -> anyhow::Result<()> {
205    use zsfm_mitra::infer::BurnMitraModel;
206
207    let buf = zsfm_core::read_stdin_limited()?;
208    let req: serde_json::Value = serde_json::from_str(&buf).context("parse JSON input")?;
209
210    let x_support = zsfm_core::parse_matrix(&req["x_support"])?;
211    let x_query = zsfm_core::parse_matrix(&req["x_query"])?;
212
213    let config = match task {
214        TaskArg::Classification => MitraConfig::classifier(),
215        TaskArg::Regression => MitraConfig::regressor(),
216    };
217
218    eprintln!("Loading Burn model from {} …", gguf.display());
219    let model = BurnMitraModel::load(gguf, config).context("load model")?;
220
221    match task {
222        TaskArg::Classification => {
223            let y_support: Vec<usize> = serde_json::from_value(req["y_support"].clone())
224                .context("expected a 1D JSON integer array for `y_support`")?;
225            let n_classes = req
226                .get("n_classes")
227                .and_then(|v| v.as_u64())
228                .unwrap_or_else(|| {
229                    y_support
230                        .iter()
231                        .copied()
232                        .max()
233                        .map(|m| m as u64 + 1)
234                        .unwrap_or(1)
235                }) as usize;
236
237            eprintln!(
238                "Running classification ({} support / {} query rows, {n_classes} classes) …",
239                x_support.len(),
240                x_query.len()
241            );
242            let logits = model
243                .predict_classification(&x_support, &y_support, &x_query, n_classes)
244                .context("predict")?;
245            let probabilities: Vec<Vec<f32>> =
246                logits.iter().map(|row| zsfm_core::softmax(row)).collect();
247
248            #[derive(Serialize)]
249            struct Resp<'a> {
250                task: &'static str,
251                logits: &'a [Vec<f32>],
252                probabilities: Vec<Vec<f32>>,
253            }
254            println!(
255                "{}",
256                serde_json::to_string_pretty(&Resp {
257                    task: "classification",
258                    logits: &logits,
259                    probabilities
260                })?
261            );
262        }
263        TaskArg::Regression => {
264            let y_support: Vec<f32> = serde_json::from_value(req["y_support"].clone())
265                .context("expected a 1D JSON numeric array for `y_support`")?;
266            eprintln!(
267                "Running regression ({} support / {} query rows) …",
268                x_support.len(),
269                x_query.len()
270            );
271            let predictions = model
272                .predict_regression(&x_support, &y_support, &x_query)
273                .context("predict")?;
274            #[derive(Serialize)]
275            struct Resp {
276                task: &'static str,
277                predictions: Vec<f32>,
278            }
279            println!(
280                "{}",
281                serde_json::to_string_pretty(&Resp {
282                    task: "regression",
283                    predictions
284                })?
285            );
286        }
287    }
288    Ok(())
289}
zsfm/mitra.rs

zsfm/
mitra.rs
