add sample_native.py

sample_native.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+"""Generate side-by-side samples from the trained native twins."""
+import argparse
+import torch
+import torch.nn.functional as F
+from tokenizers import Tokenizer
+from native_kolm import TinyLM, DEV, TOK_JSON
+
+
+@torch.no_grad()
+def generate(model, tok, prompt, max_new=120, temperature=0.8, top_k=40, seed=0):
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    ids = tok.encode(prompt).ids
+    x = torch.tensor([ids], device=DEV)
+    for _ in range(max_new):
+        logits, _ = model(x[:, -model.ctx:])
+        lg = logits[0, -1] / temperature
+        if top_k:
+            v, _ = torch.topk(lg, top_k)
+            lg[lg < v[-1]] = float("-inf")
+        p = F.softmax(lg, dim=-1).cpu()
+        nxt = torch.multinomial(p, 1, generator=g).item()
+        x = torch.cat([x, torch.tensor([[nxt]], device=DEV)], dim=1)
+    return tok.decode(x[0].tolist())
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--prompt", default="Once upon a time")
+    ap.add_argument("--seed", type=int, default=0)
+    args = ap.parse_args()
+    tok = Tokenizer.from_file(TOK_JSON)
+    for arch in ["kolm", "transformer"]:
+        m = TinyLM(tok.get_vocab_size(), ctx=256, arch=arch).to(DEV)
+        m.load_state_dict(torch.load(f"native_{arch}.pt", map_location=DEV))
+        m.eval()
+        print(f"\n=== {arch} ===")
+        print(generate(m, tok, args.prompt, seed=args.seed))
+
+
+if __name__ == "__main__":
+    main()