Skip to main content

zsfm_timesfm/
config.rs

1/// Hardcoded architecture constants for TimesFM 2.5 200M.
2///
3/// These match `TimesFM_2p5_200M_Definition` in the Python source exactly.
4/// No config.json is needed — the architecture is fixed for this model.
5#[derive(Debug, Clone)]
6pub struct TimesFMConfig {
7    /// Input patch length (p = 32).
8    pub input_patch_len: usize,
9    /// Output patch length (o = 128).
10    pub output_patch_len: usize,
11    /// Output quantile length for the continuous quantile head (os = 1024).
12    pub output_quantile_len: usize,
13    /// Number of stacked transformer layers.
14    pub num_layers: usize,
15    /// Transformer model dimension.
16    pub d_model: usize,
17    /// Feed-forward hidden dimension.
18    pub d_ff: usize,
19    /// Number of attention heads.
20    pub num_heads: usize,
21    /// Head dimension = d_model / num_heads.
22    pub head_dim: usize,
23    /// Quantile levels predicted (excludes the implicit point forecast at index 0).
24    pub quantiles: Vec<f32>,
25    /// Total number of per-timestep outputs = len(quantiles) + 1 (point).
26    pub n_outputs: usize,
27    /// Index used to extract point/AR forecast from the 10-dim output.
28    pub decode_index: usize,
29    /// Maximum context tokens the model accepts (in individual timesteps).
30    pub context_limit: usize,
31    /// RMS norm epsilon.
32    pub rms_norm_eps: f64,
33    /// RoPE base frequency (theta = 10000).
34    pub rope_theta: f64,
35}
36
37impl Default for TimesFMConfig {
38    fn default() -> Self {
39        let quantiles = vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9];
40        let n_outputs = quantiles.len() + 1; // 10
41        Self {
42            input_patch_len:    32,
43            output_patch_len:   128,
44            output_quantile_len: 1024,
45            num_layers:         20,
46            d_model:            1280,
47            d_ff:               1280,
48            num_heads:          16,
49            head_dim:           80,
50            n_outputs,
51            decode_index:       5,
52            context_limit:      16384,
53            rms_norm_eps:       1e-6,
54            rope_theta:         10000.0,
55            quantiles,
56        }
57    }
58}
59
60impl TimesFMConfig {
61    pub fn new() -> Self {
62        Self::default()
63    }
64}