import argparse
import torch
from tokenizers import Tokenizer

from model import GPT, sample

ap = argparse.ArgumentParser()
ap.add_argument("ckpt", help="чекпоинт: kit_*.pt или last.pt")
ap.add_argument("prompt", nargs="?", default="Москва — это")
ap.add_argument("--tokens", type=int, default=150)
ap.add_argument("--temp", type=float, default=0.8)
ap.add_argument("--top_k", type=int, default=40)
args = ap.parse_args()

device = "cuda"
tok = Tokenizer.from_file("data/tokenizer.json")
model = GPT().to(device)
if args.ckpt.endswith(".safetensors"):
    from safetensors.torch import load_file
    # head.weight связан с tok_emb.weight и в файле отсутствует
    model.load_state_dict({k: v.float() for k, v in load_file(args.ckpt).items()}, strict=False)
else:
    sd = torch.load(args.ckpt, map_location=device)
    model.load_state_dict(sd["model"] if "model" in sd else sd)
print(sample(model, tok, args.prompt, device, args.tokens, args.temp, args.top_k))