compress_bert.py

"""
BERT Compression Example

Demonstrates compressing real BERT models from HuggingFace.
"""

import sys
import os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import torch
from utils.huggingface_loader import load_bert_model, HuggingFaceModelLoader
from pipeline.transformer_pipeline import (
    create_default_transformer_pipeline,
    create_aggressive_transformer_pipeline,
    create_accuracy_preserving_transformer_pipeline,
)


def compress_bert_base():
    """
    Compress BERT-base-uncased using default compression strategy.
    """
    print("=" * 80)
    print("BERT-Base Compression Example")
    print("=" * 80)
    
    # Load BERT model
    print("\n[1] Loading BERT-base-uncased...")
    model, tokenizer, config = load_bert_model('bert-base-uncased', task='base')
    
    print(f"\nModel Configuration:")
    print(f"  Name: {config['model_name']}")
    print(f"  Parameters: {config['num_parameters']:,}")
    print(f"  Layers: {config['num_layers']}")
    print(f"  Hidden size: {config['hidden_size']}")
    print(f"  Attention heads: {config['num_attention_heads']}")
    print(f"  Vocab size: {config['vocab_size']}")
    
    # Extract just the encoder (BERT core without task-specific heads)
    bert_encoder = model.bert if hasattr(model, 'bert') else model
    
    # Create calibration data
    print("\n[2] Creating calibration data...")
    loader = HuggingFaceModelLoader()
    calib_data = loader.create_calibration_data(
        tokenizer, 
        model_type='bert',
        batch_size=4,
        seq_length=128,
        num_samples=4
    )
    print(f"  Calibration data shape: {calib_data.shape}")
    
    # Compress using default pipeline
    print("\n[3] Compressing with default transformer pipeline...")
    pipeline = create_default_transformer_pipeline()
    
    try:
        compressed_mir, stats = pipeline.compress(bert_encoder, calib_data)
    except Exception as e:
        print(f"  [ERROR] Error during compression: {e}")
        import traceback
        traceback.print_exc()
        return
    
    print(f"\n[OK] Compression completed!")
    print(f"\nCompression Statistics:")
    print(f"  Original size: {stats.get('original_size_mb', 0):.2f} MB")
    print(f"  Compressed size: {stats.get('compressed_size_mb', 0):.2f} MB")
    print(f"  Compression ratio: {stats.get('compression_ratio', 0):.2f}x")
    
    if 'original_nodes' in stats:
        print(f"  Original nodes: {stats['original_nodes']}")
    if 'compressed_nodes' in stats:
        print(f"  Compressed nodes: {stats['compressed_nodes']}")
    
    if 'pruning_stats' in stats:
        pstats = stats['pruning_stats']
        print(f"\nPruning Statistics:")
        print(f"  Attention heads pruned: {pstats.get('heads_pruned', 0)}")
        print(f"  Heads remaining: {pstats.get('heads_remaining', 0)}")
    
    if 'ffn_stats' in stats:
        fstats = stats['ffn_stats']
        print(f"\nFFN Compression Statistics:")
        print(f"  Layers compressed: {fstats.get('layers_compressed', 0)}")
        print(f"  Average reduction: {fstats.get('avg_reduction', 0)*100:.1f}%")
    
    if 'quantization_stats' in stats:
        qstats = stats['quantization_stats']
        print(f"\nQuantization Statistics:")
        print(f"  Layers quantized: {qstats.get('layers_quantized', 0)}")
        print(f"  Quantization method: {qstats.get('method', 'N/A')}")
    
    return compressed_mir, stats


