File size: 1,448 Bytes
6b02537 | 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 | #!/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()
|