zsfm/
moment.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_moment::config::MomentConfig;
10use zsfm_moment::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/moment-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 MOMENT-1-large 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 = "AutonLab/MOMENT-1-large")]
33        model: String,
34        #[arg(short, long, default_value = "gguf/moment-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 MOMENT forecasting from a GGUF file. Univariate, point-forecast only. No
50    /// --config flag: the architecture is fixed (MomentConfig::default()), unlike
51    /// chronos/flowstate/toto/ttm, which vary by checkpoint and need config.json.
52    ///
53    /// Reads a JSON request from stdin: {"context": [...], "horizon": N}
54    /// Outputs a JSON forecast in OpenAI-compatible format.
55    Infer {
56        #[arg(short, long, default_value = "gguf/moment-f32.gguf")]
57        gguf: 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 = "AutonLab/MOMENT-1-large")]
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 = MomentConfig::default();
113            let f32_opts = ConvertOptions {
114                output_dtype: GGMLType::F32,
115            };
116            convert(&files.safetensors_shards, &config, &f32_opts, &canonical)?;
117            eprintln!("Wrote canonical F32 GGUF to {} …", canonical.display());
118            files.cleanup_weights();
119
120            zsfm_checkpoint::recast(&canonical, &output, dtype.into())?;
121            eprintln!("Wrote {}", output.display());
122        }
123
124        Command::InspectTensors { path } => crate::common::inspect_tensors(&path)?,
125
126        Command::Infer { gguf } => {
127            return run_infer_burn(&gguf).await;
128        }
129    }
130
131    Ok(())
132}
133
134async fn run_infer_burn(gguf: &std::path::PathBuf) -> anyhow::Result<()> {
135    use zsfm_moment::infer::BurnMomentModel;
136
137    let buf = zsfm_core::read_stdin_limited()?;
138    let req: serde_json::Value = serde_json::from_str(&buf).context("parse JSON input")?;
139    let contexts = zsfm_core::parse_mv_contexts(req["context"].clone())?;
140    let horizon = zsfm_core::parse_horizon(&req)?;
141    let config = MomentConfig::default();
142
143    eprintln!("Loading Burn model from {} …", gguf.display());
144    let model = BurnMomentModel::load(gguf, config).context("load model")?;
145
146    let mut fc_outputs = Vec::new();
147    let mut total_ctx = 0usize;
148    for raw_variates in &contexts {
149        anyhow::ensure!(
150            raw_variates.len() == 1,
151            "MOMENT only supports univariate forecasting (1 variate per context)"
152        );
153        let ctx = &raw_variates[0];
154        anyhow::ensure!(!ctx.is_empty(), "context series must not be empty");
155        total_ctx += ctx.len();
156        let point = model.forecast(ctx, horizon).context("forecast")?;
157        fc_outputs.push(zsfm_core::ForecastOutput::Univariate {
158            point,
159            quantiles: Default::default(),
160        });
161    }
162    println!(
163        "{}",
164        zsfm_core::forecast_response_json("moment", total_ctx, horizon, fc_outputs)?
165    );
166    Ok(())
167}
zsfm/moment.rs

zsfm/
moment.rs
