Update src/streamlit_app.py

src/streamlit_app.py

@@ -14,10 +14,12 @@ class BanglaPoliticalNet(nn.Module):
         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)
             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),
@@ -29,32 +31,39 @@ class BanglaPoliticalNet(nn.Module):
 
     def forward(self, input_ids, attention_mask=None):
         bert_out = self.banglabert(input_ids, attention_mask=attention_mask).last_hidden_state
+
         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))
+
         cnn_concat = torch.cat(cnn_features, dim=-1)
         proj = nn.Linear(384, self.hidden_size).to(input_ids.device)
         attn_input = proj(cnn_concat)
         attn_out, _ = self.attention(attn_input, attn_input, attn_input)
         attn_pooled = attn_out[:, 0, :]
+
         logits = self.classifier(attn_pooled)
         return logits
 
 st.markdown("""
 <style>
 @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700&display=swap');
+
 html, body, [class*="css"] {
     font-family: 'Inter', sans-serif !important;
     color: #1f2937 !important;
 }
+
 .stApp {
     background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
 }
+
 h1, h2, h3 {
     color: #ffffff !important;
     text-shadow: 0 2px 4px rgba(0,0,0,0.3);
 }
+
 .stTextArea textarea {
     background-color: #ffffff !important;
     color: #1f2937 !important;
@@ -63,10 +72,12 @@ h1, h2, h3 {
     padding: 16px !important;
     font-size: 16px !important;
 }
+
 .stTextArea label {
     color: #ffffff !important;
     font-weight: 700 !important;
 }
+
 .main-card {
     background: linear-gradient(145deg, #ffffff 0%, #f8fafc 100%);
     padding: 35px;
@@ -77,6 +88,7 @@ h1, h2, h3 {
     border: 1px solid rgba(255,255,255,0.3);
     backdrop-filter: blur(10px);
 }
+
 .result-title {
     color: #475569 !important;
     font-size: 16px;
@@ -85,12 +97,14 @@ h1, h2, h3 {
     margin-bottom: 12px;
     font-weight: 700;
 }
+
 .result-value {
     font-size: 52px;
     font-weight: 800;
     margin: 0;
     text-shadow: 0 2px 4px rgba(0,0,0,0.1);
 }
+
 .section-header {
     font-size: 22px;
     font-weight: 700;
@@ -103,6 +117,7 @@ h1, h2, h3 {
     border-radius: 10px;
     box-shadow: 0 4px 12px rgba(0,0,0,0.1);
 }
+
 .model-card {
     background: linear-gradient(145deg, #ffffff 0%, #f1f5f9 100%);
     padding: 25px;
@@ -112,10 +127,12 @@ h1, h2, h3 {
     border: 1px solid rgba(255,255,255,0.5);
     transition: all 0.3s ease;
 }
+
 .model-card:hover {
     transform: translateY(-5px);
     box-shadow: 0 20px 40px rgba(0,0,0,0.2);
 }
+
 .model-name {
     color: #334155 !important;
     font-size: 15px;
@@ -124,6 +141,7 @@ h1, h2, h3 {
     border-bottom: 3px solid #e2e8f0;
     padding-bottom: 8px;
 }
+
 .prob-row {
     margin-bottom: 18px;
     background: rgba(255,255,255,0.9);
@@ -131,6 +149,7 @@ h1, h2, h3 {
     border-radius: 12px;
     box-shadow: 0 2px 8px rgba(0,0,0,0.05);
 }
+
 .prob-label {
     font-size: 15px;
     color: #1e293b !important;
@@ -140,6 +159,7 @@ h1, h2, h3 {
     justify-content: space-between;
     align-items: center;
 }
+
 .prob-bar-bg {
     width: 100%;
     height: 14px;
@@ -148,12 +168,14 @@ h1, h2, h3 {
     overflow: hidden;
     box-shadow: inset 0 2px 4px rgba(0,0,0,0.05);
 }
+
 .prob-bar-fill {
     height: 100%;
     border-radius: 7px;
     transition: width 0.8s ease;
     box-shadow: 0 0 20px rgba(0,0,0,0.2);
 }
+
 .stButton > button {
     background: linear-gradient(45deg, #3b82f6, #1d4ed8) !important;
     color: white !important;
@@ -165,18 +187,22 @@ h1, h2, h3 {
     box-shadow: 0 8px 25px rgba(59,130,246,0.4) !important;
     transition: all 0.3s ease !important;
 }
+
 .stButton > button:hover {
     transform: translateY(-2px) !important;
     box-shadow: 0 12px 35px rgba(59,130,246,0.6) !important;
 }
+
 .stRadio > div > label {
     color: #ffffff !important;
     font-weight: 600 !important;
 }
+
 .stSelectbox > label {
     color: #ffffff !important;
     font-weight: 600 !important;
 }
+
 .stExpander {
     background: rgba(255,255,255,0.1) !important;
     border-radius: 12px !important;
@@ -194,63 +220,60 @@ label_colors = {
     'Very Positive': '#16a34a'
 }
 
-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"  
-}
-
 @st.cache_resource
 def load_models():
-    models_dict = {}
-    for key, repo in models.items():
+    models_loaded = {}
+    
+    target_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"
+    }
+
+    for name, repo in target_models.items():
         try:
             tokenizer = AutoTokenizer.from_pretrained(repo)
             model = AutoModelForSequenceClassification.from_pretrained(repo)
-            models_dict[key] = (tokenizer, model.to('cuda' if torch.cuda.is_available() else 'cpu'))
+            models_loaded[name] = (tokenizer, model.to('cuda' if torch.cuda.is_available() else 'cpu'))
         except:
             continue
-        try:
-            SA_tokenizer = AutoTokenizer.from_pretrained("rocky250/bangla-political")
-            model_SA = BanglaPoliticalNet(num_classes=5)
-            model_SA.load_state_dict(torch.load("rocky250/bangla-political/pytorch_model.bin", map_location='cpu'))
-            model_SA = model_SA.to('cuda' if torch.cuda.is_available() else 'cpu')
-            models_dict["Creative Model"] = (SA_tokenizer, model_SA)
-        except:
-            pass
-    return models_dict
+            
+    return models_loaded
 
 models_dict = load_models()
 
 def predict_single_model(text, model_name):
     clean_text = normalize(text)
     tokenizer, model = models_dict[model_name]
+
     device = next(model.parameters()).device
     inputs = tokenizer(clean_text, return_tensors="pt", truncation=True, padding=True, max_length=128).to(device)
+
     with torch.no_grad():
-        if "Creative" in model_name:
-            logits = model(input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'])
-        else:
-            outputs = model(**inputs)
-            logits = outputs.logits
+        outputs = model(**inputs)
+        logits = outputs.logits
+
     probs = F.softmax(logits, dim=1).cpu().numpy()[0]
     pred_id = np.argmax(probs)
     prediction = id2label[pred_id]
+
     return prediction, probs
 
 def predict_ensemble(text):
     clean_text = normalize(text)
     all_probs = []
     all_predictions = []
-    for name in list(models_dict.keys())[:4]:
+
+    for name in models_dict.keys():
         try:
             pred, probs = predict_single_model(clean_text, name)
             all_probs.append(probs)
             all_predictions.append(pred)
         except:
             continue
+
     if all_probs:
         avg_probs = np.mean(all_probs, axis=0)
         final_pred = id2label[np.argmax(avg_probs)]
@@ -259,50 +282,49 @@ def predict_ensemble(text):
 
 st.markdown("""
 <div style='
-text-align: center;
-background: rgba(255,255,255,0.1);
-padding: 30px;
-border-radius: 20px;
-margin-bottom: 30px;
-backdrop-filter: blur(20px);
+    text-align: center;
+    background: rgba(255,255,255,0.1);
+    padding: 30px;
+    border-radius: 20px;
+    margin-bottom: 30px;
+    backdrop-filter: blur(20px);
 '>
-<h1 style='font-size: 3.5rem; margin: 0; background: linear-gradient(45deg, #ffffff, #e2e8f0); -webkit-background-clip: text; -webkit-text-fill-color: transparent; font-weight: 800;'>Political Sentiment Analysis</h1>
+    <h1 style='font-size: 3.5rem; margin: 0; background: linear-gradient(45deg, #ffffff, #e2e8f0); -webkit-background-clip: text; -webkit-text-fill-color: transparent; font-weight: 800;'>Political Sentiment Analysis</h1>
 </div>
 """, unsafe_allow_html=True)
 
 col1, col2 = st.columns([3, 1])
 with col1:
     user_input = st.text_area("Enter Bengali political text:", height=140,
-                              placeholder="এই বক্সে বাংলা রাজনৈতিক মন্তব্য লিখুন...",
-                              help="Type or paste Bengali political text for sentiment analysis")
+        placeholder="এই বক্সে বাংলা রাজনৈতিক মন্তব্য লিখুন...",
+        help="Type or paste Bengali political text for sentiment analysis")
 
 with col2:
     st.markdown("<div style='height: 20px'></div>", unsafe_allow_html=True)
     mode = st.radio("Analysis Mode:",
-                    ["Single Model", "Ensemble"],
-                    horizontal=True)
+        ["Single Model", "Ensemble"],
+        horizontal=True)
 
-if mode == "Single Model":
-    model_options = {f"({i+1}) {name}": name for i, name in enumerate(models_dict.keys())}
-    selected_model = st.selectbox("Select Model:", list(model_options.keys()), index=0)
-else:
-    st.markdown("<div style='height: 50px'></div>", unsafe_allow_html=True)
+    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)
 
 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 = model_options[selected_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 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)
 
@@ -315,13 +337,13 @@ if analyze_btn and user_input.strip():
 
                     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="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)
 
@@ -333,8 +355,8 @@ if analyze_btn and user_input.strip():
             with main_col:
                 st.markdown(f"""
                 <div class="main-card" style="border-top: 8px solid {label_colors[final_res]}; box-shadow: 0 25px 50px rgba(0,0,0,0.2);">
-                <div class="result-title" style="font-size: 18px;">ENSEMBLE CONSENSUS</div>
-                <div class="result-value" style="color: {label_colors[final_res]}; font-size: 60px;">{final_res}</div>
+                    <div class="result-title" style="font-size: 18px;">ENSEMBLE CONSENSUS</div>
+                    <div class="result-value" style="color: {label_colors[final_res]}; font-size: 60px;">{final_res}</div>
                 </div>
                 """, unsafe_allow_html=True)
 
@@ -347,26 +369,26 @@ if analyze_btn and user_input.strip():
 
                     st.markdown(f"""
                     <div class="prob-row">
-                    <div class="prob-label">
-                    <span>{label}</span>
-                    <span style="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="prob-label">
+                            <span>{label}</span>
+                            <span style="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)
 
             with details_col:
                 st.markdown('<div class="section-header">Individual Model Votes</div>', unsafe_allow_html=True)
                 model_cols = st.columns(2)
-                for idx, (name, vote) in enumerate(zip(list(models_dict.keys())[:4], all_votes[:4])):
+                for idx, (name, vote) in enumerate(zip(list(models_dict.keys()), all_votes)):
                     with model_cols[idx % 2]:
                         color = label_colors[vote]
                         st.markdown(f"""
                         <div class="model-card">
-                        <div class="model-name">{name}</div>
-                        <div style="color: {color}; font-weight: 800; font-size: 24px; margin-top: 8px;">{vote}</div>
+                            <div class="model-name">{name}</div>
+                            <div style="color: {color}; font-weight: 800; font-size: 24px; margin-top: 8px;">{vote}</div>
                         </div>
                         """, unsafe_allow_html=True)
 
@@ -382,7 +404,7 @@ with st.expander("Example Political Texts", expanded=False):
     example_cols = st.columns(3)
     for idx, example in enumerate(examples):
         with example_cols[idx]:
-            if st.button(example[:40] + "..." if len(example) > 40 else example,
-                        use_container_width=True):
+            if st.button(example[:40] + "..." if len(example) > 40 else example, 
+                       use_container_width=True):
                 st.session_state.user_input = example
-                st.rerun()
+                st.rerun()