"""Generate text with the standalone Python model with the fixed evaluation sampling policy."""

import argparse
import json
from pathlib import Path

import torch
from model_100m.model import GPT, ModelConfig
from safetensors.torch import load_file
from tokenizers import Tokenizer

p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--prompt", default="The solar system consists of")
p.add_argument("--tokens", type=int, default=128)
p.add_argument("--seed", type=int, default=42)
p.add_argument("--temperature", type=float, default=0.8, help="0 selects greedy decoding")
p.add_argument("--top-k", type=int, default=50, help="0 samples the full vocabulary")
p.add_argument("--device", default="cpu", choices=["cpu", "cuda"])
a = p.parse_args()
if a.tokens < 0 or a.temperature < 0 or a.top_k < 0:
    p.error("Generation settings must be nonnegative")
root = Path(__file__).resolve().parent
torch.set_num_threads(2)
torch.manual_seed(a.seed)
model = GPT(ModelConfig(**json.loads((root / "config.json").read_text()))).to(a.device).eval()
model.load_state_dict(load_file(str(root / "model.safetensors")), strict=True)
tokenizer = Tokenizer.from_file(str(root / "tokenizer/tokenizer.json"))
ids = [1] + tokenizer.encode(a.prompt, add_special_tokens=False).ids
generator = torch.Generator(device=a.device).manual_seed(a.seed)
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
with torch.inference_mode():
    for _ in range(a.tokens):
        logits = model(torch.tensor([ids[-model.config.block_size :]], device=a.device))[
            0, -1
        ].float()
        if not torch.isfinite(logits).all():
            raise ValueError("Non-finite generation logits")
        logits[[0, 1, 3]] = -torch.inf  # PAD/BOS/UNK; EOS remains eligible.
        if a.temperature == 0:
            token = int(logits.argmax())
        elif a.top_k:
            values, indices = torch.topk(logits / a.temperature, min(a.top_k, logits.numel()))
            token = int(indices[torch.multinomial(values.softmax(-1), 1, generator=generator)])
        else:
            token = int(
                torch.multinomial((logits / a.temperature).softmax(-1), 1, generator=generator)
            )
        ids.append(token)
        if token == 2:
            break
print(tokenizer.decode(ids, skip_special_tokens=True))