from transformers import GPT2Model, GPT2Config

# --- 1. Load a small open-weight GPT-2 model ---
print("Loading GPT-2 model...")
# You can choose 'distilgpt2' (≈82M params) or 'gpt2' (124M params)
model_name = "distilgpt2"
model = GPT2Model.from_pretrained(model_name)
model.eval()

# --- 2. Compress its weights ---
print("\nCompressing GPT-2 weights...")
compressed_data_gpt, ratio_gpt = compress_weights(model, num_levels=256)

# Save and check disk usage
with open("gpt2_compressed.pkl", "wb") as f:
    pickle.dump(compressed_data_gpt, f)
disk_size_gpt = os.path.getsize("gpt2_compressed.pkl") / (1024 * 1024)
print(f"GPT-2 compression ratio: {ratio_gpt:.4f}")
print(f"GPT-2 compressed file on disk: {disk_size_gpt:.2f} MB")