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)

tests/test_sample.py

@@ -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()