#!/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()