From 8915d274f50d01205ca1abe2edbb4f20cd9890a4 Mon Sep 17 00:00:00 2001
From: rzen <rzenetl@gmail.com>
Date: Sun, 12 Jul 2026 11:34:45 -0400
Subject: [PATCH] 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)
---
 src/sample.py        | 143 +++++++++++++++++++++++++++++++++++++++++++
 tests/test_sample.py |  86 ++++++++++++++++++++++++++
 2 files changed, 229 insertions(+)
 create mode 100644 src/sample.py
 create mode 100644 tests/test_sample.py

diff --git a/src/sample.py b/src/sample.py
new file mode 100644
index 0000000..6982665
--- /dev/null
+++ b/src/sample.py
@@ -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()
diff --git a/tests/test_sample.py b/tests/test_sample.py
new file mode 100644
index 0000000..03ef014
--- /dev/null
+++ b/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()