reuploaded after fixing max output bug for samplers

complete.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+"""Completion CLI: continue a prompt with a trained Notio checkpoint.
+
+Usage:
+    src/complete.py "Once upon a time" --checkpoint out/notio.pt --n 2
+    src/complete.py "The little cat" --temp 0.8 --top-k 40 --max-tokens 200
+"""
+import argparse
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+import sampler
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("prompt", nargs="?", default="", help="text to complete (default: story start)")
+    ap.add_argument("--checkpoint", default="out/notio.pt")
+    ap.add_argument("--temp", type=float, default=0.9)
+    ap.add_argument("--top-k", type=int, default=0, help="0 = off")
+    ap.add_argument("--max-tokens", type=int, default=900)
+    ap.add_argument("--n", type=int, default=1, help="number of completions")
+    ap.add_argument("--device", default=None)
+    a = ap.parse_args()
+
+    m, device, ck = sampler.load_model(a.checkpoint, a.device)
+    print(f"checkpoint: {a.checkpoint} | step {ck['step']} | "
+          f"train {ck.get('loss', float('nan')):.3f} | val {ck.get('val_loss', float('nan')):.3f} | {device}")
+
+    prompt = a.prompt or "Once upon a time"
+    prompt_ids = sampler.encode_prompt(prompt)
+    block = m.cfg.layer1.block_size
+    if len(prompt_ids) >= block:
+        sys.exit(f"prompt too long: {len(prompt_ids)} ids, block_size {block}")
+    print(f"prompt: {prompt!r} ({len(prompt_ids)} ids) | temp {a.temp} | top-k {a.top_k}")
+    for i in range(a.n):
+        ids = sampler.generate(m, device, prompt_ids, a.max_tokens, a.temp, a.top_k)
+        print(f"\n=== completion {i} ({len(ids)} ids) ===\n{sampler.ids_to_display(ids)}")
+
+
+if __name__ == "__main__":
+    main()

sampler.py

@@ -0,0 +1,85 @@
+"""Sampler: load a Notio checkpoint and generate completions statefully.
+
+Prompt is fed once (positions 0..P-1), then generation continues token by
+token carrying the RNN state with advancing position offsets. Decodes via
+the runtime decoder: tags -> effects, display text for humans.
+"""
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+PROJ = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJ))
+sys.path.insert(0, str(PROJ / "scripts"))
+import tokenizer as tk
+from model import Notio
+
+tk.load_vocab()
+
+
+def detach_states(states):
+    out = []
+    for st in states:
+        out.append(tuple(s.detach() for s in st) if isinstance(st, tuple) else st.detach())
+    return out
+
+
+def load_model(checkpoint, device=None):
+    """Load checkpoint -> (model, device, ck_dict). Weights-only disabled:
+    our checkpoints embed NotioConfig."""
+    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+    ck = torch.load(checkpoint, map_location=device, weights_only=False)
+    m = Notio(ck["cfg"])
+    m.load_state_dict(ck["state_dict"])
+    m.to(device)
+    m.eval()
+    return m, device, ck
+
+
+def encode_prompt(text):
+    """Text -> id list, prefixed with <bos>. Accepts literal whitespace."""
+    if isinstance(text, str):
+        text = text.encode()
+    return [1] + list(tk.encode_bytes(text))   # <bos>
+
+
+def ids_to_display(ids):
+    """id list -> human text via the runtime decoder (tags mapped/dropped)."""
+    out = bytes(np.asarray(ids, dtype=np.uint8)).translate(tk.TABLE_DEC)
+    for sent, tok in tk.SENT_DEC:
+        out = out.replace(sent, tok)
+    return tk.display(out).decode()
+
+
+@torch.no_grad()
+def generate(m, device, prompt_ids, max_tokens=900, temperature=0.9, top_k=0):
+    # NOTE: learned story length is ~700-900 ids; too-small max_tokens chops
+    # stories off before the model's natural <eos> (see src/diag_eos.py)
+    """Stateful completion. Stops at <eos> or max_tokens. Returns id list."""
+    block = m.cfg.layer1.block_size
+    if len(prompt_ids) >= block:
+        return list(prompt_ids)          # nothing left to generate
+    if len(prompt_ids) + max_tokens > block:
+        max_tokens = block - len(prompt_ids)   # wpe capacity is the ceiling
+    ids = torch.tensor([prompt_ids], dtype=torch.long, device=device)
+    logits, states = m(ids)              # one pass over the prompt, pos 0..P-1
+    states = detach_states(states)
+    pos = len(prompt_ids)
+    for _ in range(max_tokens):
+        logits, states = m(ids[:, -1:], states, pos_offset=pos)
+        states = detach_states(states)
+        pos += 1
+        logits = logits[:, -1, :] / max(temperature, 1e-6)
+        if top_k > 0:
+            v, _ = torch.topk(logits, top_k)
+            logits = torch.where(logits < v[:, -1:],
+                                 torch.full_like(logits, -float("inf")), logits)
+        probs = F.softmax(logits, dim=-1)
+        nxt = torch.multinomial(probs, 1)
+        ids = torch.cat([ids, nxt], dim=1)
+        if nxt.item() == 2:              # <eos>
+            break
+    return ids[0].tolist()