"""
Generate Portuguese text from the ode · Pessoana char-GPT.

  pip install tinygrad numpy
  python sample.py --prompt "A cidade" --tokens 500 --temperature 0.8
"""
import argparse, json, pickle
from pathlib import Path
from tinygrad import Tensor
from tinygrad.nn.state import safe_load, load_state_dict
from model import GPT, GPTConfig

HERE = Path(__file__).resolve().parent


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--dir", default=str(HERE), help="folder with model.safetensors/config.json/meta.pkl")
    p.add_argument("--prompt", default="\n")
    p.add_argument("--tokens", type=int, default=500)
    p.add_argument("--temperature", type=float, default=0.8)
    p.add_argument("--top_k", type=int, default=200)
    p.add_argument("--seed", type=int, default=1337)
    a = p.parse_args()
    Tensor.manual_seed(a.seed)

    d = Path(a.dir)
    cfg = GPTConfig(**json.load(open(d / "config.json")))
    model = GPT(cfg)
    load_state_dict(model, safe_load(str(d / "model.safetensors")))
    meta = pickle.load(open(d / "meta.pkl", "rb"))
    stoi, itos = meta["stoi"], meta["itos"]

    Tensor.training = False
    idx = Tensor([[stoi[c] for c in a.prompt]])
    out = model.generate(idx, a.tokens, temperature=a.temperature, top_k=a.top_k or None)
    print("".join(itos[i] for i in out[0].tolist()))


if __name__ == "__main__":
    main()