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

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