1use crate::{FormatError, Result};
4
5#[derive(Debug, Clone)]
10pub struct ModelMetadata {
11 pub architecture: String,
14 pub hidden_size: usize,
16 pub intermediate_size: usize,
18 pub num_hidden_layers: usize,
20 pub num_attention_heads: usize,
22 pub num_key_value_heads: usize,
24 pub vocab_size: usize,
26 pub max_position_embeddings: usize,
28 pub rms_norm_eps: f64,
30 pub rope_theta: f64,
32 pub tie_word_embeddings: bool,
34 pub head_dim: usize,
36 pub attention_bias: bool,
38 pub bos_token_id: Option<u32>,
40 pub eos_token_ids: Vec<u32>,
42 pub vision: Option<VisionConfig>,
45 pub attention_pattern: AttentionPattern,
48 pub activation: Activation,
50 pub rope_scaling: RopeScaling,
52}
53
54#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
58pub enum Activation {
59 #[default]
60 Silu,
61 GeluTanh,
62 Gelu,
63}
64
65impl Activation {
66 fn parse(name: Option<&str>) -> Self {
67 match name {
68 Some("gelu_pytorch_tanh" | "gelu_new" | "gelu_fast") => Activation::GeluTanh,
69 Some("gelu") => Activation::Gelu,
70 _ => Activation::Silu,
72 }
73 }
74}
75
76#[derive(Debug, Clone, PartialEq, Default)]
80pub enum RopeScaling {
81 #[default]
82 None,
83 Linear {
84 factor: f64,
85 },
86 Llama3 {
88 factor: f64,
89 low_freq_factor: f64,
90 high_freq_factor: f64,
91 original_max_position_embeddings: usize,
92 },
93 Yarn {
95 factor: f64,
96 original_max_position_embeddings: usize,
97 beta_fast: f64,
98 beta_slow: f64,
99 attention_factor: Option<f64>,
102 },
103 LongRope {
106 short_factor: Vec<f64>,
107 long_factor: Vec<f64>,
108 original_max_position_embeddings: usize,
109 factor: f64,
113 attention_factor: Option<f64>,
116 },
117}
118
119impl RopeScaling {
120 fn parse(config: &serde_json::Value) -> Result<Self> {
121 let Some(rs) = config.get("rope_scaling").filter(|v| !v.is_null()) else {
122 return Ok(RopeScaling::None);
123 };
124 let kind = rs
125 .get("rope_type")
126 .or_else(|| rs.get("type"))
127 .and_then(|v| v.as_str())
128 .unwrap_or("default");
129 let f = |key: &str, default: f64| rs.get(key).and_then(|v| v.as_f64()).unwrap_or(default);
130 let factor = f("factor", 1.0);
131 match kind {
132 "default" => Ok(RopeScaling::None),
133 "linear" => Ok(RopeScaling::Linear { factor }),
134 "llama3" => Ok(RopeScaling::Llama3 {
135 factor,
136 low_freq_factor: f("low_freq_factor", 1.0),
137 high_freq_factor: f("high_freq_factor", 4.0),
138 original_max_position_embeddings: rs
139 .get("original_max_position_embeddings")
140 .and_then(|v| v.as_u64())
141 .unwrap_or(8192) as usize,
142 }),
143 "yarn" => Ok(RopeScaling::Yarn {
144 factor,
145 original_max_position_embeddings: rs
146 .get("original_max_position_embeddings")
147 .and_then(|v| v.as_u64())
148 .unwrap_or(32768) as usize,
149 beta_fast: f("beta_fast", 32.0),
150 beta_slow: f("beta_slow", 1.0),
151 attention_factor: rs.get("attention_factor").and_then(|v| v.as_f64()),
152 }),
153 "longrope" => {
154 let factors = |key: &str| -> Result<Vec<f64>> {
155 rs.get(key)
156 .and_then(|v| v.as_array())
157 .map(|a| a.iter().filter_map(|x| x.as_f64()).collect())
158 .ok_or_else(|| {
159 FormatError::MissingField(format!("rope_scaling.{key}"))
160 })
161 };
162 let original = rs
167 .get("original_max_position_embeddings")
168 .or_else(|| config.get("original_max_position_embeddings"))
169 .and_then(|v| v.as_u64())
170 .ok_or_else(|| {
171 FormatError::MissingField(
172 "original_max_position_embeddings (longrope)".to_string(),
173 )
174 })? as usize;
175 let max_pos = config
176 .get("max_position_embeddings")
177 .and_then(|v| v.as_u64())
178 .unwrap_or(original as u64) as usize;
179 Ok(RopeScaling::LongRope {
180 short_factor: factors("short_factor")?,
181 long_factor: factors("long_factor")?,
182 original_max_position_embeddings: original,
183 factor: max_pos as f64 / original as f64,
184 attention_factor: rs.get("attention_factor").and_then(|v| v.as_f64()),
185 })
186 }
187 other => Err(FormatError::MissingField(format!(
188 "unsupported rope_scaling type {other:?} (supported: linear, llama3, yarn, longrope)"
189 ))),
190 }
191 }
192}
193
194#[derive(Debug, Clone)]
197pub struct VisionConfig {
198 pub image_size: usize,
200 pub patch_size: usize,
202 pub hidden_size: usize,
204 pub intermediate_size: usize,
206 pub num_hidden_layers: usize,
208 pub num_attention_heads: usize,
210 pub layer_norm_eps: f64,
212 pub scale_factor: usize,
215 pub image_token_id: u32,
217}
218
219#[derive(Debug, Clone)]
223pub struct AttentionPattern {
224 pub sliding_window: Option<usize>,
226 pub pattern: usize,
228 pub rope_local_theta: f64,
230 pub query_pre_attn_scalar: Option<f64>,
233 pub max_window_layers: Option<usize>,
239}
240
241impl Default for AttentionPattern {
242 fn default() -> Self {
243 AttentionPattern {
244 sliding_window: None,
245 pattern: 6,
246 rope_local_theta: 10000.0,
247 query_pre_attn_scalar: None,
248 max_window_layers: None,
249 }
250 }
251}
252
253impl AttentionPattern {
254 pub fn is_global_layer(&self, i: usize) -> bool {
256 self.sliding_window.is_none() || (i + 1) % self.pattern == 0
257 }
258}
259
260impl VisionConfig {
261 pub fn image_seq_len(&self) -> usize {
264 let per_side = self.image_size / self.patch_size;
265 (per_side * per_side) / (self.scale_factor * self.scale_factor)
266 }
267
268 pub fn head_dim(&self) -> usize {
270 self.hidden_size / self.num_attention_heads
271 }
272
273 fn from_hf_config(config: &serde_json::Value) -> Result<Option<Self>> {
276 let Some(v) = config.get("vision_config").filter(|v| v.is_object()) else {
277 return Ok(None);
278 };
279 let get = |key: &str| v.get(key).and_then(|x| x.as_u64()).map(|x| x as usize);
280 let image_token_id = config
281 .get("image_token_id")
282 .and_then(|x| x.as_u64())
283 .ok_or_else(|| FormatError::MissingField("image_token_id".to_string()))?
284 as u32;
285 Ok(Some(VisionConfig {
286 image_size: get("image_size").unwrap_or(512),
287 patch_size: get("patch_size")
288 .ok_or_else(|| FormatError::MissingField("vision_config.patch_size".to_string()))?,
289 hidden_size: get("hidden_size")
290 .ok_or_else(|| FormatError::MissingField("vision_config.hidden_size".to_string()))?,
291 intermediate_size: get("intermediate_size")
292 .ok_or_else(|| FormatError::MissingField("vision_config.intermediate_size".to_string()))?,
293 num_hidden_layers: get("num_hidden_layers")
294 .ok_or_else(|| FormatError::MissingField("vision_config.num_hidden_layers".to_string()))?,
295 num_attention_heads: get("num_attention_heads")
296 .ok_or_else(|| FormatError::MissingField("vision_config.num_attention_heads".to_string()))?,
297 layer_norm_eps: v
298 .get("layer_norm_eps")
299 .and_then(|x| x.as_f64())
300 .unwrap_or(1e-12),
301 scale_factor: config
302 .get("scale_factor")
303 .and_then(|x| x.as_u64())
304 .unwrap_or(2) as usize,
305 image_token_id,
306 }))
307 }
308}
309
310fn get_u64(v: &serde_json::Value, key: &str) -> Result<u64> {
311 v.get(key)
312 .and_then(|x| x.as_u64())
313 .ok_or_else(|| FormatError::MissingField(key.to_string()))
314}
315
316fn get_f64(v: &serde_json::Value, key: &str, default: f64) -> f64 {
317 v.get(key).and_then(|x| x.as_f64()).unwrap_or(default)
318}
319
320fn token_ids(v: Option<&serde_json::Value>) -> Vec<u32> {
322 match v {
323 Some(serde_json::Value::Array(arr)) => arr
324 .iter()
325 .filter_map(|x| x.as_u64().map(|n| n as u32))
326 .collect(),
327 Some(x) => x.as_u64().map(|n| vec![n as u32]).unwrap_or_default(),
328 None => Vec::new(),
329 }
330}
331
332impl ModelMetadata {
333 pub fn diffusion_placeholder(architecture: &str) -> Self {
336 Self {
337 architecture: architecture.to_string(),
338 hidden_size: 0,
339 intermediate_size: 0,
340 num_hidden_layers: 0,
341 num_attention_heads: 0,
342 num_key_value_heads: 0,
343 vocab_size: 0,
344 max_position_embeddings: 0,
345 rms_norm_eps: 1e-6,
346 rope_theta: 10_000.0,
347 tie_word_embeddings: false,
348 head_dim: 0,
349 attention_bias: false,
350 bos_token_id: None,
351 eos_token_ids: Vec::new(),
352 vision: None,
353 attention_pattern: AttentionPattern::default(),
354 activation: Activation::default(),
355 rope_scaling: RopeScaling::default(),
356 }
357 }
358
359 pub fn from_hf_config(
363 config: &serde_json::Value,
364 generation_config: Option<&serde_json::Value>,
365 ) -> Result<Self> {
366 if config.get("model_type").and_then(|x| x.as_str()) == Some("whisper")
372 && config.get("hidden_size").is_none()
373 {
374 let mut remapped = config.clone();
375 let obj = remapped
376 .as_object_mut()
377 .ok_or_else(|| FormatError::MissingField("config object".to_string()))?;
378 for (from, to) in [
379 ("d_model", "hidden_size"),
380 ("encoder_attention_heads", "num_attention_heads"),
381 ("encoder_ffn_dim", "intermediate_size"),
382 ("encoder_layers", "num_hidden_layers"),
383 ("max_target_positions", "max_position_embeddings"),
384 ] {
385 if let Some(v) = config.get(from).cloned() {
386 obj.insert(to.to_string(), v);
387 }
388 }
389 obj.insert("tie_word_embeddings".to_string(), serde_json::json!(true));
390 return Self::from_hf_config(&remapped, generation_config);
391 }
392 let text = config.get("text_config").unwrap_or(config);
395 let hidden_size = get_u64(text, "hidden_size")? as usize;
396 let num_attention_heads = get_u64(text, "num_attention_heads")? as usize;
397 let num_key_value_heads = text
398 .get("num_key_value_heads")
399 .and_then(|x| x.as_u64())
400 .map(|x| x as usize)
401 .unwrap_or(num_attention_heads);
402
403 let mut eos_token_ids = token_ids(text.get("eos_token_id"));
404 if let Some(gc) = generation_config {
405 for id in token_ids(gc.get("eos_token_id")) {
406 if !eos_token_ids.contains(&id) {
407 eos_token_ids.push(id);
408 }
409 }
410 }
411
412 let bos_token_id = generation_config
413 .and_then(|gc| gc.get("bos_token_id"))
414 .and_then(|x| x.as_u64())
415 .map(|x| x as u32)
416 .or_else(|| {
417 text.get("bos_token_id")
418 .and_then(|x| x.as_u64())
419 .map(|x| x as u32)
420 });
421
422 let architecture = config
423 .get("model_type")
424 .and_then(|x| x.as_str())
425 .ok_or_else(|| FormatError::MissingField("model_type".to_string()))?
426 .to_string();
427 let tie_default = matches!(architecture.as_str(), "gemma3" | "gemma3_text");
430
431 if hidden_size % num_attention_heads != 0 {
432 return Err(FormatError::MissingField(format!(
433 "hidden_size ({hidden_size}) not divisible by num_attention_heads ({num_attention_heads})"
434 )));
435 }
436
437 Ok(ModelMetadata {
438 architecture,
439 hidden_size,
440 intermediate_size: get_u64(text, "intermediate_size")? as usize,
441 num_hidden_layers: get_u64(text, "num_hidden_layers")? as usize,
442 num_attention_heads,
443 num_key_value_heads,
444 vocab_size: get_u64(text, "vocab_size")? as usize,
445 max_position_embeddings: text
446 .get("max_position_embeddings")
447 .and_then(|x| x.as_u64())
448 .unwrap_or(2048) as usize,
449 rms_norm_eps: get_f64(text, "rms_norm_eps", 1e-5),
450 rope_theta: get_f64(text, "rope_theta", 10000.0),
451 tie_word_embeddings: config
452 .get("tie_word_embeddings")
453 .or_else(|| text.get("tie_word_embeddings"))
454 .and_then(|x| x.as_bool())
455 .unwrap_or(tie_default),
456 head_dim: text
457 .get("head_dim")
458 .and_then(|x| x.as_u64())
459 .map(|x| x as usize)
460 .unwrap_or(hidden_size / num_attention_heads),
461 attention_bias: text
462 .get("attention_bias")
463 .and_then(|x| x.as_bool())
464 .unwrap_or(false),
465 bos_token_id,
466 eos_token_ids,
467 vision: VisionConfig::from_hf_config(config)?,
468 attention_pattern: AttentionPattern {
469 sliding_window: text
474 .get("sliding_window")
475 .and_then(|x| x.as_u64())
476 .map(|x| x as usize)
477 .filter(|_| {
478 text.get("use_sliding_window").and_then(|x| x.as_bool())
479 != Some(false)
480 }),
481 pattern: get_u64(text, "sliding_window_pattern").unwrap_or(6) as usize,
482 rope_local_theta: get_f64(text, "rope_local_base_freq", 10000.0),
483 query_pre_attn_scalar: text
484 .get("query_pre_attn_scalar")
485 .and_then(|x| x.as_f64()),
486 max_window_layers: text
487 .get("max_window_layers")
488 .and_then(|x| x.as_u64())
489 .map(|x| x as usize),
490 },
491 activation: Activation::parse(
492 text.get("hidden_act")
493 .or_else(|| text.get("hidden_activation"))
494 .and_then(|x| x.as_str()),
495 ),
496 rope_scaling: RopeScaling::parse(text)?,
497 })
498 }
499}
500
501#[cfg(test)]
502mod tests {
503 use super::*;
504
505 #[test]
506 fn parses_smollm2_style_config() {
507 let config = serde_json::json!({
508 "model_type": "llama",
509 "hidden_size": 576,
510 "intermediate_size": 1536,
511 "num_hidden_layers": 30,
512 "num_attention_heads": 9,
513 "num_key_value_heads": 3,
514 "vocab_size": 49152,
515 "max_position_embeddings": 8192,
516 "rms_norm_eps": 1e-5,
517 "rope_theta": 100000,
518 "tie_word_embeddings": true,
519 "eos_token_id": 0,
520 "bos_token_id": 0
521 });
522 let meta = ModelMetadata::from_hf_config(&config, None).unwrap();
523 assert_eq!(meta.architecture, "llama");
524 assert_eq!(meta.head_dim, 64);
525 assert_eq!(meta.num_key_value_heads, 3);
526 assert_eq!(meta.rope_theta, 100000.0);
527 assert!(meta.tie_word_embeddings);
528 assert_eq!(meta.eos_token_ids, vec![0]);
529 }
530
531 #[test]
532 fn merges_generation_config_eos_array() {
533 let config = serde_json::json!({
534 "model_type": "llama", "hidden_size": 8, "intermediate_size": 16,
535 "num_hidden_layers": 1, "num_attention_heads": 2, "vocab_size": 10,
536 "eos_token_id": 1
537 });
538 let gen = serde_json::json!({ "eos_token_id": [1, 2] });
539 let meta = ModelMetadata::from_hf_config(&config, Some(&gen)).unwrap();
540 assert_eq!(meta.eos_token_ids, vec![1, 2]);
541 assert_eq!(meta.num_key_value_heads, 2);
543 }
544
545 #[test]
546 fn qwen2_use_sliding_window_false_nulls_the_window() {
547 let base = serde_json::json!({
548 "model_type": "qwen2", "hidden_size": 8, "intermediate_size": 16,
549 "num_hidden_layers": 2, "num_attention_heads": 2, "vocab_size": 10,
550 "sliding_window": 131072, "use_sliding_window": false,
551 "max_window_layers": 28
552 });
553 let meta = ModelMetadata::from_hf_config(&base, None).unwrap();
554 assert_eq!(meta.attention_pattern.sliding_window, None);
555 assert_eq!(meta.attention_pattern.max_window_layers, Some(28));
556
557 let mut on = base.clone();
560 on["use_sliding_window"] = serde_json::json!(true);
561 let meta = ModelMetadata::from_hf_config(&on, None).unwrap();
562 assert_eq!(meta.attention_pattern.sliding_window, Some(131072));
563
564 let mut absent = base.clone();
566 absent.as_object_mut().unwrap().remove("use_sliding_window");
567 let meta = ModelMetadata::from_hf_config(&absent, None).unwrap();
568 assert_eq!(meta.attention_pattern.sliding_window, Some(131072));
569 }
570
571 #[test]
572 fn gemma3_defaults_to_tied_embeddings() {
573 let config = serde_json::json!({
577 "model_type": "gemma3_text", "hidden_size": 8, "intermediate_size": 16,
578 "num_hidden_layers": 1, "num_attention_heads": 2, "vocab_size": 10
579 });
580 let meta = ModelMetadata::from_hf_config(&config, None).unwrap();
581 assert!(meta.tie_word_embeddings);
582
583 let mut untied = config.clone();
585 untied["tie_word_embeddings"] = serde_json::json!(false);
586 let meta = ModelMetadata::from_hf_config(&untied, None).unwrap();
587 assert!(!meta.tie_word_embeddings);
588
589 let mut llama = config.clone();
591 llama["model_type"] = serde_json::json!("llama");
592 let meta = ModelMetadata::from_hf_config(&llama, None).unwrap();
593 assert!(!meta.tie_word_embeddings);
594 }
595
596 #[test]
597 fn parses_phi3_longrope_with_toplevel_original_max() {
598 let config = serde_json::json!({
602 "model_type": "phi3", "hidden_size": 3072, "intermediate_size": 8192,
603 "num_hidden_layers": 32, "num_attention_heads": 32, "vocab_size": 32064,
604 "max_position_embeddings": 131072,
605 "original_max_position_embeddings": 4096,
606 "rope_scaling": {
607 "type": "longrope",
608 "short_factor": [1.0, 1.05, 1.1],
609 "long_factor": [2.0, 2.5, 3.0]
610 }
611 });
612 let meta = ModelMetadata::from_hf_config(&config, None).unwrap();
613 match &meta.rope_scaling {
614 RopeScaling::LongRope {
615 short_factor,
616 long_factor,
617 original_max_position_embeddings,
618 factor,
619 attention_factor,
620 } => {
621 assert_eq!(short_factor, &[1.0, 1.05, 1.1]);
622 assert_eq!(long_factor, &[2.0, 2.5, 3.0]);
623 assert_eq!(*original_max_position_embeddings, 4096);
624 assert_eq!(*factor, 32.0);
625 assert_eq!(*attention_factor, None);
626 }
627 other => panic!("expected LongRope, got {other:?}"),
628 }
629 }
630
631 #[test]
632 fn parses_nested_idefics3_config() {
633 let config = serde_json::json!({
634 "model_type": "idefics3",
635 "image_token_id": 49190,
636 "scale_factor": 4,
637 "tie_word_embeddings": false,
638 "text_config": {
639 "hidden_size": 576,
640 "intermediate_size": 1536,
641 "num_hidden_layers": 30,
642 "num_attention_heads": 9,
643 "num_key_value_heads": 3,
644 "vocab_size": 49280,
645 "max_position_embeddings": 8192,
646 "rms_norm_eps": 1e-5,
647 "rope_theta": 100000,
648 "eos_token_id": 2
649 },
650 "vision_config": {
651 "hidden_size": 768,
652 "intermediate_size": 3072,
653 "num_hidden_layers": 12,
654 "num_attention_heads": 12,
655 "image_size": 512,
656 "patch_size": 16,
657 "layer_norm_eps": 1e-6
658 }
659 });
660 let meta = ModelMetadata::from_hf_config(&config, None).unwrap();
661 assert_eq!(meta.architecture, "idefics3");
662 assert_eq!(meta.hidden_size, 576);
663 assert_eq!(meta.eos_token_ids, vec![2]);
664 let v = meta.vision.expect("vision config parsed");
665 assert_eq!(v.hidden_size, 768);
666 assert_eq!(v.image_token_id, 49190);
667 assert_eq!(v.scale_factor, 4);
668 assert_eq!(v.head_dim(), 64);
669 assert_eq!(v.image_seq_len(), 64);
671 }
672}