zsfm/
toto.rs

1use std::path::PathBuf;
2
3use crate::common::DtypeArg;
4use anyhow::Context;
5use clap::Subcommand;
6
7use zsfm_gguf::GGMLType;
8use zsfm_toto::config::TotoConfig;
9use zsfm_toto::convert::{convert, ConvertOptions};
10
11#[derive(Subcommand)]
12pub enum Command {
13    /// Download a Toto model from HuggingFace and convert it to GGUF.
14    ///
15    /// The first conversion downloads the model and writes a canonical F32 GGUF to
16    /// `<model_dir>/<owner>__<name>/model-f32.gguf`, deleting the (large) downloaded
17    /// weight files afterward. Later conversions (any --dtype) recast from that
18    /// cached F32 GGUF instead of re-downloading — pass --redownload to force a
19    /// fresh download anyway (e.g. the repo was updated).
20    Convert {
21        #[arg(short, long, default_value = "Datadog/Toto-2.0-2.5B")]
22        model: String,
23        #[arg(short, long, default_value = "gguf/toto-2.5b-f16.gguf")]
24        output: PathBuf,
25        #[arg(long, default_value = "f16")]
26        dtype: DtypeArg,
27        #[arg(long, default_value = "models")]
28        model_dir: PathBuf,
29        #[arg(long, env = "HF_TOKEN")]
30        token: Option<String>,
31        /// Force a fresh download even if a cached F32 GGUF already exists.
32        #[arg(long)]
33        redownload: bool,
34    },
35    /// Print all tensor names in a local checkpoint file (safetensors, PyTorch
36    /// pickle, GGUF, or npy/npz — format auto-detected).
37    InspectTensors { path: PathBuf },
38    /// Run Toto-2 time-series forecasting from a GGUF file.
39    ///
40    /// Reads a JSON request from stdin.
41    ///
42    /// Univariate batch:
43    ///   {"context": [[t0, t1, ...], [t0, t1, ...]], "horizon": N}
44    ///
45    /// Multivariate batch (n_var variates per series):
46    ///   {"context": [[[v0_t0, v0_t1, ...], [v1_t0, v1_t1, ...]], ...], "horizon": N}
47    Infer {
48        /// Path to the GGUF file.
49        #[arg(short, long, default_value = "gguf/toto-2.5b-f16.gguf")]
50        gguf: PathBuf,
51
52        /// Path to config.json from the original HuggingFace model.
53        #[arg(long, default_value = "models/Datadog__Toto-2.0-2.5B/config.json")]
54        config: PathBuf,
55
56        /// Context length fed to the model (must be divisible by patch_size=32).
57        /// Defaults to the last 4096 timesteps (or all if shorter).
58        #[arg(long)]
59        context_length: Option<usize>,
60
61        /// Run the forward pass in F64 (double precision) to match PyTorch
62        /// numerical accuracy. Uses ~2× memory. Default: F32.
63        #[arg(long = "f64", default_value_t = false)]
64        use_f64: bool,
65    },
66    /// Upload source + GGUF files to HuggingFace Hub.
67    Upload {
68        /// HuggingFace repo to upload to (owner/name).
69        #[arg(short, long, default_value = "amaye15/toto-gguf")]
70        repo: String,
71        /// HuggingFace API token (or set HF_TOKEN env var).
72        #[arg(long, env = "HF_TOKEN")]
73        token: String,
74        /// Directory tree to upload. Defaults to the current directory.
75        #[arg(long, default_value = ".")]
76        root: PathBuf,
77    },
78
79    /// Delete cached model files for this model.
80    ///
81    /// Removes the canonical F32 GGUF and config cache at
82    /// `<model_dir>/<owner>__<name>/` created by `convert`. By default only
83    /// the cache is removed; pass `--output <path>` to also delete a
84    /// converted GGUF file (e.g. `gguf/...`).
85    Delete {
86        /// HuggingFace repo id (must match the `convert` --model you used).
87        #[arg(short, long, default_value = "Datadog/Toto-2.0-2.5B")]
88        model: String,
89        /// Directory where the cache was written (must match `convert` --model-dir).
90        #[arg(long, default_value = "models")]
91        model_dir: PathBuf,
92        /// Also delete this output GGUF file if it exists.
93        #[arg(short, long)]
94        output: Option<PathBuf>,
95    },
96}
97
98pub async fn run(command: Command) -> anyhow::Result<()> {
99    match command {
100        Command::Delete {
101            model,
102            model_dir,
103            output,
104        } => {
105            let canonical = zsfm_hub::canonical_gguf_path(&model_dir, &model);
106            crate::common::delete_cached_model(&canonical, output.as_deref())?;
107        }
108        Command::Upload { repo, token, root } => {
109            zsfm_hub::upload_repo(&repo, &token, &root).await?;
110        }
111
112        Command::Convert {
113            model,
114            output,
115            dtype,
116            model_dir,
117            token,
118            redownload,
119        } => {
120            let canonical = zsfm_hub::canonical_gguf_path(&model_dir, &model);
121            if crate::common::try_recast_from_cache(&canonical, &output, dtype.into(), redownload)?
122            {
123                return Ok(());
124            }
125
126            eprintln!("Downloading {model} into {} …", model_dir.display());
127            let files = zsfm_hub::download_model(&model, token.as_deref(), &model_dir)
128                .await
129                .context("download failed")?;
130
131            let config_str =
132                std::fs::read_to_string(&files.config_json).context("read config.json")?;
133            let config = TotoConfig::from_json(&config_str).context("parse config.json")?;
134
135            println!(
136                "Config: {} layers, hidden={}, heads={}",
137                config.num_hidden_layers, config.hidden_size, config.num_attention_heads
138            );
139
140            let f32_opts = ConvertOptions {
141                output_dtype: GGMLType::F32,
142            };
143            convert(&model, &files, &config, &f32_opts, &canonical)?;
144            eprintln!("Wrote canonical F32 GGUF to {} …", canonical.display());
145            files.cleanup_weights();
146
147            zsfm_checkpoint::recast(&canonical, &output, dtype.into())?;
148            eprintln!("Wrote {}", output.display());
149        }
150
151        Command::InspectTensors { path } => crate::common::inspect_tensors(&path)?,
152
153        Command::Infer {
154            gguf,
155            config,
156            context_length,
157            use_f64,
158        } => {
159            return run_infer_burn(&gguf, &config, context_length, use_f64).await;
160        }
161    }
162
163    Ok(())
164}
165
166#[allow(clippy::too_many_arguments)]
167async fn run_infer_burn(
168    gguf: &std::path::PathBuf,
169    config: &std::path::PathBuf,
170    context_length: Option<usize>,
171    use_f64: bool,
172) -> anyhow::Result<()> {
173    use zsfm_toto::infer::BurnTotoModel;
174    use zsfm_toto::infer::InferConfig;
175
176    if use_f64 {
177        anyhow::bail!("toto Burn engine supports compute_f64=false only");
178    }
179    let buf = zsfm_core::read_stdin_limited()?;
180    let req: serde_json::Value = serde_json::from_str(&buf).context("parse JSON input")?;
181    let contexts = zsfm_core::parse_mv_contexts(req["context"].clone())?;
182    let horizon = zsfm_core::parse_horizon(&req)?;
183
184    let config_str =
185        std::fs::read_to_string(config).with_context(|| format!("read {}", config.display()))?;
186    let toto_config: serde_json::Value = serde_json::from_str(&config_str)?;
187
188    let max_ctx = context_length.unwrap_or(4096);
189    let quantile_levels = [0.1f32, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9];
190    let median_idx = 4;
191
192    eprintln!("Loading Burn model from {} …", gguf.display());
193    let cfg = InferConfig::from_json_value(&toto_config);
194    let model = BurnTotoModel::load(gguf, cfg).context("load model")?;
195    let patch_size = model.patch_size();
196
197    let mut fc_outputs = Vec::new();
198    let mut total_ctx: usize = 0;
199    for raw_variates in &contexts {
200        let n_var = raw_variates.len();
201        anyhow::ensure!(n_var > 0, "each context must have at least one variate");
202        total_ctx += raw_variates.iter().map(|v| v.len()).sum::<usize>();
203        let v0_len = raw_variates[0].len();
204        let ctx_len = (v0_len.min(max_ctx as usize) / patch_size) * patch_size;
205        anyhow::ensure!(
206            ctx_len > 0,
207            "context too short — need at least {patch_size} timesteps, got {v0_len}"
208        );
209        let mut data: Vec<Vec<f32>> = Vec::with_capacity(n_var);
210        let mut mask: Vec<Vec<bool>> = Vec::with_capacity(n_var);
211        for (vi, raw_ctx) in raw_variates.iter().enumerate() {
212            anyhow::ensure!(
213                raw_ctx.len() >= ctx_len,
214                "variate {vi} has fewer than {ctx_len} timesteps (needed to match variate 0)"
215            );
216            let start = raw_ctx.len() - ctx_len;
217            data.push(raw_ctx[start..].to_vec());
218            mask.push(vec![true; ctx_len]);
219        }
220        eprintln!("Running forecast ({ctx_len} context steps, {n_var} variate(s) → {horizon} future steps) …");
221        let quantile_mat = model.forecast(&data, &mask, horizon).context("forecast")?;
222        fc_outputs.push(zsfm_core::quantile_matrix_to_output(
223            &quantile_mat,
224            &quantile_levels,
225            median_idx,
226        ));
227    }
228
229    println!(
230        "{}",
231        zsfm_core::forecast_response_json("toto", total_ctx, horizon, fc_outputs)?
232    );
233    Ok(())
234}
zsfm/toto.rs

zsfm/
toto.rs
