zsfm/
tabdpt.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_tabdpt::config::TabDptConfig;
10use zsfm_tabdpt::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 TabDPT (one checkpoint serves both classification and regression) from
21    /// HuggingFace and convert to GGUF.
22    ///
23    /// The first conversion downloads the model and writes a canonical F32 GGUF
24    /// alongside it, deleting the downloaded safetensors file afterward. Later
25    /// conversions (any --dtype) recast from that cached F32 GGUF instead of
26    /// re-downloading — pass --redownload to force a fresh download anyway (e.g.
27    /// the repo was updated).
28    Convert {
29        #[arg(short, long, default_value = "Layer6/TabDPT")]
30        model: String,
31        #[arg(long, default_value = "tabdpt1_2.safetensors")]
32        filename: String,
33        #[arg(short, long)]
34        output: Option<PathBuf>,
35        #[arg(long, default_value = "f32")]
36        dtype: DtypeArg,
37        #[arg(long, default_value = "models")]
38        model_dir: PathBuf,
39        #[arg(long, env = "HF_TOKEN")]
40        token: Option<String>,
41        /// Force a fresh download even if a cached F32 GGUF already exists.
42        #[arg(long)]
43        redownload: bool,
44    },
45    /// Print all tensor names in a local checkpoint file (safetensors, PyTorch
46    /// pickle, GGUF, or npy/npz — format auto-detected).
47    InspectTensors { path: PathBuf },
48    /// Upload source + GGUF files to HuggingFace Hub.
49    Upload {
50        #[arg(short, long, default_value = "amaye15/tabdpt-gguf")]
51        repo: String,
52        #[arg(long, env = "HF_TOKEN")]
53        token: String,
54        #[arg(long, default_value = ".")]
55        root: PathBuf,
56    },
57    /// Zero-shot classification/regression from a GGUF file (single forward pass — no
58    /// class-permutation ensembling; see the crate's module docs for why).
59    ///
60    /// Reads a JSON request from stdin:
61    ///   {"x_support": [[...]], "y_support": [...], "x_query": [[...]], "n_classes": N}
62    /// `n_classes` is only used for classification (ignored, may be omitted, for regression).
63    /// Outputs JSON: classification -> {"task":"classification","probabilities":[[...]]};
64    /// regression -> {"task":"regression","predictions":[...]}.
65    Infer {
66        #[arg(short, long)]
67        gguf: PathBuf,
68        #[arg(long, value_enum, default_value = "classification")]
69        task: TaskArg,
70    },
71
72    /// Delete cached model files for this model.
73    ///
74    /// Removes the canonical F32 GGUF and config cache at
75    /// `<model_dir>/<owner>__<name>/` created by `convert`. By default only
76    /// the cache is removed; pass `--output <path>` to also delete a
77    /// converted GGUF file (e.g. `gguf/...`).
78    Delete {
79        /// HuggingFace repo id (must match the `convert` --model you used).
80        #[arg(short, long, default_value = "Layer6/TabDPT")]
81        model: String,
82        /// Directory where the cache was written (must match `convert` --model-dir).
83        #[arg(long, default_value = "models")]
84        model_dir: PathBuf,
85        /// Also delete this output GGUF file if it exists.
86        #[arg(short, long)]
87        output: Option<PathBuf>,
88    },
89}
90
91fn dtype_name(d: &DtypeArg) -> &'static str {
92    match d {
93        DtypeArg::F32 => "f32",
94        DtypeArg::F16 => "f16",
95        DtypeArg::Q8 => "q8",
96    }
97}
98
99pub async fn run(command: Command) -> anyhow::Result<()> {
100    match command {
101        Command::Delete {
102            model,
103            model_dir,
104            output,
105        } => {
106            let canonical = zsfm_hub::canonical_gguf_path(&model_dir, &model);
107            crate::common::delete_cached_model(&canonical, output.as_deref())?;
108        }
109        Command::Upload { repo, token, root } => {
110            zsfm_hub::upload_repo(&repo, &token, &root).await?;
111        }
112
113        Command::Convert {
114            model,
115            filename,
116            output,
117            dtype,
118            model_dir,
119            token,
120            redownload,
121        } => {
122            let output = output.unwrap_or_else(|| {
123                PathBuf::from(format!("gguf/tabdpt-{}.gguf", dtype_name(&dtype)))
124            });
125            let canonical = zsfm_hub::canonical_gguf_path(&model_dir, &model);
126
127            if crate::common::try_recast_from_cache(&canonical, &output, dtype.into(), redownload)?
128            {
129                return Ok(());
130            }
131            // download_file() only creates model_dir itself, not the namespaced
132            // subdirectory canonical_gguf_path() computes — unlike download_model(),
133            // which creates that exact directory as a side effect.
134            if let Some(parent) = canonical.parent() {
135                std::fs::create_dir_all(parent).context("create canonical GGUF dir")?;
136            }
137
138            eprintln!(
139                "Downloading {model}/{filename} into {} …",
140                model_dir.display()
141            );
142            let path = zsfm_hub::download_file(&model, &filename, token.as_deref(), &model_dir)
143                .await
144                .context("download failed")?;
145
146            let config = TabDptConfig::default_v1_2();
147            let f32_opts = ConvertOptions {
148                output_dtype: GGMLType::F32,
149            };
150            convert(std::slice::from_ref(&path), &config, &f32_opts, &canonical)?;
151            eprintln!("Wrote canonical F32 GGUF to {} …", canonical.display());
152            let _ = std::fs::remove_file(&path);
153
154            zsfm_checkpoint::recast(&canonical, &output, dtype.into())?;
155            eprintln!("Wrote {}", output.display());
156        }
157
158        Command::InspectTensors { path } => crate::common::inspect_tensors(&path)?,
159
160        Command::Infer { gguf, task } => {
161            run_infer_burn(&gguf, &task).await?;
162        }
163    }
164
165    Ok(())
166}
167
168async fn run_infer_burn(gguf: &std::path::PathBuf, task: &TaskArg) -> anyhow::Result<()> {
169    use zsfm_tabdpt::infer::BurnTabDptModel;
170
171    let buf = zsfm_core::read_stdin_limited()?;
172    let req: serde_json::Value = serde_json::from_str(&buf).context("parse JSON input")?;
173
174    let x_support = zsfm_core::parse_matrix(&req["x_support"])?;
175    let x_query = zsfm_core::parse_matrix(&req["x_query"])?;
176
177    let config = TabDptConfig::default_v1_2();
178    eprintln!("Loading Burn model from {} …", gguf.display());
179    let model = BurnTabDptModel::load(gguf, config).context("load model")?;
180
181    match task {
182        TaskArg::Classification => {
183            let y_support: Vec<usize> = serde_json::from_value(req["y_support"].clone())
184                .context("expected a 1D JSON integer array for `y_support`")?;
185            let n_classes = req
186                .get("n_classes")
187                .and_then(|v| v.as_u64())
188                .unwrap_or_else(|| {
189                    y_support
190                        .iter()
191                        .copied()
192                        .max()
193                        .map(|m| m as u64 + 1)
194                        .unwrap_or(1)
195                }) as usize;
196
197            eprintln!(
198                "Running classification ({} support / {} query rows, {n_classes} classes) …",
199                x_support.len(),
200                x_query.len()
201            );
202            let probabilities = model
203                .predict_classification(&x_support, &y_support, &x_query, n_classes)
204                .context("predict")?;
205
206            #[derive(Serialize)]
207            struct Resp {
208                task: &'static str,
209                probabilities: Vec<Vec<f32>>,
210            }
211            println!(
212                "{}",
213                serde_json::to_string_pretty(&Resp {
214                    task: "classification",
215                    probabilities
216                })?
217            );
218        }
219        TaskArg::Regression => {
220            let y_support: Vec<f32> = serde_json::from_value(req["y_support"].clone())
221                .context("expected a 1D JSON numeric array for `y_support`")?;
222            eprintln!(
223                "Running regression ({} support / {} query rows) …",
224                x_support.len(),
225                x_query.len()
226            );
227            let predictions = model
228                .predict_regression(&x_support, &y_support, &x_query)
229                .context("predict")?;
230            #[derive(Serialize)]
231            struct Resp {
232                task: &'static str,
233                predictions: Vec<f32>,
234            }
235            println!(
236                "{}",
237                serde_json::to_string_pretty(&Resp {
238                    task: "regression",
239                    predictions
240                })?
241            );
242        }
243    }
244    Ok(())
245}
zsfm/tabdpt.rs

zsfm/
tabdpt.rs
