zsfm/
chronos.rs

1use std::path::PathBuf;
2
3use crate::common::DtypeArg;
4use anyhow::Context;
5use clap::Subcommand;
6
7use zsfm_chronos::config::Chronos2Config;
8use zsfm_chronos::convert::{convert, ConvertOptions};
9use zsfm_gguf::GGMLType;
10
11#[derive(Subcommand)]
12pub enum Command {
13    /// Download a Chronos-2 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 = "amazon/chronos-2")]
22        model: String,
23        #[arg(short, long, default_value = "gguf/chronos-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 Chronos-2 forecasting from a GGUF file.
39    ///
40    /// Reads a JSON request from stdin: {"context": [...], "horizon": N}
41    /// Outputs a JSON forecast in OpenAI-compatible format with all quantile levels.
42    /// The `point` field contains the median quantile (q0.5). Univariate only.
43    Infer {
44        #[arg(short, long, default_value = "gguf/chronos-f16.gguf")]
45        gguf: PathBuf,
46        #[arg(long, default_value = "models/amazon__chronos-2/config.json")]
47        config: PathBuf,
48    },
49    /// Upload source + GGUF files to HuggingFace Hub.
50    Upload {
51        #[arg(short, long, default_value = "amaye15/chronos-gguf")]
52        repo: String,
53        #[arg(long, env = "HF_TOKEN")]
54        token: String,
55        /// Directory tree to upload. Defaults to the current directory.
56        #[arg(long, default_value = ".")]
57        root: PathBuf,
58    },
59
60    /// Delete cached model files for this model.
61    ///
62    /// Removes the canonical F32 GGUF and config cache at
63    /// `<model_dir>/<owner>__<name>/` created by `convert`. By default only
64    /// the cache is removed; pass `--output <path>` to also delete a
65    /// converted GGUF file (e.g. `gguf/...`).
66    Delete {
67        /// HuggingFace repo id (must match the `convert` --model you used).
68        #[arg(short, long, default_value = "amazon/chronos-2")]
69        model: String,
70        /// Directory where the cache was written (must match `convert` --model-dir).
71        #[arg(long, default_value = "models")]
72        model_dir: PathBuf,
73        /// Also delete this output GGUF file if it exists.
74        #[arg(short, long)]
75        output: Option<PathBuf>,
76    },
77}
78
79pub async fn run(command: Command) -> anyhow::Result<()> {
80    match command {
81        Command::Delete {
82            model,
83            model_dir,
84            output,
85        } => {
86            let canonical = zsfm_hub::canonical_gguf_path(&model_dir, &model);
87            crate::common::delete_cached_model(&canonical, output.as_deref())?;
88        }
89        Command::Upload { repo, token, root } => {
90            zsfm_hub::upload_repo(&repo, &token, &root).await?;
91        }
92
93        Command::Convert {
94            model,
95            output,
96            dtype,
97            model_dir,
98            token,
99            redownload,
100        } => {
101            let canonical = zsfm_hub::canonical_gguf_path(&model_dir, &model);
102            if crate::common::try_recast_from_cache(&canonical, &output, dtype.into(), redownload)?
103            {
104                return Ok(());
105            }
106
107            eprintln!("Downloading {model} into {} …", model_dir.display());
108            let files = zsfm_hub::download_model(&model, token.as_deref(), &model_dir)
109                .await
110                .context("download failed")?;
111
112            let config_str =
113                std::fs::read_to_string(&files.config_json).context("read config.json")?;
114            let config = Chronos2Config::from_json(&config_str).context("parse config.json")?;
115
116            println!(
117                "Config: {} layers, d_model={}, d_ff={}, {} heads, patch_size={}",
118                config.num_layers,
119                config.d_model,
120                config.d_ff,
121                config.num_heads,
122                config.chronos_config.input_patch_size,
123            );
124
125            let f32_opts = ConvertOptions {
126                output_dtype: GGMLType::F32,
127            };
128            convert(&model, &files, &config, &f32_opts, &canonical)?;
129            eprintln!("Wrote canonical F32 GGUF to {} …", canonical.display());
130            files.cleanup_weights();
131
132            zsfm_checkpoint::recast(&canonical, &output, dtype.into())?;
133            eprintln!("Wrote {}", output.display());
134        }
135
136        Command::InspectTensors { path } => crate::common::inspect_tensors(&path)?,
137
138        Command::Infer { gguf, config } => {
139            return run_infer_burn(&gguf, &config).await;
140        }
141    }
142
143    Ok(())
144}
145
146async fn run_infer_burn(
147    gguf: &std::path::PathBuf,
148    config: &std::path::PathBuf,
149) -> anyhow::Result<()> {
150    use zsfm_chronos::infer::BurnChronosModel;
151
152    let buf = zsfm_core::read_stdin_limited()?;
153    let req: serde_json::Value = serde_json::from_str(&buf).context("parse JSON input")?;
154
155    let contexts = zsfm_core::parse_mv_contexts(req["context"].clone())?;
156    let horizon = zsfm_core::parse_horizon(&req)?;
157
158    let config_str =
159        std::fs::read_to_string(config).with_context(|| format!("read {}", config.display()))?;
160    let c2_config = Chronos2Config::from_json(&config_str).context("parse config.json")?;
161
162    eprintln!("Loading Burn model from {} …", gguf.display());
163    let model = BurnChronosModel::load(gguf, zsfm_chronos::infer::InferConfig::from(&c2_config))
164        .context("load model")?;
165
166    let patch_size = model.patch_size();
167    let patch_stride = model.patch_stride();
168    let max_ctx = model.context_length();
169    let quantile_levels = model.quantiles().to_vec();
170    let median_idx = crate::common::median_index(&quantile_levels);
171
172    let mut fc_outputs = Vec::new();
173    let mut total_ctx: usize = 0;
174    for raw_variates in &contexts {
175        anyhow::ensure!(
176            raw_variates.len() == 1,
177            "Chronos-2 only supports univariate forecasting (1 variate per context)"
178        );
179        let mut ctx = raw_variates[0].clone();
180        anyhow::ensure!(!ctx.is_empty(), "context series must not be empty");
181        total_ctx += ctx.len();
182
183        let total_len = ctx.len();
184        let usable = total_len.min(max_ctx);
185        let ctx_len = (usable / patch_stride) * patch_stride;
186        anyhow::ensure!(
187            ctx_len >= patch_size,
188            "context too short — need at least {patch_size} timesteps, got {ctx_len}"
189        );
190        let start = total_len.saturating_sub(ctx_len);
191        ctx = ctx[start..].to_vec();
192
193        eprintln!(
194            "Running forecast ({} context steps → {horizon} future steps) …",
195            ctx.len()
196        );
197        let quantile_mat = model.forecast(&ctx, horizon).context("forecast")?;
198        let qmat: zsfm_core::QuantileMatrix =
199            quantile_mat.into_iter().map(|row| vec![row]).collect();
200        fc_outputs.push(zsfm_core::quantile_matrix_to_output(
201            &qmat,
202            &quantile_levels,
203            median_idx,
204        ));
205    }
206
207    println!(
208        "{}",
209        zsfm_core::forecast_response_json("chronos", total_ctx, horizon, fc_outputs)?
210    );
211    Ok(())
212}
zsfm/chronos.rs

zsfm/
chronos.rs
