1#[derive(Clone)]
3pub struct MoiraiConfig {
4 pub d_model: usize, pub n_layers: usize, pub n_heads: usize, pub head_dim: usize, pub d_ff: usize, pub max_seq_len: usize, pub patch_sizes: Vec<usize>, pub max_patch_size: usize, }
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 pub fn patch_idx(&self, patch_size: usize) -> usize {
30 self.patch_sizes.iter().position(|&p| p == patch_size)
31 .unwrap_or(2) }
33}