| 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)) |
|
|