"""Run text generation with the from-scratch TinyStories Transformer.

Setup (uv):
    uv venv --python 3.12   # torch~=2.6.0 has no wheels yet for newer Pythons
    uv pip install "cs336_basics @ git+https://github.com/zhL-d/llm-from-scratch.git@2b1593c1bc049dde93d332b9f5a6cf9d57575744"
    uv pip install huggingface_hub
    uv run generate.py "Once upon a time"

This script only depends on the `cs336_basics` package (the from-scratch
implementation, pinned to the exact commit this checkpoint was produced
under) and `huggingface_hub` (to pull the tokenizer files from the
dataset repo). The model weights (`model.pt`) and `config.json` in this
repo are loaded directly from disk.
"""
import json
import sys
from pathlib import Path

import torch
from huggingface_hub import hf_hub_download

from cs336_basics.transformer_lm import TransformerLM
from cs336_basics.tokenizer import Tokenizer
from cs336_basics.decoding import Decoding

HERE = Path(__file__).parent


def load_tokenizer(config: dict) -> Tokenizer:
    vocab_path = hf_hub_download(
        repo_id=config["tokenizer_repo"],
        filename=config["vocab_file"],
        repo_type="dataset",
    )
    merges_path = hf_hub_download(
        repo_id=config["tokenizer_repo"],
        filename=config["merges_file"],
        repo_type="dataset",
    )
    return Tokenizer.from_files(
        vocab_filepath=vocab_path,
        merges_filepath=merges_path,
        special_tokens=config["special_tokens"],
    )


def load_model(config: dict) -> TransformerLM:
    model = TransformerLM(
        vocab_size=config["vocab_size"],
        context_length=config["context_length"],
        num_layers=config["num_layers"],
        d_model=config["d_model"],
        num_heads=config["num_heads"],
        d_ff=config["d_ff"],
        rope_theta=config["rope_theta"],
    )
    state = torch.load(HERE / "model.pt", map_location="cpu")
    model.load_state_dict(state["model"])
    model.eval()
    return model


def generate(
    prompt: str,
    max_tokens: int = 256,
    temperature: float = 0.7,
    topp_threshold: float = 0.9,
) -> str:
    config = json.loads((HERE / "config.json").read_text())

    tokenizer = load_tokenizer(config)
    model = load_model(config)

    prompt_ids = tokenizer.encode(prompt)
    eos_id = tokenizer.encode(config["special_tokens"][0])[0]

    generated_ids = Decoding(
        model=model,
        prompt_ids=prompt_ids,
        max_tokens=max_tokens,
        temperature=temperature,
        topp_threshold=topp_threshold,
        eos_id=eos_id,
        context_length=config["context_length"],
    )

    return tokenizer.decode(generated_ids[0].tolist())


if __name__ == "__main__":
    prompt = sys.argv[1] if len(sys.argv) > 1 else "Once upon a time"
    print(generate(prompt))