"""Load the model from this folder (or the Hub) and generate text."""

import argparse

import torch
from tokenizers import Tokenizer
import transformers
from transformers import AutoModelForCausalLM


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default=".", help="Local folder or Hugging Face model ID")
    parser.add_argument("--prompt", default="The purpose of education is")
    parser.add_argument("--max-new-tokens", type=int, default=128)
    parser.add_argument("--temperature", type=float, default=0.8)
    parser.add_argument("--top-k", type=int, default=40)
    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
    dtype_arg = (
        {"dtype": dtype}
        if int(transformers.__version__.split(".", 1)[0]) >= 5
        else {"torch_dtype": dtype}
    )
    model = AutoModelForCausalLM.from_pretrained(
        args.model, trust_remote_code=True, **dtype_arg
    ).to(device).eval()
    tokenizer_path = args.model + "/tokenizer.json" if args.model != "." else "tokenizer.json"
    if not tokenizer_path.startswith(("http://", "https://")) and "/" in args.model and not __import__("os").path.exists(tokenizer_path):
        from huggingface_hub import hf_hub_download
        tokenizer_path = hf_hub_download(args.model, "tokenizer.json")
    tokenizer = Tokenizer.from_file(tokenizer_path)
    eos_id = tokenizer.token_to_id("<eos>")
    token_ids = tokenizer.encode(args.prompt).ids

    generator = torch.Generator(device=device).manual_seed(42)
    with torch.inference_mode():
        for _ in range(args.max_new_tokens):
            context = token_ids[-model.config.context_length :]
            input_ids = torch.tensor([context], device=device, dtype=torch.long)
            logits = model(input_ids=input_ids).logits[0, -1].float()
            k = min(args.top_k, logits.numel())
            values, indexes = torch.topk(logits, k=k)
            probabilities = torch.softmax(values / args.temperature, dim=-1)
            next_id = int(indexes[torch.multinomial(probabilities, 1, generator=generator)].item())
            if next_id == eos_id:
                break
            token_ids.append(next_id)

    print(tokenizer.decode(token_ids, skip_special_tokens=True))


if __name__ == "__main__":
    main()