File size: 1,296 Bytes
e4345ad | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 | """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()]))
|