"""
Example: Compress a HuggingFace model for CPU deployment.

This example demonstrates how to use CyMeta's CPU-optimized compression
to make LLM models runnable on CPU without GPU.
"""

import torch
from cymeta import (
    ConversionConfig,
    convert_pretrained_model,
    ensure_cpu_model,
    estimate_model_size_mb,
    get_compression_summary,
)


def main():
    print("=" * 70)
    print("CyMeta CPU-Optimized Compression Example")
    print("=" * 70)
    
    # Step 1: Configure for CPU deployment
    print("\n1. Configuring CPU-optimized compression...")
    config = ConversionConfig(
        cpu_optimized=True,  # Enable aggressive compression for CPU
        verbose=True,        # Show progress
    )
    print(f"   - Dictionary size: {config.dict_size} (reduced for CPU)")
    print(f"   - Gating hidden: {config.gating_hidden} (reduced for CPU)")
    print(f"   - Compression iterations: {config.compress_iterations}")
    
    # Step 2: Load and compress model
    print("\n2. Loading and compressing model...")
    print("   This may take several minutes depending on model size...")
    
    model_name = "distilbert-base-uncased"  # Start with a smaller model
    compressed_model = convert_pretrained_model(
        model_name,
        config=config,
    )
    
    # Step 3: Ensure model is on CPU
    print("\n3. Ensuring model is CPU-ready...")
    compressed_model = ensure_cpu_model(compressed_model)
    print("   ✓ Model is on CPU and ready for inference")
    
    # Step 4: Get compression statistics
    print("\n4. Compression Statistics:")
    summary = get_compression_summary(compressed_model)
    print(f"   - Total layers: {summary['total_layers']}")
    print(f"   - Compressed layers: {summary['compressed_layers']}")
    print(f"   - Average compression ratio: {summary['average_compression_ratio']:.2f}x")
    print(f"   - Compression coverage: {summary['compression_coverage']:.1%}")
    
    # Step 5: Estimate model size
    print("\n5. Model Size:")
    size_mb = estimate_model_size_mb(compressed_model)
    print(f"   - Estimated size: {size_mb:.2f} MB")
    print(f"   - Suitable for CPU deployment: {'Yes' if size_mb < 2000 else 'Consider larger compression'}")
    
    # Step 6: Test inference on CPU
    print("\n6. Testing CPU inference...")
    from transformers import AutoTokenizer
    
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    text = "Hello, this is a test of CPU inference with CyMeta compression!"
    
    inputs = tokenizer(text, return_tensors="pt")
    
    with torch.no_grad():
        outputs = compressed_model(**inputs)
    
    print(f"   ✓ Inference successful on CPU!")
    print(f"   - Input text: {text}")
    print(f"   - Output shape: {outputs.last_hidden_state.shape}")
    
    # Step 7: Save compressed model
    print("\n7. Saving compressed model...")
    save_path = "./compressed_model_cpu"
    compressed_model.save_pretrained(save_path)
    tokenizer.save_pretrained(save_path)
    print(f"   ✓ Model saved to: {save_path}")
    print(f"   You can now load it with: AutoModel.from_pretrained('{save_path}')")
    
    print("\n" + "=" * 70)
    print("CPU compression complete! Model is ready for CPU deployment.")
    print("=" * 70)


if __name__ == "__main__":
    main()

