zsfm/
tirex.rs

1use std::collections::BTreeMap;
2use std::path::PathBuf;
3
4use crate::common::DtypeArg;
5use anyhow::Context;
6use clap::Subcommand;
7
8use zsfm_gguf::GGMLType;
9use zsfm_tirex::config::TiRexConfig;
10use zsfm_tirex::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/tirex-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 TiRex from HuggingFace and convert to GGUF.
25    ///
26    /// Downloads model.ckpt and reads it directly via the in-workspace
27    /// pickle reader — no Python or intermediate extraction step required.
28    ///
29    /// The first conversion downloads the model and writes a canonical F32 GGUF to
30    /// `<model_dir>/<owner>__<name>/model-f32.gguf`, deleting the downloaded .ckpt
31    /// afterward. Later conversions (any --dtype) recast from that cached F32 GGUF
32    /// instead of re-downloading — pass --redownload to force a fresh download
33    /// anyway (e.g. the repo was updated).
34    Convert {
35        #[arg(short, long, default_value = "NX-AI/TiRex")]
36        model: String,
37        #[arg(short, long, default_value = "gguf/tirex-f32.gguf")]
38        output: PathBuf,
39        #[arg(long, default_value = "f32")]
40        dtype: DtypeArg,
41        #[arg(long, default_value = "models")]
42        model_dir: PathBuf,
43        #[arg(long, env = "HF_TOKEN")]
44        token: Option<String>,
45        /// Force a fresh download even if a cached F32 GGUF already exists.
46        #[arg(long)]
47        redownload: bool,
48    },
49    /// Print all tensor names in a local checkpoint file (safetensors, PyTorch
50    /// pickle, GGUF, or npy/npz — format auto-detected).
51    InspectTensors { path: PathBuf },
52    /// Run TiRex forecasting from a GGUF file. Univariate only. No --config flag: the
53    /// architecture is fixed (TiRexConfig::default_from_ckpt()), unlike
54    /// chronos/flowstate/toto/ttm, which vary by checkpoint and need config.json.
55    ///
56    /// Reads a JSON request from stdin: {"context": [...], "horizon": N}
57    /// Outputs a JSON forecast in OpenAI-compatible format with all quantile levels.
58    /// The `point` field contains the median quantile (q0.5).
59    Infer {
60        #[arg(short, long, default_value = "gguf/tirex-f32.gguf")]
61        gguf: PathBuf,
62    },
63
64    /// Delete cached model files for this model.
65    ///
66    /// Removes the canonical F32 GGUF and config cache at
67    /// `<model_dir>/<owner>__<name>/` created by `convert`. By default only
68    /// the cache is removed; pass `--output <path>` to also delete a
69    /// converted GGUF file (e.g. `gguf/...`).
70    Delete {
71        /// HuggingFace repo id (must match the `convert` --model you used).
72        #[arg(short, long, default_value = "NX-AI/TiRex")]
73        model: String,
74        /// Directory where the cache was written (must match `convert` --model-dir).
75        #[arg(long, default_value = "models")]
76        model_dir: PathBuf,
77        /// Also delete this output GGUF file if it exists.
78        #[arg(short, long)]
79        output: Option<PathBuf>,
80    },
81}
82
83pub async fn run(command: Command) -> anyhow::Result<()> {
84    match command {
85        Command::Delete {
86            model,
87            model_dir,
88            output,
89        } => {
90            let canonical = zsfm_hub::canonical_gguf_path(&model_dir, &model);
91            crate::common::delete_cached_model(&canonical, output.as_deref())?;
92        }
93        Command::Upload { repo, token, root } => {
94            zsfm_hub::upload_repo(&repo, &token, &root).await?;
95        }
96
97        Command::Convert {
98            model,
99            output,
100            dtype,
101            model_dir,
102            token,
103            redownload,
104        } => {
105            let canonical = zsfm_hub::canonical_gguf_path(&model_dir, &model);
106            if crate::common::try_recast_from_cache(&canonical, &output, dtype.into(), redownload)?
107            {
108                return Ok(());
109            }
110            // download_file() only creates model_dir itself, not the namespaced
111            // subdirectory canonical_gguf_path() computes — unlike download_model(),
112            // which creates that exact directory as a side effect.
113            if let Some(parent) = canonical.parent() {
114                std::fs::create_dir_all(parent).context("create canonical GGUF dir")?;
115            }
116
117            eprintln!("Downloading {model} into {} …", model_dir.display());
118            let ckpt_path =
119                zsfm_hub::download_file(&model, "model.ckpt", token.as_deref(), &model_dir)
120                    .await
121                    .context("download failed")?;
122
123            let config = TiRexConfig::default_from_ckpt();
124            let f32_opts = ConvertOptions {
125                output_dtype: GGMLType::F32,
126            };
127            convert(&ckpt_path, &config, &f32_opts, &canonical)?;
128            eprintln!("Wrote canonical F32 GGUF to {} …", canonical.display());
129            let _ = std::fs::remove_file(&ckpt_path);
130
131            zsfm_checkpoint::recast(&canonical, &output, dtype.into())?;
132            eprintln!("Wrote {}", output.display());
133        }
134
135        Command::InspectTensors { path } => crate::common::inspect_tensors(&path)?,
136
137        Command::Infer { gguf } => {
138            return run_infer_burn(&gguf).await;
139        }
140    }
141
142    Ok(())
143}
144
145async fn run_infer_burn(gguf: &std::path::PathBuf) -> anyhow::Result<()> {
146    use std::collections::BTreeMap;
147    use zsfm_tirex::infer::BurnTiRexModel;
148
149    let buf = zsfm_core::read_stdin_limited()?;
150    let req: serde_json::Value = serde_json::from_str(&buf).context("parse JSON input")?;
151    let contexts = zsfm_core::parse_mv_contexts(req["context"].clone())?;
152    let horizon = zsfm_core::parse_horizon(&req)?;
153    let config = TiRexConfig::default_from_ckpt();
154
155    eprintln!("Loading Burn model from {} …", gguf.display());
156    let model = BurnTiRexModel::load(gguf, config.clone()).context("load model")?;
157
158    let mut fc_outputs = Vec::new();
159    let mut total_ctx = 0usize;
160    for raw_variates in &contexts {
161        anyhow::ensure!(
162            raw_variates.len() == 1,
163            "TiRex only supports univariate forecasting (1 variate per context)"
164        );
165        let ctx = &raw_variates[0];
166        anyhow::ensure!(!ctx.is_empty(), "context series must not be empty");
167        total_ctx += ctx.len();
168        let (quantiles, median) = model.forecast(ctx, horizon).context("forecast")?;
169        let mut q_map: BTreeMap<String, Vec<f32>> = BTreeMap::new();
170        for (i, q) in config.quantiles.iter().enumerate() {
171            q_map.insert(format!("{q:.2}"), quantiles[i].clone());
172        }
173        fc_outputs.push(zsfm_core::ForecastOutput::Univariate {
174            point: median,
175            quantiles: q_map,
176        });
177    }
178    println!(
179        "{}",
180        zsfm_core::forecast_response_json("tirex", total_ctx, horizon, fc_outputs)?
181    );
182    Ok(())
183}
zsfm/tirex.rs

zsfm/
tirex.rs
