"""
Verify this model is true BitNet (96%+ ternary weights)
"""
import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("Chris4K/bitnet-gpt2-1.58bit")

total = 0
ternary = 0

for name, param in model.named_parameters():
    if 'weight' in name:
        flat = param.data.flatten()
        is_ternary = (
            torch.isclose(flat, torch.tensor(-1.0), atol=1e-3) |
            torch.isclose(flat, torch.tensor(0.0), atol=1e-3) |
            torch.isclose(flat, torch.tensor(1.0), atol=1e-3)
        )
        ternary += is_ternary.sum().item()
        total += len(flat)

print(f"Ternary percentage: {ternary/total*100:.2f}%")
print("PASS: Real BitNet!" if ternary/total > 0.8 else "FAIL: Fake!")