"""Standalone sampling from a downloaded Hugging Face Shakespeare model folder."""

import argparse
import json
from pathlib import Path

import torch
from model import GPT, ModelConfig
from safetensors.torch import load_file


def load_model(folder, device="cpu"):
    folder = Path(folder)
    config = json.loads((folder / "config.json").read_text())
    tokenizer = json.loads((folder / "tokenizer.json").read_text())
    model = GPT(ModelConfig(**config)).to(device).eval()
    model.load_state_dict(load_file(str(folder / "model.safetensors")), strict=True)
    return model, tokenizer["characters"]


@torch.no_grad()
def generate(model, characters, prompt, tokens=300, seed=1339):
    indices = {character: index for index, character in enumerate(characters)}
    if set(prompt) - indices.keys():
        raise ValueError("prompt contains characters absent from the training vocabulary")
    context = [indices[c] for c in prompt] or [0]
    random = torch.Generator().manual_seed(seed)
    device = next(model.parameters()).device
    output = []
    for _ in range(tokens):
        logits = model(torch.tensor([context[-model.config.block_size :]], device=device))
        probabilities = logits[0, -1].softmax(-1).cpu()
        token = torch.multinomial(probabilities, 1, generator=random).item()
        context.append(token)
        output.append(characters[token])
    return prompt + "".join(output)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--folder", default=".")
    parser.add_argument("--prompt", default="ROMEO:\n")
    parser.add_argument("--tokens", type=int, default=300)
    parser.add_argument("--seed", type=int, default=1339)
    parser.add_argument("--device", default="cpu", choices=("cpu", "cuda", "mps"))
    args = parser.parse_args()
    if args.tokens < 0:
        parser.error("tokens must be nonnegative")
    model, characters = load_model(args.folder, args.device)
    print(generate(model, characters, args.prompt, args.tokens, args.seed))