Views
No views yet
1import torch
2import torch.nn as nn
3import time
4
5VOCAB_SIZE = 256
6EMBED_DIM = 64
7NUM_HEADS = 4
8NUM_LAYERS = 2
9SEQ_LEN = 32
10DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
11
12class MeowGPT(nn.Module):
13 def __init__(self):
14 super().__init__()
15 self.token_emb = nn.Embedding(VOCAB_SIZE, EMBED_DIM)
16 self.pos_emb = nn.Embedding(SEQ_LEN, EMBED_DIM)
17 layer = nn.TransformerEncoderLayer(d_model=EMBED_DIM, nhead=NUM_HEADS, dim_feedforward=EMBED_DIM*4, batch_first=True)
18 self.transformer = nn.TransformerEncoder(layer, num_layers=NUM_LAYERS)
19 self.head = nn.Linear(EMBED_DIM, VOCAB_SIZE)
20
21 def forward(self, x):
22 B, T = x.shape
23 positions = torch.arange(T, device=x.device)
24 x = self.token_emb(x) + self.pos_emb(positions)
25 mask = torch.triu(torch.ones(T, T, device=DEVICE) * float('-inf'), diagonal=1)
26 x = self.transformer(x, mask=mask)
27 return self.head(x)
28
29def load_bot():
30 model = MeowGPT().to(DEVICE)
31 try:
32 model.load_state_dict(torch.load("meow_model.bin", map_location=DEVICE))
33 print("Loaded meow_model.bin successfully!")
34 except FileNotFoundError:
35 print("Error: meow_model.bin not found. Run train.py first!")
36 exit()
37 model.eval()
38 return model
39
40def chat():
41 model = load_bot()
42 print("\n--- MeowGPT is listening (type 'exit' to stop) ---")
43
44 while True:
45 text = input("You: ")
46 if text.lower() in ['exit', 'quit']: break
47
48 # Prepare input
49 input_ids = [ord(c) for c in text]
50 x = torch.tensor([input_ids], device=DEVICE)
51
52 generated_text = ""
53 for _ in range(10): # Max prediction length
54 if x.size(1) > SEQ_LEN: x = x[:, -SEQ_LEN:] # Truncate if too long
55
56 with torch.no_grad():
57 logits = model(x)
58
59 # Greedy decoding: pick the highest probability token
60 next_token_id = logits[0, -1].argmax().item()
61
62 if next_token_id == 0: # EOS token
63 break
64
65 generated_char = chr(next_token_id)
66 generated_text += generated_char
67
68 # Auto-regressive: append output to input
69 x = torch.cat([x, torch.tensor([[next_token_id]], device=DEVICE)], dim=1)
70
71 print(f"Bot: {generated_text}")
72
73if __name__ == "__main__":
74 chat()