zsfm_core/
json.rs

1use std::collections::BTreeMap;
2use std::time::{SystemTime, UNIX_EPOCH};
3
4use serde::Serialize;
5
6use crate::forecast::QuantileMatrix;
7
8#[derive(Serialize)]
9#[serde(untagged)]
10pub enum ForecastOutput {
11    Univariate {
12        point: Vec<f32>,
13        quantiles: BTreeMap<String, Vec<f32>>,
14    },
15    Multivariate {
16        variates: Vec<VariateForecast>,
17    },
18}
19
20#[derive(Serialize)]
21pub struct VariateForecast {
22    pub point: Vec<f32>,
23    pub quantiles: BTreeMap<String, Vec<f32>>,
24}
25
26/// Turn one forecast's `QuantileMatrix` into a `ForecastOutput`, choosing the
27/// univariate (flat `point`/`quantiles`) or multivariate (`variates: [...]`)
28/// shape based on how many variates the matrix actually has.
29///
30/// `quantile_levels[i]` is the probability level for `qmat[i]`; `median_idx`
31/// is the row holding the point forecast (usually the index of level 0.5).
32pub fn quantile_matrix_to_output(
33    qmat: &QuantileMatrix,
34    quantile_levels: &[f32],
35    median_idx: usize,
36) -> ForecastOutput {
37    let n_var = qmat.first().map(|q| q.len()).unwrap_or(0);
38
39    let quantiles_for = |vi: usize| -> BTreeMap<String, Vec<f32>> {
40        let mut quantiles = BTreeMap::new();
41        for (qi, &level) in quantile_levels.iter().enumerate() {
42            if let Some(row) = qmat.get(qi) {
43                if let Some(q) = row.get(vi) {
44                    quantiles.insert(format!("{level:.2}"), q.clone());
45                }
46            }
47        }
48        quantiles
49    };
50    let point_for = |vi: usize| -> Vec<f32> {
51        qmat.get(median_idx)
52            .and_then(|v| v.get(vi))
53            .cloned()
54            .unwrap_or_default()
55    };
56
57    if n_var <= 1 {
58        ForecastOutput::Univariate {
59            point: point_for(0),
60            quantiles: quantiles_for(0),
61        }
62    } else {
63        let variates = (0..n_var)
64            .map(|vi| VariateForecast {
65                point: point_for(vi),
66                quantiles: quantiles_for(vi),
67            })
68            .collect();
69        ForecastOutput::Multivariate { variates }
70    }
71}
72
73/// Wrap one or more `ForecastOutput`s in the shared response envelope every
74/// model's `infer` CLI subcommand (and later, Python binding) returns.
75pub fn forecast_response_json(
76    model_name: &str,
77    context_length: usize,
78    forecast_length: usize,
79    outputs: Vec<ForecastOutput>,
80) -> anyhow::Result<String> {
81    #[derive(Serialize)]
82    struct ForecastResponse {
83        id: String,
84        object: &'static str,
85        created: u64,
86        model: String,
87        choices: Vec<Choice>,
88        usage: Usage,
89    }
90    #[derive(Serialize)]
91    struct Choice {
92        index: usize,
93        forecast: ForecastOutput,
94        finish_reason: &'static str,
95    }
96    #[derive(Serialize)]
97    struct Usage {
98        context_length: usize,
99        forecast_length: usize,
100    }
101
102    let created = SystemTime::now()
103        .duration_since(UNIX_EPOCH)
104        .unwrap_or_default()
105        .as_secs();
106
107    let resp = ForecastResponse {
108        id: format!("forecast-{created:016x}"),
109        object: "forecast",
110        created,
111        model: model_name.to_string(),
112        choices: outputs
113            .into_iter()
114            .enumerate()
115            .map(|(i, forecast)| Choice {
116                index: i,
117                forecast,
118                finish_reason: "stop",
119            })
120            .collect(),
121        usage: Usage {
122            context_length,
123            forecast_length,
124        },
125    };
126
127    Ok(serde_json::to_string_pretty(&resp)?)
128}
zsfm_core/json.rs

zsfm_core/
json.rs
