ohms_adaptq/novaq/
progress.rs

1use indicatif::{ProgressBar, ProgressStyle, MultiProgress};
2use std::time::{Duration, Instant};
3use nu_ansi_term::Color::{Green, Blue, Yellow, Red, Cyan};
4
5/// Comprehensive quantization progress tracking system
6/// Replaces verbose iteration logs with clean, meaningful progress indicators
7#[derive(Debug)]
8pub struct QuantizationProgressTracker {
9    multi_progress: MultiProgress,
10    main_progress: ProgressBar,
11    phase_progress: ProgressBar,
12    current_phase: QuantizationPhase,
13    start_time: Instant,
14    total_phases: u64,
15    completed_phases: u64,
16    quality_metrics: QualityMetrics,
17    verbosity_level: VerbosityLevel,
18}
19
20#[derive(Debug, Clone, Copy)]
21pub enum QuantizationPhase {
22    ModelLoading,
23    CodebookInitialization,
24    Level1Refinement,
25    Level2Refinement,
26    QualityValidation,
27    ModelSaving,
28    Complete,
29}
30
31#[derive(Debug, Clone, Copy)]
32pub enum VerbosityLevel {
33    Silent,     // No output except errors
34    Minimal,    // Progress bars only
35    Standard,   // Progress bars + phase summaries
36    Detailed,   // Standard + quality metrics
37}
38
39#[derive(Debug, Clone)]
40pub struct QualityMetrics {
41    pub mse: f32,
42    pub accuracy: f32,
43    pub compression_ratio: f32,
44    pub recovery_count: u32,
45    pub nan_issues: u32,
46    pub inf_issues: u32,
47}
48
49impl Default for QualityMetrics {
50    fn default() -> Self {
51        Self {
52            mse: 0.0,
53            accuracy: 0.0,
54            compression_ratio: 0.0,
55            recovery_count: 0,
56            nan_issues: 0,
57            inf_issues: 0,
58        }
59    }
60}
61
62impl QuantizationProgressTracker {
63    pub fn new(verbosity: VerbosityLevel) -> Self {
64        let multi_progress = MultiProgress::new();
65        
66        // Main progress bar for overall quantization
67        let main_progress = multi_progress.add(ProgressBar::new(6));
68        main_progress.set_style(
69            ProgressStyle::default_bar()
70                .template("{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} {msg}")
71                .unwrap()
72                .progress_chars("#>-")
73        );
74        main_progress.set_message("Initializing NOVAQ Quantization");
75        
76        // Phase-specific progress bar
77        let phase_progress = multi_progress.add(ProgressBar::new(100));
78        phase_progress.set_style(
79            ProgressStyle::default_bar()
80                .template("  └─ {spinner:.yellow} [{bar:30.yellow/dim}] {percent}% {msg}")
81                .unwrap()
82                .progress_chars("█▉▊▋▌▍▎▏ ")
83        );
84        
85        Self {
86            multi_progress,
87            main_progress,
88            phase_progress,
89            current_phase: QuantizationPhase::ModelLoading,
90            start_time: Instant::now(),
91            total_phases: 6,
92            completed_phases: 0,
93            quality_metrics: QualityMetrics::default(),
94            verbosity_level: verbosity,
95        }
96    }
97    
98    /// Start a new quantization phase
99    pub fn start_phase(&mut self, phase: QuantizationPhase, max_iterations: Option<u64>) {
100        self.current_phase = phase;
101        
102        match self.verbosity_level {
103            VerbosityLevel::Silent => return,
104            _ => {}
105        }
106        
107        let (phase_name, description) = self.get_phase_info(phase);
108        
109        // Update main progress
110        self.main_progress.set_position(self.completed_phases);
111        self.main_progress.set_message(format!("{} {}", phase_name, description));
112        
113        // Reset and configure phase progress
114        let iterations = max_iterations.unwrap_or(100);
115        self.phase_progress.reset();
116        self.phase_progress.set_length(iterations);
117        self.phase_progress.set_message(description.to_string());
118        
119        if matches!(self.verbosity_level, VerbosityLevel::Standard | VerbosityLevel::Detailed) {
120            println!("{}", Blue.bold().paint(format!("🔄 Starting: {}", phase_name)));
121        }
122    }
123    
124    /// Update iteration progress within current phase
125    pub fn update_iteration(&mut self, iteration: u64, metrics: Option<&QualityMetrics>) {
126        if matches!(self.verbosity_level, VerbosityLevel::Silent) {
127            return;
128        }
129        
130        self.phase_progress.set_position(iteration);
131        
132        if let Some(m) = metrics {
133            self.quality_metrics = m.clone();
134            
135            // Update message with key metrics
136            let message = match self.current_phase {
137                QuantizationPhase::Level1Refinement | QuantizationPhase::Level2Refinement => {
138                    if m.recovery_count > 0 {
139                        format!("MSE: {:.4}, Acc: {:.2}% [Recovered: {}]", m.mse, m.accuracy * 100.0, m.recovery_count)
140                    } else {
141                        format!("MSE: {:.4}, Accuracy: {:.2}%", m.mse, m.accuracy * 100.0)
142                    }
143                },
144                _ => format!("Quality: {:.2}%", m.accuracy * 100.0),
145            };
146            self.phase_progress.set_message(message);
147        }
148    }
149    
150    /// Complete current phase
151    pub fn complete_phase(&mut self) {
152        self.completed_phases += 1;
153        
154        match self.verbosity_level {
155            VerbosityLevel::Silent => return,
156            _ => {}
157        }
158        
159        self.phase_progress.finish_and_clear();
160        
161        let (phase_name, _) = self.get_phase_info(self.current_phase);
162        let elapsed = self.start_time.elapsed();
163        
164        if matches!(self.verbosity_level, VerbosityLevel::Standard | VerbosityLevel::Detailed) {
165            let status_icon = if self.quality_metrics.recovery_count > 0 { "🛡️" } else { "✅" };
166            println!("{} Completed: {} ({})", 
167                     status_icon, 
168                     phase_name, 
169                     self.format_duration(elapsed));
170                     
171            if matches!(self.verbosity_level, VerbosityLevel::Detailed) {
172                self.print_phase_metrics();
173            }
174        }
175        
176        self.main_progress.set_position(self.completed_phases);
177    }
178    
179    /// Complete entire quantization process
180    pub fn complete(&mut self) {
181        match self.verbosity_level {
182            VerbosityLevel::Silent => return,
183            _ => {}
184        }
185        
186        self.phase_progress.finish_and_clear();
187        self.main_progress.finish_and_clear();
188        
189        let total_time = self.start_time.elapsed();
190        
191        println!();
192        println!("{}", Green.bold().paint("🎉 NOVAQ Quantization Complete!"));
193        println!();
194        
195        if matches!(self.verbosity_level, VerbosityLevel::Standard | VerbosityLevel::Detailed) {
196            println!("{}", Cyan.bold().paint("Summary:"));
197            println!("  Duration: {}", self.format_duration(total_time));
198            println!("  Final Accuracy: {:.2}%", self.quality_metrics.accuracy * 100.0);
199            
200            if self.quality_metrics.compression_ratio > 0.0 {
201                println!("  Compression Ratio: {:.1}x", self.quality_metrics.compression_ratio);
202            }
203            
204            if self.quality_metrics.recovery_count > 0 {
205                println!("  {} Recoveries: {} (NaN: {}, Inf: {})", 
206                         Yellow.paint("🛡️"),
207                         self.quality_metrics.recovery_count,
208                         self.quality_metrics.nan_issues,
209                         self.quality_metrics.inf_issues);
210            } else {
211                println!("  {} No numerical issues detected", Green.paint("✨"));
212            }
213            println!();
214        }
215    }
216    
217    /// Report error and cleanup
218    pub fn error(&mut self, error_msg: &str) {
219        self.phase_progress.abandon_with_message(format!("❌ {}", error_msg));
220        self.main_progress.abandon_with_message("❌ Quantization failed");
221        
222        if !matches!(self.verbosity_level, VerbosityLevel::Silent) {
223            println!("{}", Red.bold().paint(format!("❌ Error: {}", error_msg)));
224        }
225    }
226    
227    /// Set verbosity level
228    pub fn set_verbosity(&mut self, level: VerbosityLevel) {
229        self.verbosity_level = level;
230    }
231    
232    /// Check if we should show detailed output
233    pub fn should_log_iteration(&self, iteration: u64) -> bool {
234        match self.verbosity_level {
235            VerbosityLevel::Silent => false,
236            VerbosityLevel::Minimal => false,
237            VerbosityLevel::Standard => iteration % 25 == 0,
238            VerbosityLevel::Detailed => iteration % 10 == 0,
239        }
240    }
241    
242    fn get_phase_info(&self, phase: QuantizationPhase) -> (&str, &str) {
243        match phase {
244            QuantizationPhase::ModelLoading => ("Model Loading", "Loading and preprocessing model weights"),
245            QuantizationPhase::CodebookInitialization => ("Codebook Init", "Initializing L1 and L2 codebooks"),
246            QuantizationPhase::Level1Refinement => ("L1 Refinement", "Optimizing level-1 codebook centroids"),
247            QuantizationPhase::Level2Refinement => ("L2 Refinement", "Optimizing level-2 residual codebooks"),
248            QuantizationPhase::QualityValidation => ("Quality Check", "Validating compression quality"),
249            QuantizationPhase::ModelSaving => ("Model Saving", "Serializing compressed model"),
250            QuantizationPhase::Complete => ("Complete", "Quantization finished successfully"),
251        }
252    }
253    
254    fn format_duration(&self, duration: Duration) -> String {
255        let total_secs = duration.as_secs();
256        if total_secs < 60 {
257            format!("{:.1}s", duration.as_secs_f32())
258        } else if total_secs < 3600 {
259            format!("{}m {}s", total_secs / 60, total_secs % 60)
260        } else {
261            format!("{}h {}m {}s", total_secs / 3600, (total_secs % 3600) / 60, total_secs % 60)
262        }
263    }
264    
265    fn print_phase_metrics(&self) {
266        if self.quality_metrics.accuracy > 0.0 {
267            println!("    Accuracy: {:.3}%", self.quality_metrics.accuracy * 100.0);
268        }
269        if self.quality_metrics.mse > 0.0 {
270            println!("    MSE Loss: {:.6}", self.quality_metrics.mse);
271        }
272        if self.quality_metrics.recovery_count > 0 {
273            println!("    Numerical Recoveries: {}", self.quality_metrics.recovery_count);
274        }
275    }
276}
277
278/// Create progress tracker with environment-based verbosity
279pub fn create_progress_tracker() -> QuantizationProgressTracker {
280    let verbosity = match std::env::var("NOVAQ_VERBOSITY").as_deref() {
281        Ok("silent") => VerbosityLevel::Silent,
282        Ok("minimal") => VerbosityLevel::Minimal,
283        Ok("detailed") => VerbosityLevel::Detailed,
284        _ => VerbosityLevel::Standard,
285    };
286    
287    QuantizationProgressTracker::new(verbosity)
288}
289
290#[cfg(test)]
291mod tests {
292    use super::*;
293    
294    #[test]
295    fn test_progress_tracker_creation() {
296        let tracker = QuantizationProgressTracker::new(VerbosityLevel::Standard);
297        assert_eq!(tracker.completed_phases, 0);
298        assert_eq!(tracker.total_phases, 6);
299    }
300    
301    #[test]
302    fn test_phase_transitions() {
303        let mut tracker = QuantizationProgressTracker::new(VerbosityLevel::Minimal);
304        
305        tracker.start_phase(QuantizationPhase::CodebookInitialization, Some(50));
306        assert!(matches!(tracker.current_phase, QuantizationPhase::CodebookInitialization));
307        
308        tracker.complete_phase();
309        assert_eq!(tracker.completed_phases, 1);
310    }
311    
312    #[test]
313    fn test_metrics_update() {
314        let mut tracker = QuantizationProgressTracker::new(VerbosityLevel::Standard);
315        let metrics = QualityMetrics {
316            mse: 0.001,
317            accuracy: 0.95,
318            compression_ratio: 4.2,
319            recovery_count: 2,
320            nan_issues: 1,
321            inf_issues: 1,
322        };
323        
324        tracker.update_iteration(10, Some(&metrics));
325        assert!((tracker.quality_metrics.accuracy - 0.95).abs() < f32::EPSILON);
326        assert_eq!(tracker.quality_metrics.recovery_count, 2);
327    }
328    
329    #[test]
330    fn test_verbosity_levels() {
331        let silent_tracker = QuantizationProgressTracker::new(VerbosityLevel::Silent);
332        assert!(!silent_tracker.should_log_iteration(10));
333        
334        let detailed_tracker = QuantizationProgressTracker::new(VerbosityLevel::Detailed);
335        assert!(detailed_tracker.should_log_iteration(10));
336        assert!(!detailed_tracker.should_log_iteration(15));
337    }
338}
ohms_adaptq/novaq/progress.rs

ohms_adaptq/novaq/
progress.rs
