zsfm/
tabpfn.rs

1use std::path::PathBuf;
2
3use anyhow::Context;
4use clap::Subcommand;
5use serde::Serialize;
6
7use zsfm_tabpfn::config::TabPfnConfig;
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/tabpfn-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 (classification only — regression's
24    /// bar-distribution head is out of scope, see the crate's module docs).
25    ///
26    /// NON-COMMERCIAL WEIGHTS LICENSE (TabPFN-3 Non-Commercial License v1.0, Prior Labs GmbH):
27    /// research, internal benchmarking, and evaluation only — no production, commercial, or
28    /// hosted-service use without a separate commercial license from Prior Labs.
29    ///
30    /// Convert the checkpoint first with the generic converter (names pass through unchanged,
31    /// so there's no `tabpfn convert` subcommand):
32    ///   zsfm convert --repo Prior-Labs/tabpfn_3 --file classifier --format ckpt -o gguf/tabpfn-v3-f32.gguf
33    ///
34    /// Reads a JSON request from stdin:
35    ///   {"x_support": [[...]], "y_support": [...], "x_query": [[...]], "n_classes": N}
36    /// Outputs JSON: {"task":"classification","probabilities":[[...]]}.
37    Infer {
38        #[arg(short, long)]
39        gguf: PathBuf,
40    },
41
42    /// Delete cached model files for this model.
43    ///
44    /// Removes the canonical F32 GGUF and config cache at
45    /// `<model_dir>/<owner>__<name>/` created by `convert`. By default only
46    /// the cache is removed; pass `--output <path>` to also delete a
47    /// converted GGUF file (e.g. `gguf/...`).
48    Delete {
49        /// HuggingFace repo id (must match the `convert` --model you used).
50        #[arg(short, long, default_value = "Prior-Labs/tabpfn_3")]
51        model: String,
52        /// Directory where the cache was written (must match `convert` --model-dir).
53        #[arg(long, default_value = "models")]
54        model_dir: PathBuf,
55        /// Also delete this output GGUF file if it exists.
56        #[arg(short, long)]
57        output: Option<PathBuf>,
58    },
59}
60
61pub async fn run(command: Command) -> anyhow::Result<()> {
62    match command {
63        Command::Delete {
64            model,
65            model_dir,
66            output,
67        } => {
68            let canonical = zsfm_hub::canonical_gguf_path(&model_dir, &model);
69            crate::common::delete_cached_model(&canonical, output.as_deref())?;
70        }
71        Command::Upload { repo, token, root } => {
72            zsfm_hub::upload_repo(&repo, &token, &root).await?;
73        }
74
75        Command::InspectTensors { path } => crate::common::inspect_tensors(&path)?,
76
77        Command::Infer { gguf } => {
78            run_infer_burn(&gguf).await?;
79        }
80    }
81
82    Ok(())
83}
84
85async fn run_infer_burn(gguf: &std::path::PathBuf) -> anyhow::Result<()> {
86    use zsfm_tabpfn::infer::BurnTabPfnModel;
87
88    let buf = zsfm_core::read_stdin_limited()?;
89    let req: serde_json::Value = serde_json::from_str(&buf).context("parse JSON input")?;
90
91    let x_support = zsfm_core::parse_matrix(&req["x_support"])?;
92    let x_query = zsfm_core::parse_matrix(&req["x_query"])?;
93    let y_support: Vec<usize> = serde_json::from_value(req["y_support"].clone())
94        .context("expected a 1D JSON integer array for `y_support`")?;
95    let n_classes = req
96        .get("n_classes")
97        .and_then(|v| v.as_u64())
98        .unwrap_or_else(|| {
99            y_support
100                .iter()
101                .copied()
102                .max()
103                .map(|m| m as u64 + 1)
104                .unwrap_or(1)
105        }) as usize;
106
107    let config = TabPfnConfig::v3_default();
108    eprintln!("Loading Burn model from {} …", gguf.display());
109    let model = BurnTabPfnModel::load(gguf, config).context("load model")?;
110
111    eprintln!(
112        "Running classification ({} support / {} query rows, {n_classes} classes) …",
113        x_support.len(),
114        x_query.len()
115    );
116    let probabilities = model
117        .predict_classification(&x_support, &y_support, &x_query, n_classes)
118        .context("predict")?;
119
120    #[derive(Serialize)]
121    struct Resp {
122        task: &'static str,
123        probabilities: Vec<Vec<f32>>,
124    }
125    println!(
126        "{}",
127        serde_json::to_string_pretty(&Resp {
128            task: "classification",
129            probabilities
130        })?
131    );
132    Ok(())
133}
zsfm/tabpfn.rs

zsfm/
tabpfn.rs
