| import torch |
| import json |
| from tokenizers import Tokenizer |
| from safetensors.torch import load_file |
| from model.simbot import SIMGPT |
|
|
| |
| |
| |
| tokenizer = Tokenizer.from_file("tokenizer.json") |
|
|
| with open("config.json") as f: |
| cfg = json.load(f) |
|
|
| |
| |
| |
| model = SIMGPT( |
| vocab_size=cfg["vocab_size"], |
| block_size=cfg["block_size"], |
| n_layers=cfg["n_layers"], |
| n_heads=cfg["n_heads"], |
| d_model=cfg["d_model"] |
| ) |
|
|
| state_dict = load_file("simbot.safetensors") |
| model.load_state_dict(state_dict) |
| model.eval() |
|
|
| print("SimBot GPT ready. Type 'exit' to quit.\n") |
|
|
| |
| |
| |
| while True: |
| user_input = input("User: ").strip() |
| if user_input.lower() in {"exit", "quit"}: |
| break |
|
|
| prompt = f"<bos>\nUser: {user_input}\nAssistant:" |
| ids = tokenizer.encode(prompt).ids |
| x = torch.tensor(ids).unsqueeze(0) |
|
|
| with torch.no_grad(): |
| for _ in range(80): |
| logits = model(x) |
| next_id = torch.argmax(logits[:, -1, :], dim=-1).item() |
| x = torch.cat([x, torch.tensor([[next_id]])], dim=1) |
|
|
| output = tokenizer.decode(x[0].tolist()) |
| print("\nAssistant:", output.split("Assistant:")[-1].strip(), "\n") |
|
|
