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.iter().position(|&p| p == patch_size)
31            .unwrap_or(2) // default to index 2 = 32
32    }
33}