"""Continue a text prompt with MP-01 100M base.

Usage (from this folder):
  python run.py "The history of the printing press"
  python run.py "Once upon a time" --tokens 200 --temperature 0.7

Runs on a CUDA GPU when one is available, otherwise on the CPU.
"""
import argparse
import json
import sys
from pathlib import Path

import torch
from safetensors.torch import load_file
from tokenizers import Tokenizer

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
from vbq.model import VBQ, VBQConfig  # noqa: E402


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("prompt", nargs="?", default="The history of the printing press")
    p.add_argument("--tokens", type=int, default=150, help="new tokens to generate")
    p.add_argument("--temperature", type=float, default=0.8)
    p.add_argument("--top-k", type=int, default=50)
    p.add_argument("--top-p", type=float, default=0.95)
    p.add_argument("--repetition-penalty", type=float, default=1.1)
    p.add_argument("--seed", type=int, default=None)
    a = p.parse_args()
    if a.seed is not None:
        torch.manual_seed(a.seed)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    cfg = VBQConfig(**json.loads((HERE / "config.json").read_text())["model_cfg"])
    model = VBQ(cfg)
    model.load_state_dict(load_file(str(HERE / "model.safetensors")))
    model.to(device=device, dtype=torch.bfloat16 if device == "cuda" else torch.float32).eval()

    tok = Tokenizer.from_file(str(HERE / "tokenizer.json"))
    ids = tok.encode(a.prompt, add_special_tokens=False).ids
    idx = torch.tensor([ids], device=device)
    print(a.prompt, end="", flush=True)
    new, shown = [], ""
    with torch.no_grad():
        for t in model.generate_cached(idx, max_new_tokens=a.tokens, temperature=a.temperature, top_k=a.top_k,
                                       top_p=a.top_p, eos_id=0, repetition_penalty=a.repetition_penalty):
            new.append(t)
            text = tok.decode(new)          # decode the whole tail so multi-byte characters print correctly
            print(text[len(shown):], end="", flush=True)
            shown = text
    print()


if __name__ == "__main__":
    main()