ferrum_types/
models.rs

1//! Model-related types and configurations
2
3use crate::{devices::*, ids::ModelId, FerrumError, Result};
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6
7/// Portable repository-relative path shared by GGUF selection and its evidence.
8pub fn validate_gguf_filename(path: &str) -> Result<()> {
9    let valid = path.to_ascii_lowercase().ends_with(".gguf")
10        && !path.chars().any(|c| {
11            c.is_control() || matches!(c, '\\' | '%' | '?' | '#' | ':' | '*' | '<' | '>' | '|')
12        })
13        && path
14            .split('/')
15            .all(|part| !matches!(part, "" | "." | "..") && !part.ends_with([' ', '.']));
16    if !valid {
17        return Err(FerrumError::config("--gguf-file must be a portable repository-relative .gguf path without URL metacharacters"));
18    }
19    Ok(())
20}
21
22/// Model type enumeration
23#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
24pub enum ModelType {
25    /// LLaMA family models
26    Llama,
27    /// Mistral family models
28    Mistral,
29    /// Qwen family models
30    Qwen,
31    /// Phi family models
32    Phi,
33    /// Gemma family models
34    Gemma,
35    /// Code-specific models
36    Code(String),
37    /// Embedding models (BERT, etc.)
38    Embedding,
39    /// CLIP vision-language models
40    Clip,
41    /// Custom model implementation
42    Custom(String),
43}
44
45impl std::fmt::Display for ModelType {
46    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
47        match self {
48            ModelType::Llama => write!(f, "llama"),
49            ModelType::Mistral => write!(f, "mistral"),
50            ModelType::Qwen => write!(f, "qwen"),
51            ModelType::Phi => write!(f, "phi"),
52            ModelType::Gemma => write!(f, "gemma"),
53            ModelType::Embedding => write!(f, "embedding"),
54            ModelType::Clip => write!(f, "clip"),
55            ModelType::Code(name) => write!(f, "code-{}", name),
56            ModelType::Custom(name) => write!(f, "custom-{}", name),
57        }
58    }
59}
60
61/// Model information and metadata
62#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct ModelInfo {
64    /// Model identifier
65    pub model_id: ModelId,
66    /// Model type/architecture
67    pub model_type: ModelType,
68    /// Number of parameters
69    pub num_parameters: u64,
70    /// Hidden dimension size
71    pub hidden_size: usize,
72    /// Number of transformer layers
73    pub num_layers: usize,
74    /// Number of attention heads
75    pub num_heads: usize,
76    /// Number of key-value heads (for GQA)
77    pub num_kv_heads: usize,
78    /// Vocabulary size
79    pub vocab_size: usize,
80    /// Maximum sequence length
81    pub max_sequence_length: usize,
82    /// Data type used by the model
83    pub dtype: DataType,
84    /// Device where model is loaded
85    pub device: Device,
86    /// Model version or revision
87    pub version: Option<String>,
88    /// Model license
89    pub license: Option<String>,
90    /// Additional model metadata
91    pub metadata: HashMap<String, serde_json::Value>,
92}
93
94impl ModelInfo {
95    /// Calculate approximate model size in bytes
96    pub fn estimated_size_bytes(&self) -> u64 {
97        // Rough estimation: parameters * dtype size + some overhead
98        let param_size = self.num_parameters * self.dtype.size_bytes() as u64;
99        // Add ~20% overhead for embeddings, activations, etc.
100        (param_size as f64 * 1.2) as u64
101    }
102
103    /// Check if model supports a specific sequence length
104    pub fn supports_sequence_length(&self, length: usize) -> bool {
105        length <= self.max_sequence_length
106    }
107
108    /// Get memory requirements for inference
109    pub fn memory_requirements(
110        &self,
111        batch_size: usize,
112        sequence_length: usize,
113    ) -> ModelMemoryRequirements {
114        let param_memory = self.estimated_size_bytes();
115
116        // Estimate KV cache size: layers * heads * seq_len * head_dim * 2 (key + value) * dtype * batch_size
117        let head_dim = self.hidden_size / self.num_heads;
118        let kv_cache_per_token =
119            self.num_layers * self.num_kv_heads * head_dim * 2 * self.dtype.size_bytes();
120        let kv_cache_memory = (kv_cache_per_token * sequence_length * batch_size) as u64;
121
122        // Estimate activation memory (rough approximation)
123        let activation_memory =
124            (self.hidden_size * sequence_length * batch_size * self.dtype.size_bytes()) as u64 * 4;
125
126        ModelMemoryRequirements {
127            parameter_memory: param_memory,
128            kv_cache_memory,
129            activation_memory,
130            total_estimated: param_memory + kv_cache_memory + activation_memory,
131        }
132    }
133}
134
135/// Memory requirements for model inference
136#[derive(Debug, Clone, Serialize, Deserialize)]
137pub struct ModelMemoryRequirements {
138    /// Memory required for model parameters
139    pub parameter_memory: u64,
140    /// Memory required for KV cache
141    pub kv_cache_memory: u64,
142    /// Memory required for activations
143    pub activation_memory: u64,
144    /// Total estimated memory requirement
145    pub total_estimated: u64,
146}
147
148/// Model configuration for runtime
149#[derive(Debug, Clone, Serialize, Deserialize)]
150pub struct ModelConfig {
151    /// Model identifier
152    pub model_id: ModelId,
153    /// Path to model files
154    pub model_path: String,
155    /// Model type/architecture
156    pub model_type: ModelType,
157    /// Data type to use for inference
158    pub dtype: DataType,
159    /// Target device
160    pub device: Device,
161    /// Maximum batch size
162    pub max_batch_size: usize,
163    /// Maximum sequence length
164    pub max_sequence_length: usize,
165    /// Tensor parallelism size
166    pub tensor_parallel_size: Option<usize>,
167    /// Pipeline parallelism size  
168    pub pipeline_parallel_size: Option<usize>,
169    /// Quantization configuration
170    pub quantization: Option<QuantizationConfig>,
171    /// Use flash attention if available
172    pub use_flash_attention: bool,
173    /// Use paged attention for KV cache
174    pub use_paged_attention: bool,
175    /// Enable CUDA graphs for optimization
176    pub enable_cuda_graphs: bool,
177    /// Additional configuration parameters
178    pub extra_config: HashMap<String, serde_json::Value>,
179}
180
181impl ModelConfig {
182    /// Create a new model configuration
183    pub fn new(model_id: impl Into<ModelId>, model_path: impl Into<String>) -> Self {
184        Self {
185            model_id: model_id.into(),
186            model_path: model_path.into(),
187            model_type: ModelType::Custom("unknown".to_string()),
188            dtype: DataType::FP16,
189            device: Device::CPU,
190            max_batch_size: 1,
191            max_sequence_length: 2048,
192            tensor_parallel_size: None,
193            pipeline_parallel_size: None,
194            quantization: None,
195            use_flash_attention: false,
196            use_paged_attention: false,
197            enable_cuda_graphs: false,
198            extra_config: HashMap::new(),
199        }
200    }
201
202    /// Validate the configuration
203    pub fn validate(&self) -> Result<()> {
204        if self.model_path.is_empty() {
205            return Err(FerrumError::config("Model path cannot be empty"));
206        }
207
208        if self.max_batch_size == 0 {
209            return Err(FerrumError::config("Max batch size must be positive"));
210        }
211
212        if self.max_sequence_length == 0 {
213            return Err(FerrumError::config("Max sequence length must be positive"));
214        }
215
216        if let Some(tp_size) = self.tensor_parallel_size {
217            if tp_size == 0 {
218                return Err(FerrumError::config("Tensor parallel size must be positive"));
219            }
220        }
221
222        if let Some(pp_size) = self.pipeline_parallel_size {
223            if pp_size == 0 {
224                return Err(FerrumError::config(
225                    "Pipeline parallel size must be positive",
226                ));
227            }
228        }
229
230        Ok(())
231    }
232}
233
234/// Quantization configuration
235#[derive(Debug, Clone, Serialize, Deserialize)]
236pub enum QuantizationConfig {
237    /// GPTQ quantization
238    GPTQ {
239        bits: u8,
240        group_size: usize,
241        desc_act: bool,
242    },
243    /// AWQ quantization
244    AWQ {
245        bits: u8,
246        zero_point: bool,
247        version: String,
248    },
249    /// FP8 quantization
250    FP8 { e4m3: bool, kv_cache: bool },
251    /// INT8 quantization
252    INT8 { symmetric: bool, per_channel: bool },
253    /// INT4 quantization
254    INT4 { symmetric: bool, group_size: usize },
255    /// SmoothQuant
256    SmoothQuant { alpha: f32, calibration_size: usize },
257}
258
259impl QuantizationConfig {
260    /// Get the number of bits used by this quantization method
261    pub fn bits(&self) -> u8 {
262        match self {
263            QuantizationConfig::GPTQ { bits, .. } => *bits,
264            QuantizationConfig::AWQ { bits, .. } => *bits,
265            QuantizationConfig::FP8 { .. } => 8,
266            QuantizationConfig::INT8 { .. } => 8,
267            QuantizationConfig::INT4 { .. } => 4,
268            QuantizationConfig::SmoothQuant { .. } => 8,
269        }
270    }
271
272    /// Check if this quantization preserves accuracy well
273    pub fn is_high_accuracy(&self) -> bool {
274        match self {
275            QuantizationConfig::FP8 { .. } => true,
276            QuantizationConfig::INT8 { .. } => true,
277            QuantizationConfig::SmoothQuant { .. } => true,
278            _ => false,
279        }
280    }
281}
282
283/// Token usage statistics
284#[derive(Debug, Clone, Serialize, Deserialize)]
285pub struct TokenUsage {
286    /// Number of tokens in the prompt
287    pub prompt_tokens: usize,
288    /// Number of tokens generated
289    pub completion_tokens: usize,
290    /// Total tokens processed
291    pub total_tokens: usize,
292}
293
294impl TokenUsage {
295    /// Create new token usage
296    pub fn new(prompt_tokens: usize, completion_tokens: usize) -> Self {
297        Self {
298            prompt_tokens,
299            completion_tokens,
300            total_tokens: prompt_tokens + completion_tokens,
301        }
302    }
303
304    /// Add completion tokens
305    pub fn add_completion_tokens(&mut self, tokens: usize) {
306        self.completion_tokens += tokens;
307        self.total_tokens = self.prompt_tokens + self.completion_tokens;
308    }
309}
310
311/// RoPE (Rotary Position Embedding) scaling configuration
312#[derive(Debug, Clone, Serialize, Deserialize)]
313pub struct RopeScaling {
314    /// Type of scaling: "linear", "dynamic", etc.
315    pub scaling_type: String,
316    /// Scaling factor
317    pub factor: f32,
318}
319
320/// Normalization type used in the model
321#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
322pub enum NormType {
323    /// Layer Normalization
324    LayerNorm,
325    /// Root Mean Square Normalization
326    RMSNorm,
327}
328
329/// Activation function type
330#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
331pub enum Activation {
332    /// Gaussian Error Linear Unit
333    GELU,
334    /// GELU tanh approximation (HF `gelu_pytorch_tanh`, Gemma family)
335    GeluTanh,
336    /// Sigmoid Linear Unit
337    SiLU,
338    /// Rectified Linear Unit
339    ReLU,
340    /// Swish activation
341    Swish,
342}
343
344/// Attention configuration for model architecture
345#[derive(Debug, Clone, Serialize, Deserialize, Default)]
346pub struct AttentionConfig {
347    /// Whether attention uses bias
348    pub attention_bias: bool,
349    /// Sliding window size (None for full attention)
350    pub sliding_window: Option<usize>,
351}
352
353/// Model loading source specification
354#[derive(Debug, Clone, Serialize, Deserialize)]
355pub enum ModelSource {
356    /// Local file path
357    Local(String),
358    /// Hugging Face Hub model
359    HuggingFace {
360        repo_id: String,
361        revision: Option<String>,
362        cache_dir: Option<String>,
363    },
364    /// URL download
365    Url {
366        url: String,
367        headers: HashMap<String, String>,
368    },
369    /// S3-compatible storage
370    S3 {
371        bucket: String,
372        key: String,
373        region: Option<String>,
374        endpoint: Option<String>,
375    },
376}
ferrum_types/models.rs

ferrum_types/
models.rs
