Skip to main content

zsfm_moirai/
config.rs

1/// Moirai-1.0-R-large configuration.
2#[derive(Clone)]
3pub struct MoiraiConfig {
4    pub d_model: usize,          // 1024
5    pub n_layers: usize,         // 24
6    pub n_heads: usize,          // 16
7    pub head_dim: usize,         // 64
8    pub d_ff: usize,             // 2736
9    pub max_seq_len: usize,      // 512 (max context length in timesteps)
10    pub patch_sizes: Vec<usize>, // [8, 16, 32, 64, 128]
11    pub max_patch_size: usize,   // 128 (output head dimension)
12}
13
14impl MoiraiConfig {
15    pub fn default() -> Self {
16        Self {
17            d_model: 1024,
18            n_layers: 24,
19            n_heads: 16,
20            head_dim: 64,
21            d_ff: 2736,
22            max_seq_len: 512,
23            patch_sizes: vec![8, 16, 32, 64, 128],
24            max_patch_size: 128,
25        }
26    }
27
28    /// Index of a given patch size in patch_sizes.
29    pub fn patch_idx(&self, patch_size: usize) -> usize {
30        self.patch_sizes
31            .iter()
32            .position(|&p| p == patch_size)
33            .unwrap_or(2) // default to index 2 = 32
34    }
35}