File size: 2,732 Bytes
8e0f023 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | #!/usr/bin/env python3
"""Generate text with host/ref.py - the exact-integer reference the ROM is
verified against - and detokenise it.
Greedy argmax with ties to the lowest index, exactly as rom/nn.s does, so the
strings printed here are the strings the cartridge produces. The ROM free-runs
from a single seed token; --prompt is a host-only convenience for inspecting
conditional behaviour and is labelled as such wherever it is used.
"""
import argparse
import json
import os
import sys
import numpy as np
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "host"))
import ref
def load_vocab(path="data/vocab.json"):
v = json.load(open(path))
return v["vocab"], v["base"]
def detok(ids, vocab):
return "".join(vocab[i] for i in ids)
def encode(text, vocab, base):
"""Longest-match tokenisation against the 64-symbol vocabulary."""
order = sorted(range(len(vocab)), key=lambda i: -len(vocab[i]))
out, i = [], 0
while i < len(text):
for t in order:
s = vocab[t]
if s and text.startswith(s, i):
out.append(t)
i += len(s)
break
else:
i += 1
return out
def run(model, prompt_ids, nsteps):
"""Teacher-force the prompt, then free-run greedy. Returns the full id list."""
r = ref.Runner(model)
ids = list(prompt_ids)
cur = ids[0]
p = 0
while p < nsteps:
nxt = r.step(cur, p)
p += 1
if p < len(ids):
cur = ids[p] # still inside the prompt
else:
ids.append(nxt)
cur = nxt
return ids
def main():
ap = argparse.ArgumentParser()
ap.add_argument("npz")
ap.add_argument("--vocab", default="data/vocab.json")
ap.add_argument("--seeds", default=None, help="comma separated seed token ids")
ap.add_argument("--prompt", default=None, help="host-only text prompt")
ap.add_argument("--n", type=int, default=ref.T - 1)
ap.add_argument("--survey", action="store_true", help="all 64 seed tokens")
a = ap.parse_args()
vocab, base = load_vocab(a.vocab)
m = ref.Model.from_npz(a.npz)
if a.prompt is not None:
ids = encode(a.prompt, vocab, base)
out = run(m, ids, a.n)
print("prompt (host-only) %-20r -> %r"
% (detok(ids, vocab), detok(out, vocab)))
return 0
seeds = (list(range(len(vocab))) if a.survey else
[int(s) for s in (a.seeds or "1").split(",")])
for s in seeds:
out = run(m, [s], a.n)
print("seed %2d %-8r -> %r" % (s, vocab[s], detok(out, vocab)))
return 0
if __name__ == "__main__":
sys.exit(main())
|