"""Load the released SAN x-small weights and run a forward pass.

Requires: torch, transformers (for the tokenizer). san_model.py + san_triton.py
must be in the same directory (or on PYTHONPATH).
"""
import torch
from san_model import SimpleAttentionNetwork, SANConfig
from transformers import AutoTokenizer

# 1. Build the architecture — must match the released checkpoint (12 layers).
cfg = SANConfig(
    vocab_size=49152, d_model=384, num_heads=8, num_kv_heads=4,
    num_layers=12, max_seq_len=2048, dtype="bfloat16",
)
model = SimpleAttentionNetwork(cfg).to(torch.bfloat16).eval()

# 2. Load weights. The raw san_latest.pt also works:
#    sd = torch.load("san_latest.pt", map_location="cpu")["model_state_dict"]
sd = torch.load("pytorch_model.bin", map_location="cpu")
model.load_state_dict(sd)
print("loaded", round(sum(p.numel() for p in model.parameters()) / 1e6, 2), "M params")

# 3. Tokenizer (SmolLM2-135M vocab 49152) — local copy shipped in this repo.
tok = AutoTokenizer.from_pretrained("tokenizer")

# 4. Greedy forward over the vocab.
ids = tok("The capital of France is", return_tensors="pt")["input_ids"]
with torch.no_grad():
    logits = model(ids)               # (1, T, 49152)
    next_id = logits[0, -1].argmax(-1)
print("next token:", tok.decode([next_id.item()]))