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
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = GPT(GPTConfig(**json.loads((root/'config.json').read_text())))
load_model(model, root/'model.safetensors')
model.to(device).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)
tokens = torch.tensor(ids, device=device).unsqueeze(0)
out = model.generate(tokens, 64, 0.8, 100)[0].tolist()
print(tokenizer.decode(out))