zsfm/
lib.rs

1use pyo3::prelude::*;
2use std::path::PathBuf;
3
4/// Python bindings for zsfm — zero-shot forecasting and tabular foundation models.
5///
6/// Mirrors the Rust CLI (`zsfm <model> convert / infer / delete`) via `uv + pyo3 + maturin`.
7/// Every model exposes a `Model` class (e.g. `TtmModel`, `ChronosModel`) with
8/// `forecast`/`predict`, plus top-level `convert`/`delete` helpers that dispatch
9/// by model name. See `list_models()` for the full registry.
10#[pymodule]
11fn zsfm(m: &Bound<'_, PyModule>) -> PyResult<()> {
12    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
13    m.add_function(wrap_pyfunction!(list_models, m)?)?;
14    m.add_function(wrap_pyfunction!(list_forecasters, m)?)?;
15    m.add_function(wrap_pyfunction!(list_tabular, m)?)?;
16    m.add_function(wrap_pyfunction!(convert, m)?)?;
17    m.add_function(wrap_pyfunction!(delete, m)?)?;
18
19    // forecasters
20    m.add_class::<TotoModel>()?;
21    m.add_class::<ChronosModel>()?;
22    m.add_class::<TimesFmModel>()?;
23    m.add_class::<SundialModel>()?;
24    m.add_class::<TtmModel>()?;
25    m.add_class::<LagLlamaModel>()?;
26    m.add_class::<MomentModel>()?;
27    m.add_class::<MoiraiModel>()?;
28    m.add_class::<Moirai2Model>()?;
29    m.add_class::<FlowStateModel>()?;
30    m.add_class::<TirexModel>()?;
31
32    // tabular
33    m.add_class::<MitraModel>()?;
34    m.add_class::<TabDptModel>()?;
35    m.add_class::<TabIclModel>()?;
36    m.add_class::<TabPfnModel>()?;
37    m.add_class::<TabFmModel>()?;
38    Ok(())
39}
40
41#[pyfunction]
42fn list_models() -> Vec<&'static str> {
43    vec![
44        "toto", "chronos", "timesfm", "sundial", "ttm", "lag_llama", "moment", "moirai",
45        "moirai2", "flowstate", "tirex", "mitra", "tabdpt", "tabicl", "tabpfn", "tabfm",
46    ]
47}
48#[pyfunction]
49fn list_forecasters() -> Vec<&'static str> {
50    vec![
51        "toto", "chronos", "timesfm", "sundial", "ttm", "lag_llama", "moment", "moirai",
52        "moirai2", "flowstate", "tirex",
53    ]
54}
55#[pyfunction]
56fn list_tabular() -> Vec<&'static str> {
57    vec!["mitra", "tabdpt", "tabicl", "tabpfn", "tabfm"]
58}
59
60// ---------------------------------------------------------------------------
61// Helpers
62// ---------------------------------------------------------------------------
63
64fn dtype_from_str(s: &str) -> PyResult<zsfm_gguf::GGMLType> {
65    match s {
66        "f32" => Ok(zsfm_gguf::GGMLType::F32),
67        "f16" => Ok(zsfm_gguf::GGMLType::F16),
68        "q8" => Ok(zsfm_gguf::GGMLType::Q8_0),
69        "bf16" => Ok(zsfm_gguf::GGMLType::BF16),
70        _ => Err(pyo3::exceptions::PyValueError::new_err(format!(
71            "unknown dtype {s:?}: expected one of f32, f16, q8, bf16"
72        ))),
73    }
74}
75fn delete_cached_model(canonical: &std::path::Path, output: Option<&std::path::Path>) -> anyhow::Result<()> {
76    let mut deleted_any = false;
77    if let Some(cache_dir) = canonical.parent() {
78        if cache_dir.exists() {
79            let count = walk_file_count(cache_dir);
80            std::fs::remove_dir_all(cache_dir)?;
81            println!("Deleted cache directory {} ({} files)", cache_dir.display(), count);
82            deleted_any = true;
83        } else if canonical.exists() {
84            std::fs::remove_file(canonical)?;
85            println!("Deleted {}", canonical.display());
86            deleted_any = true;
87        } else {
88            println!("No cache found at {} (already deleted?)", cache_dir.display());
89        }
90    } else if canonical.exists() {
91        std::fs::remove_file(canonical)?;
92        println!("Deleted {}", canonical.display());
93        deleted_any = true;
94    }
95    if let Some(out) = output {
96        if out.exists() {
97            std::fs::remove_file(out)?;
98            println!("Deleted output {}", out.display());
99            deleted_any = true;
100        } else {
101            println!("Output file not found: {} (already deleted?)", out.display());
102        }
103    }
104    if !deleted_any {
105        println!("Nothing to delete.");
106    }
107    Ok(())
108}
109
110fn walk_file_count(dir: &std::path::Path) -> usize {
111    let mut count = 0;
112    if let Ok(entries) = std::fs::read_dir(dir) {
113        for entry in entries.flatten() {
114            let path = entry.path();
115            if path.is_dir() {
116                count += walk_file_count(&path);
117            } else {
118                count += 1;
119            }
120        }
121    }
122    count
123}
124
125
126// Generic convert/delete that dispatch by model name — mirrors `zsfm <model> convert/delete`.
127#[pyfunction]
128#[pyo3(signature = (model, output=None, dtype="f16", model_dir="models", token=None, redownload=false, task=None, filename=None))]
129fn convert(
130    model: &str,
131    output: Option<String>,
132    dtype: &str,
133    model_dir: &str,
134    token: Option<String>,
135    redownload: bool,
136    task: Option<String>,
137    filename: Option<String>,
138) -> PyResult<()> {
139    let dtype_ty = dtype_from_str(dtype)?;
140    let model_dir = PathBuf::from(model_dir);
141    let rt = tokio::runtime::Runtime::new().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
142    rt.block_on(async {
143        match model {
144            "toto" => {
145                let repo = "Datadog/Toto-2.0-2.5B";
146                let out = output.map(PathBuf::from).unwrap_or_else(|| PathBuf::from("gguf/toto-2.5b-f16.gguf"));
147                let canonical = zsfm_hub::canonical_gguf_path(&model_dir, repo);
148                if canonical.exists() && !redownload {
149                    zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
150                    return Ok(());
151                }
152                let files = zsfm_hub::download_model(repo, token.as_deref(), &model_dir).await.map_err(|e| anyhow::anyhow!(e.to_string()))?;
153                let s = std::fs::read_to_string(&files.config_json).map_err(|e| anyhow::anyhow!(e.to_string()))?;
154                let cfg = zsfm_toto::config::TotoConfig::from_json(&s).map_err(|e| anyhow::anyhow!(e.to_string()))?;
155                zsfm_toto::convert::convert(repo, &files, &cfg, &zsfm_toto::convert::ConvertOptions { output_dtype: zsfm_gguf::GGMLType::F32 }, &canonical).map_err(|e| anyhow::anyhow!(e.to_string()))?;
156                files.cleanup_weights();
157                zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
158                Ok::<(), anyhow::Error>(())
159            }
160            "chronos" => {
161                let repo = "amazon/chronos-2";
162                let out = output.map(PathBuf::from).unwrap_or_else(|| PathBuf::from("gguf/chronos-f16.gguf"));
163                let canonical = zsfm_hub::canonical_gguf_path(&model_dir, repo);
164                if canonical.exists() && !redownload {
165                    zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
166                    return Ok(());
167                }
168                let files = zsfm_hub::download_model(repo, token.as_deref(), &model_dir).await.map_err(|e| anyhow::anyhow!(e.to_string()))?;
169                let s = std::fs::read_to_string(&files.config_json).map_err(|e| anyhow::anyhow!(e.to_string()))?;
170                let cfg = zsfm_chronos::config::Chronos2Config::from_json(&s).map_err(|e| anyhow::anyhow!(e.to_string()))?;
171                zsfm_chronos::convert::convert(repo, &files, &cfg, &zsfm_chronos::convert::ConvertOptions { output_dtype: zsfm_gguf::GGMLType::F32 }, &canonical).map_err(|e| anyhow::anyhow!(e.to_string()))?;
172                files.cleanup_weights();
173                zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
174                Ok(())
175            }
176            "timesfm" => {
177                let repo = "google/timesfm-2.5-200m-pytorch";
178                let out = output.map(PathBuf::from).unwrap_or_else(|| PathBuf::from("gguf/timesfm.gguf"));
179                let canonical = zsfm_hub::canonical_gguf_path(&model_dir, repo);
180                if canonical.exists() && !redownload {
181                    zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
182                    return Ok(());
183                }
184                let files = zsfm_hub::download_model(repo, token.as_deref(), &model_dir).await.map_err(|e| anyhow::anyhow!(e.to_string()))?;
185                let cfg = zsfm_timesfm::config::TimesFMConfig::new();
186                zsfm_timesfm::convert::convert(repo, &files, &cfg, &zsfm_timesfm::convert::ConvertOptions { output_dtype: zsfm_gguf::GGMLType::F32 }, &canonical).map_err(|e| anyhow::anyhow!(e.to_string()))?;
187                files.cleanup_weights();
188                zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
189                Ok(())
190            }
191            "sundial" => {
192                let repo = "thuml/sundial-base-128m";
193                let out = output.map(PathBuf::from).unwrap_or_else(|| PathBuf::from("gguf/sundial-f16.gguf"));
194                let canonical = zsfm_hub::canonical_gguf_path(&model_dir, repo);
195                if canonical.exists() && !redownload {
196                    zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
197                    return Ok(());
198                }
199                let files = zsfm_hub::download_model(repo, token.as_deref(), &model_dir).await.map_err(|e| anyhow::anyhow!(e.to_string()))?;
200                let s = std::fs::read_to_string(&files.config_json).map_err(|e| anyhow::anyhow!(e.to_string()))?;
201                let cfg: zsfm_sundial::config::SundialConfig = serde_json::from_str(&s).map_err(|e| anyhow::anyhow!(e.to_string()))?;
202                zsfm_sundial::convert::convert(repo, &files, &cfg, &zsfm_sundial::convert::ConvertOptions { output_dtype: zsfm_gguf::GGMLType::F32 }, &canonical).map_err(|e| anyhow::anyhow!(e.to_string()))?;
203                files.cleanup_weights();
204                zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
205                Ok(())
206            }
207            "ttm" => {
208                let repo = "ibm-granite/granite-timeseries-ttm-r2";
209                let out = output.map(PathBuf::from).unwrap_or_else(|| PathBuf::from("gguf/ttm-f32.gguf"));
210                let canonical = zsfm_hub::canonical_gguf_path(&model_dir, repo);
211                if canonical.exists() && !redownload {
212                    zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
213                    return Ok(());
214                }
215                let files = zsfm_hub::download_model(repo, token.as_deref(), &model_dir).await.map_err(|e| anyhow::anyhow!(e.to_string()))?;
216                let s = std::fs::read_to_string(&files.config_json).map_err(|e| anyhow::anyhow!(e.to_string()))?;
217                let cfg = zsfm_ttm::config::TtmConfig::from_json(&s).map_err(|e| anyhow::anyhow!(e.to_string()))?;
218                zsfm_ttm::convert::convert(repo, &files, &cfg, &zsfm_ttm::convert::ConvertOptions { output_dtype: zsfm_gguf::GGMLType::F32 }, &canonical).map_err(|e| anyhow::anyhow!(e.to_string()))?;
219                files.cleanup_weights();
220                zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
221                Ok(())
222            }
223            "lag_llama" | "lag-llama" => {
224                let repo = "time-series-foundation-models/Lag-Llama";
225                let out = output.map(PathBuf::from).unwrap_or_else(|| PathBuf::from("gguf/lag_llama-f32.gguf"));
226                let canonical = zsfm_hub::canonical_gguf_path(&model_dir, repo);
227                if canonical.exists() && !redownload {
228                    zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
229                    return Ok(());
230                }
231                if let Some(p) = canonical.parent() { std::fs::create_dir_all(p).map_err(|e| anyhow::anyhow!(e.to_string()))?; }
232                let ckpt = zsfm_hub::download_file(repo, "lag-llama.ckpt", token.as_deref(), &model_dir).await.map_err(|e| anyhow::anyhow!(e.to_string()))?;
233                let cfg = zsfm_lag_llama::config::LagLlamaConfig::default_from_ckpt();
234                zsfm_lag_llama::convert::convert(&ckpt, &cfg, &zsfm_lag_llama::convert::ConvertOptions { output_dtype: zsfm_gguf::GGMLType::F32 }, &canonical).map_err(|e| anyhow::anyhow!(e.to_string()))?;
235                let _ = std::fs::remove_file(&ckpt);
236                zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
237                Ok(())
238            }
239            "moment" => {
240                let repo = "AutonLab/MOMENT-1-large";
241                let out = output.map(PathBuf::from).unwrap_or_else(|| PathBuf::from("gguf/moment-f32.gguf"));
242                let canonical = zsfm_hub::canonical_gguf_path(&model_dir, repo);
243                if canonical.exists() && !redownload {
244                    zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
245                    return Ok(());
246                }
247                let files = zsfm_hub::download_model(repo, token.as_deref(), &model_dir).await.map_err(|e| anyhow::anyhow!(e.to_string()))?;
248                let cfg = zsfm_moment::config::MomentConfig::default();
249                zsfm_moment::convert::convert(&files.safetensors_shards, &cfg, &zsfm_moment::convert::ConvertOptions { output_dtype: zsfm_gguf::GGMLType::F32 }, &canonical).map_err(|e| anyhow::anyhow!(e.to_string()))?;
250                files.cleanup_weights();
251                zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
252                Ok(())
253            }
254            "moirai" => {
255                let repo = "Salesforce/moirai-1.0-R-large";
256                let out = output.map(PathBuf::from).unwrap_or_else(|| PathBuf::from("gguf/moirai-f32.gguf"));
257                let canonical = zsfm_hub::canonical_gguf_path(&model_dir, repo);
258                if canonical.exists() && !redownload {
259                    zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
260                    return Ok(());
261                }
262                let files = zsfm_hub::download_model(repo, token.as_deref(), &model_dir).await.map_err(|e| anyhow::anyhow!(e.to_string()))?;
263                let cfg = zsfm_moirai::config::MoiraiConfig::default();
264                zsfm_moirai::convert::convert(&files.safetensors_shards, &cfg, &zsfm_moirai::convert::ConvertOptions { output_dtype: zsfm_gguf::GGMLType::F32 }, &canonical).map_err(|e| anyhow::anyhow!(e.to_string()))?;
265                files.cleanup_weights();
266                zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
267                Ok(())
268            }
269            "moirai2" | "moirai-2" => {
270                let repo = "Salesforce/moirai-2.0-R-small";
271                let out = output.map(PathBuf::from).unwrap_or_else(|| PathBuf::from("gguf/moirai2-f32.gguf"));
272                let canonical = zsfm_hub::canonical_gguf_path(&model_dir, repo);
273                if canonical.exists() && !redownload {
274                    zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
275                    return Ok(());
276                }
277                let files = zsfm_hub::download_model(repo, token.as_deref(), &model_dir).await.map_err(|e| anyhow::anyhow!(e.to_string()))?;
278                let cfg = zsfm_moirai2::config::Moirai2Config::default();
279                zsfm_moirai2::convert::convert(&files.safetensors_shards, &cfg, &zsfm_moirai2::convert::ConvertOptions { output_dtype: zsfm_gguf::GGMLType::F32 }, &canonical).map_err(|e| anyhow::anyhow!(e.to_string()))?;
280                files.cleanup_weights();
281                zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
282                Ok(())
283            }
284            "flowstate" | "flowstate-r1" => {
285                let repo = "ibm-granite/granite-timeseries-flowstate-r1";
286                let out = output.map(PathBuf::from).unwrap_or_else(|| PathBuf::from("gguf/flowstate-r1-f16.gguf"));
287                let canonical = zsfm_hub::canonical_gguf_path(&model_dir, repo);
288                if canonical.exists() && !redownload {
289                    zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
290                    return Ok(());
291                }
292                let files = zsfm_hub::download_model(repo, token.as_deref(), &model_dir).await.map_err(|e| anyhow::anyhow!(e.to_string()))?;
293                let s = std::fs::read_to_string(&files.config_json).map_err(|e| anyhow::anyhow!(e.to_string()))?;
294                let cfg = zsfm_flowstate::config::FlowStateConfig::from_json(&s).map_err(|e| anyhow::anyhow!(e.to_string()))?;
295                zsfm_flowstate::convert::convert(repo, &files, &cfg, &zsfm_flowstate::convert::ConvertOptions { output_dtype: zsfm_gguf::GGMLType::F32 }, &canonical).map_err(|e| anyhow::anyhow!(e.to_string()))?;
296                files.cleanup_weights();
297                zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
298                Ok(())
299            }
300            "tirex" => {
301                let repo = "NX-AI/TiRex";
302                let out = output.map(PathBuf::from).unwrap_or_else(|| PathBuf::from("gguf/tirex-f32.gguf"));
303                let canonical = zsfm_hub::canonical_gguf_path(&model_dir, repo);
304                if canonical.exists() && !redownload {
305                    zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
306                    return Ok(());
307                }
308                if let Some(p) = canonical.parent() { std::fs::create_dir_all(p).map_err(|e| anyhow::anyhow!(e.to_string()))?; }
309                let ckpt = zsfm_hub::download_file(repo, "model.ckpt", token.as_deref(), &model_dir).await.map_err(|e| anyhow::anyhow!(e.to_string()))?;
310                let cfg = zsfm_tirex::config::TiRexConfig::default_from_ckpt();
311                zsfm_tirex::convert::convert(&ckpt, &cfg, &zsfm_tirex::convert::ConvertOptions { output_dtype: zsfm_gguf::GGMLType::F32 }, &canonical).map_err(|e| anyhow::anyhow!(e.to_string()))?;
312                let _ = std::fs::remove_file(&ckpt);
313                zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
314                Ok(())
315            }
316            "mitra" => {
317                let task_str = task.as_deref().unwrap_or("classification");
318                let repo = if task_str == "regression" { "autogluon/mitra-regressor" } else { "autogluon/mitra-classifier" };
319                let variant_dir = model_dir.join(format!("mitra-{task_str}"));
320                let out = output.map(PathBuf::from).unwrap_or_else(|| PathBuf::from(format!("gguf/mitra-{task_str}-{dtype}.gguf")));
321                let canonical = zsfm_hub::canonical_gguf_path(&variant_dir, repo);
322                if canonical.exists() && !redownload {
323                    zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
324                    return Ok(());
325                }
326                let files = zsfm_hub::download_model(repo, token.as_deref(), &variant_dir).await.map_err(|e| anyhow::anyhow!(e.to_string()))?;
327                let cfg = if task_str == "regression" { zsfm_mitra::config::MitraConfig::regressor() } else { zsfm_mitra::config::MitraConfig::classifier() };
328                let p = files.safetensors_shards.first().ok_or_else(|| anyhow::anyhow!("no shard"))?;
329                zsfm_mitra::convert::convert(std::slice::from_ref(p), &cfg, &zsfm_mitra::convert::ConvertOptions { output_dtype: zsfm_gguf::GGMLType::F32 }, &canonical).map_err(|e| anyhow::anyhow!(e.to_string()))?;
330                files.cleanup_weights();
331                zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
332                Ok(())
333            }
334            "tabdpt" => {
335                let repo = "Layer6/TabDPT";
336                let fname = filename.as_deref().unwrap_or("tabdpt1_2.safetensors");
337                let out = output.map(PathBuf::from).unwrap_or_else(|| PathBuf::from(format!("gguf/tabdpt-{dtype}.gguf")));
338                let canonical = zsfm_hub::canonical_gguf_path(&model_dir, repo);
339                if canonical.exists() && !redownload {
340                    zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
341                    return Ok(());
342                }
343                if let Some(p) = canonical.parent() { std::fs::create_dir_all(p).map_err(|e| anyhow::anyhow!(e.to_string()))?; }
344                let p = zsfm_hub::download_file(repo, fname, token.as_deref(), &model_dir).await.map_err(|e| anyhow::anyhow!(e.to_string()))?;
345                let cfg = zsfm_tabdpt::config::TabDptConfig::default_v1_2();
346                zsfm_tabdpt::convert::convert(std::slice::from_ref(&p), &cfg, &zsfm_tabdpt::convert::ConvertOptions { output_dtype: zsfm_gguf::GGMLType::F32 }, &canonical).map_err(|e| anyhow::anyhow!(e.to_string()))?;
347                let _ = std::fs::remove_file(&p);
348                zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
349                Ok(())
350            }
351            "tabfm" => {
352                let task_str = task.as_deref().unwrap_or("classification");
353                let repo = "google/tabfm-1.0.0-pytorch";
354                let variant_dir = model_dir.join(format!("tabfm-{task_str}"));
355                let out = output.map(PathBuf::from).unwrap_or_else(|| PathBuf::from(format!("gguf/tabfm-{task_str}-{dtype}.gguf")));
356                let canonical = zsfm_hub::canonical_gguf_path(&variant_dir, repo);
357                if canonical.exists() && !redownload {
358                    zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
359                    return Ok(());
360                }
361                let files = zsfm_hub::download_model_prefixed(repo, task_str, token.as_deref(), &variant_dir).await.map_err(|e| anyhow::anyhow!(e.to_string()))?;
362                let s = std::fs::read_to_string(&files.config_json).map_err(|e| anyhow::anyhow!(e.to_string()))?;
363                let cfg = zsfm_tabfm::config::TabFMConfig::from_json(&s).map_err(|e| anyhow::anyhow!(e.to_string()))?;
364                let p = files.safetensors_shards.first().ok_or_else(|| anyhow::anyhow!("no shard"))?;
365                zsfm_tabfm::convert::convert(repo, p, &cfg, &zsfm_tabfm::convert::ConvertOptions { output_dtype: zsfm_gguf::GGMLType::F32 }, &canonical).map_err(|e| anyhow::anyhow!(e.to_string()))?;
366                files.cleanup_weights();
367                zsfm_checkpoint::recast(&canonical, &out, dtype_ty).map_err(|e| anyhow::anyhow!(e.to_string()))?;
368                Ok(())
369            }
370            other => Err(anyhow::anyhow!("unknown model {other:?}; expected one of {}", list_models().join(", ")))
371        }
372    }).map_err(|e: anyhow::Error| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
373    Ok(())
374}
375
376#[pyfunction]
377#[pyo3(signature = (model, model_dir="models", output=None))]
378fn delete(model: &str, model_dir: &str, output: Option<String>) -> PyResult<()> {
379    // Map friendly names to canonical repo ids for per-model delete
380    let (repo, variant_dir): (String, Option<String>) = match model {
381        "toto" => ("Datadog/Toto-2.0-2.5B".into(), None),
382        "chronos" => ("amazon/chronos-2".into(), None),
383        "timesfm" => ("google/timesfm-2.5-200m-pytorch".into(), None),
384        "sundial" => ("thuml/sundial-base-128m".into(), None),
385        "ttm" => ("ibm-granite/granite-timeseries-ttm-r2".into(), None),
386        "lag_llama" | "lag-llama" => ("time-series-foundation-models/Lag-Llama".into(), None),
387        "moment" => ("AutonLab/MOMENT-1-large".into(), None),
388        "moirai" => ("Salesforce/moirai-1.0-R-large".into(), None),
389        "moirai2" | "moirai-2" => ("Salesforce/moirai-2.0-R-small".into(), None),
390        "flowstate" | "flowstate-r1" => ("ibm-granite/granite-timeseries-flowstate-r1".into(), None),
391        "tirex" => ("NX-AI/TiRex".into(), None),
392        "mitra" | "mitra-classification" => ("autogluon/mitra-classifier".into(), Some("mitra-classification".into())),
393        "mitra-regression" => ("autogluon/mitra-regressor".into(), Some("mitra-regression".into())),
394        "tabdpt" => ("Layer6/TabDPT".into(), None),
395        "tabicl" => ("jingang/TabICL".into(), None),
396        "tabpfn" => ("Prior-Labs/tabpfn_3".into(), None),
397        "tabfm" | "tabfm-classification" => ("google/tabfm-1.0.0-pytorch".into(), Some("tabfm-classification".into())),
398        "tabfm-regression" => ("google/tabfm-1.0.0-pytorch".into(), Some("tabfm-regression".into())),
399        // also accept raw repo ids
400        _ if model.contains('/') => (model.to_string(), None),
401        _ => return Err(pyo3::exceptions::PyValueError::new_err(format!("unknown model {model:?}"))),
402    };
403    let model_dir = PathBuf::from(model_dir);
404    let canonical = if let Some(v) = variant_dir {
405        let vd = model_dir.join(v);
406        zsfm_hub::canonical_gguf_path(&vd, &repo)
407    } else {
408        zsfm_hub::canonical_gguf_path(&model_dir, &repo)
409    };
410    let out = output.map(PathBuf::from);
411    delete_cached_model(&canonical, out.as_deref()).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
412    Ok(())
413}
414
415// ---------------------------------------------------------------------------
416// Forecasters — one #[pyclass] per model (point or quantile)
417// ---------------------------------------------------------------------------
418
419use zsfm_toto::config::TotoConfig;
420use zsfm_toto::infer::TotoModel as RustTotoModel;
421
422#[pyclass]
423struct TotoModel {
424    inner: RustTotoModel,
425    _gguf: String,
426}
427
428#[pymethods]
429impl TotoModel {
430    #[new]
431    #[pyo3(signature = (gguf, config=None, context_length=None, use_f64=false))]
432    fn new(gguf: String, config: Option<String>, context_length: Option<usize>, use_f64: bool) -> PyResult<Self> {
433        let gguf_path = PathBuf::from(&gguf);
434        let cfg_path = config.unwrap_or_else(|| "models/Datadog__Toto-2.0-2.5B/config.json".into());
435        let s = std::fs::read_to_string(&cfg_path).map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("{cfg_path}: {e}")))?;
436        let v: serde_json::Value = serde_json::from_str(&s).map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
437        let max_ctx = context_length.unwrap_or(4096);
438        let inner = RustTotoModel::builder(&gguf_path)
439            .config_json(&v)
440            .with_compute_f64(use_f64)
441            .build()
442            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("load Toto {gguf}: {e}")))?;
443        // Store max_ctx via a wrapper? For now keep simple: the Rust model already encodes patch_size; we trim in forecast.
444        let _ = max_ctx;
445        Ok(Self { inner, _gguf: gguf })
446    }
447    /// Forecast from a univariate context. For batch/multivariate, call
448    /// `forecast_batch` (or loop). Returns the point forecast (median).
449    fn forecast(&self, context: Vec<f32>, horizon: usize) -> PyResult<Vec<f32>> {
450        let data = vec![context.clone()];
451        let mask = vec![vec![true; context.len()]];
452        let qmat = self.inner.forecast(&data, &mask, horizon).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
453        let median_idx = 4; // q0.5
454        Ok(qmat[median_idx][0].clone())
455    }
456
457    /// Batch forecast: context is List[List[float]] with shape [batch][time].
458    fn forecast_batch(&self, contexts: Vec<Vec<f32>>, horizon: usize) -> PyResult<Vec<Vec<f32>>> {
459        let mut out = Vec::new();
460        for ctx in contexts {
461            let data = vec![ctx.clone()];
462            let mask = vec![vec![true; ctx.len()]];
463            let qmat = self.inner.forecast(&data, &mask, horizon).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
464            out.push(qmat[4][0].clone());
465        }
466        Ok(out)
467    }
468
469    /// Full quantile matrix for a univariate context: one row per quantile level
470    /// (see `quantiles()` for the levels, in the same order), each `horizon` long.
471    fn forecast_quantiles(&self, context: Vec<f32>, horizon: usize) -> PyResult<Vec<Vec<f32>>> {
472        let data = vec![context.clone()];
473        let mask = vec![vec![true; context.len()]];
474        let qmat = self.inner.forecast(&data, &mask, horizon).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
475        Ok(qmat.into_iter().map(|q| q[0].clone()).collect())
476    }
477
478    /// The 9 quantile levels each row of `forecast_quantiles()` corresponds to.
479    fn quantiles(&self) -> Vec<f32> {
480        vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
481    }
482}
483
484use zsfm_timesfm::infer::TimesFMModel as RustTimesFmModel;
485#[pyclass]
486struct TimesFmModel { inner: RustTimesFmModel }
487#[pymethods]
488impl TimesFmModel {
489    #[new]
490    fn new(gguf: String) -> PyResult<Self> {
491        let p = PathBuf::from(&gguf);
492        let inner = RustTimesFmModel::load(&p).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("load TimesFM {gguf}: {e}")))?;
493        Ok(Self { inner })
494    }
495    fn forecast(&self, context: Vec<f32>, horizon: usize) -> PyResult<Vec<f32>> {
496        let out = self.inner.forecast(&context, horizon).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
497        Ok(out.into_iter().next().unwrap_or_default())
498    }
499
500    /// Full quantile matrix: one row per quantile level (see `quantiles()`), each `horizon` long.
501    /// TimesFM's point forecast (from `forecast()`) is a separate dedicated model output, not
502    /// derived from these quantiles.
503    fn forecast_quantiles(&self, context: Vec<f32>, horizon: usize) -> PyResult<Vec<Vec<f32>>> {
504        let mut out = self.inner.forecast(&context, horizon).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
505        if !out.is_empty() {
506            out.remove(0); // drop the dedicated point-forecast row, keep q0.1..q0.9
507        }
508        Ok(out)
509    }
510
511    /// The 9 quantile levels each row of `forecast_quantiles()` corresponds to.
512    fn quantiles(&self) -> Vec<f32> {
513        vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
514    }
515}
516
517use zsfm_sundial::infer::SundialModel as RustSundialModel;
518#[pyclass]
519struct SundialModel { inner: RustSundialModel }
520#[pymethods]
521impl SundialModel {
522    #[new]
523    fn new(gguf: String) -> PyResult<Self> {
524        let p = PathBuf::from(&gguf);
525        let inner = RustSundialModel::builder(&p).build().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("load Sundial {gguf}: {e}")))?;
526        Ok(Self { inner })
527    }
528    fn forecast(&self, context: Vec<f32>, horizon: usize) -> PyResult<Vec<f32>> {
529        let raw = self.inner.forecast(&context, &candle_core::Device::Cpu).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
530        Ok(raw.into_iter().take(horizon).collect())
531    }
532}
533
534use zsfm_ttm::config::TtmConfig;
535use zsfm_ttm::infer::TtmModel as RustTtmModel;
536
537#[pyclass]
538struct TtmModel {
539    inner: RustTtmModel,
540    _config_path: Option<String>,
541}
542
543#[pymethods]
544impl TtmModel {
545    #[new]
546    #[pyo3(signature = (gguf, config=None))]
547    fn new(gguf: String, config: Option<String>) -> PyResult<Self> {
548        let gguf_path = PathBuf::from(&gguf);
549        let ttm_config = if let Some(cfg_path) = &config {
550            let s = std::fs::read_to_string(cfg_path).map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("read config {cfg_path}: {e}")))?;
551            TtmConfig::from_json(&s).map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?
552        } else {
553            let canonical = PathBuf::from("models/ibm-granite__granite-timeseries-ttm-r2/config.json");
554            if canonical.exists() {
555                let s = std::fs::read_to_string(&canonical).map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))?;
556                TtmConfig::from_json(&s).map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?
557            } else {
558                return Err(pyo3::exceptions::PyFileNotFoundError::new_err(format!("config not found at {canonical:?} and no `config` arg given")));
559            }
560        };
561        let inner = RustTtmModel::builder(&gguf_path).config(ttm_config).build().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("load TTM {gguf}: {e}")))?;
562        Ok(Self { inner, _config_path: config })
563    }
564    fn forecast(&self, context: Vec<f32>, horizon: usize) -> PyResult<Vec<f32>> {
565        let out = self.inner.forecast(&context).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
566        Ok(out.into_iter().take(horizon).collect())
567    }
568}
569
570use zsfm_lag_llama::config::LagLlamaConfig;
571use zsfm_lag_llama::infer::LagLlamaModel as RustLagLlamaModel;
572#[pyclass]
573struct LagLlamaModel { inner: RustLagLlamaModel }
574#[pymethods]
575impl LagLlamaModel {
576    #[new]
577    fn new(gguf: String) -> PyResult<Self> {
578        let p = PathBuf::from(&gguf);
579        let cfg = LagLlamaConfig::default_from_ckpt();
580        let inner = RustLagLlamaModel::load(&p, cfg).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("load LagLlama {gguf}: {e}")))?;
581        Ok(Self { inner })
582    }
583    fn forecast(&self, context: Vec<f32>, horizon: usize) -> PyResult<Vec<f32>> {
584        let out = self.inner.forecast(&context, horizon).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
585        Ok(out)
586    }
587}
588
589use zsfm_moment::config::MomentConfig;
590use zsfm_moment::infer::MomentModel as RustMomentModel;
591#[pyclass]
592struct MomentModel { inner: RustMomentModel }
593#[pymethods]
594impl MomentModel {
595    #[new]
596    fn new(gguf: String) -> PyResult<Self> {
597        let p = PathBuf::from(&gguf);
598        let cfg = MomentConfig::default();
599        let inner = RustMomentModel::load(&p, cfg).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("load Moment {gguf}: {e}")))?;
600        Ok(Self { inner })
601    }
602    fn forecast(&self, context: Vec<f32>, horizon: usize) -> PyResult<Vec<f32>> {
603        let out = self.inner.forecast(&context, horizon).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
604        Ok(out)
605    }
606}
607
608use zsfm_moirai::config::MoiraiConfig;
609use zsfm_moirai::infer::MoiraiModel as RustMoiraiModel;
610#[pyclass]
611struct MoiraiModel { inner: RustMoiraiModel }
612#[pymethods]
613impl MoiraiModel {
614    #[new]
615    fn new(gguf: String) -> PyResult<Self> {
616        let p = PathBuf::from(&gguf);
617        let cfg = MoiraiConfig::default();
618        let inner = RustMoiraiModel::load(&p, cfg).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("load Moirai {gguf}: {e}")))?;
619        Ok(Self { inner })
620    }
621    fn forecast(&self, context: Vec<f32>, horizon: usize) -> PyResult<Vec<f32>> {
622        let out = self.inner.forecast(&context, horizon).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
623        Ok(out)
624    }
625}
626
627use zsfm_moirai2::config::Moirai2Config;
628use zsfm_moirai2::infer::Moirai2Model as RustMoirai2Model;
629#[pyclass]
630struct Moirai2Model { inner: RustMoirai2Model }
631#[pymethods]
632impl Moirai2Model {
633    #[new]
634    fn new(gguf: String) -> PyResult<Self> {
635        let p = PathBuf::from(&gguf);
636        let cfg = Moirai2Config::default();
637        let inner = RustMoirai2Model::load(&p, cfg).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("load Moirai2 {gguf}: {e}")))?;
638        Ok(Self { inner })
639    }
640    fn forecast(&self, context: Vec<f32>, horizon: usize) -> PyResult<Vec<f32>> {
641        let out = self.inner.forecast(&context, horizon).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
642        Ok(out)
643    }
644}
645
646use zsfm_flowstate::config::FlowStateConfig;
647use zsfm_flowstate::infer::FlowStateModel as RustFlowStateModel;
648#[pyclass]
649struct FlowStateModel { inner: RustFlowStateModel }
650#[pymethods]
651impl FlowStateModel {
652    #[new]
653    #[pyo3(signature = (gguf, config=None))]
654    fn new(gguf: String, config: Option<String>) -> PyResult<Self> {
655        let p = PathBuf::from(&gguf);
656        let cfg_path = config.map(PathBuf::from).unwrap_or_else(|| PathBuf::from("models/ibm-granite__granite-timeseries-flowstate-r1/config.json"));
657        let s = std::fs::read_to_string(&cfg_path).map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("read {}: {e}", cfg_path.display())))?;
658        let cfg = FlowStateConfig::from_json(&s).map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
659        let inner = RustFlowStateModel::builder(&p).config_from(&cfg).build().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("load FlowState {gguf}: {e}")))?;
660        Ok(Self { inner })
661    }
662    fn forecast(&self, context: Vec<f32>, horizon: usize) -> PyResult<Vec<f32>> {
663        let qmat = self.inner.forecast(&context, horizon).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
664        let median_idx = self.inner.config.median_index();
665        Ok(qmat[median_idx].clone())
666    }
667
668    /// Full quantile matrix: one row per quantile level (see `quantiles()`), each `horizon` long.
669    fn forecast_quantiles(&self, context: Vec<f32>, horizon: usize) -> PyResult<Vec<Vec<f32>>> {
670        self.inner.forecast(&context, horizon).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
671    }
672
673    /// The quantile levels each row of `forecast_quantiles()` corresponds to (from config.json).
674    fn quantiles(&self) -> Vec<f32> {
675        self.inner.config.quantiles().to_vec()
676    }
677}
678
679use zsfm_tirex::config::TiRexConfig;
680use zsfm_tirex::infer::TiRexModel as RustTirexModel;
681#[pyclass]
682struct TirexModel { inner: RustTirexModel }
683#[pymethods]
684impl TirexModel {
685    #[new]
686    fn new(gguf: String) -> PyResult<Self> {
687        let p = PathBuf::from(&gguf);
688        let cfg = TiRexConfig::default_from_ckpt();
689        let inner = RustTirexModel::load(&p, cfg).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("load TiRex {gguf}: {e}")))?;
690        Ok(Self { inner })
691    }
692    /// Returns the median (q0.5) forecast. Despite its old internal name, the second
693    /// element of the underlying Rust `forecast()` tuple is the median quantile row,
694    /// not a separate mean statistic — `zsfm_tirex::infer::TiRexModel::forecast` docs it
695    /// explicitly.
696    fn forecast(&self, context: Vec<f32>, horizon: usize) -> PyResult<Vec<f32>> {
697        let (_quantiles, median) = self.inner.forecast(&context, horizon).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
698        Ok(median)
699    }
700
701    /// Full quantile matrix: one row per quantile level (see `quantiles()`), each `horizon` long.
702    fn forecast_quantiles(&self, context: Vec<f32>, horizon: usize) -> PyResult<Vec<Vec<f32>>> {
703        let (quantiles, _median) = self.inner.forecast(&context, horizon).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
704        Ok(quantiles)
705    }
706
707    /// The quantile levels each row of `forecast_quantiles()` corresponds to.
708    fn quantiles(&self) -> Vec<f32> {
709        TiRexConfig::default_from_ckpt().quantiles
710    }
711}
712
713use zsfm_chronos::config::Chronos2Config;
714use zsfm_chronos::infer::ChronosModel as RustChronosModel;
715
716#[pyclass]
717struct ChronosModel { inner: RustChronosModel }
718#[pymethods]
719impl ChronosModel {
720    #[new]
721    #[pyo3(signature = (gguf, config=None))]
722    fn new(gguf: String, config: Option<String>) -> PyResult<Self> {
723        let gguf_path = PathBuf::from(&gguf);
724        let cfg_path = config.map(PathBuf::from).unwrap_or_else(|| PathBuf::from("models/amazon__chronos-2/config.json"));
725        let s = std::fs::read_to_string(&cfg_path).map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("read {}: {e}", cfg_path.display())))?;
726        let cfg = Chronos2Config::from_json(&s).map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
727        let inner = RustChronosModel::builder(&gguf_path).config_from(&cfg).build().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("load Chronos {gguf}: {e}")))?;
728        Ok(Self { inner })
729    }
730    fn forecast(&self, context: Vec<f32>, horizon: usize) -> PyResult<Vec<f32>> {
731        let qmat = self.inner.forecast(&context, horizon).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
732        let levels = self.inner.config.quantiles();
733        let median_idx = levels.iter().position(|&q| (q - 0.5).abs() < 1e-6).unwrap_or(levels.len()/2);
734        Ok(qmat[median_idx].clone())
735    }
736    fn forecast_quantiles(&self, context: Vec<f32>, horizon: usize) -> PyResult<Vec<Vec<f32>>> {
737        let qmat = self.inner.forecast(&context, horizon).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
738        Ok(qmat)
739    }
740    fn quantiles(&self) -> Vec<f32> { self.inner.config.quantiles().to_vec() }
741}
742
743// ---------------------------------------------------------------------------
744// Tabular models
745// ---------------------------------------------------------------------------
746
747use zsfm_mitra::config::MitraConfig;
748use zsfm_mitra::MitraModel as RustMitraModel;
749#[pyclass]
750struct MitraModel { inner: RustMitraModel, is_classifier: bool }
751#[pymethods]
752impl MitraModel {
753    #[new]
754    #[pyo3(signature = (gguf, task="classification"))]
755    fn new(gguf: String, task: &str) -> PyResult<Self> {
756        let p = PathBuf::from(&gguf);
757        let cfg = match task {
758            "classification" => MitraConfig::classifier(),
759            "regression" => MitraConfig::regressor(),
760            _ => return Err(pyo3::exceptions::PyValueError::new_err("task must be 'classification' or 'regression'"))
761        };
762        let is_classifier = task == "classification";
763        let inner = RustMitraModel::load(&p, cfg).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("load Mitra {gguf}: {e}")))?;
764        Ok(Self { inner, is_classifier })
765    }
766    fn predict_classification(&self, x_support: Vec<Vec<f32>>, y_support: Vec<usize>, x_query: Vec<Vec<f32>>, n_classes: usize) -> PyResult<Vec<Vec<f32>>> {
767        if !self.is_classifier { return Err(pyo3::exceptions::PyValueError::new_err("model was loaded as regressor, not classifier")) }
768        let logits = self.inner.predict_classification(&x_support, &y_support, &x_query, n_classes).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
769        Ok(logits)
770    }
771    fn predict_regression(&self, x_support: Vec<Vec<f32>>, y_support: Vec<f32>, x_query: Vec<Vec<f32>>) -> PyResult<Vec<f32>> {
772        if self.is_classifier { return Err(pyo3::exceptions::PyValueError::new_err("model was loaded as classifier, not regressor")) }
773        let out = self.inner.predict_regression(&x_support, &y_support, &x_query).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
774        Ok(out)
775    }
776}
777
778use zsfm_tabdpt::config::TabDptConfig;
779use zsfm_tabdpt::TabDptModel as RustTabDptModel;
780#[pyclass]
781struct TabDptModel { inner: RustTabDptModel }
782#[pymethods]
783impl TabDptModel {
784    #[new]
785    fn new(gguf: String) -> PyResult<Self> {
786        let p = PathBuf::from(&gguf);
787        let cfg = TabDptConfig::default_v1_2();
788        let inner = RustTabDptModel::load(&p, cfg).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("load TabDPT {gguf}: {e}")))?;
789        Ok(Self { inner })
790    }
791    fn predict_classification(&self, x_support: Vec<Vec<f32>>, y_support: Vec<usize>, x_query: Vec<Vec<f32>>, n_classes: usize) -> PyResult<Vec<Vec<f32>>> {
792        let out = self.inner.predict_classification(&x_support, &y_support, &x_query, n_classes).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
793        Ok(out)
794    }
795    fn predict_regression(&self, x_support: Vec<Vec<f32>>, y_support: Vec<f32>, x_query: Vec<Vec<f32>>) -> PyResult<Vec<f32>> {
796        let out = self.inner.predict_regression(&x_support, &y_support, &x_query).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
797        Ok(out)
798    }
799}
800
801use zsfm_tabicl::config::TabIclConfig;
802use zsfm_tabicl::TabIclModel as RustTabIclModel;
803#[pyclass]
804struct TabIclModel { inner: RustTabIclModel }
805#[pymethods]
806impl TabIclModel {
807    #[new]
808    fn new(gguf: String) -> PyResult<Self> {
809        let p = PathBuf::from(&gguf);
810        let cfg = TabIclConfig::v2();
811        let inner = RustTabIclModel::load(&p, cfg).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("load TabICL {gguf}: {e}")))?;
812        Ok(Self { inner })
813    }
814    fn predict_classification(&self, x_support: Vec<Vec<f32>>, y_support: Vec<usize>, x_query: Vec<Vec<f32>>, n_classes: usize) -> PyResult<Vec<Vec<f32>>> {
815        let out = self.inner.predict_classification(&x_support, &y_support, &x_query, n_classes).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
816        Ok(out)
817    }
818}
819
820use zsfm_tabpfn::config::TabPfnConfig;
821use zsfm_tabpfn::TabPfnModel as RustTabPfnModel;
822#[pyclass]
823struct TabPfnModel { inner: RustTabPfnModel }
824#[pymethods]
825impl TabPfnModel {
826    #[new]
827    fn new(gguf: String) -> PyResult<Self> {
828        let p = PathBuf::from(&gguf);
829        let cfg = TabPfnConfig::v3_default();
830        let inner = RustTabPfnModel::load(&p, cfg).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("load TabPFN {gguf}: {e}")))?;
831        Ok(Self { inner })
832    }
833    fn predict_classification(&self, x_support: Vec<Vec<f32>>, y_support: Vec<usize>, x_query: Vec<Vec<f32>>, n_classes: usize) -> PyResult<Vec<Vec<f32>>> {
834        let out = self.inner.predict_classification(&x_support, &y_support, &x_query, n_classes).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
835        Ok(out)
836    }
837}
838
839use zsfm_tabfm::config::TabFMConfig;
840use zsfm_tabfm::TabFMModel as RustTabFMModel;
841#[pyclass]
842struct TabFmModel { inner: RustTabFMModel, is_classifier: bool }
843#[pymethods]
844impl TabFmModel {
845    #[new]
846    #[pyo3(signature = (gguf, config=None))]
847    fn new(gguf: String, config: Option<String>) -> PyResult<Self> {
848        let p = PathBuf::from(&gguf);
849        let cfg_path = if let Some(c) = config { PathBuf::from(c) } else {
850            // Auto-detect: try classification then regression config
851            let cand = PathBuf::from("models/tabfm-classification/google__tabfm-1.0.0-pytorch/classification_config.json");
852            if cand.exists() { cand } else { PathBuf::from("models/tabfm-regression/google__tabfm-1.0.0-pytorch/regression_config.json") }
853        };
854        let s = std::fs::read_to_string(&cfg_path).map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("read {}: {e}", cfg_path.display())))?;
855        let cfg = TabFMConfig::from_json(&s).map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?;
856        let is_classifier = cfg.is_classifier;
857        let inner = RustTabFMModel::builder(&p).config_from(&cfg).build().map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("load TabFM {gguf}: {e}")))?;
858        Ok(Self { inner, is_classifier })
859    }
860    fn predict(&self, x: Vec<Vec<f32>>, y: Vec<f32>, train_size: usize) -> PyResult<Vec<Vec<f32>>> {
861        let out = self.inner.predict(&x, &y, train_size, None, None).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
862        Ok(out)
863    }
864}
zsfm/lib.rs

zsfm/
lib.rs
