zsfm/
tabicl.rs

1use std::path::PathBuf;
2
3use anyhow::Context;
4use clap::Subcommand;
5use serde::Serialize;
6
7use zsfm_tabicl::config::TabIclConfig;
8
9#[derive(Subcommand)]
10pub enum Command {
11    /// Print all tensor names in a local checkpoint file (safetensors, PyTorch
12    /// pickle, GGUF, or npy/npz — format auto-detected).
13    InspectTensors { path: PathBuf },
14    /// Upload source + GGUF files to HuggingFace Hub.
15    Upload {
16        #[arg(short, long, default_value = "amaye15/tabicl-gguf")]
17        repo: String,
18        #[arg(long, env = "HF_TOKEN")]
19        token: String,
20        #[arg(long, default_value = ".")]
21        root: PathBuf,
22    },
23    /// Zero-shot classification from a GGUF file (v2 checkpoint; classification only — see the
24    /// crate's module docs for what's out of scope).
25    ///
26    /// Convert the checkpoint first with the generic converter (names pass through unchanged,
27    /// so there's no `tabicl convert` subcommand):
28    ///   zsfm convert --repo jingang/TabICL --file v2 --format ckpt -o gguf/tabicl-v2-f32.gguf
29    ///
30    /// Reads a JSON request from stdin:
31    ///   {"x_support": [[...]], "y_support": [...], "x_query": [[...]], "n_classes": N}
32    /// `n_classes` must be <= 10. Outputs JSON: {"task":"classification","probabilities":[[...]]}.
33    Infer {
34        #[arg(short, long)]
35        gguf: PathBuf,
36    },
37
38    /// Delete cached model files for this model.
39    ///
40    /// Removes the canonical F32 GGUF and config cache at
41    /// `<model_dir>/<owner>__<name>/` created by `convert`. By default only
42    /// the cache is removed; pass `--output <path>` to also delete a
43    /// converted GGUF file (e.g. `gguf/...`).
44    Delete {
45        /// HuggingFace repo id (must match the `convert` --model you used).
46        #[arg(short, long, default_value = "jingang/TabICL")]
47        model: String,
48        /// Directory where the cache was written (must match `convert` --model-dir).
49        #[arg(long, default_value = "models")]
50        model_dir: PathBuf,
51        /// Also delete this output GGUF file if it exists.
52        #[arg(short, long)]
53        output: Option<PathBuf>,
54    },
55}
56
57pub async fn run(command: Command) -> anyhow::Result<()> {
58    match command {
59        Command::Delete {
60            model,
61            model_dir,
62            output,
63        } => {
64            let canonical = zsfm_hub::canonical_gguf_path(&model_dir, &model);
65            crate::common::delete_cached_model(&canonical, output.as_deref())?;
66        }
67        Command::Upload { repo, token, root } => {
68            zsfm_hub::upload_repo(&repo, &token, &root).await?;
69        }
70
71        Command::InspectTensors { path } => crate::common::inspect_tensors(&path)?,
72
73        Command::Infer { gguf } => {
74            run_infer_burn(&gguf).await?;
75        }
76    }
77
78    Ok(())
79}
80
81async fn run_infer_burn(gguf: &std::path::PathBuf) -> anyhow::Result<()> {
82    use zsfm_tabicl::infer::BurnTabIclModel;
83
84    let buf = zsfm_core::read_stdin_limited()?;
85    let req: serde_json::Value = serde_json::from_str(&buf).context("parse JSON input")?;
86
87    let x_support = zsfm_core::parse_matrix(&req["x_support"])?;
88    let x_query = zsfm_core::parse_matrix(&req["x_query"])?;
89    let y_support: Vec<usize> = serde_json::from_value(req["y_support"].clone())
90        .context("expected a 1D JSON integer array for `y_support`")?;
91    let n_classes = req
92        .get("n_classes")
93        .and_then(|v| v.as_u64())
94        .unwrap_or_else(|| {
95            y_support
96                .iter()
97                .copied()
98                .max()
99                .map(|m| m as u64 + 1)
100                .unwrap_or(1)
101        }) as usize;
102
103    let config = TabIclConfig::v2();
104    eprintln!("Loading Burn model from {} …", gguf.display());
105    let model = BurnTabIclModel::load(gguf, config).context("load model")?;
106
107    eprintln!(
108        "Running classification ({} support / {} query rows, {n_classes} classes) …",
109        x_support.len(),
110        x_query.len()
111    );
112    let probabilities = model
113        .predict_classification(&x_support, &y_support, &x_query, n_classes)
114        .context("predict")?;
115
116    #[derive(Serialize)]
117    struct Resp {
118        task: &'static str,
119        probabilities: Vec<Vec<f32>>,
120    }
121    println!(
122        "{}",
123        serde_json::to_string_pretty(&Resp {
124            task: "classification",
125            probabilities
126        })?
127    );
128    Ok(())
129}
zsfm/tabicl.rs

zsfm/
tabicl.rs
