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