Update app.py

app.py

@@ -6,35 +6,42 @@ import torch.optim as optim
 class TinyTextAI(nn.Module):
     def __init__(self, vocab_size, num_classes):
         super(TinyTextAI, self).__init__()
-        self.embedding = nn.EmbeddingBag(vocab_size, 8, sparse=False)
-        self.fc = nn.Linear(8, num_classes)
+        self.embedding = nn.EmbeddingBag(vocab_size, 16, sparse=False)
+        self.fc1 = nn.Linear(16, 12)
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Linear(12, num_classes)
 
     def forward(self, text, offsets):
         embedded = self.embedding(text, offsets)
-        return self.fc(embedded)
+        x = self.relu(self.fc1(embedded))
+        return self.fc2(x)
 
 data = {
-    "hello": 0,
-    "hi": 0,
-    "greetings": 0,
-    "how are you": 1,
-    "how is it going": 1,
-    "what is your name": 2,
-    "who are you": 2
+    "hello": 0, "hi": 0, "hey": 0,
+    "how are you": 1, "status": 1, "up": 1,
+    "name": 2, "who": 2, "identity": 2,
+    "bye": 3, "goodbye": 3, "exit": 3
+}
+
+responses = {
+    0: "Hello! I am feeling smarter now.",
+    1: "Systems are nominal. My parameters are tuned!",
+    2: "I am a Version 2 Tiny AI.",
+    3: "Goodbye! Come back to train me more soon."
 }
-responses = {0: "Hello there!", 1: "I am doing great, thank you!", 2: "I am a Tiny AI built by you."}
 
 vocab = {word: i for i, word in enumerate(set(" ".join(data.keys()).split()))}
+vocab["<UNK>"] = len(vocab)
 vocab_size = len(vocab)
 num_classes = len(responses)
 
 model = TinyTextAI(vocab_size, num_classes)
-optimizer = optim.Adam(model.parameters(), lr=0.1)
+optimizer = optim.Adam(model.parameters(), lr=0.05)
 criterion = nn.CrossEntropyLoss()
 
 def prepare_text(text):
-    tokens = [vocab[w] for w in text.lower().split() if w in vocab]
-    if not tokens: return torch.tensor([0]), torch.tensor([0])
+    tokens = [vocab.get(w, vocab["<UNK>"]) for w in text.lower().split()]
+    if not tokens: return torch.tensor([vocab["<UNK>"]]), torch.tensor([0])
     return torch.tensor(tokens, dtype=torch.int64), torch.tensor([0], dtype=torch.int64)
 
 def train_model(epochs):
@@ -51,7 +58,7 @@ def train_model(epochs):
             loss.backward()
             optimizer.step()
             total_loss += loss.item()
-        if epoch % 10 == 0:
+        if epoch % 20 == 0:
             log.append(f"Epoch {epoch} - Error: {total_loss:.4f}")
     return "\n".join(log)
 
@@ -64,16 +71,16 @@ def chat(user_input):
     return responses[prediction]
 
 with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 Tiny Text Learner")
+    gr.Markdown("# 🚀 Tiny AI v2: The Hidden Layer")
     with gr.Row():
-        epochs = gr.Number(label="Training Rounds", value=100)
-        btn_train = gr.Button("Train AI")
-    status = gr.Textbox(label="Training Status")
-    chat_input = gr.Textbox(label="Ask me something (e.g., 'Hi' or 'Who are you?')")
+        epochs = gr.Number(label="Training Rounds", value=200)
+        btn_train = gr.Button("Re-Train AI")
+    status = gr.Textbox(label="Neural Progress")
+    chat_input = gr.Textbox(label="Talk to the AI")
     btn_chat = gr.Button("Send")
-    chat_output = gr.Textbox(label="AI Response")
+    chat_output = gr.Textbox(label="Response")
     
     btn_train.click(train_model, inputs=epochs, outputs=status)
     btn_chat.click(chat, inputs=chat_input, outputs=chat_output)
 
-demo.launch()
+demo.launch()