ferrotensor_hub/
models.rs

1//! Supported model types and auto-detection
2
3use crate::config::ModelConfig;
4
5/// Supported model architectures
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum ModelType {
8    /// Llama family (Llama 2, Llama 3, etc.)
9    Llama,
10    /// Mistral models
11    Mistral,
12    /// Microsoft Phi models
13    Phi,
14    /// Qwen models (Alibaba)
15    Qwen,
16    /// Google Gemma models
17    Gemma,
18    /// Unknown/unsupported
19    Unknown,
20}
21
22impl ModelType {
23    /// Detect model type from config.
24    pub fn from_config(config: &ModelConfig) -> Self {
25        // Check architectures field
26        for arch in &config.architectures {
27            let arch_lower = arch.to_lowercase();
28            if arch_lower.contains("llama") {
29                return Self::Llama;
30            }
31            if arch_lower.contains("mistral") {
32                return Self::Mistral;
33            }
34            if arch_lower.contains("phi") {
35                return Self::Phi;
36            }
37            if arch_lower.contains("qwen") {
38                return Self::Qwen;
39            }
40            if arch_lower.contains("gemma") {
41                return Self::Gemma;
42            }
43        }
44
45        // Check model_type field
46        let model_type = config.model_type.to_lowercase();
47        if model_type.contains("llama") {
48            return Self::Llama;
49        }
50        if model_type.contains("mistral") {
51            return Self::Mistral;
52        }
53        if model_type.contains("phi") {
54            return Self::Phi;
55        }
56        if model_type.contains("qwen") {
57            return Self::Qwen;
58        }
59        if model_type.contains("gemma") {
60            return Self::Gemma;
61        }
62
63        Self::Unknown
64    }
65
66    /// Check if the model type is supported by Ferrotensor.
67    pub fn is_supported(&self) -> bool {
68        matches!(
69            self,
70            Self::Llama | Self::Mistral | Self::Phi | Self::Qwen | Self::Gemma
71        )
72    }
73
74    /// Get the display name.
75    pub fn name(&self) -> &'static str {
76        match self {
77            Self::Llama => "Llama",
78            Self::Mistral => "Mistral",
79            Self::Phi => "Phi",
80            Self::Qwen => "Qwen",
81            Self::Gemma => "Gemma",
82            Self::Unknown => "Unknown",
83        }
84    }
85}
86
87impl std::fmt::Display for ModelType {
88    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
89        write!(f, "{}", self.name())
90    }
91}
92
93/// Supported quantization formats
94#[derive(Debug, Clone, Copy, PartialEq, Eq)]
95pub enum QuantFormat {
96    /// Full precision FP32
97    FP32,
98    /// Half precision FP16
99    FP16,
100    /// Brain float 16
101    BF16,
102    /// 8-bit integer (GPTQ, etc.)
103    Int8,
104    /// 4-bit integer
105    Int4,
106}
107
108impl QuantFormat {
109    /// Bytes per parameter.
110    pub fn bytes_per_param(&self) -> f32 {
111        match self {
112            Self::FP32 => 4.0,
113            Self::FP16 | Self::BF16 => 2.0,
114            Self::Int8 => 1.0,
115            Self::Int4 => 0.5,
116        }
117    }
118
119    /// Estimate model size in GB.
120    pub fn estimate_size_gb(&self, num_params: usize) -> f32 {
121        (num_params as f32 * self.bytes_per_param()) / 1_000_000_000.0
122    }
123}
ferrotensor_hub/models.rs

ferrotensor_hub/
models.rs
