Skip to main content

combs_formats/
metadata.rs

1//! Model metadata parsed from HuggingFace `config.json` (+ `generation_config.json`).
2
3use crate::{FormatError, Result};
4
5/// Architecture + hyperparameter description of a model, format-agnostic.
6///
7/// Field names follow the HuggingFace Llama config convention; other
8/// architecture families remap their config onto this struct in their adapter.
9#[derive(Debug, Clone)]
10pub struct ModelMetadata {
11    /// Architecture identifier, e.g. `"llama"`, `"smollm2"` (from
12    /// `config.json::model_type`). The model registry keys on this string.
13    pub architecture: String,
14    /// Hidden size (model dimension).
15    pub hidden_size: usize,
16    /// MLP intermediate size.
17    pub intermediate_size: usize,
18    /// Number of transformer layers.
19    pub num_hidden_layers: usize,
20    /// Number of attention query heads.
21    pub num_attention_heads: usize,
22    /// Number of key/value heads (GQA). Equal to `num_attention_heads` for MHA.
23    pub num_key_value_heads: usize,
24    /// Vocabulary size.
25    pub vocab_size: usize,
26    /// Maximum positional embeddings the model was built for.
27    pub max_position_embeddings: usize,
28    /// RMSNorm epsilon.
29    pub rms_norm_eps: f64,
30    /// RoPE base frequency (theta).
31    pub rope_theta: f64,
32    /// Whether lm_head is tied to the embedding matrix.
33    pub tie_word_embeddings: bool,
34    /// Per-head dimension, derived: `hidden_size / num_attention_heads`.
35    pub head_dim: usize,
36    /// Whether attention projections carry biases.
37    pub attention_bias: bool,
38    /// Beginning-of-sequence token id, if defined.
39    pub bos_token_id: Option<u32>,
40    /// End-of-sequence token ids (merged from config + generation_config).
41    pub eos_token_ids: Vec<u32>,
42    /// Vision-tower hyperparameters for multimodal models (Idefics3/SmolVLM
43    /// today); `None` for text-only models.
44    pub vision: Option<VisionConfig>,
45    /// Layer-type attention pattern (Gemma sliding-window interleave);
46    /// defaults to all-global (Llama-family behavior).
47    pub attention_pattern: AttentionPattern,
48    /// MLP activation (`hidden_act` / `hidden_activation`).
49    pub activation: Activation,
50    /// RoPE frequency scaling (`rope_scaling`); `None` variant when absent.
51    pub rope_scaling: RopeScaling,
52}
53
54/// MLP activation function, parsed from `hidden_act`/`hidden_activation`.
55/// The tanh-approximation family (`gelu_pytorch_tanh`, `gelu_new`,
56/// `gelu_fast`) all map to [`Activation::GeluTanh`].
57#[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            // silu/swish and anything unknown: the llama-family default.
71            _ => Activation::Silu,
72        }
73    }
74}
75
76/// RoPE frequency scaling, parsed from HF `rope_scaling` (accepts both the
77/// modern `rope_type` and the legacy `type` key). Table math lives in
78/// `combs-models::rope`; this is parse-only.
79#[derive(Debug, Clone, PartialEq, Default)]
80pub enum RopeScaling {
81    #[default]
82    None,
83    Linear {
84        factor: f64,
85    },
86    /// Llama-3.1+ piecewise frequency scaling.
87    Llama3 {
88        factor: f64,
89        low_freq_factor: f64,
90        high_freq_factor: f64,
91        original_max_position_embeddings: usize,
92    },
93    /// YaRN NTK-by-parts interpolation (+ attention temperature).
94    Yarn {
95        factor: f64,
96        original_max_position_embeddings: usize,
97        beta_fast: f64,
98        beta_slow: f64,
99        /// Explicit attention scaling; `None` = the YaRN default
100        /// `0.1·ln(factor) + 1`.
101        attention_factor: Option<f64>,
102    },
103    /// Phi-3 LongRoPE: per-dimension frequency divisors, one set for
104    /// prompts within the pretraining context and one beyond it.
105    LongRope {
106        short_factor: Vec<f64>,
107        long_factor: Vec<f64>,
108        original_max_position_embeddings: usize,
109        /// Context extension ratio `max_position / original_max` (phi
110        /// derives the attention temperature from it, not from a config
111        /// `factor` key).
112        factor: f64,
113        /// Explicit attention scaling; `None` = the LongRoPE default
114        /// `sqrt(1 + ln(factor)/ln(original_max))` (1.0 when factor ≤ 1).
115        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                // Phi-3 keeps the pretraining context length top-level (not
163                // inside rope_scaling), with max_position_embeddings already
164                // raised to the extended value; their ratio drives the
165                // attention temperature.
166                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/// Vision-encoder hyperparameters parsed from `config.json::vision_config`
195/// (plus top-level `scale_factor` / `image_token_id`).
196#[derive(Debug, Clone)]
197pub struct VisionConfig {
198    /// Square input image size (pixels).
199    pub image_size: usize,
200    /// Patch size (pixels) of the patch embedding.
201    pub patch_size: usize,
202    /// Vision hidden size.
203    pub hidden_size: usize,
204    /// Vision MLP intermediate size.
205    pub intermediate_size: usize,
206    /// Vision transformer layers.
207    pub num_hidden_layers: usize,
208    /// Vision attention heads (MHA — kv heads == q heads).
209    pub num_attention_heads: usize,
210    /// LayerNorm epsilon (SigLIP: 1e-6).
211    pub layer_norm_eps: f64,
212    /// Pixel-shuffle scale factor of the connector (scale² patches are
213    /// merged into one visual token).
214    pub scale_factor: usize,
215    /// Token id whose span in the prompt is replaced by visual embeddings.
216    pub image_token_id: u32,
217}
218
219/// Attention RoPE/scale settings that vary by layer type (Gemma2/3):
220/// `pattern`-th layers are "global" (full attention, `rope_theta`); the
221/// rest are "local" (sliding-window attention, `rope_local_theta`).
222#[derive(Debug, Clone)]
223pub struct AttentionPattern {
224    /// Sliding-window span for local layers; `None` = all layers global.
225    pub sliding_window: Option<usize>,
226    /// Every Nth layer is global (HF `sliding_window_pattern`, default 6).
227    pub pattern: usize,
228    /// RoPE base frequency for local layers (`rope_local_base_freq`).
229    pub rope_local_theta: f64,
230    /// Attention logit scale divisor (`query_pre_attn_scalar`); when
231    /// `None`, the scale is `1/sqrt(head_dim)`.
232    pub query_pre_attn_scalar: Option<f64>,
233    /// Qwen2-style partition, stored raw: the first N layers are global and
234    /// layers >= N slide — the inverse of `pattern`'s every-Nth-global.
235    /// `ArchSpec::resolve` turns it into the per-layer layout. All shipped
236    /// qwen2.5 checkpoints disable sliding anyway (`use_sliding_window:
237    /// false` nulls `sliding_window` at parse).
238    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    /// Whether layer `i` uses global attention (vs sliding-window local).
255    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    /// Visual tokens per image:
262    /// `(image_size / patch_size)² / scale_factor²`.
263    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    /// Vision per-head dimension.
269    pub fn head_dim(&self) -> usize {
270        self.hidden_size / self.num_attention_heads
271    }
272
273    /// Parses the vision section of a multimodal `config.json` (returns
274    /// `None` when no `vision_config` object is present).
275    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
320/// Extracts token ids from a config value that may be a single id or an array.
321fn 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    /// Minimal placeholder metadata for diffusion components that do not
334    /// carry a language-model `config.json` (UNet, VAE, etc.).
335    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    /// Parses metadata from a HuggingFace `config.json` value, optionally
360    /// merged with a `generation_config.json` value (which can override/add
361    /// bos/eos ids).
362    pub fn from_hf_config(
363        config: &serde_json::Value,
364        generation_config: Option<&serde_json::Value>,
365    ) -> Result<Self> {
366        // Whisper names its dimensions differently (d_model, encoder_layers,
367        // …). Remap them onto the shared keys and re-enter; the guard on
368        // `hidden_size` keeps the second pass out of this branch. Whisper-only
369        // geometry (mel bins, audio positions) is derived from tensor shapes
370        // at model load, so nothing else needs a slot here.
371        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        // Multimodal configs (Idefics3/SmolVLM) nest the text hyperparameters
393        // under `text_config`; the architecture id stays at the root.
394        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        // HF config classes carry per-family defaults the JSON omits:
428        // gemma3 ties word embeddings unless stated otherwise.
429        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                // Qwen2-family configs carry `sliding_window` even when
470                // sliding is off (`use_sliding_window: false`); an explicit
471                // false must null the window or the layer-pattern math would
472                // invent gemma-style local layers.
473                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        // GQA default: kv heads == q heads.
542        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        // Explicitly enabled (rare long-context qwen) keeps the window, so
558        // the registry guard can reject it loudly instead of running wrong.
559        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        // Absent key (gemma/mistral style) keeps the window too.
565        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        // Gemma3 config.json omits tie_word_embeddings entirely; the HF
574        // config class default (true) must apply — the checkpoints ship no
575        // lm_head.weight.
576        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        // An explicit false still wins (hypothetical untied variant)…
584        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        // …and non-gemma architectures keep the false default.
590        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        // Phi-3-mini-128k shape: `original_max_position_embeddings` lives at
599        // the TOP level (not inside rope_scaling), max_position already
600        // extended; factor derives from the ratio.
601        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        // (512/16)² / 4² = 1024/16 = 64 visual tokens per image.
670        assert_eq!(v.image_seq_len(), 64);
671    }
672}