Skip to main content

zsfm_moirai2/
config.rs

1/// Moirai-2.0-R-small configuration.
2#[derive(Clone)]
3pub struct Moirai2Config {
4    pub d_model: usize,           // 384
5    pub n_layers: usize,          // 6
6    pub n_heads: usize,           // 6
7    pub head_dim: usize,          // 64
8    pub d_ff: usize,              // 1024
9    pub patch_size: usize,        // 16
10    pub num_predict_token: usize, // 4
11    pub num_quantiles: usize,     // 9
12    pub max_seq_len: usize,       // 512 (in timesteps)
13    pub rope_dim: usize,          // 32 (partial_factor=(0.0, 0.5) of head_dim=64)
14    pub median_quantile: usize,   // 4 (0-indexed: 0.5 is index 4 of 9 levels)
15}
16
17impl Moirai2Config {
18    pub fn default() -> Self {
19        Self {
20            d_model: 384,
21            n_layers: 6,
22            n_heads: 6,
23            head_dim: 64,
24            d_ff: 1024,
25            patch_size: 16,
26            num_predict_token: 4,
27            num_quantiles: 9,
28            max_seq_len: 512,
29            rope_dim: 32,
30            median_quantile: 4,
31        }
32    }
33
34    pub fn max_ctx_tokens(&self) -> usize {
35        self.max_seq_len / self.patch_size
36    }
37}