Update app.py

app.py

@@ -6,203 +6,169 @@ import torch.nn as nn
 import torch.nn.functional as F
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.ensemble import RandomForestClassifier
-from sklearn.linear_model import LogisticRegression
-from sklearn.svm import SVC
-from sklearn.naive_bayes import MultinomialNB
 from sklearn.preprocessing import LabelEncoder
 import kagglehub
 import time
 import random
 
-# --- 1. CORE MODEL LOGIC ---
-# Using the H3MOS and EpisodicMemory architecture provided in the benchmarks
+# --- 1. ARCHITECTURE (From your benchmark) ---
 
 class EpisodicMemory:
-    """Mimics Hippocampal retention and retrieval"""
     def __init__(self, capacity=2000):
         self.memory_x, self.memory_y = [], []
         self.capacity = capacity
-
     def store(self, x, y):
         curr_x, curr_y = x.detach().cpu(), y.detach().cpu()
         for i in range(curr_x.size(0)):
             if len(self.memory_x) >= self.capacity:
                 self.memory_x.pop(0); self.memory_y.pop(0)
             self.memory_x.append(curr_x[i]); self.memory_y.append(curr_y[i])
-
     def retrieve(self, query_x, k=5):
-        if len(self.memory_x) < k: return None
+        if not self.memory_x: return None
         mem_tensor = torch.stack(self.memory_x).to(query_x.device)
         distances = torch.cdist(query_x, mem_tensor)
         top_k_indices = torch.topk(distances, k, largest=False).indices
-        retrieved_y = [torch.stack([self.memory_y[idx] for idx in sample_indices]) for sample_indices in top_k_indices]
-        return torch.stack(retrieved_y).to(query_x.device)
-
-class ExecutiveCore(nn.Module):
-    def __init__(self, input_dim, hidden_dim):
-        super().__init__()
-        self.net = nn.Sequential(
-            nn.Linear(input_dim, hidden_dim), 
-            nn.LayerNorm(hidden_dim), 
-            nn.GELU(), 
-            nn.Dropout(0.2), 
-            nn.Linear(hidden_dim, hidden_dim), 
-            nn.GELU()
-        )
-    def forward(self, x): return self.net(x)
-
-class MotorPolicy(nn.Module):
-    def __init__(self, hidden_dim, output_dim):
-        super().__init__()
-        self.fc = nn.Linear(hidden_dim, output_dim)
-    def forward(self, x): return self.fc(x)
+        return torch.stack([torch.stack([self.memory_y[idx] for idx in s_idx]) for s_idx in top_k_indices]).to(query_x.device)
 
 class H3MOS(nn.Module):
-    """The DISTIL-H3MOS model architecture"""
     def __init__(self, input_dim, hidden_dim, output_dim):
         super().__init__()
