Skip to main content

ohms_adaptq/
lib.rs

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
19// NOVAQ Democratic Access - No Restrictions
20// Core compression technology available to everyone
21// No admin controls - completely open access
22
23/// Public NOVAQ compression interface - no restrictions
24pub 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    /// Create with specific verbosity level
49    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    /// Compress any model with NOVAQ - no restrictions
59    /// Uses automatic recovery if enabled (default: enabled)
60    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    /// Compress model without automatic recovery (original behavior)
70    pub fn compress_model_basic(&mut self, weights: Vec<WeightMatrix>) -> Result<NOVAQModel> {
71        self.engine.quantize_model(weights)
72    }
73    
74    /// Enable or disable automatic recovery
75    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    /// Get recovery statistics
85    pub fn get_recovery_stats(&self) -> &RecoveryStats {
86        self.recovery_manager.get_stats()
87    }
88    
89    /// Print recovery statistics summary
90    pub fn print_recovery_summary(&self) {
91        self.recovery_manager.print_recovery_summary();
92    }
93    
94    /// Reset recovery statistics
95    pub fn reset_recovery_stats(&mut self) {
96        self.recovery_manager.reset_stats();
97    }
98    
99    /// Validate NOVAQ model quality
100    pub fn validate_model(&self, model: &NOVAQModel) -> Result<ValidationReport> {
101        let mut issues = Vec::new();
102        
103        // Realistic validation thresholds based on bit depth
104        let min_compression_ratio = 2.0; // At least 2x compression
105        let min_bit_accuracy = match model.config.target_bits {
106            b if b <= 1.0 => 0.85,  // 1-bit: 85% accuracy is excellent
107            b if b <= 2.0 => 0.90,  // 2-bit: 90% accuracy is excellent  
108            b if b <= 4.0 => 0.95,  // 4-bit: 95% accuracy is excellent
109            _ => 0.98,              // Higher bits: 98% accuracy expected
110        };
111        
112        // Check compression ratio
113        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        // Check bit accuracy
119        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        // Quality score calculation
125        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    /// Fetch and compress model from Hugging Face
139    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    /// Fetch and compress model from Ollama
152    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    /// Fetch and compress model from URL
164    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    /// Compress local model file
177    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    /// Get compression statistics
189    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}