def compress_distilbert():
    """
    Compress DistilBERT (smaller BERT variant) using aggressive compression.
    """
    print("\n\n" + "=" * 80)
    print("DistilBERT Compression Example (Aggressive)")
    print("=" * 80)
    
    # Load DistilBERT model
    print("\n[1] Loading distilbert-base-uncased...")
    try:
        model, tokenizer, config = load_bert_model('distilbert-base-uncased', task='base')
    except Exception as e:
        print(f"  [ERROR] Error loading DistilBERT: {e}")
        print(f"  Skipping DistilBERT example...")
        return None, None
    
    print(f"\nModel Configuration:")
    print(f"  Name: {config['model_name']}")
    print(f"  Parameters: {config['num_parameters']:,}")
    print(f"  Layers: {config['num_layers']}")
    print(f"  Hidden size: {config['hidden_size']}")
    print(f"  Attention heads: {config['num_attention_heads']}")
    
    # Create calibration data
    print("\n[2] Creating calibration data...")
    loader = HuggingFaceModelLoader()
    calib_data = loader.create_calibration_data(
        tokenizer,
        model_type='bert',
        batch_size=4,
        seq_length=128,
        num_samples=4
    )
    
    # Parse to HIR
    print("\n[3] Compressing DistilBERT...")
    
    # Extract encoder
    distilbert_encoder = model.distilbert if hasattr(model, 'distilbert') else model
    
    # Compress with aggressive pipeline
    pipeline = create_aggressive_transformer_pipeline()
    
    try:
        compressed_mir, stats = pipeline.compress(distilbert_encoder, calib_data)
    except Exception as e:
        print(f"  [ERROR] Error during compression: {e}")
        import traceback
        traceback.print_exc()
        return None, None
    
    print(f"\n[OK] Compression completed!")
    print(f"\nCompression Statistics:")
    print(f"  Original size: {stats.get('original_size_mb', 0):.2f} MB")
    print(f"  Compressed size: {stats.get('compressed_size_mb', 0):.2f} MB")
    print(f"  Compression ratio: {stats.get('compression_ratio', 0):.2f}x")
    if 'original_nodes' in stats and 'compressed_nodes' in stats:
        print(f"  Node reduction: {stats['original_nodes']} -> {stats['compressed_nodes']}")
    
    return compressed_mir, stats


def compare_compression_strategies():
    """
    Compare different compression strategies on BERT.
    """
    print("\n\n" + "=" * 80)
    print("BERT Compression Strategy Comparison")
    print("=" * 80)
    
    # Load model once
    print("\n[1] Loading BERT-base-uncased...")
    model, tokenizer, config = load_bert_model('bert-base-uncased', task='base')
    
    # Create calibration data
    loader = HuggingFaceModelLoader()
    calib_data = loader.create_calibration_data(tokenizer, model_type='bert', num_samples=4)
    
    # Extract encoder
    print("\n[2] Preparing model...")
    bert_encoder = model.bert if hasattr(model, 'bert') else model
    
    # Test each strategy
    strategies = {
        'Accuracy-Preserving': create_accuracy_preserving_transformer_pipeline(),
        'Default': create_default_transformer_pipeline(),
        'Aggressive': create_aggressive_transformer_pipeline(),
    }
    
    results = {}
    
    print("\n[3] Testing compression strategies...")
    for strategy_name, pipeline in strategies.items():
        print(f"\n  Testing {strategy_name} strategy...")
        try:
            compressed_mir, stats = pipeline.compress(bert_encoder, calib_data)
            results[strategy_name] = stats
            print(f"    [OK] {stats['compression_ratio']:.2f}x compression")
        except Exception as e:
            print(f"    [ERROR] Error: {e}")
            results[strategy_name] = None
    
    # Print comparison table
    print("\n" + "=" * 80)
    print("Compression Strategy Comparison Results")
    print("=" * 80)
    print(f"\n{'Strategy':<15} {'Original (MB)':<15} {'Compressed (MB)':<18} {'Ratio':<10}")
    print("-" * 80)
    
    for strategy_name, stats in results.items():
        if stats:
            print(f"{strategy_name:<15} {stats['original_size_mb']:<15.2f} "
                  f"{stats['compressed_size_mb']:<18.2f} {stats['compression_ratio']:<10.2f}x")
        else:
            print(f"{strategy_name:<15} {'N/A':<15} {'N/A':<18} {'N/A':<10}")
    
    print("\n" + "=" * 80)


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="BERT Compression Examples")
    parser.add_argument(
        '--example',
        type=str,
        choices=['bert', 'distilbert', 'compare', 'all'],
        default='all',
        help='Which example to run'
    )
    
    args = parser.parse_args()
    
    try:
        if args.example == 'bert' or args.example == 'all':
            compress_bert_base()
        
        if args.example == 'distilbert' or args.example == 'all':
            compress_distilbert()
        
        if args.example == 'compare' or args.example == 'all':
            compare_compression_strategies()
        
        print("\n" + "=" * 80)
        print("[OK] All examples completed successfully!")
        print("=" * 80)
        
    except KeyboardInterrupt:
        print("\n\n[WARNING] Interrupted by user")
    except Exception as e:
        print(f"\n\n[ERROR] Error: {e}")
        import traceback
        traceback.print_exc()