Phase 5: autoregressive sampler (temperature + top-k)
- src/sample.py: load checkpoint + tokenizer, generate token-by-token with temperature/top-k (or greedy), stop at <|endoftext|>, crop to context window. --suite samples a fixed prompt suite into reports/samples-<config>.md. - tests/test_sample.py: length, context-window safety, early stop, greedy determinism, str output (5 tests, tiny untrained model)
This commit is contained in:
+143
@@ -0,0 +1,143 @@
|
||||
"""Autoregressive sampling from a trained checkpoint.
|
||||
|
||||
Loads the tokenizer and a checkpoint, then generates text one token at a time,
|
||||
feeding each sampled token back in. Supports temperature and top-k sampling.
|
||||
|
||||
Usage:
|
||||
python src/sample.py --prompt "Once upon a time"
|
||||
python src/sample.py --prompt "The dog" --temperature 0.8 --top-k 200 -n 3
|
||||
python src/sample.py --suite # sample a fixed prompt suite -> reports/
|
||||
|
||||
`--suite` samples a fixed set of story openings and writes a gallery to
|
||||
reports/samples-<config>.md — useful to eyeball qualitative progress.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(ROOT / "src"))
|
||||
|
||||
from model import GPT # noqa: E402
|
||||
from tokenizer import BPETokenizer # noqa: E402
|
||||
|
||||
EOT = "<|endoftext|>"
|
||||
TOK_PATH = ROOT / "artifacts" / "tokenizer.json"
|
||||
|
||||
SUITE_PROMPTS = [
|
||||
"Once upon a time,",
|
||||
"One day, a little girl named Lily",
|
||||
"Tom was very happy because",
|
||||
"The dog and the cat were",
|
||||
"In a big forest, there lived",
|
||||
]
|
||||
|
||||
|
||||
def pick_device() -> str:
|
||||
return "mps" if torch.backends.mps.is_available() else "cpu"
|
||||
|
||||
|
||||
def load_model(config: str, device: str) -> GPT:
|
||||
ckpt_path = ROOT / "checkpoints" / config / "ckpt.pt"
|
||||
if not ckpt_path.exists():
|
||||
sys.exit(f"no checkpoint at {ckpt_path} — train first")
|
||||
ckpt = torch.load(ckpt_path, map_location=device)
|
||||
cfg = SimpleNamespace(**ckpt["model_cfg"])
|
||||
model = GPT(cfg).to(device)
|
||||
model.load_state_dict(ckpt["model"])
|
||||
model.eval()
|
||||
print(f"loaded {ckpt_path} (iter {ckpt['iter']}, best val {ckpt.get('best_val', float('nan')):.4f})")
|
||||
return model
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def generate(
|
||||
model: GPT,
|
||||
idx: torch.Tensor,
|
||||
max_new_tokens: int,
|
||||
temperature: float = 0.8,
|
||||
top_k: int | None = None,
|
||||
stop_id: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Extend a (1, T) sequence by up to max_new_tokens, stopping early at stop_id."""
|
||||
ctx_len = model.cfg.context_length
|
||||
for _ in range(max_new_tokens):
|
||||
idx_cond = idx[:, -ctx_len:] # never feed more than the context window
|
||||
logits, _ = model(idx_cond)
|
||||
logits = logits[:, -1, :] # last position's predictions
|
||||
|
||||
if temperature <= 0.0: # greedy
|
||||
next_id = logits.argmax(dim=-1, keepdim=True)
|
||||
else:
|
||||
logits = logits / temperature
|
||||
if top_k is not None:
|
||||
kth = torch.topk(logits, min(top_k, logits.size(-1)))[0][:, [-1]]
|
||||
logits = logits.masked_fill(logits < kth, float("-inf"))
|
||||
probs = F.softmax(logits, dim=-1)
|
||||
next_id = torch.multinomial(probs, num_samples=1)
|
||||
|
||||
idx = torch.cat([idx, next_id], dim=1)
|
||||
if stop_id is not None and next_id.item() == stop_id:
|
||||
break
|
||||
return idx
|
||||
|
||||
|
||||
def sample_text(model, tok, prompt, device, max_new_tokens, temperature, top_k) -> str:
|
||||
eot_id = tok.special_tokens[EOT]
|
||||
ids = tok.encode(prompt) or [eot_id]
|
||||
idx = torch.tensor([ids], dtype=torch.long, device=device)
|
||||
out = generate(model, idx, max_new_tokens, temperature, top_k, stop_id=eot_id)
|
||||
ids_out = [i for i in out[0].tolist() if i != eot_id]
|
||||
return tok.decode(ids_out)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--config", default="v1")
|
||||
ap.add_argument("--prompt", default="Once upon a time,")
|
||||
ap.add_argument("--max-new-tokens", type=int, default=200)
|
||||
ap.add_argument("--temperature", type=float, default=0.8)
|
||||
ap.add_argument("--top-k", type=int, default=200)
|
||||
ap.add_argument("-n", "--num-samples", type=int, default=1)
|
||||
ap.add_argument("--seed", type=int, default=None)
|
||||
ap.add_argument("--suite", action="store_true")
|
||||
args = ap.parse_args()
|
||||
|
||||
if args.seed is not None:
|
||||
torch.manual_seed(args.seed)
|
||||
|
||||
device = pick_device()
|
||||
tok = BPETokenizer.load(TOK_PATH)
|
||||
model = load_model(args.config, device)
|
||||
|
||||
def one(prompt: str) -> str:
|
||||
return sample_text(model, tok, prompt, device,
|
||||
args.max_new_tokens, args.temperature, args.top_k)
|
||||
|
||||
if args.suite:
|
||||
lines = [f"# Sample gallery — {args.config}", ""]
|
||||
for prompt in SUITE_PROMPTS:
|
||||
print(f"\n=== {prompt!r} ===")
|
||||
text = one(prompt)
|
||||
print(text)
|
||||
lines += [f"**Prompt:** {prompt}", "", f"> {text}", ""]
|
||||
out_path = ROOT / "reports" / f"samples-{args.config}.md"
|
||||
out_path.parent.mkdir(exist_ok=True)
|
||||
out_path.write_text("\n".join(lines))
|
||||
print(f"\nsaved {out_path}")
|
||||
else:
|
||||
for i in range(args.num_samples):
|
||||
if args.num_samples > 1:
|
||||
print(f"\n=== sample {i + 1} ===")
|
||||
print(one(args.prompt))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,86 @@
|
||||
"""Sampling tests: generation shape, early stop, greedy determinism.
|
||||
|
||||
Uses a tiny untrained model (output is gibberish — we only check mechanics).
|
||||
Run: `.venv/bin/python tests/test_sample.py`
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
ROOT = Path(__file__).resolve().parent.parent
|
||||
sys.path.insert(0, str(ROOT / "src"))
|
||||
|
||||
from model import GPT # noqa: E402
|
||||
from sample import generate, sample_text # noqa: E402
|
||||
from tokenizer import BPETokenizer # noqa: E402
|
||||
|
||||
EOT = "<|endoftext|>"
|
||||
|
||||
|
||||
def _tiny_model():
|
||||
cfg = SimpleNamespace(vocab_size=300, context_length=32, n_layers=2,
|
||||
n_heads=2, d_model=32, d_ff=64)
|
||||
torch.manual_seed(0)
|
||||
return GPT(cfg).eval()
|
||||
|
||||
|
||||
def _tiny_tokenizer():
|
||||
corpus = ("The cat sat. Tom ran. Lily played happily.\n" + EOT + "\n") * 30
|
||||
return BPETokenizer.train(corpus, vocab_size=300, special_tokens=[EOT])
|
||||
|
||||
|
||||
def test_generate_appends_expected_length():
|
||||
model = _tiny_model()
|
||||
idx = torch.zeros((1, 3), dtype=torch.long)
|
||||
out = generate(model, idx, max_new_tokens=10, temperature=0.8, top_k=50)
|
||||
# Without an early stop, output grows by exactly max_new_tokens.
|
||||
assert out.shape[1] == 3 + 10
|
||||
|
||||
|
||||
def test_generate_respects_context_window():
|
||||
model = _tiny_model() # context_length = 32
|
||||
idx = torch.zeros((1, 40), dtype=torch.long) # longer than the window
|
||||
out = generate(model, idx, max_new_tokens=5, temperature=0.8, top_k=50)
|
||||
assert out.shape[1] == 40 + 5 # must not error on over-long context
|
||||
|
||||
|
||||
def test_generate_stops_at_stop_id():
|
||||
model = _tiny_model()
|
||||
idx = torch.zeros((1, 2), dtype=torch.long)
|
||||
# Greedy with stop_id equal to whatever greedy picks first -> stops immediately.
|
||||
with torch.no_grad():
|
||||
first = model(idx)[0][:, -1, :].argmax().item()
|
||||
out = generate(model, idx, max_new_tokens=20, temperature=0.0, stop_id=first)
|
||||
assert out.shape[1] == 3 # original 2 + the single stop token
|
||||
|
||||
|
||||
def test_greedy_is_deterministic():
|
||||
model = _tiny_model()
|
||||
idx = torch.zeros((1, 4), dtype=torch.long)
|
||||
a = generate(model, idx, max_new_tokens=8, temperature=0.0)
|
||||
b = generate(model, idx, max_new_tokens=8, temperature=0.0)
|
||||
assert torch.equal(a, b)
|
||||
|
||||
|
||||
def test_sample_text_returns_str_without_eot():
|
||||
model = _tiny_model()
|
||||
tok = _tiny_tokenizer()
|
||||
text = sample_text(model, tok, "The cat", "cpu",
|
||||
max_new_tokens=15, temperature=0.8, top_k=50)
|
||||
assert isinstance(text, str)
|
||||
assert EOT not in text # the special token is stripped from output
|
||||
|
||||
|
||||
def main() -> None:
|
||||
tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
|
||||
for t in tests:
|
||||
t()
|
||||
print(f"ok {t.__name__}")
|
||||
print(f"\n{len(tests)} passed")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user