zsfm/
flowstate.rs

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

zsfm/
flowstate.rs
