1pub mod manifest;
2pub mod verification;
3pub mod universal_loader;
4pub mod model_fetcher;
5pub mod novaq;
6pub mod real_model_loader;
7pub mod streaming_loader;
8
9pub use manifest::*;
10pub use verification::*;
11pub use universal_loader::{UniversalModel, UniversalLoader, load_any_model, find_model};
12pub use model_fetcher::{ModelFetcher, ModelSource, parse_model_source, FetchResult, ModelMetadata, ModelFormat};
13pub use novaq::{NOVAQEngine, NOVAQConfig, NOVAQModel, WeightMatrix, QuantizationRecoveryManager, RecoveryStats, QuantizationProgressTracker, VerbosityLevel};
14pub use real_model_loader::{RealModelLoader, ModelStats};
15pub use streaming_loader::{StreamingModelLoader};
16
17pub type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;
18
19pub struct PublicNOVAQ {
25 engine: NOVAQEngine,
26 recovery_manager: QuantizationRecoveryManager,
27 auto_recovery_enabled: bool,
28 verbosity_level: VerbosityLevel,
29}
30
31impl PublicNOVAQ {
32 pub fn new(config: NOVAQConfig) -> Self {
33 let verbosity = match std::env::var("NOVAQ_VERBOSITY").as_deref() {
34 Ok("silent") => VerbosityLevel::Silent,
35 Ok("minimal") => VerbosityLevel::Minimal,
36 Ok("detailed") => VerbosityLevel::Detailed,
37 _ => VerbosityLevel::Standard,
38 };
39
40 Self {
41 engine: NOVAQEngine::new(config.clone()),
42 recovery_manager: QuantizationRecoveryManager::new(config),
43 auto_recovery_enabled: true,
44 verbosity_level: verbosity,
45 }
46 }
47
48 pub fn new_with_verbosity(config: NOVAQConfig, verbosity: VerbosityLevel) -> Self {
50 Self {
51 engine: NOVAQEngine::new(config.clone()),
52 recovery_manager: QuantizationRecoveryManager::new(config),
53 auto_recovery_enabled: true,
54 verbosity_level: verbosity,
55 }
56 }
57
58 pub fn compress_model(&mut self, weights: Vec<WeightMatrix>) -> Result<NOVAQModel> {
61 if self.auto_recovery_enabled {
62 self.recovery_manager.quantize_with_recovery_and_progress(weights, self.verbosity_level)
63 } else {
64 let mut progress = QuantizationProgressTracker::new(self.verbosity_level);
65 self.engine.quantize_model_with_progress(weights, &mut progress)
66 }
67 }
68
69 pub fn compress_model_basic(&mut self, weights: Vec<WeightMatrix>) -> Result<NOVAQModel> {
71 self.engine.quantize_model(weights)
72 }
73
74 pub fn set_auto_recovery(&mut self, enabled: bool) {
76 self.auto_recovery_enabled = enabled;
77 if enabled {
78 println!("🛡️ Automatic recovery enabled - quantization will attempt to recover from failures");
79 } else {
80 println!("⚠️ Automatic recovery disabled - quantization will fail immediately on errors");
81 }
82 }
83
84 pub fn get_recovery_stats(&self) -> &RecoveryStats {
86 self.recovery_manager.get_stats()
87 }
88
89 pub fn print_recovery_summary(&self) {
91 self.recovery_manager.print_recovery_summary();
92 }
93
94 pub fn reset_recovery_stats(&mut self) {
96 self.recovery_manager.reset_stats();
97 }
98
99 pub fn validate_model(&self, model: &NOVAQModel) -> Result<ValidationReport> {
101 let mut issues = Vec::new();
102
103 let min_compression_ratio = 2.0; let min_bit_accuracy = match model.config.target_bits {
106 b if b <= 1.0 => 0.85, b if b <= 2.0 => 0.90, b if b <= 4.0 => 0.95, _ => 0.98, };
111
112 if model.compression_ratio < min_compression_ratio {
114 issues.push(format!("Compression ratio {:.1}x below minimum {:.1}x",
115 model.compression_ratio, min_compression_ratio));
116 }
117
118 if model.bit_accuracy < min_bit_accuracy {
120 issues.push(format!("Bit accuracy {:.1}% below minimum {:.1}% for {:.1}-bit quantization",
121 model.bit_accuracy * 100.0, min_bit_accuracy * 100.0, model.config.target_bits));
122 }
123
124 let quality_score = (model.compression_ratio / 100.0 + model.bit_accuracy) / 2.0;
126
127 let passed_validation = issues.is_empty();
128
129 Ok(ValidationReport {
130 compression_ratio: model.compression_ratio,
131 bit_accuracy: model.bit_accuracy,
132 quality_score,
133 passed_validation,
134 issues,
135 })
136 }
137
138 pub fn compress_hf_model(&mut self, repo: &str, file: Option<&str>) -> Result<NOVAQModel> {
140 let source = ModelSource::HuggingFace {
141 repo: repo.to_string(),
142 file: file.map(|f| f.to_string())
143 };
144
145 let fetch_result = ModelFetcher::fetch(&source)?;
146 let weights = RealModelLoader::load_model(&fetch_result)?;
147
148 self.compress_model(weights)
149 }
150
151 pub fn compress_ollama_model(&mut self, model: &str) -> Result<NOVAQModel> {
153 let source = ModelSource::Ollama {
154 model: model.to_string()
155 };
156
157 let fetch_result = ModelFetcher::fetch(&source)?;
158 let weights = RealModelLoader::load_model(&fetch_result)?;
159
160 self.compress_model(weights)
161 }
162
163 pub fn compress_url_model(&mut self, url: &str, filename: Option<&str>) -> Result<NOVAQModel> {
165 let source = ModelSource::Url {
166 url: url.to_string(),
167 filename: filename.map(|f| f.to_string())
168 };
169
170 let fetch_result = ModelFetcher::fetch(&source)?;
171 let weights = RealModelLoader::load_model(&fetch_result)?;
172
173 self.compress_model(weights)
174 }
175
176 pub fn compress_local_model(&mut self, path: &str) -> Result<NOVAQModel> {
178 let source = ModelSource::LocalPath {
179 path: std::path::PathBuf::from(path)
180 };
181
182 let fetch_result = ModelFetcher::fetch(&source)?;
183 let weights = RealModelLoader::load_model(&fetch_result)?;
184
185 self.compress_model(weights)
186 }
187
188 pub fn get_compression_stats(&self, model: &NOVAQModel) -> CompressionStats {
190 CompressionStats {
191 compression_ratio: model.compression_ratio,
192 bit_accuracy: model.bit_accuracy,
193 quality_score: (model.compression_ratio / 100.0 + model.bit_accuracy) / 2.0,
194 target_bits: model.config.target_bits,
195 num_subspaces: model.config.num_subspaces,
196 codebook_size_l1: model.config.codebook_size_l1,
197 codebook_size_l2: model.config.codebook_size_l2,
198 }
199 }
200}
201
202#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
203pub struct ValidationReport {
204 pub compression_ratio: f32,
205 pub bit_accuracy: f32,
206 pub quality_score: f32,
207 pub passed_validation: bool,
208 pub issues: Vec<String>,
209}
210
211#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
212pub struct CompressionStats {
213 pub compression_ratio: f32,
214 pub bit_accuracy: f32,
215 pub quality_score: f32,
216 pub target_bits: f32,
217 pub num_subspaces: usize,
218 pub codebook_size_l1: usize,
219 pub codebook_size_l2: usize,
220}