1use serde::Deserialize;
2
3#[derive(Debug, Clone, Deserialize)]
4pub struct SundialConfig {
5 pub hidden_size: usize,
6 pub intermediate_size: usize,
7 pub num_hidden_layers: usize,
8 pub num_attention_heads: usize,
9 pub input_token_len: usize,
10 #[serde(default = "default_output_token_lens")]
11 pub output_token_lens: Vec<usize>,
12 pub rope_theta: f64,
13 #[serde(default = "default_flow_depth")]
14 pub flow_loss_depth: usize,
15 #[serde(default = "default_sampling_steps")]
16 pub num_sampling_steps: usize,
17}
18
19fn default_output_token_lens() -> Vec<usize> {
20 vec![720]
21}
22fn default_flow_depth() -> usize {
23 3
24}
25fn default_sampling_steps() -> usize {
26 50
27}
28
29impl SundialConfig {
30 pub fn head_dim(&self) -> usize {
31 self.hidden_size / self.num_attention_heads
32 }
33 pub fn output_token_len(&self) -> usize {
34 self.output_token_lens[0]
35 }
36}
37
38impl Default for SundialConfig {
39 fn default() -> Self {
40 Self {
41 hidden_size: 768,
42 intermediate_size: 3072,
43 num_hidden_layers: 12,
44 num_attention_heads: 12,
45 input_token_len: 16,
46 output_token_lens: vec![720],
47 rope_theta: 10000.0,
48 flow_loss_depth: 3,
49 num_sampling_steps: 50,
50 }
51 }
52}