ferrotensor_hub/
config.rs

1//! Model configuration parsing
2
3use crate::HubResult;
4use serde::{Deserialize, Serialize};
5use std::fs;
6use std::path::Path;
7
8/// Model configuration from config.json
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct ModelConfig {
11    /// Model architecture type
12    #[serde(default)]
13    pub architectures: Vec<String>,
14
15    /// Hidden size / embedding dimension
16    #[serde(default)]
17    pub hidden_size: usize,
18
19    /// Number of attention heads
20    #[serde(default)]
21    pub num_attention_heads: usize,
22
23    /// Number of key-value heads (for GQA)
24    #[serde(default)]
25    pub num_key_value_heads: Option<usize>,
26
27    /// Number of transformer layers
28    #[serde(default)]
29    pub num_hidden_layers: usize,
30
31    /// Intermediate size (MLP)
32    #[serde(default)]
33    pub intermediate_size: usize,
34
35    /// Vocabulary size
36    #[serde(default)]
37    pub vocab_size: usize,
38
39    /// Maximum context length
40    #[serde(default)]
41    pub max_position_embeddings: usize,
42
43    /// RMS norm epsilon
44    #[serde(default = "default_rms_eps")]
45    pub rms_norm_eps: f32,
46
47    /// RoPE theta
48    #[serde(default = "default_rope_theta")]
49    pub rope_theta: f32,
50
51    /// Tie word embeddings
52    #[serde(default)]
53    pub tie_word_embeddings: bool,
54
55    /// Model type string
56    #[serde(default)]
57    pub model_type: String,
58}
59
60fn default_rms_eps() -> f32 {
61    1e-5
62}
63fn default_rope_theta() -> f32 {
64    10000.0
65}
66
67impl ModelConfig {
68    /// Load config from a config.json file.
69    pub fn from_file(path: &Path) -> HubResult<Self> {
70        let content = fs::read_to_string(path)?;
71        Ok(serde_json::from_str(&content)?)
72    }
73
74    /// Load config from a model directory.
75    pub fn from_model_dir(model_dir: &Path) -> HubResult<Self> {
76        Self::from_file(&model_dir.join("config.json"))
77    }
78
79    /// Get the head dimension.
80    pub fn head_dim(&self) -> usize {
81        self.hidden_size / self.num_attention_heads
82    }
83
84    /// Get number of KV heads (for GQA).
85    pub fn kv_heads(&self) -> usize {
86        self.num_key_value_heads.unwrap_or(self.num_attention_heads)
87    }
88
89    /// Check if using Grouped Query Attention.
90    pub fn uses_gqa(&self) -> bool {
91        self.kv_heads() != self.num_attention_heads
92    }
93}
94
95impl Default for ModelConfig {
96    fn default() -> Self {
97        Self {
98            architectures: vec![],
99            hidden_size: 4096,
100            num_attention_heads: 32,
101            num_key_value_heads: None,
102            num_hidden_layers: 32,
103            intermediate_size: 11008,
104            vocab_size: 32000,
105            max_position_embeddings: 4096,
106            rms_norm_eps: 1e-5,
107            rope_theta: 10000.0,
108            tie_word_embeddings: false,
109            model_type: "llama".to_string(),
110        }
111    }
112}
ferrotensor_hub/config.rs

ferrotensor_hub/
config.rs
