compress_gpt2.py

"""
GPT-2 Compression Example

Demonstrates compressing GPT-2 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_gpt2_model, HuggingFaceModelLoader
from pipeline.transformer_pipeline import (
    create_default_transformer_pipeline,
    create_aggressive_transformer_pipeline,
    create_accuracy_preserving_transformer_pipeline,
)


def compress_gpt2_base():
    """
    Compress GPT-2 base model using default compression strategy.
    """
    print("=" * 80)
    print("GPT-2 Base Compression Example")
    print("=" * 80)
    
    # Load GPT-2 model
    print("\n[1] Loading GPT-2 base model...")
    model, tokenizer, config = load_gpt2_model('gpt2', task='causal-lm')
    
    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']}")
    
    # Create calibration data
    print("\n[2] Creating calibration data...")
    loader = HuggingFaceModelLoader()
    calib_data = loader.create_calibration_data(
        tokenizer, 
        model_type='gpt2',
        batch_size=4,
        seq_length=128,
        num_samples=4
    )
    print(f"  Calibration data shape: {calib_data.shape}")
    
    # Extract transformer (GPT-2 has 'transformer' attribute)
    print("\n[3] Compressing with default transformer pipeline...")
    gpt2_transformer = model.transformer if hasattr(model, 'transformer') else model
    
    # Compress using default pipeline
    pipeline = create_default_transformer_pipeline()
    
    try:
        compressed_mir, stats = pipeline.compress(gpt2_transformer, 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}%")
    
    return compressed_mir, stats


def compress_distilgpt2():
    """
    Compress DistilGPT-2 (smaller GPT-2 variant) using aggressive compression.
    """
    print("\n\n" + "=" * 80)
    print("DistilGPT-2 Compression Example (Aggressive)")
    print("=" * 80)
    
    # Load DistilGPT-2 model
    print("\n[1] Loading distilgpt2...")
    try:
        model, tokenizer, config = load_gpt2_model('distilgpt2', task='causal-lm')
    except Exception as e:
        print(f"  [ERROR] Error loading DistilGPT-2: {e}")
        print(f"  Skipping DistilGPT-2 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='gpt2',
        batch_size=4,
        seq_length=128,
        num_samples=4
    )
    
    # Compress with aggressive pipeline
    print("\n[3] Compressing DistilGPT-2...")
    
    # Extract transformer
    gpt2_transformer = model.transformer if hasattr(model, 'transformer') else model
    
    # Compress with aggressive pipeline
    pipeline = create_aggressive_transformer_pipeline()
    
    try:
        compressed_mir, stats = pipeline.compress(gpt2_transformer, 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 GPT-2.
    """
    print("\n\n" + "=" * 80)
    print("GPT-2 Compression Strategy Comparison")
    print("=" * 80)
    
    # Load model once
    print("\n[1] Loading GPT-2 base model...")
    model, tokenizer, config = load_gpt2_model('gpt2', task='causal-lm')
    
    # Create calibration data
    loader = HuggingFaceModelLoader()
    calib_data = loader.create_calibration_data(tokenizer, model_type='gpt2', num_samples=4)
    
    # Extract transformer
    print("\n[2] Preparing model...")
    gpt2_transformer = model.transformer if hasattr(model, 'transformer') 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(gpt2_transformer, 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':<25} {'Original (MB)':<15} {'Compressed (MB)':<18} {'Ratio':<10}")
    print("-" * 80)
    
    for strategy_name, stats in results.items():
        if stats:
            print(f"{strategy_name:<25} {stats['original_size_mb']:<15.2f} "
                  f"{stats['compressed_size_mb']:<18.2f} {stats['compression_ratio']:<10.2f}x")
        else:
            print(f"{strategy_name:<25} {'N/A':<15} {'N/A':<18} {'N/A':<10}")
    
    print("\n" + "=" * 80)


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="GPT-2 Compression Examples")
    parser.add_argument(
        '--example',
        type=str,
        choices=['gpt2', 'distilgpt2', 'compare', 'all'],
        default='all',
        help='Which example to run'
    )
    
    args = parser.parse_args()
    
    try:
        if args.example == 'gpt2' or args.example == 'all':
            compress_gpt2_base()
        
        if args.example == 'distilgpt2' or args.example == 'all':
            compress_distilgpt2()
        
        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()