zsfm/
timesfm.rs

1use std::collections::BTreeMap;
2use std::path::PathBuf;
3
4use crate::common::DtypeArg;
5use anyhow::{bail, Context, Result};
6use clap::Subcommand;
7
8use zsfm_core::ForecastOutput;
9use zsfm_gguf::GGMLType;
10use zsfm_timesfm::config::TimesFMConfig;
11use zsfm_timesfm::convert::{convert, ConvertOptions};
12
13#[derive(Subcommand)]
14pub enum Command {
15    /// Upload source + GGUF files to HuggingFace Hub.
16    Upload {
17        #[arg(short, long, default_value = "amaye15/timesfm-gguf")]
18        repo: String,
19        #[arg(long, env = "HF_TOKEN")]
20        token: String,
21        /// Directory tree to upload. Defaults to the current directory.
22        #[arg(long, default_value = ".")]
23        root: PathBuf,
24    },
25    /// Download and convert model weights to GGUF format.
26    ///
27    /// The first conversion downloads the model and writes a canonical F32 GGUF to
28    /// `<model_dir>/<owner>__<name>/model-f32.gguf`, deleting the (large) downloaded
29    /// weight files afterward. Later conversions (any --dtype) recast from that
30    /// cached F32 GGUF instead of re-downloading — pass --redownload to force a
31    /// fresh download anyway (e.g. the repo was updated).
32    Convert {
33        #[arg(short, long, default_value = "google/timesfm-2.5-200m-pytorch")]
34        model: String,
35        #[arg(short, long, default_value = "gguf/timesfm.gguf")]
36        output: PathBuf,
37        #[arg(long, default_value = "f16")]
38        dtype: DtypeArg,
39        #[arg(long, env = "HF_TOKEN")]
40        token: Option<String>,
41        /// Local directory to cache downloaded files.
42        #[arg(long, default_value = "models")]
43        model_dir: PathBuf,
44        /// Force a fresh download even if a cached F32 GGUF already exists.
45        #[arg(long)]
46        redownload: bool,
47    },
48    /// Print all tensor names in a local checkpoint file (safetensors, PyTorch
49    /// pickle, GGUF, or npy/npz — format auto-detected).
50    InspectTensors { path: PathBuf },
51    /// Run inference on a time series. Univariate only; the architecture is fixed, so no
52    /// `--config` is needed.
53    ///
54    /// Reads a JSON request from stdin: {"context": [...], "horizon": N}
55    /// Outputs a JSON forecast in OpenAI-compatible format with point and quantile forecasts.
56    Infer {
57        #[arg(short, long, default_value = "gguf/timesfm.gguf")]
58        gguf: PathBuf,
59    },
60
61    /// Delete cached model files for this model.
62    ///
63    /// Removes the canonical F32 GGUF and config cache at
64    /// `<model_dir>/<owner>__<name>/` created by `convert`. By default only
65    /// the cache is removed; pass `--output <path>` to also delete a
66    /// converted GGUF file (e.g. `gguf/...`).
67    Delete {
68        /// HuggingFace repo id (must match the `convert` --model you used).
69        #[arg(short, long, default_value = "google/timesfm-2.5-200m-pytorch")]
70        model: String,
71        /// Directory where the cache was written (must match `convert` --model-dir).
72        #[arg(long, default_value = "models")]
73        model_dir: PathBuf,
74        /// Also delete this output GGUF file if it exists.
75        #[arg(short, long)]
76        output: Option<PathBuf>,
77    },
78}
79
80pub async fn run(command: Command) -> Result<()> {
81    match command {
82        Command::Delete {
83            model,
84            model_dir,
85            output,
86        } => {
87            let canonical = zsfm_hub::canonical_gguf_path(&model_dir, &model);
88            crate::common::delete_cached_model(&canonical, output.as_deref())?;
89        }
90        Command::Upload { repo, token, root } => {
91            zsfm_hub::upload_repo(&repo, &token, &root).await?;
92        }
93        Command::Convert {
94            model,
95            output,
96            dtype,
97            token,
98            model_dir,
99            redownload,
100        } => {
101            cmd_convert(
102                &model,
103                &output,
104                dtype.into(),
105                token.as_deref(),
106                &model_dir,
107                redownload,
108            )
109            .await?;
110        }
111        Command::InspectTensors { path } => crate::common::inspect_tensors(&path)?,
112        Command::Infer { gguf } => cmd_infer_burn(&gguf)?,
113    }
114    Ok(())
115}
116
117async fn cmd_convert(
118    model: &str,
119    output: &PathBuf,
120    output_dtype: GGMLType,
121    token: Option<&str>,
122    model_dir: &PathBuf,
123    redownload: bool,
124) -> Result<()> {
125    let canonical = zsfm_hub::canonical_gguf_path(model_dir, model);
126    if crate::common::try_recast_from_cache(&canonical, output, output_dtype, redownload)? {
127        return Ok(());
128    }
129
130    eprintln!("Downloading {model} into {} …", model_dir.display());
131    let files = zsfm_hub::download_model(model, token, model_dir)
132        .await
133        .context("download failed")?;
134    let config = TimesFMConfig::new();
135
136    convert(
137        model,
138        &files,
139        &config,
140        &ConvertOptions {
141            output_dtype: GGMLType::F32,
142        },
143        &canonical,
144    )?;
145    eprintln!("Wrote canonical F32 GGUF to {} …", canonical.display());
146    files.cleanup_weights();
147
148    if let Some(parent) = output.parent() {
149        if !parent.as_os_str().is_empty() {
150            std::fs::create_dir_all(parent)
151                .with_context(|| format!("create output dir {}", parent.display()))?;
152        }
153    }
154    zsfm_checkpoint::recast(&canonical, output, output_dtype)?;
155    eprintln!("Wrote {}", output.display());
156    Ok(())
157}
158
159fn cmd_infer_burn(gguf_path: &PathBuf) -> Result<()> {
160    use zsfm_timesfm::infer::BurnTimesFMModel;
161
162    let buf = zsfm_core::read_stdin_limited()?;
163    let req: serde_json::Value = serde_json::from_str(&buf).context("parse JSON input")?;
164    let contexts = zsfm_core::parse_mv_contexts(req["context"].clone())?;
165    let horizon = zsfm_core::parse_horizon(&req)?;
166
167    eprintln!("Loading Burn model from {} …", gguf_path.display());
168    let model = BurnTimesFMModel::load(gguf_path).context("load model")?;
169
170    let quantile_labels = [
171        "0.10", "0.20", "0.30", "0.40", "0.50", "0.60", "0.70", "0.80", "0.90",
172    ];
173
174    let mut fc_outputs = Vec::new();
175    let mut total_ctx = 0usize;
176    for raw_variates in &contexts {
177        anyhow::ensure!(
178            raw_variates.len() == 1,
179            "TimesFM only supports univariate forecasting (1 variate per context)"
180        );
181        let ctx = &raw_variates[0];
182        if ctx.is_empty() {
183            bail!("context series must not be empty");
184        }
185        total_ctx += ctx.len();
186        eprintln!(
187            "Running forecast (context={}, horizon={}) …",
188            ctx.len(),
189            horizon
190        );
191        let outputs = model.forecast(ctx, horizon)?;
192
193        let point = outputs.first().cloned().unwrap_or_default();
194        let mut quantiles = BTreeMap::new();
195        for (i, label) in quantile_labels.iter().enumerate() {
196            if let Some(q) = outputs.get(i + 1) {
197                quantiles.insert(label.to_string(), q.clone());
198            }
199        }
200        fc_outputs.push(ForecastOutput::Univariate { point, quantiles });
201    }
202    println!(
203        "{}",
204        zsfm_core::forecast_response_json("timesfm", total_ctx, horizon, fc_outputs)?
205    );
206    Ok(())
207}
zsfm/timesfm.rs

zsfm/
timesfm.rs
