Spaces:
Sleeping
Sleeping
File size: 2,887 Bytes
e908071 824268b e908071 824268b 649d462 92923aa 824268b 649d462 92923aa 824268b 649d462 824268b 92923aa 824268b 649d462 824268b 92923aa 824268b 649d462 824268b 92923aa 824268b 649d462 824268b 92923aa 824268b 649d462 824268b e908071 824268b e908071 649d462 824268b 649d462 824268b 649d462 92923aa 824268b 63b238e 92923aa 649d462 | 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | import gradio as gr
import torch
import torch.nn as nn
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, 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)
x = self.relu(self.fc1(embedded))
return self.fc2(x)
data = {
"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."
}
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.05)
criterion = nn.CrossEntropyLoss()
def prepare_text(text):
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):
model.train()
log = []
for epoch in range(int(epochs)):
total_loss = 0
for text, label in data.items():
input_tensor, offsets = prepare_text(text)
target = torch.tensor([label], dtype=torch.int64)
optimizer.zero_grad()
output = model(input_tensor, offsets)
loss = criterion(output, target)
loss.backward()
optimizer.step()
total_loss += loss.item()
if epoch % 20 == 0:
log.append(f"Epoch {epoch} - Error: {total_loss:.4f}")
return "\n".join(log)
def chat(user_input):
model.eval()
input_tensor, offsets = prepare_text(user_input)
with torch.no_grad():
output = model(input_tensor, offsets)
prediction = torch.argmax(output, dim=1).item()
return responses[prediction]
with gr.Blocks() as demo:
gr.Markdown("# 🚀 Tiny AI v2: The Hidden Layer")
with gr.Row():
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="Response")
btn_train.click(train_model, inputs=epochs, outputs=status)
btn_chat.click(chat, inputs=chat_input, outputs=chat_output)
demo.launch()
|