compress.py

import os
import torch
import time
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from env import device, prepare_model


def get_model_metadata(model_path):
    if not os.path.exists(model_path):
        return None

    state_dict = torch.load(model_path, map_location='cpu')

    num_params = sum(p.numel() for p in state_dict.values())

    size_mb = num_params * 4 / (1024 * 1024)

    return {
        "parameters": f"{num_params / 1_000_000:.2f}M",
        "estimated_size_mb": round(size_mb, 2)
    }

def compress_model(model_path, mode="auto"):
    model_name = "bert-base-uncased"
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)

    print(f"[compress] Loading state_dict from: {model_path}")
    state_dict = torch.load(model_path, map_location=device)
    model.load_state_dict(state_dict)

    model, applied_mode = prepare_model(model, mode=mode)
    model.to(device).eval()

    base_name = os.path.basename(model_path).split('.')[0]
    save_path = f"{base_name}_{applied_mode}.pt"
    torch.save(model.state_dict(), save_path)
    size_mb = os.path.getsize(save_path) / (1024 * 1024)

    inputs = tokenizer("This movie was fantastic!", return_tensors="pt").to(device)
    start = time.time()
    with torch.no_grad():
        _ = model(**inputs)
    end = time.time()
    latency = (end - start) * 1000

    return {
        "source_model": model_path,
        "compressed_model_path": save_path,
        "size_mb": round(size_mb, 2),
        "latency_ms": round(latency, 2),
        "compression_type": applied_mode
    }