compression_pipeline.py

"""
Integrated Model Compression Pipeline

Combines quantization, attention head pruning, and FFN neuron pruning
into a single, easy-to-use compression system.
"""

import torch
import torch.nn as nn
from typing import Dict, Optional, Tuple, List
import numpy as np
import copy
import time

# Import compression modules
from passes.enhanced_head_importance import compute_head_importance_enhanced
from passes.head_pruner import prune_model_by_importance as prune_heads
from passes.ffn_importance import compute_ffn_importance
from passes.ffn_pruner import prune_ffn_by_importance


class CompressionPipeline:
    """
    End-to-end compression pipeline for transformer models.
    
    Applies compression techniques in optimal order:
    1. Structural pruning (heads + FFN)
    2. Quantization
    """
    
    def __init__(self, model: nn.Module, tokenizer, model_type: str = 'auto'):
        """
        Initialize compression pipeline.
        
        Args:
            model: HuggingFace transformer model
            tokenizer: Corresponding tokenizer
            model_type: 'bert', 'gpt2', or 'auto'
        """
        self.original_model = model
        self.tokenizer = tokenizer
        self.model_type = self._detect_model_type(model) if model_type == 'auto' else model_type
        self.compressed_model = None
        self.compression_stats = {}
        
    def _detect_model_type(self, model) -> str:
        """Detect model architecture."""
        model_class = type(model).__name__.lower()
        if 'bert' in model_class or 'roberta' in model_class:
            return 'bert'
        elif 'gpt' in model_class:
            return 'gpt2'
        else:
            return 'unknown'
    
    def _count_parameters(self, model) -> int:
        """Count model parameters."""
        return sum(p.numel() for p in model.parameters())
    
    def _prepare_calibration_data(self, texts: Optional[List[str]] = None, max_length: int = 128) -> Dict:
        """Prepare calibration data for importance computation."""
        if texts is None:
            # Default calibration texts
            texts = [
                "The transformer architecture has revolutionized natural language processing.",
                "Attention mechanisms enable models to focus on relevant information.",
                "Feed-forward networks process information in each layer.",
                "Model compression reduces computational requirements significantly.",
                "Quantization and pruning are complementary compression techniques.",
                "Structural pruning physically removes model components.",
                "Efficient inference is critical for deployment at scale.",
                "Deep learning models continue to grow in size and capability.",
            ]
        
        inputs = self.tokenizer(
            texts,
            return_tensors='pt',
            padding=True,
            truncation=True,
            max_length=max_length
        )
        return inputs
    
    def compress(
        self,
        head_prune_ratio: float = 0.25,
        ffn_prune_ratio: float = 0.25,
        quantize: bool = True,
        calibration_texts: Optional[List[str]] = None,
        num_calibration_samples: int = 8,
        verbose: bool = True
    ) -> Tuple[nn.Module, Dict]:
        """
        Apply full compression pipeline.
        
        Args:
            head_prune_ratio: Fraction of attention heads to prune (0.0 to 1.0)
            ffn_prune_ratio: Fraction of FFN neurons to prune (0.0 to 1.0)
            quantize: Whether to apply INT8 quantization
            calibration_texts: Custom calibration texts (optional)
            num_calibration_samples: Number of samples for importance computation
            verbose: Print progress messages
        
        Returns:
            Tuple of (compressed_model, compression_statistics)
        """
        start_time = time.time()
        
        if verbose:
            print("\n" + "=" * 70)
            print("INTEGRATED MODEL COMPRESSION PIPELINE")
            print("=" * 70)
            print(f"\nModel: {type(self.original_model).__name__}")
            print(f"Type: {self.model_type.upper()}")
            print(f"Compression strategy:")
            print(f"  - Attention heads: {head_prune_ratio*100:.0f}% pruning")
            print(f"  - FFN neurons: {ffn_prune_ratio*100:.0f}% pruning")
            print(f"  - Quantization: {'Enabled' if quantize else 'Disabled'}")
        
        # Track original size
        original_params = self._count_parameters(self.original_model)
        self.compression_stats['original_params'] = original_params
        
        if verbose:
            print(f"\nOriginal parameters: {original_params:,}")
        
        # Create working copy
        model = copy.deepcopy(self.original_model)
        model.eval()
        
        # Prepare calibration data
        if verbose:
            print("\n[1/4] Preparing calibration data...")
        calibration_data = self._prepare_calibration_data(calibration_texts)
        
        # Step 1: Prune attention heads
        if head_prune_ratio > 0:
            if verbose:
                print(f"\n[2/4] Pruning {head_prune_ratio*100:.0f}% attention heads...")
            
            try:
                # Compute head importance
                head_importance = compute_head_importance_enhanced(
                    model,
                    calibration_data['input_ids'],
                    method='entropy',
                    num_samples=num_calibration_samples
                )
                
                # Prune heads
                model, head_stats = prune_heads(
                    model,
                    head_importance,
                    prune_ratio=head_prune_ratio,
                    model_type=self.model_type
                )
                
                self.compression_stats['heads_pruned'] = head_stats.get('total_heads_pruned', 0)
                
                if verbose:
                    print(f"  Pruned {head_stats.get('total_heads_pruned', 0)} attention heads")
                    params_after_heads = self._count_parameters(model)
                    print(f"  Parameters: {params_after_heads:,}")
            
            except Exception as e:
                if verbose:
                    print(f"  Warning: Head pruning failed: {e}")
                self.compression_stats['heads_pruned'] = 0
        else:
            if verbose:
                print("\n[2/4] Skipping attention head pruning")
            self.compression_stats['heads_pruned'] = 0
        
        # Step 2: Prune FFN neurons
        if ffn_prune_ratio > 0:
            if verbose:
                print(f"\n[3/4] Pruning {ffn_prune_ratio*100:.0f}% FFN neurons...")
            
            try:
                # Compute FFN importance
                ffn_importance = compute_ffn_importance(
                    model,
                    calibration_data['input_ids'],
                    method='weight_magnitude',
                    num_samples=num_calibration_samples,
                    model_type=self.model_type
                )
                
                # Prune FFN
                model = prune_ffn_by_importance(
                    model,
                    ffn_importance,
                    prune_ratio=ffn_prune_ratio,
                    model_type=self.model_type,
                    inplace=True
                )
                
                if verbose:
                    params_after_ffn = self._count_parameters(model)
                    print(f"  FFN pruning complete")
                    print(f"  Parameters: {params_after_ffn:,}")
            
            except Exception as e:
                if verbose:
                    print(f"  Warning: FFN pruning failed: {e}")
        else:
            if verbose:
                print("\n[3/4] Skipping FFN pruning")
        
        # Calculate structural compression
        params_after_pruning = self._count_parameters(model)
        structural_compression = original_params / params_after_pruning
        self.compression_stats['params_after_pruning'] = params_after_pruning
        self.compression_stats['structural_compression'] = structural_compression
        
        # Step 3: Quantization
        if quantize:
            if verbose:
                print("\n[4/4] Applying INT8 quantization...")
            
            try:
                # Dynamic quantization
                model = torch.quantization.quantize_dynamic(
                    model,
                    {nn.Linear},
                    dtype=torch.qint8
                )
                
                self.compression_stats['quantized'] = True
                
                if verbose:
                    print("  Quantization complete")
            
            except Exception as e:
                if verbose:
                    print(f"  Warning: Quantization failed: {e}")
                self.compression_stats['quantized'] = False
        else:
            if verbose:
                print("\n[4/4] Skipping quantization")
            self.compression_stats['quantized'] = False
        
        # Calculate total compression
        quantization_factor = 4.0 if quantize else 1.0
        total_compression = structural_compression * quantization_factor
        
        self.compression_stats['params_removed'] = original_params - params_after_pruning
        self.compression_stats['total_compression'] = total_compression
        self.compression_stats['compression_time'] = time.time() - start_time
        
        # Test forward pass
        if verbose:
            print("\n" + "=" * 70)
            print("COMPRESSION SUMMARY")
            print("=" * 70)
            print(f"Original parameters: {original_params:,}")
            print(f"After pruning: {params_after_pruning:,}")
            print(f"Parameters removed: {self.compression_stats['params_removed']:,} ({self.compression_stats['params_removed']/original_params*100:.1f}%)")
            print(f"Structural compression: {structural_compression:.2f}x")
            if quantize:
                print(f"With quantization (4x): {total_compression:.2f}x")
            print(f"Compression time: {self.compression_stats['compression_time']:.2f}s")
            
            # Test inference
            print("\nTesting compressed model...")
            try:
                with torch.no_grad():
                    test_input = calibration_data['input_ids'][:1]
                    outputs = model(test_input)
                print("  [OK] Model inference successful")
            except Exception as e:
                print(f"  [WARNING] Inference test failed: {e}")
        
        self.compressed_model = model
        return model, self.compression_stats
    
    def save_compressed_model(self, output_path: str):
        """Save compressed model to disk."""
        if self.compressed_model is None:
            raise ValueError("No compressed model available. Run compress() first.")
        
        torch.save({
            'model_state_dict': self.compressed_model.state_dict(),
            'compression_stats': self.compression_stats,
            'model_type': self.model_type,
        }, output_path)
        print(f"\nCompressed model saved to: {output_path}")
    
    def get_compression_report(self) -> str:
        """Generate detailed compression report."""
        if not self.compression_stats:
            return "No compression performed yet."
        
        report = []
        report.append("\n" + "=" * 70)
        report.append("COMPRESSION REPORT")
        report.append("=" * 70)
        report.append(f"\nOriginal Parameters: {self.compression_stats['original_params']:,}")
        report.append(f"Compressed Parameters: {self.compression_stats['params_after_pruning']:,}")
        report.append(f"Parameters Removed: {self.compression_stats['params_removed']:,}")
        report.append(f"\nStructural Compression: {self.compression_stats['structural_compression']:.2f}x")
        
        if self.compression_stats.get('quantized'):
            report.append(f"Total Compression (with quantization): {self.compression_stats['total_compression']:.2f}x")
        
        report.append(f"\nAttention Heads Pruned: {self.compression_stats.get('heads_pruned', 0)}")
        report.append(f"Quantized: {'Yes' if self.compression_stats.get('quantized') else 'No'}")
        report.append(f"Compression Time: {self.compression_stats['compression_time']:.2f}s")
        report.append("=" * 70)
        
        return "\n".join(report)


