Update src/streamlit_app.py

src/streamlit_app.py

@@ -2,10 +2,10 @@ import streamlit as st
 import torch
 import torch.nn.functional as F
 import numpy as np
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModel
 from normalizer import normalize
 import torch.nn as nn
-from transformers import AutoModel
+from huggingface_hub import hf_hub_download
 
 st.set_page_config(page_title="Political Sentiment", layout="wide")
 
@@ -13,38 +13,45 @@ class BanglaPoliticalNet(nn.Module):
     def __init__(self, num_classes=5):
         super().__init__()
         self.banglabert = AutoModel.from_pretrained("csebuetnlp/banglabert")
+        
         self.hidden_size = self.banglabert.config.hidden_size
-
+        
         self.cnn_layers = nn.ModuleList([
-            nn.Conv1d(self.hidden_size, 128, kernel_size=k, padding=k//2)
+            nn.Conv1d(self.hidden_size, 128, kernel_size=k, padding=k//2) 
             for k in [3,5,7]
         ])
-
+        
         self.attention = nn.MultiheadAttention(self.hidden_size, 8, batch_first=True)
         self.classifier = nn.Sequential(
-            nn.Dropout(0.3),
+            nn.LayerNorm(self.hidden_size),
+            nn.Dropout(0.4),
             nn.Linear(self.hidden_size, 512),
-            nn.ReLU(),
-            nn.Dropout(0.2),
-            nn.Linear(512, num_classes)
+            nn.GELU(),
+            nn.Dropout(0.3),
+            nn.Linear(512, 256),
+            nn.GELU(),
+            nn.Linear(256, num_classes)
         )
-
+        
+        self.explainability_weights = nn.Parameter(torch.ones(num_classes) * 0.1)
+    
     def forward(self, input_ids, attention_mask=None):
         bert_out = self.banglabert(input_ids, attention_mask=attention_mask).last_hidden_state
+        
+        cnn_outs = [F.relu(cnn(bert_out.transpose(1,2)).transpose(1,2)) for cnn in self.cnn_layers]
+        cnn_concat = torch.cat(cnn_outs, dim=-1)
 
-        cnn_features = []
-        for cnn in self.cnn_layers:
-            cnn_out = cnn(bert_out.transpose(1,2)).transpose(1,2)
-            cnn_features.append(F.relu(cnn_out))
+        if not hasattr(self, 'cnn_proj'):
 
-        cnn_concat = torch.cat(cnn_features, dim=-1)
-        proj = nn.Linear(384, self.hidden_size).to(input_ids.device)
-        attn_input = proj(cnn_concat)
+            self.cnn_proj = nn.Linear(384, self.hidden_size).to(input_ids.device)
+        
+        attn_input = self.cnn_proj(cnn_concat)
         attn_out, _ = self.attention(attn_input, attn_input, attn_input)
-        attn_pooled = attn_out[:, 0, :]
+        pooled = attn_out[:, 0, :]
+        
+        logits = self.classifier(pooled)
+        return logits, self.explainability_weights
 
-        logits = self.classifier(attn_pooled)
-        return logits
 
 st.markdown("""
 <style>
@@ -223,23 +230,44 @@ label_colors = {
 @st.cache_resource
 def load_models():
     models_loaded = {}
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
     
-    target_models = {
+    standard_models = {
         "model_banglabert": "rocky250/Sentiment-banglabert",
         "model_mbert": "rocky250/Sentiment-mbert",
         "model_bbase": "rocky250/Sentiment-bbase",
-        "model_xlmr": "rocky250/Sentiment-xlmr",
-        "bangla_political": "rocky250/bangla-political"
+        "model_xlmr": "rocky250/Sentiment-xlmr"
     }
 
-    for name, repo in target_models.items():
+    for name, repo in standard_models.items():
         try:
             tokenizer = AutoTokenizer.from_pretrained(repo)
             model = AutoModelForSequenceClassification.from_pretrained(repo)
-            models_loaded[name] = (tokenizer, model.to('cuda' if torch.cuda.is_available() else 'cpu'))
-        except:
+            models_loaded[name] = (tokenizer, model.to(device))
+        except Exception as e:
+            print(f"Skipped {name}: {e}")
             continue
-            
+
+    try:
+
+        model_path = hf_hub_download(repo_id="rocky250/bangla-political", filename="pytorch_model.bin")
+        
+
+        tokenizer = AutoTokenizer.from_pretrained("rocky250/bangla-political")
+        
+ 
+        model = BanglaPoliticalNet(num_classes=5)
+        
+
+        if not hasattr(model, 'cnn_proj'):
+             model.cnn_proj = nn.Linear(384, model.hidden_size)
+             
+        model.load_state_dict(torch.load(model_path, map_location=device), strict=False)
+        
+        models_loaded["bangla_political"] = (tokenizer, model.to(device))
+    except Exception as e:
+        print(f"Skipped bangla_political: {e}")
+
     return models_loaded
 
 models_dict = load_models()
@@ -252,8 +280,11 @@ def predict_single_model(text, model_name):
     inputs = tokenizer(clean_text, return_tensors="pt", truncation=True, padding=True, max_length=128).to(device)
 
     with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
+        if isinstance(model, BanglaPoliticalNet):
+             logits, _ = model(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'])
+        else:
+            outputs = model(**inputs)
+            logits = outputs.logits
 
     probs = F.softmax(logits, dim=1).cpu().numpy()[0]
     pred_id = np.argmax(probs)
@@ -280,6 +311,7 @@ def predict_ensemble(text):
         return final_pred, all_predictions, avg_probs
     return "Error", [], np.zeros(5)
 
+
 st.markdown("""
 <div style='
     text-align: center;
@@ -308,45 +340,50 @@ with col2:
     selected_model = None
     if mode == "Single Model":
         model_options = {name: name for name in models_dict.keys()}
-        selected_model = st.selectbox("Select Model:", list(model_options.keys()), index=0)
+        if model_options:
+            selected_model = st.selectbox("Select Model:", list(model_options.keys()), index=0)
+        else:
+            st.warning("No models loaded.")
 
 analyze_btn = st.button("ANALYZE SENTIMENT", type="primary", use_container_width=True)
 
 if analyze_btn and user_input.strip():
     with st.spinner('Processing with models...'):
         if mode == "Single Model":
-            model_name = selected_model
-            final_res, probs = predict_single_model(user_input, model_name)
-
-            col1, col2 = st.columns([1, 2])
-            with col1:
-                st.markdown(f"""
-                <div class="main-card" style="border-top: 8px solid {label_colors[final_res]}">
-                    <div class="result-title">{model_name}</div>
-                    <div class="result-value" style="color: {label_colors[final_res]}">{final_res}</div>
-                    <div style="font-size: 18px; color: #64748b; margin-top: 15px;">Confidence: {max(probs)*100:.1f}%</div>
-                </div>
-                """, unsafe_allow_html=True)
-
-            with col2:
-                st.markdown('<div class="section-header">Confidence Scores</div>', unsafe_allow_html=True)
-                for i in range(5):
-                    label = id2label[i]
-                    prob = probs[i] * 100
-                    color = label_colors[label]
+            if selected_model:
+                final_res, probs = predict_single_model(user_input, selected_model)
 
+                col1, col2 = st.columns([1, 2])
+                with col1:
                     st.markdown(f"""
-                    <div class="prob-row">
-                        <div class="prob-label">
-                            <span style="font-weight: 700;">{label}</span>
-                            <span style="font-weight: 700; color: {color};">{prob:.1f}%</span>
-                        </div>
-                        <div class="prob-bar-bg">
-                            <div class="prob-bar-fill" style="width: {min(prob, 100)}%; background: linear-gradient(90deg, {color}, {color}cc);"></div>
-                        </div>
+                    <div class="main-card" style="border-top: 8px solid {label_colors[final_res]}">
+                        <div class="result-title">{selected_model}</div>
+                        <div class="result-value" style="color: {label_colors[final_res]}">{final_res}</div>
+                        <div style="font-size: 18px; color: #64748b; margin-top: 15px;">Confidence: {max(probs)*100:.1f}%</div>
                     </div>
                     """, unsafe_allow_html=True)
 
+                with col2:
+                    st.markdown('<div class="section-header">Confidence Scores</div>', unsafe_allow_html=True)
+                    for i in range(5):
+                        label = id2label[i]
+                        prob = probs[i] * 100
+                        color = label_colors[label]
+
+                        st.markdown(f"""
+                        <div class="prob-row">
+                            <div class="prob-label">
+                                <span style="font-weight: 700;">{label}</span>
+                                <span style="font-weight: 700; color: {color};">{prob:.1f}%</span>
+                            </div>
+                            <div class="prob-bar-bg">
+                                <div class="prob-bar-fill" style="width: {min(prob, 100)}%; background: linear-gradient(90deg, {color}, {color}cc);"></div>
+                            </div>
+                        </div>
+                        """, unsafe_allow_html=True)
+            else:
+                 st.error("Model not selected or failed to load.")
+
         else:
             final_res, all_votes, avg_probs = predict_ensemble(user_input)