Update app.py

app.py

@@ -1,701 +0,0 @@
-# app.py
-"""
-Research-grade KV cache compression benchmark application.
-RocketKV-enhanced SPG with 450x compression capability.
-FIXED: CUDA assert errors, safer default parameters, GPT-2 sequence limits.
-"""
-
-import gradio as gr
-import torch
-import numpy as np
-import matplotlib.pyplot as plt
-import seaborn as sns
-from datetime import datetime
-import json
-import pandas as pd
-import tempfile
-import os
-import logging
-from typing import Dict, List, Any, Tuple
-
-from config import (
-    CompressionConfig, CompressionType, EnhancedSPGConfig, 
-    ProvingConfig, ResearchConstants, SUPPORTED_MODELS, BENCHMARK_CONFIGS
-)
-from benchmark import (
-    run_research_benchmark, export_proof_bundle, verify_proof_bundle,
-    BenchmarkMetrics
-)
-from compression import detect_model_layers
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Set style for plots
-plt.style.use('seaborn-v0_8-darkgrid')
-sns.set_palette("husl")
-
-# Global state for results
-current_results = {}
-
-
-def run_benchmark(model_key, compression_type, benchmark_type, dataset_subset,
-                 eval_samples, n_seeds, seq_length, generation_length,
-                 base_decay_rate, sink_tokens, recent_window,
-                 enable_adaptive, target_perplexity_delta, 
-                 enable_progressive, progressive_quality_threshold,
-                 initial_compression_ratio, max_compression_ratio,
-                 sequence_compression_ratio, head_compression_ratio,
-                 head_retention_mode, magnitude_threshold_mode,
-                 min_tokens_for_stability, recent_boost_factor,
-                 fail_on_cpu):
-    """Run comprehensive benchmark with all compression methods."""
-    
-    # Enable synchronous CUDA for debugging
-    if torch.cuda.is_available():
-        os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
-    
-    # Validate sequence length for GPT-2
-    if model_key == "gpt2" and seq_length > 1024:
-        logger.warning(f"Reducing sequence length from {seq_length} to 1024 for GPT-2")
-        seq_length = 1024
-    
-    try:
-        # Create base configuration
-        base_config = CompressionConfig(
-            model_key=model_key,
-            compression_type=CompressionType[compression_type.upper()],
-            benchmark_type=benchmark_type,
-            benchmark_subset=dataset_subset if benchmark_type == "longbench" else None,
-            eval_samples=int(eval_samples),
-            n_seeds=int(n_seeds),
-            prefill_length=int(seq_length),
-            generation_length=int(generation_length),
-            fail_on_cpu_fallback=fail_on_cpu
-        )
-        
-        # Configure Enhanced SPG with safer parameters
-        base_config.enhanced_spg_config = EnhancedSPGConfig(
-            base_decay_rate=float(base_decay_rate),
-            sink_tokens=int(sink_tokens),
-            recent_window=int(recent_window),
-            enable_adaptive=enable_adaptive,
-            target_perplexity_delta=float(target_perplexity_delta),
-            enable_progressive=enable_progressive,
-            quality_threshold=float(progressive_quality_threshold),
-            initial_compression_ratio=float(initial_compression_ratio),
-            max_compression_ratio=float(max_compression_ratio),
-            target_compression_ratio=float(max_compression_ratio),
-            sequence_compression_ratio=float(sequence_compression_ratio),
-            head_compression_ratio=float(head_compression_ratio),
-            head_retention_mode=head_retention_mode,
-            magnitude_threshold_mode=magnitude_threshold_mode,
-            min_tokens_for_stability=int(min_tokens_for_stability),
-            recent_boost_factor=float(recent_boost_factor),
-            enable_two_stage=True,
-            use_hybrid_sparse_attention=True,
-            use_snapkv_plus_plus=True,
-            stage1_compression_ratio=20.0,  # Safer default
-            stage2_compression_ratio=20.0   # For 400x total
-        )
-        
-        # Store results
-        results = {}
-        model_name = base_config.model_name
-        
-        # Run benchmark for selected compression type
-        logger.info(f"Running {compression_type} benchmark...")
-        metrics, summary, records, fingerprints = run_research_benchmark(
-            model_name, base_config
-        )
-        
-        results[compression_type] = {
-            'metrics': metrics,
-            'summary': summary,
-            'records': records
-        }
-        
-        # Also run NONE compression for baseline comparison
-        if compression_type != "none":
-            logger.info("Running baseline (no compression) benchmark...")
-            baseline_config = CompressionConfig(
-                model_key=model_key,
-                compression_type=CompressionType.NONE,
-                benchmark_type=benchmark_type,
-                benchmark_subset=dataset_subset if benchmark_type == "longbench" else None,
-                eval_samples=int(eval_samples),
-                n_seeds=int(n_seeds),
-                prefill_length=int(seq_length),
-                generation_length=int(generation_length),
-                fail_on_cpu_fallback=fail_on_cpu
-            )
-            
-            try:
-                baseline_metrics, baseline_summary, baseline_records, _ = run_research_benchmark(
-                    model_name, baseline_config
-                )
-                
-                results['none'] = {
-                    'metrics': baseline_metrics,
-                    'summary': baseline_summary,
-                    'records': baseline_records
-                }
-            except Exception as e:
-                logger.error(f"Baseline benchmark failed: {e}")
-                # Continue without baseline
-        
-        # Store globally for export
-        global current_results
-        current_results = results
-        
-        # Create visualizations
-        plots = create_visualizations(results, benchmark_type)
-        
-        # Create summary text
-        summary_text = create_summary_text(results, benchmark_type)
-        
-        # Export proof bundle
-        with tempfile.TemporaryDirectory() as tmpdir:
-            bundle_path = export_proof_bundle(
-                tmpdir, base_config, metrics, summary, records, fingerprints
-            )
-            
-            # Verify the bundle
-            verification = verify_proof_bundle(
-                tmpdir, base_config, base_config.proving
-            )
-            
-            verification_text = f"Proof verification: {'PASSED ✓' if verification['ok'] else 'FAILED ✗'}"
-            if not verification['ok']:
-                verification_text += f"\nFailures: {verification['failures']}"
-        
-        return plots, summary_text, verification_text
-        
-    except Exception as e:
-        logger.error(f"Benchmark failed: {e}", exc_info=True)
-        return [], f"Error: {str(e)}", "Verification failed due to error"
-
-
-def create_visualizations(results: Dict, benchmark_type: str) -> List:
-    """Create comprehensive visualizations from benchmark results."""
-    plots = []
-    
-    # 1. Compression Ratio Comparison
-    fig, ax = plt.subplots(figsize=(10, 6))
-    methods = []
-    ratios = []
-    errors = []
-    
-    for method, data in results.items():
-        if 'metrics' in data and hasattr(data['metrics'], 'compression_ratio_mean'):
-            methods.append(method.upper())
-            ratios.append(data['metrics'].compression_ratio_mean)
-            errors.append(data['metrics'].compression_ratio_std)
-    
-    if methods:
-        bars = ax.bar(methods, ratios, yerr=errors, capsize=5)
-        ax.set_ylabel('Compression Ratio')
-        ax.set_title('KV Cache Compression Ratios')
-        ax.grid(True, alpha=0.3)
-        
-        # Add value labels on bars
-        for bar, ratio in zip(bars, ratios):
-            height = bar.get_height()
-            ax.text(bar.get_x() + bar.get_width()/2., height,
-                   f'{ratio:.1f}x', ha='center', va='bottom')
-    
-    plt.tight_layout()
-    plots.append(fig)
-    
-    # 2. Memory Usage Comparison
-    fig, ax = plt.subplots(figsize=(10, 6))
-    memories = []
-    memory_errors = []
-    
-    for method, data in results.items():
-        if 'metrics' in data and hasattr(data['metrics'], 'kv_cache_memory_mb'):
-            memories.append(data['metrics'].kv_cache_memory_mb)
-            memory_errors.append(0)  # No std for memory in current implementation
-    
-    if methods and memories:
-        bars = ax.bar(methods, memories, yerr=memory_errors, capsize=5, color='coral')
-        ax.set_ylabel('Memory Usage (MB)')
-        ax.set_title('KV Cache Memory Footprint')
-        ax.grid(True, alpha=0.3)
-        
-        for bar, mem in zip(bars, memories):
-            height = bar.get_height()
-            ax.text(bar.get_x() + bar.get_width()/2., height,
-                   f'{mem:.1f}', ha='center', va='bottom')
-    
-    plt.tight_layout()
-    plots.append(fig)
-    
-    # 3. Benchmark-specific metrics
-    if benchmark_type == "wikitext":
-        # Perplexity comparison
-        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))
-        
-        # Prefill perplexity
-        prefill_ppls = []
-        prefill_errors = []
-        gen_ppls = []
-        gen_errors = []
-        
-        for method, data in results.items():
-            if 'metrics' in data:
-                metrics = data['metrics']
-                if hasattr(metrics, 'prefill_perplexity_mean'):
-                    prefill_ppls.append(metrics.prefill_perplexity_mean)
-                    prefill_errors.append(metrics.prefill_perplexity_std)
-                if hasattr(metrics, 'generation_perplexity_mean'):
-                    gen_ppls.append(metrics.generation_perplexity_mean)
-                    gen_errors.append(metrics.generation_perplexity_std)
-        
-        if prefill_ppls:
-            ax1.bar(methods[:len(prefill_ppls)], prefill_ppls, yerr=prefill_errors, capsize=5, color='skyblue')
-            ax1.set_ylabel('Perplexity')
-            ax1.set_title('Prefill Perplexity')
-            ax1.grid(True, alpha=0.3)
-        
-        if gen_ppls:
-            ax2.bar(methods[:len(gen_ppls)], gen_ppls, yerr=gen_errors, capsize=5, color='lightgreen')
-            ax2.set_ylabel('Perplexity')
-            ax2.set_title('Generation Perplexity')
-            ax2.grid(True, alpha=0.3)
-        
-        plt.suptitle('Quality Metrics: Perplexity Comparison')
-        plt.tight_layout()
-        plots.append(fig)
-        
-    elif benchmark_type in ["niah", "ruler", "scbench"]:
-        # Accuracy metrics
-        fig, ax = plt.subplots(figsize=(10, 6))
-        accuracies = []
-        
-        for method, data in results.items():
-            if 'summary' in data:
-                if benchmark_type == "niah" and 'niah_accuracy' in data['summary']:
-                    accuracies.append(data['summary']['niah_accuracy'])
-                elif benchmark_type == "ruler" and 'ruler_exact_match' in data['summary']:
-                    accuracies.append(data['summary']['ruler_exact_match'])
-                elif benchmark_type == "scbench" and 'scbench_accuracy' in data['summary']:
-                    accuracies.append(data['summary']['scbench_accuracy'])
-        
-        if accuracies:
-            bars = ax.bar(methods[:len(accuracies)], accuracies, color='gold')
-            ax.set_ylabel('Accuracy')
-            ax.set_ylim(0, 1.1)
-            ax.set_title(f'{benchmark_type.upper()} Accuracy')
-            ax.grid(True, alpha=0.3)
-            
-            for bar, acc in zip(bars, accuracies):
-                height = bar.get_height()
-                ax.text(bar.get_x() + bar.get_width()/2., height,
-                       f'{acc:.2%}', ha='center', va='bottom')
-        
-        plt.tight_layout()
-        plots.append(fig)
-    
-    # 4. Speed comparison
-    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))
-    
-    prefill_times = []
-    decode_times = []
-    
-    for method, data in results.items():
-        if 'metrics' in data:
-            metrics = data['metrics']
-            if hasattr(metrics, 'prefill_time_mean'):
-                prefill_times.append(metrics.prefill_time_mean * 1000)  # Convert to ms
-            if hasattr(metrics, 'decode_time_per_token_mean_ms'):
-                decode_times.append(metrics.decode_time_per_token_mean_ms)
-    
-    if prefill_times:
-        ax1.bar(methods[:len(prefill_times)], prefill_times, color='purple', alpha=0.7)
-        ax1.set_ylabel('Time (ms)')
-        ax1.set_title('Prefill Time')
-        ax1.grid(True, alpha=0.3)
-    
-    if decode_times:
-        ax2.bar(methods[:len(decode_times)], decode_times, color='orange', alpha=0.7)
-        ax2.set_ylabel('Time per Token (ms)')
-        ax2.set_title('Decode Time')
-        ax2.grid(True, alpha=0.3)
-    
-    plt.suptitle('Performance Metrics: Speed Comparison')
-    plt.tight_layout()
-    plots.append(fig)
-    
-    return plots
-
-
-def create_summary_text(results: Dict, benchmark_type: str) -> str:
-    """Create detailed summary text from results."""
-    summary_lines = []
-    summary_lines.append("=" * 60)
-    summary_lines.append("BENCHMARK RESULTS SUMMARY")
-    summary_lines.append("=" * 60)
-    summary_lines.append(f"Benchmark Type: {benchmark_type.upper()}")
-    summary_lines.append(f"Timestamp: {datetime.now().isoformat()}")
-    summary_lines.append("")
-    
-    for method, data in results.items():
-        if 'summary' not in data:
-            continue
-            
-        summary = data['summary']
-        metrics = data['metrics'] if 'metrics' in data else None
-        
-        summary_lines.append(f"Method: {method.upper()}")
-        summary_lines.append("-" * 40)
-        
-        # Compression metrics
-        if 'compression_ratio' in summary:
-            summary_lines.append(f"Compression Ratio: {summary['compression_ratio']:.1f}x")
-        if 'kv_cache_memory_mb' in summary:
-            summary_lines.append(f"KV Cache Memory: {summary['kv_cache_memory_mb']:.2f} MB")
-        
-        # Quality metrics
-        if benchmark_type == "wikitext":
-            if 'prefill_perplexity' in summary:
-                summary_lines.append(f"Prefill Perplexity: {summary['prefill_perplexity']:.2f}")
-            if 'generation_perplexity' in summary:
-                summary_lines.append(f"Generation Perplexity: {summary['generation_perplexity']:.2f}")
-        elif benchmark_type == "niah" and 'niah_accuracy' in summary:
-            summary_lines.append(f"NIAH Accuracy: {summary['niah_accuracy']:.2%}")
-        elif benchmark_type == "ruler" and 'ruler_exact_match' in summary:
-            summary_lines.append(f"RULER Exact Match: {summary['ruler_exact_match']:.2%}")
-        elif benchmark_type == "scbench" and 'scbench_accuracy' in summary:
-            summary_lines.append(f"SCBench Accuracy: {summary['scbench_accuracy']:.2%}")
-        elif benchmark_type == "longbench" and 'longbench_accuracy' in summary:
-            summary_lines.append(f"LongBench Accuracy: {summary['longbench_accuracy']:.2%}")
-        
-        # Performance metrics
-        if 'prefill_time_ms' in summary:
-            summary_lines.append(f"Prefill Time: {summary['prefill_time_ms']:.2f} ms")
-        if 'decode_time_ms' in summary:
-            summary_lines.append(f"Decode Time per Token: {summary['decode_time_ms']:.2f} ms")
-        if 'throughput_tokens_sec' in summary:
-            summary_lines.append(f"Throughput: {summary['throughput_tokens_sec']:.1f} tokens/sec")
-        if 'end_to_end_throughput' in summary:
-            summary_lines.append(f"End-to-End Throughput: {summary['end_to_end_throughput']:.1f} tokens/sec")
-        if 'peak_memory_mb' in summary:
-            summary_lines.append(f"Peak Memory: {summary['peak_memory_mb']:.2f} MB")
-        
-        summary_lines.append("")
-    
-    # Add statistical comparison if baseline is available
-    if 'none' in results and len(results) > 1:
-        summary_lines.append("COMPARISON WITH BASELINE")
-        summary_lines.append("-" * 40)
-        
-        baseline_summary = results['none']['summary']
-        
-        for method, data in results.items():
-            if method == 'none' or 'summary' not in data:
-                continue
-                
-            summary = data['summary']
-            
-            # Calculate improvements
-            if 'compression_ratio' in summary:
-                summary_lines.append(f"{method.upper()} vs Baseline:")
-                summary_lines.append(f"  Compression: {summary['compression_ratio']:.1f}x")
-            
-            if 'kv_cache_memory_mb' in summary and 'kv_cache_memory_mb' in baseline_summary:
-                baseline_mem = baseline_summary['kv_cache_memory_mb']
-                method_mem = summary['kv_cache_memory_mb']
-                if baseline_mem > 0:
-                    reduction = (1 - method_mem / baseline_mem) * 100
-                    summary_lines.append(f"  Memory Reduction: {reduction:.1f}%")
-            
-            # Quality degradation for WikiText
-            if benchmark_type == "wikitext":
-                if 'generation_perplexity' in summary and 'generation_perplexity' in baseline_summary:
-                    baseline_ppl = baseline_summary['generation_perplexity']
-                    method_ppl = summary['generation_perplexity']
-                    if baseline_ppl > 0:
-                        degradation = ((method_ppl - baseline_ppl) / baseline_ppl) * 100
-                        summary_lines.append(f"  Perplexity Change: {degradation:+.1f}%")
-            
-            # Accuracy comparison for other benchmarks
-            elif benchmark_type == "niah":
-                if 'niah_accuracy' in summary and 'niah_accuracy' in baseline_summary:
-                    acc_diff = summary['niah_accuracy'] - baseline_summary['niah_accuracy']
-                    summary_lines.append(f"  Accuracy Difference: {acc_diff:+.2%}")
-            
-            summary_lines.append("")
-    
-    return "\n".join(summary_lines)
-
-
-def export_results(format_type):
-    """Export current results in specified format."""
-    if not current_results:
-        return "No results to export. Please run a benchmark first."
-    
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    
-    if format_type == "JSON":
-        filename = f"results_{timestamp}.json"
-        
-        # Convert numpy types to Python types for JSON serialization
-        def convert_numpy(obj):
-            if isinstance(obj, np.ndarray):
-                return obj.tolist()
-            elif isinstance(obj, (np.integer, np.int64, np.int32)):
-                return int(obj)
-            elif isinstance(obj, (np.floating, np.float64, np.float32)):
-                return float(obj)
-            elif isinstance(obj, BenchmarkMetrics):
-                return obj.__dict__
-            return obj
-        
-        serializable_results = json.loads(
-            json.dumps(current_results, default=convert_numpy)
-        )
-        
-        with open(filename, 'w') as f:
-            json.dump(serializable_results, f, indent=2)
-        
-        return f"Results exported to {filename}"
-        
-    elif format_type == "CSV":
-        filename = f"results_{timestamp}.csv"
-        
-        # Flatten results for CSV
-        rows = []
-        for method, data in current_results.items():
-            if 'summary' in data:
-                row = {'method': method}
-                row.update(data['summary'])
-                rows.append(row)
-        
-        if rows:
-            df = pd.DataFrame(rows)
-            df.to_csv(filename, index=False)
-            return f"Results exported to {filename}"
-        else:
-            return "No summary data to export"
-            
-    elif format_type == "LaTeX":
-        filename = f"results_{timestamp}.tex"
-        
-        # Create LaTeX table
-        latex_lines = [
-            "\\begin{table}[h]",
-            "\\centering",
-            "\\caption{KV Cache Compression Results}",
-            "\\begin{tabular}{lccc}",
-            "\\hline",
-            "Method & Compression & Memory (MB) & Throughput (tok/s) \\\\",
-            "\\hline"
-        ]
-        
-        for method, data in current_results.items():
-            if 'summary' in data:
-                s = data['summary']
-                comp = f"{s.get('compression_ratio', 1.0):.1f}x"
-                mem = f"{s.get('kv_cache_memory_mb', 0):.1f}"
-                thr = f"{s.get('throughput_tokens_sec', 0):.1f}"
-                latex_lines.append(f"{method.upper()} & {comp} & {mem} & {thr} \\\\")
-        
-        latex_lines.extend([
-            "\\hline",
-            "\\end{tabular}",
-            "\\end{table}"
-        ])
-        
-        with open(filename, 'w') as f:
-            f.write('\n'.join(latex_lines))
-        
-        return f"LaTeX table exported to {filename}"
-    
-    return "Invalid export format"
-
-
-# Create Gradio interface
-def create_interface():
-    with gr.Blocks(title="RocketKV-Enhanced SPG Benchmark") as demo:
-        gr.Markdown("""
-        # 🚀 RocketKV-Enhanced SPG Compression Benchmark
-        
-        Research-grade KV cache compression with **450x compression capability**.
-        Implements Enhanced Sliding Precision Gradient with RocketKV-style optimizations.
-        
-        **Features:**
-        - Multiple compression methods (SPG, Adaptive, Enhanced, Progressive)
-        - Comprehensive benchmarks (WikiText, NIAH, RULER, SCBench, LongBench)
-        - Attestable proof generation and verification
-        - Real-time visualization and analysis
-        """)
-        
-        with gr.Tab("Configuration"):
-            with gr.Row():
-                with gr.Column():
-                    gr.Markdown("### Model & Benchmark Settings")
-                    model_dropdown = gr.Dropdown(
-                        choices=list(SUPPORTED_MODELS.keys()),
-                        value="gpt2",
-                        label="Model"
-                    )
-                    
-                    compression_dropdown = gr.Dropdown(
-                        choices=["none", "spg", "adaptive_spg", "enhanced_spg", "progressive_spg"],
-                        value="enhanced_spg",
-                        label="Compression Method"
-                    )
-                    
-                    benchmark_dropdown = gr.Dropdown(
-                        choices=["wikitext", "niah", "ruler", "scbench", "longbench"],
-                        value="wikitext",
-                        label="Benchmark Type"
-                    )
-                    
-                    dataset_subset = gr.Dropdown(
-                        choices=BENCHMARK_CONFIGS["longbench"]["subsets"],
-                        value="narrativeqa",
-                        label="LongBench Subset (if applicable)",
-                        visible=False
-                    )
-                    
-                    # Show/hide subset based on benchmark type
-                    def update_subset_visibility(benchmark_type):
-                        return gr.update(visible=(benchmark_type == "longbench"))
-                    
-                    benchmark_dropdown.change(
-                        update_subset_visibility,
-                        inputs=[benchmark_dropdown],
-                        outputs=[dataset_subset]
-                    )
-                    
-                with gr.Column():
-                    gr.Markdown("### Evaluation Parameters")
-                    eval_samples = gr.Slider(1, 100, value=20, step=1, label="Evaluation Samples")
-                    n_seeds = gr.Slider(1, 5, value=3, step=1, label="Random Seeds")
-                    seq_length = gr.Slider(128, 1024, value=512, step=128, 
-                                          label="Sequence Length (max 1024 for GPT-2)")
-                    generation_length = gr.Slider(16, 128, value=64, step=16, label="Generation Length")
-                    
-            with gr.Row():
-                with gr.Column():
-                    gr.Markdown("### SPG Core Parameters")
-                    base_decay = gr.Slider(0.8, 0.99, value=0.95, step=0.01, label="Base Decay Rate")
-                    sink_tokens = gr.Slider(0, 8, value=2, step=1, label="Sink Tokens")
-                    recent_window = gr.Slider(8, 64, value=32, step=8, label="Recent Window")
-                    
-                with gr.Column():
-                    gr.Markdown("### Adaptive SPG")
-                    enable_adaptive = gr.Checkbox(value=False, label="Enable Adaptive")
-                    target_ppl_delta = gr.Slider(0.5, 5.0, value=1.8, step=0.1, 
-                                                label="Target Perplexity Delta")
-                    
-            with gr.Row():
-                with gr.Column():
-                    gr.Markdown("### Progressive Compression")
-                    enable_progressive = gr.Checkbox(value=False, label="Enable Progressive")
-                    quality_threshold = gr.Slider(0.005, 0.05, value=0.01, step=0.005,
-                                                 label="Quality Threshold")
-                    initial_compression = gr.Slider(10.0, 200.0, value=50.0, step=5.0,
-                                                   label="Initial Compression Ratio")
-                    max_compression = gr.Slider(100.0, 500.0, value=400.0, step=25.0,
-                                               label="Max Compression Ratio")
-                    
-                with gr.Column():
-                    gr.Markdown("### Enhanced SPG (RocketKV-style)")
-                    sequence_comp_ratio = gr.Slider(0.0001, 0.001, value=0.0001, step=0.00005,
-                                                   label="Sequence Compression Ratio")
-                    head_comp_ratio = gr.Slider(0.0001, 0.001, value=0.0001, step=0.00005,
-                                               label="Head Compression Ratio")
-                    head_retention = gr.Dropdown(
-                        choices=["conservative", "aggressive"],
-                        value="aggressive",
-                        label="Head Retention Mode"
-                    )
-                    magnitude_mode = gr.Dropdown(
-                        choices=["conservative", "aggressive", "extreme"],
-                        value="aggressive",  # Changed from "extreme" for stability
-                        label="Magnitude Threshold Mode"
-                    )
-                    
-            with gr.Row():
-                with gr.Column():
-                    gr.Markdown("### Stability Parameters")
-                    min_tokens_stability = gr.Slider(4, 16, value=8, step=1,
-                                                    label="Min Tokens for Stability")
-                    recent_boost = gr.Slider(0.0, 0.5, value=0.1, step=0.05,
-                                           label="Recent Boost Factor")
-                    
-                with gr.Column():
-                    gr.Markdown("### System Settings")
-                    fail_on_cpu = gr.Checkbox(value=False, label="Fail on CPU Fallback")
-                    
-        with gr.Tab("Run Benchmark"):
-            run_button = gr.Button("🚀 Run Benchmark", variant="primary")
-            
-            with gr.Row():
-                progress_text = gr.Textbox(label="Progress", lines=10)
-                
-            with gr.Row():
-                plot_gallery = gr.Gallery(label="Results Visualization", columns=2, height="auto")
-                
-            with gr.Row():
-                summary_output = gr.Textbox(label="Summary", lines=20)
-                verification_output = gr.Textbox(label="Proof Verification", lines=5)
-                
-        with gr.Tab("Export Results"):
-            gr.Markdown("### Export Options")
-            
-            export_format = gr.Radio(
-                choices=["JSON", "CSV", "LaTeX"],
-                value="JSON",
-                label="Export Format"
-            )
-            
-            export_button = gr.Button("📥 Export Results")
-            export_status = gr.Textbox(label="Export Status")
-            
-            export_button.click(
-                export_results,
-                inputs=[export_format],
-                outputs=[export_status]
-            )
-        
-        # Connect the run button
-        run_button.click(
-            run_benchmark,
-            inputs=[
-                model_dropdown, compression_dropdown, benchmark_dropdown, dataset_subset,
-                eval_samples, n_seeds, seq_length, generation_length,
-                base_decay, sink_tokens, recent_window,
-                enable_adaptive, target_ppl_delta,
-                enable_progressive, quality_threshold,
-                initial_compression, max_compression,
-                sequence_comp_ratio, head_comp_ratio,
-                head_retention, magnitude_mode,
-                min_tokens_stability, recent_boost,
-                fail_on_cpu
-            ],
-            outputs=[plot_gallery, summary_output, verification_output]
-        )
-        
-    return demo
-
-
-if __name__ == "__main__":
-    # Set up logging
-    logging.basicConfig(
-        level=logging.INFO,
-        format='%(asctime)s - %(levelname)s - %(message)s'
-    )
-    
-    # Create and launch the interface
-    demo = create_interface()
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        show_error=True
-    )