def compress_model(
    model_name: str,
    head_prune_ratio: float = 0.25,
    ffn_prune_ratio: float = 0.25,
    quantize: bool = True,
    output_path: Optional[str] = None,
    verbose: bool = True
) -> Tuple[nn.Module, Dict]:
    """
    Convenience function to compress a HuggingFace model.
    
    Args:
        model_name: HuggingFace model name (e.g., 'bert-base-uncased')
        head_prune_ratio: Fraction of attention heads to prune
        ffn_prune_ratio: Fraction of FFN neurons to prune
        quantize: Whether to apply quantization
        output_path: Path to save compressed model (optional)
        verbose: Print progress messages
    
    Returns:
        Tuple of (compressed_model, compression_statistics)
    
    Example:
        >>> model, stats = compress_model('bert-base-uncased', 
        ...                               head_prune_ratio=0.25,
        ...                               ffn_prune_ratio=0.25,
        ...                               quantize=True)
        >>> print(f"Compression: {stats['total_compression']:.2f}x")
    """
    from transformers import AutoModel, AutoTokenizer
    
    if verbose:
        print(f"Loading model: {model_name}")
    
    # Load model and tokenizer
    try:
        model = AutoModel.from_pretrained(model_name, attn_implementation='eager')
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        
        # Handle padding token
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
    
    except Exception as e:
        raise ValueError(f"Failed to load model '{model_name}': {e}")
    
    # Create pipeline and compress
    pipeline = CompressionPipeline(model, tokenizer)
    compressed_model, stats = pipeline.compress(
        head_prune_ratio=head_prune_ratio,
        ffn_prune_ratio=ffn_prune_ratio,
        quantize=quantize,
        verbose=verbose
    )
    
    # Save if requested
    if output_path:
        pipeline.save_compressed_model(output_path)
    
    return compressed_model, stats


if __name__ == "__main__":
    print("Integrated Model Compression Pipeline")
    print("=" * 70)
    print("\nUsage:")
    print("  from compression_pipeline import compress_model")
    print("  model, stats = compress_model('bert-base-uncased',")
    print("                                head_prune_ratio=0.25,")
    print("                                ffn_prune_ratio=0.25,")
    print("                                quantize=True)")