zsfm/
sundial.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_sundial::config::SundialConfig;
10use zsfm_sundial::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/sundial-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 and convert model 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 = "thuml/sundial-base-128m")]
33        model: String,
34        #[arg(short, long, default_value = "gguf/sundial-f16.gguf")]
35        output: PathBuf,
36        #[arg(long, default_value = "f16")]
37        dtype: DtypeArg,
38        #[arg(long, env = "HF_TOKEN")]
39        token: Option<String>,
40        #[arg(long, default_value = "models")]
41        model_dir: PathBuf,
42        /// Force a fresh download even if a cached F32 GGUF already exists.
43        #[arg(long)]
44        redownload: bool,
45    },
46    /// Run inference on a GGUF model. Univariate, point-forecast only (flow-matching). No
47    /// --config flag: unlike chronos/flowstate/toto/ttm (which read a HuggingFace
48    /// config.json), Sundial's architecture is read directly from GGUF metadata written
49    /// at convert time, so nothing external is needed at inference time.
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/sundial-f16.gguf")]
56        gguf: PathBuf,
57        /// Override the GGUF ODE step count (default: use model metadata, typically 50).
58        /// 10-20 steps recommended; latency scales linearly with this value.
59        #[arg(long)]
60        steps: Option<u32>,
61    },
62    /// Print all tensor names in a local checkpoint file (safetensors, PyTorch
63    /// pickle, GGUF, or npy/npz — format auto-detected).
64    InspectTensors { path: PathBuf },
65
66    /// Delete cached model files for this model.
67    ///
68    /// Removes the canonical F32 GGUF and config cache at
69    /// `<model_dir>/<owner>__<name>/` created by `convert`. By default only
70    /// the cache is removed; pass `--output <path>` to also delete a
71    /// converted GGUF file (e.g. `gguf/...`).
72    Delete {
73        /// HuggingFace repo id (must match the `convert` --model you used).
74        #[arg(short, long, default_value = "thuml/sundial-base-128m")]
75        model: String,
76        /// Directory where the cache was written (must match `convert` --model-dir).
77        #[arg(long, default_value = "models")]
78        model_dir: PathBuf,
79        /// Also delete this output GGUF file if it exists.
80        #[arg(short, long)]
81        output: Option<PathBuf>,
82    },
83}
84
85pub async fn run(command: Command) -> anyhow::Result<()> {
86    match command {
87        Command::Delete {
88            model,
89            model_dir,
90            output,
91        } => {
92            let canonical = zsfm_hub::canonical_gguf_path(&model_dir, &model);
93            crate::common::delete_cached_model(&canonical, output.as_deref())?;
94        }
95        Command::Upload { repo, token, root } => {
96            zsfm_hub::upload_repo(&repo, &token, &root).await?;
97        }
98
99        Command::Convert {
100            model,
101            output,
102            dtype,
103            token,
104            model_dir,
105            redownload,
106        } => {
107            let output_dtype: GGMLType = dtype.into();
108            let canonical = zsfm_hub::canonical_gguf_path(&model_dir, &model);
109            if crate::common::try_recast_from_cache(&canonical, &output, output_dtype, redownload)?
110            {
111                return Ok(());
112            }
113
114            eprintln!("Downloading {model} into {} …", model_dir.display());
115            let files = zsfm_hub::download_model(&model, token.as_deref(), &model_dir)
116                .await
117                .context("download failed")?;
118
119            let config_str = std::fs::read_to_string(&files.config_json)
120                .with_context(|| format!("read {}", files.config_json.display()))?;
121            let config: SundialConfig =
122                serde_json::from_str(&config_str).context("parse config.json")?;
123
124            convert(
125                &model,
126                &files,
127                &config,
128                &ConvertOptions {
129                    output_dtype: GGMLType::F32,
130                },
131                &canonical,
132            )?;
133            eprintln!("Wrote canonical F32 GGUF to {} …", canonical.display());
134            files.cleanup_weights();
135
136            zsfm_checkpoint::recast(&canonical, &output, output_dtype)?;
137            eprintln!("Wrote {}", output.display());
138        }
139
140        Command::Infer { gguf, steps } => {
141            return run_infer_burn(&gguf, steps).await;
142        }
143
144        Command::InspectTensors { path } => crate::common::inspect_tensors(&path)?,
145    }
146    Ok(())
147}
148
149async fn run_infer_burn(gguf: &std::path::PathBuf, steps: Option<u32>) -> anyhow::Result<()> {
150    use zsfm_sundial::infer::BurnSundialModel;
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    let contexts = zsfm_core::parse_mv_contexts(req["context"].clone())?;
155    let horizon = zsfm_core::parse_horizon(&req)?;
156
157    eprintln!("Loading Burn model from {} …", gguf.display());
158    let model = BurnSundialModel::load(gguf, steps.map(|s| s as usize)).context("load model")?;
159    eprintln!("Model loaded.");
160
161    let mut fc_outputs = Vec::new();
162    let mut total_ctx = 0usize;
163    for raw_variates in &contexts {
164        anyhow::ensure!(
165            raw_variates.len() == 1,
166            "Sundial only supports univariate forecasting (1 variate per context)"
167        );
168        let ctx = &raw_variates[0];
169        anyhow::ensure!(!ctx.is_empty(), "context series must not be empty");
170        total_ctx += ctx.len();
171        eprintln!("Running forecast (context len = {}) …", ctx.len());
172        let raw = model.forecast(ctx)?;
173        let point: Vec<f32> = raw.into_iter().take(horizon).collect();
174        fc_outputs.push(zsfm_core::ForecastOutput::Univariate {
175            point,
176            quantiles: Default::default(),
177        });
178    }
179    println!(
180        "{}",
181        zsfm_core::forecast_response_json("sundial", total_ctx, horizon, fc_outputs)?
182    );
183    Ok(())
184}
zsfm/sundial.rs

zsfm/
sundial.rs
