laya/
config.rs

1//! Checkpoint configuration, validated exactly as upstream validates it.
2
3use laya_core::error::{Error, Result};
4use serde::Deserialize;
5use serde_json::{Map, Value};
6use std::collections::BTreeMap;
7
8pub const FULL_ATTENTION: &str = "full_attention";
9pub const SLIDING_ATTENTION: &str = "sliding_attention";
10
11fn default_norm_eps() -> f64 {
12    1e-5
13}
14fn default_activation() -> String {
15    "gelu".into()
16}
17fn default_model_type() -> String {
18    "modernbert".into()
19}
20fn default_local_attention() -> usize {
21    128
22}
23fn default_global_every() -> usize {
24    3
25}
26fn default_global_rope() -> f64 {
27    160_000.0
28}
29fn default_local_rope() -> f64 {
30    10_000.0
31}
32fn default_max_positions() -> usize {
33    8192
34}
35
36/// `encoder/config.json`.
37#[derive(Clone, Debug, Deserialize)]
38pub struct EncoderConfig {
39    pub vocab_size: usize,
40    pub hidden_size: usize,
41    pub intermediate_size: usize,
42    pub num_hidden_layers: usize,
43    pub num_attention_heads: usize,
44
45    #[serde(default = "default_model_type")]
46    pub model_type: String,
47    #[serde(default = "default_norm_eps")]
48    pub norm_eps: f64,
49    #[serde(default)]
50    pub norm_bias: bool,
51    #[serde(default)]
52    pub attention_bias: bool,
53    #[serde(default)]
54    pub mlp_bias: bool,
55    #[serde(default = "default_activation")]
56    pub hidden_activation: String,
57    #[serde(default = "default_local_attention")]
58    pub local_attention: usize,
59    #[serde(default = "default_global_every")]
60    pub global_attn_every_n_layers: usize,
61    #[serde(default = "default_global_rope")]
62    pub global_rope_theta: f64,
63    #[serde(default = "default_local_rope")]
64    pub local_rope_theta: f64,
65    #[serde(default = "default_max_positions")]
66    pub max_position_embeddings: usize,
67    #[serde(default)]
68    pub layer_types: Option<Vec<String>>,
69    #[serde(default)]
70    pub rope_parameters: Option<BTreeMap<String, Value>>,
71}
72
73impl EncoderConfig {
74    /// Parse and validate, mirroring `EncoderConfig.from_dict`.
75    pub fn from_value(value: &Value) -> Result<Self> {
76        let mut config: Self = serde_json::from_value(value.clone())
77            .map_err(|error| Error::Config(format!("Invalid encoder config: {error}")))?;
78
79        if config.model_type != "modernbert" {
80            return Err(Error::Config(format!(
81                "Unsupported encoder: {:?}; expected modernbert",
82                config.model_type
83            )));
84        }
85        if config.hidden_activation != "gelu" {
86            return Err(Error::Config(format!(
87                "Unsupported encoder activation: {:?}",
88                config.hidden_activation
89            )));
90        }
91        if config.num_attention_heads == 0
92            || config.hidden_size % config.num_attention_heads != 0
93            || config.head_dim() % 2 != 0
94        {
95            return Err(Error::Config(
96                "ModernBERT requires an even, integral attention head dimension".into(),
97            ));
98        }
99
100        let layer_types = match config.layer_types.take() {
101            Some(types) => types,
102            None => (0..config.num_hidden_layers)
103                .map(|index| {
104                    if config.global_attn_every_n_layers == 0
105                        || index % config.global_attn_every_n_layers == 0
106                    {
107                        FULL_ATTENTION.to_string()
108                    } else {
109                        SLIDING_ATTENTION.to_string()
110                    }
111                })
112                .collect(),
113        };
114        if layer_types.len() != config.num_hidden_layers
115            || layer_types
116                .iter()
117                .any(|kind| kind != FULL_ATTENTION && kind != SLIDING_ATTENTION)
118        {
119            return Err(Error::Config("Invalid ModernBERT layer_types".into()));
120        }
121
122        for kind in [FULL_ATTENTION, SLIDING_ATTENTION] {
123            if !layer_types.iter().any(|value| value == kind) {
124                continue;
125            }
126            let rope_type = config
127                .rope_parameters
128                .as_ref()
129                .and_then(|params| params.get(kind))
130                .and_then(|params| params.get("rope_type"))
131                .and_then(Value::as_str)
132                .unwrap_or("default");
133            if rope_type != "default" {
134                return Err(Error::Config(
135                    "Only default (unscaled) ModernBERT RoPE is supported".into(),
136                ));
137            }
138        }
139
140        config.layer_types = Some(layer_types);
141        Ok(config)
142    }
143
144    pub fn head_dim(&self) -> usize {
145        self.hidden_size / self.num_attention_heads
146    }
147
148    pub fn layer_types(&self) -> &[String] {
149        self.layer_types
150            .as_deref()
151            .expect("layer_types filled in by from_value")
152    }
153
154    /// RoPE base for an attention kind, honouring per-kind overrides.
155    pub fn rope_base(&self, kind: &str) -> f64 {
156        let fallback = if kind == FULL_ATTENTION {
157            self.global_rope_theta
158        } else {
159            self.local_rope_theta
160        };
161        self.rope_parameters
162            .as_ref()
163            .and_then(|params| params.get(kind))
164            .and_then(|params| params.get("rope_theta"))
165            .and_then(Value::as_f64)
166            .unwrap_or(fallback)
167    }
168}
169
170/// `rl_agent_config.json`.
171#[derive(Clone, Debug)]
172pub struct AgentConfig {
173    pub head_layers: usize,
174    pub max_len: usize,
175    pub head_max_len: usize,
176    /// Number of action classes, i.e. `len(act_costs) + 1`.
177    pub act_classes: usize,
178    pub temperature: [f64; 3],
179    pub temperature_by_options: BTreeMap<String, f64>,
180    /// Retained so `convert` can round-trip the file untouched.
181    pub raw: Value,
182}
183
184impl AgentConfig {
185    pub fn from_value(value: &Value, max_position_embeddings: usize) -> Result<Self> {
186        let object: &Map<String, Value> = value
187            .as_object()
188            .ok_or_else(|| Error::Config("Laya config must be an object".into()))?;
189
190        if !object.contains_key("encoder") || !object.contains_key("head_layers") {
191            return Err(Error::Config(
192                "Laya config must specify encoder and head_layers".into(),
193            ));
194        }
195
196        let head_layers = object["head_layers"]
197            .as_u64()
198            .ok_or_else(|| Error::Config("head_layers must be an integer".into()))?
199            as usize;
200        // Upstream hardcodes 512, which caps a ModernBERT checkpoint well below
201        // the 8192 positions it was trained for, and rejects an encoder shorter
202        // than 512 outright. Default to what the encoder actually supports. An
203        // explicit max_len still wins, so the parity fixtures are unaffected.
204        let max_len = object
205            .get("max_len")
206            .and_then(Value::as_u64)
207            .map_or(max_position_embeddings, |value| value as usize);
208        let head_max_len = object
209            .get("head_max_len")
210            .and_then(Value::as_u64)
211            .unwrap_or(192) as usize;
212
213        if !(head_max_len > 4 && max_len > head_max_len && max_len <= max_position_embeddings) {
214            return Err(Error::Config(
215                "Expected 4 < head_max_len < max_len <= max_position_embeddings".into(),
216            ));
217        }
218
219        let act_classes = object
220            .get("act_costs")
221            .and_then(Value::as_object)
222            .map_or(0, Map::len)
223            + 1;
224
225        let temperature: Vec<f64> = match object.get("temperature") {
226            Some(Value::Array(values)) => values.iter().filter_map(Value::as_f64).collect(),
227            None => vec![1.0, 1.0, 1.0],
228            _ => return Err(Error::Config("temperature must be a list".into())),
229        };
230        if temperature.len() != 3 {
231            return Err(Error::Config(
232                "Calibration temperatures must be finite and positive".into(),
233            ));
234        }
235
236        let mut temperature_by_options = BTreeMap::new();
237        if let Some(Value::Object(map)) = object.get("temperature_by_options") {
238            for (key, value) in map {
239                let value = value.as_f64().ok_or_else(|| {
240                    Error::Config("Calibration temperatures must be finite and positive".into())
241                })?;
242                temperature_by_options.insert(key.clone(), value);
243            }
244        }
245
246        if temperature
247            .iter()
248            .chain(temperature_by_options.values())
249            .any(|value| !value.is_finite() || *value <= 0.0)
250        {
251            return Err(Error::Config(
252                "Calibration temperatures must be finite and positive".into(),
253            ));
254        }
255
256        Ok(Self {
257            head_layers,
258            max_len,
259            head_max_len,
260            act_classes,
261            temperature: [temperature[0], temperature[1], temperature[2]],
262            temperature_by_options,
263            raw: value.clone(),
264        })
265    }
266
267    /// Calibration scale for a question type and option count.
268    pub fn temperature_for(&self, kind: laya_core::QuestionKind, k: usize) -> f64 {
269        let bucket = laya_core::calibrate::temp_bucket(kind, k);
270        self.temperature_by_options
271            .get(&bucket)
272            .copied()
273            .unwrap_or(self.temperature[kind.index()])
274    }
275}
276
277#[cfg(test)]
278mod tests {
279    use super::*;
280    use serde_json::json;
281
282    fn encoder(extra: Value) -> Value {
283        let mut base = json!({
284            "model_type": "modernbert",
285            "vocab_size": 128,
286            "hidden_size": 64,
287            "intermediate_size": 96,
288            "num_hidden_layers": 3,
289            "num_attention_heads": 1,
290            "local_attention": 16,
291            "max_position_embeddings": 256
292        });
293        if let (Value::Object(base), Value::Object(extra)) = (&mut base, extra) {
294            base.extend(extra);
295        }
296        base
297    }
298
299    #[test]
300    fn derives_alternating_layer_types() {
301        let config = EncoderConfig::from_value(&encoder(json!({}))).unwrap();
302        assert_eq!(
303            config.layer_types(),
304            [FULL_ATTENTION, SLIDING_ATTENTION, SLIDING_ATTENTION]
305        );
306    }
307
308    #[test]
309    fn rope_base_falls_back_per_kind() {
310        let config = EncoderConfig::from_value(&encoder(json!({}))).unwrap();
311        assert_eq!(config.rope_base(FULL_ATTENTION), 160_000.0);
312        assert_eq!(config.rope_base(SLIDING_ATTENTION), 10_000.0);
313    }
314
315    #[test]
316    fn rope_override_is_honoured() {
317        let config = EncoderConfig::from_value(&encoder(json!({
318            "rope_parameters": {"full_attention": {"rope_theta": 5.0}}
319        })))
320        .unwrap();
321        assert_eq!(config.rope_base(FULL_ATTENTION), 5.0);
322    }
323
324    #[test]
325    fn rejects_unsupported_encoders() {
326        assert!(EncoderConfig::from_value(&encoder(json!({"model_type": "bert"}))).is_err());
327        assert!(EncoderConfig::from_value(&encoder(json!({"hidden_activation": "silu"}))).is_err());
328        assert!(EncoderConfig::from_value(&encoder(json!({
329            "rope_parameters": {"full_attention": {"rope_type": "yarn"}}
330        })))
331        .is_err());
332    }
333
334    #[test]
335    fn rejects_odd_head_dimensions() {
336        assert!(EncoderConfig::from_value(&encoder(json!({"num_attention_heads": 5}))).is_err());
337    }
338
339    #[test]
340    fn agent_config_validates_lengths() {
341        let good = json!({"encoder": "x", "head_layers": 1, "max_len": 128, "head_max_len": 32});
342        assert!(AgentConfig::from_value(&good, 256).is_ok());
343
344        let bad = json!({"encoder": "x", "head_layers": 1, "max_len": 32, "head_max_len": 128});
345        assert!(AgentConfig::from_value(&bad, 256).is_err());
346
347        let missing = json!({"encoder": "x"});
348        assert!(AgentConfig::from_value(&missing, 256).is_err());
349    }
350
351    #[test]
352    fn max_len_defaults_to_encoder_capacity() {
353        let config = json!({"encoder": "x", "head_layers": 1});
354
355        let agent = AgentConfig::from_value(&config, 8192).unwrap();
356        assert_eq!(agent.max_len, 8192);
357        // The question budget is not a function of encoder size, so it holds.
358        assert_eq!(agent.head_max_len, 192);
359
360        // An encoder smaller than upstream's hardcoded 512 now works instead
361        // of failing the max_len <= max_position_embeddings check.
362        let small = AgentConfig::from_value(&config, 256).unwrap();
363        assert_eq!(small.max_len, 256);
364
365        // An explicit value still wins over the encoder capacity.
366        let pinned = json!({"encoder": "x", "head_layers": 1, "max_len": 1024});
367        assert_eq!(
368            AgentConfig::from_value(&pinned, 8192).unwrap().max_len,
369            1024
370        );
371    }
372
373    #[test]
374    fn agent_config_rejects_bad_temperatures() {
375        let bad = json!({
376            "encoder": "x", "head_layers": 1, "max_len": 128, "head_max_len": 32,
377            "temperature": [1.0, 0.0, 1.0]
378        });
379        assert!(AgentConfig::from_value(&bad, 256).is_err());
380    }
381}
laya/config.rs

laya/
config.rs
