"""
Example: Compress a HuggingFace transformer using CyMeta.

This script shows how to:
1. Load a HuggingFace model
2. Convert all linear layers to CyMetaLinear
3. Run inference with the compressed model
"""

import torch
from transformers import AutoTokenizer

from cymeta.exports.converter import ConversionConfig, convert_pretrained_model


def main():
    model_name = "distilbert-base-uncased"

    config = ConversionConfig(
        dict_size=32,
        shift_bits=8,
        gating_hidden=16,
        compress_iterations=50,
        verbose=True,
    )

    print(f"Loading {model_name}...")
    compressed_model = convert_pretrained_model(model_name, config=config, torch_dtype=torch.float32)
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    inputs = tokenizer("CyMeta makes model compression easy!", return_tensors="pt")
    with torch.no_grad():
        outputs = compressed_model(**inputs)

    print("Compression complete.")
    print(f"Hidden states shape: {outputs.last_hidden_state.shape}")


if __name__ == "__main__":
    main()


