needle_core/
config.rs

1/// FFN activation function (matches Python `activation` config field).
2#[derive(Clone, Debug, Default, PartialEq)]
3pub enum FfnActivation {
4    #[default]
5    DRelu,
6    SwiGLU,
7    GeGLU,
8}
9
10impl FfnActivation {
11    pub fn parse(s: &str) -> Self {
12        match s {
13            "swiglu" => Self::SwiGLU,
14            "geglu" => Self::GeGLU,
15            _ => Self::DRelu, // "drelu" and unknown → DRelu
16        }
17    }
18}
19
20/// Mirror of Python TransformerConfig — all hyperparameters for Needle SAN.
21#[derive(Clone, Debug)]
22pub struct TransformerConfig {
23    pub d_model: usize,
24    pub num_heads: usize,
25    pub num_kv_heads: usize,
26    pub num_layers: usize,
27    pub num_dec_layers: usize,
28    pub vocab_size: usize,
29    pub max_enc_len: usize,
30    pub max_dec_len: usize,
31    pub ffn_dim: usize,
32    pub no_feedforward: bool,
33    /// FFN activation function (drelu/swiglu/geglu).
34    pub activation: FfnActivation,
35    /// RoPE base frequency (Python: rope_theta, default 10000.0).
36    pub rope_theta: f32,
37    /// Special token IDs
38    pub pad_id: u32,
39    pub eos_id: u32,
40    pub bos_id: u32,
41    pub unk_id: u32,
42    pub tool_call_id: u32,
43    pub tools_id: u32,
44}
45
46impl Default for TransformerConfig {
47    fn default() -> Self {
48        Self {
49            d_model: 512,
50            num_heads: 8,
51            num_kv_heads: 4,
52            num_layers: 12,
53            num_dec_layers: 8,
54            vocab_size: 8192,
55            max_enc_len: 1024,
56            max_dec_len: 512,
57            ffn_dim: 2048,
58            no_feedforward: true,
59            activation: FfnActivation::DRelu,
60            rope_theta: 10000.0,
61            pad_id: 0,
62            eos_id: 1,
63            bos_id: 2,
64            unk_id: 3,
65            tool_call_id: 4,
66            tools_id: 5,
67        }
68    }
69}
70
71impl TransformerConfig {
72    pub fn head_dim(&self) -> usize {
73        self.d_model / self.num_heads
74    }
75
76    /// Number of times each KV head is repeated to cover all Q heads.
77    pub fn kv_repeat(&self) -> usize {
78        self.num_heads / self.num_kv_heads
79    }
80
81    /// Validate structural invariants. Returns an error string if violated.
82    /// Called by the inference engine before constructing a model.
83    pub fn validate(&self) -> Result<(), &'static str> {
84        if self.num_heads == 0 {
85            return Err("num_heads must be > 0");
86        }
87        if self.num_kv_heads == 0 {
88            return Err("num_kv_heads must be > 0");
89        }
90        if self.d_model == 0 {
91            return Err("d_model must be > 0");
92        }
93        if !self.d_model.is_multiple_of(self.num_heads) {
94            return Err("d_model must be divisible by num_heads");
95        }
96        if !self.num_heads.is_multiple_of(self.num_kv_heads) {
97            return Err("num_heads must be divisible by num_kv_heads");
98        }
99        if self.vocab_size == 0 {
100            return Err("vocab_size must be > 0");
101        }
102        Ok(())
103    }
104}
needle_core/config.rs

needle_core/
config.rs
