ai_model_vault/
quantization.rs

1//! Model quantization pipeline.
2//!
3//! Provides structured quantization profiles, batch quantization, and
4//! integration with the conversion system.
5
6use serde::{Deserialize, Serialize};
7use std::collections::BTreeMap;
8use std::path::{Path, PathBuf};
9
10use crate::error::{Result, VaultError};
11
12// ── Types ────────────────────────────────────────────────────────────────────
13
14/// Supported quantization methods.
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
16#[serde(rename_all = "snake_case")]
17pub enum QuantMethod {
18    /// 4-bit quantization (fastest inference, lowest quality).
19    Q4_0,
20    /// 4-bit K-quant medium (good balance).
21    Q4KM,
22    /// 5-bit K-quant medium.
23    Q5KM,
24    /// 8-bit quantization (highest quality quantized).
25    Q8_0,
26    /// 16-bit float (half precision).
27    F16,
28    /// 32-bit float (full precision, no quantization).
29    F32,
30}
31
32impl QuantMethod {
33    /// Approximate bits-per-weight for this method.
34    pub fn bits_per_weight(self) -> f64 {
35        match self {
36            Self::Q4_0 => 4.0,
37            Self::Q4KM => 4.5,
38            Self::Q5KM => 5.5,
39            Self::Q8_0 => 8.0,
40            Self::F16 => 16.0,
41            Self::F32 => 32.0,
42        }
43    }
44
45    /// Returns the string used in conversion options.
46    pub fn as_conversion_str(&self) -> &'static str {
47        match self {
48            Self::Q4_0 => "q4_0",
49            Self::Q4KM => "q4_k_m",
50            Self::Q5KM => "q5_k_m",
51            Self::Q8_0 => "q8_0",
52            Self::F16 => "f16",
53            Self::F32 => "f32",
54        }
55    }
56}
57
58impl std::fmt::Display for QuantMethod {
59    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
60        f.write_str(self.as_conversion_str())
61    }
62}
63
64impl std::str::FromStr for QuantMethod {
65    type Err = VaultError;
66    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
67        match s.to_lowercase().replace('-', "_").as_str() {
68            "q4_0" | "q4" => Ok(Self::Q4_0),
69            "q4_k_m" | "q4km" | "q4_k" => Ok(Self::Q4KM),
70            "q5_k_m" | "q5km" | "q5_k" => Ok(Self::Q5KM),
71            "q8_0" | "q8" => Ok(Self::Q8_0),
72            "f16" | "fp16" | "half" => Ok(Self::F16),
73            "f32" | "fp32" | "float" | "full" => Ok(Self::F32),
74            _ => Err(VaultError::InvalidInput(format!(
75                "Unknown quantization method: {s}"
76            ))),
77        }
78    }
79}
80
81/// A named quantization profile with target method and metadata.
82#[derive(Debug, Clone, Serialize, Deserialize)]
83pub struct QuantProfile {
84    /// Profile name (e.g. "edge-deploy", "server-quality").
85    pub name: String,
86    /// Target quantization method.
87    pub method: QuantMethod,
88    /// Optional description.
89    pub description: Option<String>,
90    /// Extra key-value metadata.
91    #[serde(default)]
92    pub metadata: BTreeMap<String, String>,
93}
94
95/// Result of a quantization operation.
96#[derive(Debug, Clone, Serialize, Deserialize)]
97pub struct QuantResult {
98    /// Source file path.
99    pub source: PathBuf,
100    /// Output file path.
101    pub output: PathBuf,
102    /// Method used.
103    pub method: QuantMethod,
104    /// Original size in bytes.
105    pub original_bytes: u64,
106    /// Quantized size in bytes.
107    pub quantized_bytes: u64,
108    /// Compression ratio (original / quantized).
109    pub compression_ratio: f64,
110}
111
112/// Batch quantization report.
113#[derive(Debug, Clone, Default, Serialize, Deserialize)]
114pub struct BatchQuantReport {
115    pub results: Vec<QuantResult>,
116    pub failed: Vec<(PathBuf, String)>,
117}
118
119// ── Profile Store ────────────────────────────────────────────────────────────
120
121const PROFILES_FILE: &str = "quant_profiles.json";
122
123/// Manages named quantization profiles on disk.
124#[derive(Debug)]
125pub struct QuantProfileStore {
126    path: PathBuf,
127    profiles: BTreeMap<String, QuantProfile>,
128}
129
130impl QuantProfileStore {
131    /// Load or create profile store.
132    pub fn new(vault_path: &Path) -> Result<Self> {
133        let path = vault_path.join(PROFILES_FILE);
134        let profiles = if path.exists() {
135            let data = std::fs::read_to_string(&path)
136                .map_err(|e| VaultError::StorageError(format!("read quant profiles: {e}")))?;
137            serde_json::from_str(&data).unwrap_or_default()
138        } else {
139            BTreeMap::new()
140        };
141        Ok(Self { path, profiles })
142    }
143
144    fn save(&self) -> Result<()> {
145        let data = serde_json::to_string_pretty(&self.profiles)
146            .map_err(|e| VaultError::StorageError(format!("serialize quant profiles: {e}")))?;
147        std::fs::write(&self.path, data)
148            .map_err(|e| VaultError::StorageError(format!("write quant profiles: {e}")))
149    }
150
151    /// Add or update a profile.
152    pub fn set(&mut self, profile: QuantProfile) -> Result<()> {
153        self.profiles.insert(profile.name.clone(), profile);
154        self.save()
155    }
156
157    /// Remove a profile by name.
158    pub fn remove(&mut self, name: &str) -> Result<bool> {
159        let removed = self.profiles.remove(name).is_some();
160        if removed {
161            self.save()?;
162        }
163        Ok(removed)
164    }
165
166    /// Get a profile by name.
167    pub fn get(&self, name: &str) -> Option<&QuantProfile> {
168        self.profiles.get(name)
169    }
170
171    /// List all profiles.
172    pub fn list(&self) -> Vec<&QuantProfile> {
173        self.profiles.values().collect()
174    }
175}
176
177// ── Estimation ───────────────────────────────────────────────────────────────
178
179/// Estimate the output size after quantization.
180pub fn estimate_quantized_size(original_bytes: u64, from: QuantMethod, to: QuantMethod) -> u64 {
181    let ratio = to.bits_per_weight() / from.bits_per_weight();
182    (original_bytes as f64 * ratio).ceil() as u64
183}
184
185// ── Tests ────────────────────────────────────────────────────────────────────
186
187#[cfg(test)]
188mod tests {
189    use super::*;
190    use tempfile::tempdir;
191
192    #[test]
193    fn test_quant_method_roundtrip() {
194        for method in [
195            QuantMethod::Q4_0,
196            QuantMethod::Q4KM,
197            QuantMethod::Q5KM,
198            QuantMethod::Q8_0,
199            QuantMethod::F16,
200            QuantMethod::F32,
201        ] {
202            let s = method.to_string();
203            let parsed: QuantMethod = s.parse().unwrap();
204            assert_eq!(parsed, method);
205        }
206    }
207
208    #[test]
209    fn test_quant_method_parse_aliases() {
210        assert_eq!("q4".parse::<QuantMethod>().unwrap(), QuantMethod::Q4_0);
211        assert_eq!("fp16".parse::<QuantMethod>().unwrap(), QuantMethod::F16);
212        assert_eq!("half".parse::<QuantMethod>().unwrap(), QuantMethod::F16);
213        assert_eq!("q4_k".parse::<QuantMethod>().unwrap(), QuantMethod::Q4KM);
214        assert_eq!("q8".parse::<QuantMethod>().unwrap(), QuantMethod::Q8_0);
215    }
216
217    #[test]
218    fn test_quant_method_parse_invalid() {
219        assert!("q3".parse::<QuantMethod>().is_err());
220    }
221
222    #[test]
223    fn test_bits_per_weight() {
224        assert!((QuantMethod::Q4_0.bits_per_weight() - 4.0).abs() < f64::EPSILON);
225        assert!((QuantMethod::F32.bits_per_weight() - 32.0).abs() < f64::EPSILON);
226    }
227
228    #[test]
229    fn test_estimate_quantized_size() {
230        // F32 → Q4_0 should be ~8× smaller
231        let est = estimate_quantized_size(1_000_000, QuantMethod::F32, QuantMethod::Q4_0);
232        assert_eq!(est, 125000);
233    }
234
235    #[test]
236    fn test_profile_store_crud() {
237        let dir = tempdir().unwrap();
238        let mut store = QuantProfileStore::new(dir.path()).unwrap();
239
240        let profile = QuantProfile {
241            name: "edge".into(),
242            method: QuantMethod::Q4KM,
243            description: Some("Edge deployment".into()),
244            metadata: BTreeMap::new(),
245        };
246        store.set(profile).unwrap();
247        assert_eq!(store.list().len(), 1);
248        assert_eq!(store.get("edge").unwrap().method, QuantMethod::Q4KM);
249
250        store.remove("edge").unwrap();
251        assert!(store.list().is_empty());
252    }
253
254    #[test]
255    fn test_profile_store_persistence() {
256        let dir = tempdir().unwrap();
257
258        {
259            let mut store = QuantProfileStore::new(dir.path()).unwrap();
260            store
261                .set(QuantProfile {
262                    name: "prod".into(),
263                    method: QuantMethod::Q8_0,
264                    description: None,
265                    metadata: BTreeMap::new(),
266                })
267                .unwrap();
268        }
269
270        // Re-open
271        let store = QuantProfileStore::new(dir.path()).unwrap();
272        assert_eq!(store.list().len(), 1);
273        assert_eq!(store.get("prod").unwrap().method, QuantMethod::Q8_0);
274    }
275
276    #[test]
277    fn test_batch_report_default() {
278        let report = BatchQuantReport::default();
279        assert!(report.results.is_empty());
280        assert!(report.failed.is_empty());
281    }
282}
ai_model_vault/quantization.rs

ai_model_vault/
quantization.rs
