x-small / load_example.py

Upload folder using huggingface_hub

e4345ad verified about 2 months ago

1.3 kB

"""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()]))

	"""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()]))