-        self.executive = ExecutiveCore(input_dim, hidden_dim)
-        self.motor = MotorPolicy(hidden_dim, output_dim)
+        self.executive = nn.Sequential(nn.Linear(input_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.GELU())
+        self.motor = nn.Linear(hidden_dim, output_dim)
         self.hippocampus = EpisodicMemory()
-
     def forward(self, x, training_mode=False):
         z = self.executive(x)
-        if training_mode or len(self.hippocampus.memory_x) < 10: 
-            return self.motor(z)
-        
-        past_labels = self.hippocampus.retrieve(x, k=5)
         raw_logits = self.motor(z)
-        
+        if training_mode: return raw_logits
+        past_labels = self.hippocampus.retrieve(x, k=5)
         if past_labels is None: return raw_logits
-        
         mem_votes = torch.zeros_like(raw_logits)
         for i in range(x.size(0)):
             votes = torch.bincount(past_labels[i], minlength=raw_logits.size(1)).float()
             mem_votes[i] = votes
-            
         return (0.8 * raw_logits) + (0.2 * F.softmax(mem_votes, dim=1) * 5.0)
 
-# --- 2. DATA & TRAINING ---
+# --- 2. DATA LOAD & QUICK TRAIN ---
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Loading data on {device}...")
+print(f"Initializing models on {device}...")
 
 path = kagglehub.dataset_download('dewanmukto/social-messages-and-emoji-reactions')
 df = pd.read_csv(path+"/messages_emojis.csv").dropna(subset=['content'])
 
-# Mapping logic from the original benchmark
 sent_map = {'❤️':'Pos', '👍':'Pos', '😂':'Pos', '💯':'Pos', '😢':'Neg', '😭':'Neg', '😮':'Neu'}
 intent_map = {'❤️':'Emotion', '👍':'Agreement', '😂':'Emotion', '😮':'Surprise'}
 
-tfidf = TfidfVectorizer(max_features=1000, stop_words='english')
+tfidf = TfidfVectorizer(max_features=500, stop_words='english')
 X_sparse = tfidf.fit_transform(df['content'])
 X_dense = torch.FloatTensor(X_sparse.toarray()).to(device)
 
-tasks = {
-    'emoji': df['emoji'].values,
-    'sentiment': df['emoji'].apply(lambda x: sent_map.get(x, 'Neutral')).values,
-    'intent': df['emoji'].apply(lambda x: intent_map.get(x, 'Other')).values
-}
-
 model_zoo = {}
 encoders = {}
 
-for task, y_labels in tasks.items():
+# We train H3MOS and a RandomForest for variety
+for task in ['emoji', 'sentiment', 'intent']:
+    y_labels = df['emoji'].values if task == 'emoji' else df['emoji'].apply(lambda x: sent_map.get(x, 'Neutral') if task == 'sentiment' else intent_map.get(x, 'Other')).values
     le = LabelEncoder()
     y_enc = torch.LongTensor(le.fit_transform(y_labels)).to(device)
     encoders[task] = le
     
-    # Train H3MOS
-    h3_model = H3MOS(X_dense.shape[1], 64, len(le.classes_)).to(device)
-    opt = torch.optim.Adam(h3_model.parameters(), lr=0.01)
-    for _ in range(30):
-        opt.zero_grad()
-        loss = F.cross_entropy(h3_model(X_dense, True), y_enc)
-        loss.backward(); opt.step()
+    # H3MOS
+    h3 = H3MOS(X_dense.shape[1], 64, len(le.classes_)).to(device)
+    opt = torch.optim.Adam(h3.parameters(), lr=0.01)
+    for _ in range(20):
+        opt.zero_grad(); F.cross_entropy(h3(X_dense, True), y_enc).backward(); opt.step()
     
-    # Store Sklearn models for comparison
-    rf = RandomForestClassifier(n_estimators=50).fit(X_sparse, y_labels)
-    
-    model_zoo[task] = {"H3MOS": h3_model, "RandomForest": rf}
-
-# --- 3. CHAT INTERFACE ---
+    # RandomForest
+    rf = RandomForestClassifier(n_estimators=20).fit(X_sparse, y_labels)
+    model_zoo[task] = {"H3MOS": h3, "RF": rf}
 
-def get_avatar(seed):
-    return f"https://api.dicebear.com/7.x/adventurer/svg?seed={seed}"
+# --- 3. THE UI & LOGIC ---
 
 CSS = """
-.reaction-pill { 
-    background: rgba(0, 0, 0, 0.05); 
-    border-radius: 12px; 
-    padding: 2px 10px; 
-    font-size: 16px; 
-    margin-top: 8px; 
-    display: inline-block; 
-    border: 1px solid #ddd;
+.reaction-btn { 
+    background: #f0f2f5; border: 1px solid #ddd; border-radius: 15px; 
+    padding: 2px 8px; font-size: 14px; cursor: pointer; margin-top: 5px;
 }
+.bot-header { display: flex; align-items: center; margin-bottom: 5px; }
+.bot-avatar { width: 28px; height: 28px; border-radius: 50%; margin-right: 8px; border: 1px solid #eee; }
+.bot-name { font-weight: bold; font-size: 0.9em; color: #555; }
 """
 
-def predict_all(text):
+def get_avatar_url(name):
+    return f"https://api.dicebear.com/7.x/adventurer/svg?seed={name}"
+
+def predict(text):
     vec_s = tfidf.transform([text])
     vec_t = torch.FloatTensor(vec_s.toarray()).to(device)
-    
-    results = {}
+    res = {}
     for task in ['emoji', 'sentiment', 'intent']:
-        # H3MOS Inference
         with torch.no_grad():
-            h3_out = model_zoo[task]["H3MOS"](vec_t)
-            h3_pred = encoders[task].inverse_transform([torch.argmax(h3_out).item()])[0]
-        
-        # RF Inference
-        rf_pred = model_zoo[task]["RandomForest"].predict(vec_s)[0]
-        results[task] = {"H3MOS": h3_pred, "RandomForest": rf_pred}
-    return results
+            h3_idx = torch.argmax(model_zoo[task]["H3MOS"](vec_t)).item()
+            h3_p = encoders[task].inverse_transform([h3_idx])[0]
+        rf_p = model_zoo[task]["RF"].predict(vec_s)[0]
+        res[task] = {"DISTIL-H3MOS": h3_p, "RandomForest": rf_p}
+    return res
 
-def chat_fn(message, history):
+def chat_interface(message, history):
     if not message: return "", history
     
-    preds = predict_all(message)
+    preds = predict(message)
     
-    # Create Reaction HTML (Hover shows model breakdown)
-    emoji_h3 = preds['emoji']['H3MOS']
-    emoji_rf = preds['emoji']['RandomForest']
-    reaction_html = f"<div class='reaction-pill' title='H3MOS: {emoji_h3} | RF: {emoji_rf}'>{emoji_h3} 🤖</div>"
+    # Reaction Logic
+    h3_emoji = preds['emoji']['DISTIL-H3MOS']
+    rf_emoji = preds['emoji']['RandomForest']
+    details = f"DISTIL-H3MOS: {h3_emoji} | RandomForest: {rf_emoji}"
+    reaction_html = f"<button class='reaction-btn' title='{details}'>{h3_emoji} 🤖</button>"
     
     # 1. Add User Message
     history.append({"role": "user", "content": f"{message}<br>{reaction_html}"})
     yield history
 
-    # 2. Sequential Bot Replies
-    bots = [("DISTIL-H3MOS", "H3MOS"), ("RandomForest", "RandomForest")]
-    for bot_name, key in bots:
-        time.sleep(random.uniform(0.8, 1.5)) # Simulation delay
+    # 2. Simulate Group Members Replying
+    bots = ["DISTIL-H3MOS", "RandomForest"]
+    random.shuffle(bots)
+    
+    for bot in bots:
+        # Simulate "Typing..."
+        time.sleep(random.uniform(0.5, 1.2))
+        
+        sent = preds['sentiment'][bot]
+        intent = preds['intent'][bot]
+        avatar = get_avatar_url(bot)
         
-        sent = preds['sentiment'][key]
-        intent = preds['intent'][key]
+        # Format as a social media message
+        bot_content = f"""
+        <div class="bot-header">
+            <img src="{avatar}" class="bot-avatar">
+            <span class="bot-name">{bot}</span>
+        </div>
+        <div style="padding-left: 36px;">
+            <b>Sentiment:</b> {sent}<br>
+            <b>Intent:</b> {intent}
+        </div>
+        """
         
-        bot_msg = f"**Sentiment:** {sent}  \n**Intent:** {intent}"
-        history.append({
-            "role": "assistant", 
-            "content": bot_msg,
-            "metadata": {"title": bot_name}
-        })
+        history.append({"role": "assistant", "content": bot_content})
         yield history
 
 with gr.Blocks() as demo:
-    gr.Markdown("## 💬 Social AI Group Chat")
-    gr.Markdown("Type a message. Models will react with emojis and reply with their analysis.")
+    gr.Markdown("### 📱 Model Group Chat")
+    gr.Markdown("The models below analyze your message for sentiment and intent in real-time.")
     
     chatbot = gr.Chatbot(
         elem_id="chat-window",
-        avatar_images=(None, "https://api.dicebear.com/7.x/bottts/svg?seed=User"),
-        bubble_full_width=False
+        # avatar_images takes a tuple: (user_avatar, bot_placeholder)
+        avatar_images=(get_avatar_url("User"), None),
+        height=500
     )
     
     with gr.Row():
-        msg_input = gr.Textbox(placeholder="Say something...", show_label=False, scale=4)
-        submit = gr.Button("Send", variant="primary")
+        txt = gr.Textbox(placeholder="Type a message...", show_label=False, scale=4)
+        btn = gr.Button("Send", variant="primary")
 
-    msg_input.submit(chat_fn, [msg_input, chatbot], [msg_input, chatbot])
-    submit.click(chat_fn, [msg_input, chatbot], [msg_input, chatbot])
+    txt.submit(chat_interface, [txt, chatbot], [txt, chatbot])
+    btn.click(chat_interface, [txt, chatbot], [txt, chatbot])
 
-# Corrected launch for Gradio 6.0 compatibility
+# Launch with the CSS injection
 demo.launch(css=CSS)