Update app.py

app.py

@@ -5,29 +5,31 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from sklearn.feature_extraction.text import TfidfVectorizer
-from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
+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
-from sklearn.utils import Bunch
 import kagglehub
 import time
 import random
-import threading
 
-# --- 1. CORE MODEL LOGIC (Ported from your script) ---
+# --- 1. CORE MODEL LOGIC ---
+# Using the H3MOS and EpisodicMemory architecture provided in the benchmarks
 
 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
         mem_tensor = torch.stack(self.memory_x).to(query_x.device)
@@ -39,7 +41,14 @@ class EpisodicMemory:
 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())
+        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):
@@ -49,34 +58,46 @@ class MotorPolicy(nn.Module):
     def forward(self, x): return self.fc(x)
 
 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.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)
+        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 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 SETUP ---
+# --- 2. DATA & TRAINING ---
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Loading data on {device}...")
 
-print("Downloading dataset and training models...")
 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')
-X = tfidf.fit_transform(df['content'])
+X_sparse = tfidf.fit_transform(df['content'])
+X_dense = torch.FloatTensor(X_sparse.toarray()).to(device)
 
-targets = {
+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
@@ -84,126 +105,104 @@ targets = {
 
 model_zoo = {}
 encoders = {}
-algs = {
-    "RandomForest": RandomForestClassifier(n_estimators=50),
-    "SVM": SVC(kernel='linear', probability=True),
-    "NaiveBayes": MultinomialNB(),
-    "LogReg": LogisticRegression(max_iter=500),
-    "DISTIL-H3MOS": "torch"
-}
 
-for task, y in targets.items():
+for task, y_labels in tasks.items():
     le = LabelEncoder()
-    y_enc = le.fit_transform(y)
+    y_enc = torch.LongTensor(le.fit_transform(y_labels)).to(device)
     encoders[task] = le
-    for name, clf in algs.items():
-        if name not in model_zoo: model_zoo[name] = {}
-        if name == "DISTIL-H3MOS":
-            model = H3MOS(X.shape[1], 64, len(le.classes_))
-            # (Simplified training for demo speed)
-            optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
-            X_t = torch.FloatTensor(X.toarray())
-            y_t = torch.LongTensor(y_enc)
-            for _ in range(20):
-                optimizer.zero_grad(); loss = F.cross_entropy(model(X_t, True), y_t); loss.backward(); optimizer.step()
-            model_zoo[name][task] = model
-        else:
-            clf.fit(X, y)
-            model_zoo[name][task] = clf
-
-# --- 3. GRADIO UI & CHAT LOGIC ---
-
-# Dicebear Avatars
+    
+    # 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()
+    
+    # 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 ---
+
 def get_avatar(seed):
     return f"https://api.dicebear.com/7.x/adventurer/svg?seed={seed}"
 
 CSS = """
 .reaction-pill { 
-    background: rgba(255, 255, 255, 0.8); 
+    background: rgba(0, 0, 0, 0.05); 
     border-radius: 12px; 
-    padding: 2px 8px; 
-    font-size: 14px; 
-    margin-top: 5px; 
+    padding: 2px 10px; 
+    font-size: 16px; 
+    margin-top: 8px; 
     display: inline-block; 
     border: 1px solid #ddd;
-    cursor: help;
 }
-.chat-window { border-radius: 15px; }
 """
 
-def chat_interface(message, history):
+def predict_all(text):
+    vec_s = tfidf.transform([text])
+    vec_t = torch.FloatTensor(vec_s.toarray()).to(device)
+    
+    results = {}
+    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
+
+def chat_fn(message, history):
     if not message: return "", history
     
-    # 1. Process Reactions for User Message
-    vec = tfidf.transform([message])
-    vec_t = torch.FloatTensor(vec.toarray())
+    preds = predict_all(message)
     
-    reactions = []
-    reaction_details = "Models reacted: "
-    for name in model_zoo.keys():
-        if name == "DISTIL-H3MOS":
-            with torch.no_grad():
-                res = torch.argmax(model_zoo[name]['emoji'](vec_t)).item()
-                emoji = encoders['emoji'].inverse_transform([res])[0]
-        else:
-            emoji = model_zoo[name]['emoji'].predict(vec)[0]
-        reactions.append(emoji)
-        reaction_details += f"{name} ({emoji}) "
-
-    # Most frequent emoji as primary reaction
-    primary_emoji = max(set(reactions), key=reactions.count)
-    reaction_html = f"<div class='reaction-pill' title='{reaction_details}'>{primary_emoji} {len(reactions)}</div>"
+    # 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>"
     
-    user_msg_formatted = f"{message}<br>{reaction_html}"
-    history.append({"role": "user", "content": user_msg_formatted})
+    # 1. Add User Message
+    history.append({"role": "user", "content": f"{message}<br>{reaction_html}"})
     yield history
 
-    # 2. Simulate Bot Responses
-    active_models = ["DISTIL-H3MOS", "RandomForest", "LogReg", "SVM"]
-    random.shuffle(active_models)
-
-    for bot in active_models:
-        # Simulate typing delay
-        time.sleep(random.uniform(0.5, 1.5))
+    # 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
         
-        # Predict Sentiment/Intent
-        if bot == "DISTIL-H3MOS":
-            with torch.no_grad():
-                s_idx = torch.argmax(model_zoo[bot]['sentiment'](vec_t)).item()
-                i_idx = torch.argmax(model_zoo[bot]['intent'](vec_t)).item()
-                sent = encoders['sentiment'].inverse_transform([s_idx])[0]
-                intent = encoders['intent'].inverse_transform([i_idx])[0]
-        else:
-            sent = model_zoo[bot]['sentiment'].predict(vec)[0]
-            intent = model_zoo[bot]['intent'].predict(vec)[0]
-
-        bot_content = f"**Sentiment:** {sent} | **Intent:** {intent}"
+        sent = preds['sentiment'][key]
+        intent = preds['intent'][key]
+        
+        bot_msg = f"**Sentiment:** {sent}  \n**Intent:** {intent}"
         history.append({
             "role": "assistant", 
-            "content": bot_content, 
-            "metadata": {"title": f"{bot}", "avatar": get_avatar(bot)}
+            "content": bot_msg,
+            "metadata": {"title": bot_name}
         })
         yield history
 
-with gr.Blocks(css=CSS) as demo:
-    gr.Markdown("# 🤖 Multitask Model Group Chat\n*A benchmark-turned-chat app featuring H3MOS, RF, and SVM.*")
+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.")
     
     chatbot = gr.Chatbot(
         elem_id="chat-window",
-        type="messages",
         avatar_images=(None, "https://api.dicebear.com/7.x/bottts/svg?seed=User"),
         bubble_full_width=False
     )
     
     with gr.Row():
-        txt = gr.Textbox(
-            show_label=False,
-            placeholder="Type a message to see how the models react...",
-            scale=4
-        )
-        submit_btn = gr.Button("Send", variant="primary", scale=1)
+        msg_input = gr.Textbox(placeholder="Say something...", show_label=False, scale=4)
+        submit = gr.Button("Send", variant="primary")
 
-    txt.submit(chat_interface, [txt, chatbot], [txt, chatbot])
-    submit_btn.click(chat_interface, [txt, chatbot], [txt, chatbot])
+    msg_input.submit(chat_fn, [msg_input, chatbot], [msg_input, chatbot])
+    submit.click(chat_fn, [msg_input, chatbot], [msg_input, chatbot])
 
-demo.launch()
+# Corrected launch for Gradio 6.0 compatibility
+demo.launch(css=CSS)