import json, sys, torch
from pathlib import Path
from safetensors.torch import load_model
from tokenizers import Tokenizer
from model import GPT, GPTConfig
root = Path(__file__).parent
model = GPT(GPTConfig(**json.loads((root/'config.json').read_text())))
load_model(model, root/'model.safetensors')
model.eval(); tokenizer = Tokenizer.from_file(str(root/'tokenizer.json'))
prompt = sys.argv[1] if len(sys.argv) > 1 else '文明開化とは'
ids = tokenizer.encode(prompt).ids
torch.manual_seed(1945)
out = model.generate(torch.tensor(ids).unsqueeze(0), 64, 0.8, 100)[0].tolist()
print(tokenizer.decode(out))