File size: 2,063 Bytes
dbd41fe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | import torch
from model import MiniTransformer
from config import *
device = "cuda" if torch.cuda.is_available() else "cpu"
model = MiniTransformer().to(device)
model.load_state_dict(
torch.load("mini.pt", map_location=device, weights_only=True)
)
model.eval()
print("Model loaded. Enter your prompts below. Press Ctrl+C to exit.")
context_tensor = None
while True:
try:
prompt = input("\nUser: ")
if not prompt.strip():
continue
# Append special separator for chat if desired
new_tokens = [ord(c) % 256 for c in prompt + "\nBot: "]
new_tensor = torch.tensor([new_tokens], dtype=torch.long, device=device)
if context_tensor is None:
context_tensor = new_tensor
else:
context_tensor = torch.cat([context_tensor, new_tensor], dim=1)
print("Bot: ", end="", flush=True)
for _ in range(200):
x_crop = context_tensor[:, -BLOCK_SIZE:]
with torch.no_grad():
logits = model(x_crop)
# Apply temperature scaling to soften probability distributions
temperature = 0.8
logits_scaled = logits[:, -1, :] / temperature
# Probabilistic sampling to prevent deterministic repetition
probs = torch.softmax(logits_scaled, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
context_tensor = torch.cat([context_tensor, next_token], dim=1)
# Print the generated character dynamically
char = chr(int(next_token.item()))
print(char, end="", flush=True)
# Optional: stop generating if the model outputs a newline
# if char == '\n':
# break
print() # Newline after generation completes
except KeyboardInterrupt:
print("\nExiting chat loop...")
break
except EOFError:
break
|