zsfm/
ttm.rs

1use std::path::PathBuf;
2
3use crate::common::DtypeArg;
4use anyhow::Context;
5use clap::Subcommand;
6
7use zsfm_core::ForecastOutput;
8use zsfm_gguf::GGMLType;
9use zsfm_ttm::config::TtmConfig;
10use zsfm_ttm::convert::{convert, ConvertOptions};
11
12#[derive(Subcommand)]
13pub enum Command {
14    /// Upload source + GGUF files to HuggingFace Hub.
15    Upload {
16        #[arg(short, long, default_value = "amaye15/ttm-gguf")]
17        repo: String,
18        #[arg(long, env = "HF_TOKEN")]
19        token: String,
20        /// Directory tree to upload. Defaults to the current directory.
21        #[arg(long, default_value = ".")]
22        root: PathBuf,
23    },
24    /// Download TTM from HuggingFace and convert to GGUF.
25    ///
26    /// The first conversion downloads the model and writes a canonical F32 GGUF to
27    /// `<model_dir>/<owner>__<name>/model-f32.gguf`, deleting the (large) downloaded
28    /// weight files afterward. Later conversions (any --dtype) recast from that
29    /// cached F32 GGUF instead of re-downloading — pass --redownload to force a
30    /// fresh download anyway (e.g. the repo was updated).
31    Convert {
32        #[arg(short, long, default_value = "ibm-granite/granite-timeseries-ttm-r2")]
33        model: String,
34        #[arg(short, long, default_value = "gguf/ttm-f32.gguf")]
35        output: PathBuf,
36        #[arg(long, default_value = "f32")]
37        dtype: DtypeArg,
38        #[arg(long, default_value = "models")]
39        model_dir: PathBuf,
40        #[arg(long, env = "HF_TOKEN")]
41        token: Option<String>,
42        /// Force a fresh download even if a cached F32 GGUF already exists.
43        #[arg(long)]
44        redownload: bool,
45    },
46    /// Print all tensor names in a local checkpoint file (safetensors, PyTorch
47    /// pickle, GGUF, or npy/npz — format auto-detected).
48    InspectTensors { path: PathBuf },
49    /// Run TTM forecasting from a GGUF file. Univariate, point-forecast only.
50    ///
51    /// Reads a JSON request from stdin: {"context": [...], "horizon": N}
52    /// Outputs a JSON forecast in OpenAI-compatible format.
53    Infer {
54        /// Path to the GGUF file.
55        #[arg(short, long, default_value = "gguf/ttm-f32.gguf")]
56        gguf: PathBuf,
57        /// Path to config.json (original HuggingFace model).
58        #[arg(
59            long,
60            default_value = "models/ibm-granite__granite-timeseries-ttm-r2/config.json"
61        )]
62        config: PathBuf,
63    },
64
65    /// Delete cached model files for this model.
66    ///
67    /// Removes the canonical F32 GGUF and config cache at
68    /// `<model_dir>/<owner>__<name>/` created by `convert`. By default only
69    /// the cache is removed; pass `--output <path>` to also delete a
70    /// converted GGUF file (e.g. `gguf/...`).
71    Delete {
72        /// HuggingFace repo id (must match the `convert` --model you used).
73        #[arg(short, long, default_value = "ibm-granite/granite-timeseries-ttm-r2")]
74        model: String,
75        /// Directory where the cache was written (must match `convert` --model-dir).
76        #[arg(long, default_value = "models")]
77        model_dir: PathBuf,
78        /// Also delete this output GGUF file if it exists.
79        #[arg(short, long)]
80        output: Option<PathBuf>,
81    },
82}
83
84pub async fn run(command: Command) -> anyhow::Result<()> {
85    match command {
86        Command::Delete {
87            model,
88            model_dir,
89            output,
90        } => {
91            let canonical = zsfm_hub::canonical_gguf_path(&model_dir, &model);
92            crate::common::delete_cached_model(&canonical, output.as_deref())?;
93        }
94        Command::Upload { repo, token, root } => {
95            zsfm_hub::upload_repo(&repo, &token, &root).await?;
96        }
97
98        Command::Convert {
99            model,
100            output,
101            dtype,
102            model_dir,
103            token,
104            redownload,
105        } => {
106            let canonical = zsfm_hub::canonical_gguf_path(&model_dir, &model);
107            if crate::common::try_recast_from_cache(&canonical, &output, dtype.into(), redownload)?
108            {
109                return Ok(());
110            }
111
112            eprintln!("Downloading {model} into {} …", model_dir.display());
113            let files = zsfm_hub::download_model(&model, token.as_deref(), &model_dir)
114                .await
115                .context("download failed")?;
116
117            let config_str =
118                std::fs::read_to_string(&files.config_json).context("read config.json")?;
119            let config = TtmConfig::from_json(&config_str).context("parse config.json")?;
120
121            println!(
122                "Config: context={}, pred={}, patch={}/{}, d_model={}, {} enc layers, {} dec layers, {} adaptive levels",
123                config.context_length, config.prediction_length,
124                config.patch_length, config.patch_stride,
125                config.d_model, config.num_layers, config.decoder_num_layers,
126                config.adaptive_patching_levels,
127            );
128
129            let f32_opts = ConvertOptions {
130                output_dtype: GGMLType::F32,
131            };
132            convert(&model, &files, &config, &f32_opts, &canonical)?;
133            eprintln!("Wrote canonical F32 GGUF to {} …", canonical.display());
134            files.cleanup_weights();
135
136            zsfm_checkpoint::recast(&canonical, &output, dtype.into())?;
137            eprintln!("Wrote {}", output.display());
138        }
139
140        Command::InspectTensors { path } => crate::common::inspect_tensors(&path)?,
141
142        Command::Infer { gguf, config } => {
143            return run_infer_burn(&gguf, &config).await;
144        }
145    }
146
147    Ok(())
148}
149
150async fn run_infer_burn(
151    gguf: &std::path::PathBuf,
152    config: &std::path::PathBuf,
153) -> anyhow::Result<()> {
154    use zsfm_ttm::infer::BurnTtmModel;
155
156    let buf = zsfm_core::read_stdin_limited()?;
157    let req: serde_json::Value = serde_json::from_str(&buf).context("parse JSON input")?;
158    let contexts = zsfm_core::parse_mv_contexts(req["context"].clone())?;
159    let horizon = zsfm_core::parse_horizon(&req)?;
160
161    let config_str =
162        std::fs::read_to_string(config).with_context(|| format!("read {}", config.display()))?;
163    let ttm_config = TtmConfig::from_json(&config_str).context("parse config.json")?;
164    let min_ctx = ttm_config.patch_length;
165
166    eprintln!("Loading Burn model from {} …", gguf.display());
167    let model = BurnTtmModel::load(gguf, ttm_config).context("load model")?;
168
169    let mut fc_outputs = Vec::new();
170    let mut total_ctx = 0usize;
171    for raw_variates in &contexts {
172        anyhow::ensure!(
173            raw_variates.len() == 1,
174            "TTM only supports univariate forecasting (1 variate per context)"
175        );
176        let ctx = &raw_variates[0];
177        anyhow::ensure!(
178            ctx.len() >= min_ctx,
179            "context too short — need at least {min_ctx} timesteps, got {}",
180            ctx.len()
181        );
182        total_ctx += ctx.len();
183        let raw = model.forecast(ctx).context("forecast")?;
184        let point: Vec<f32> = raw.into_iter().take(horizon).collect();
185        fc_outputs.push(ForecastOutput::Univariate {
186            point,
187            quantiles: Default::default(),
188        });
189    }
190    println!(
191        "{}",
192        zsfm_core::forecast_response_json("ttm", total_ctx, horizon, fc_outputs)?
193    );
194    Ok(())
195}
zsfm/ttm.rs

zsfm/
ttm.rs
