Update app.py

app.py

@@ -6,6 +6,7 @@ from huggingface_hub import hf_hub_download
 from importlib import import_module
 import shutil
 import os
+import numpy as np
 
 # Load inference.py and model
 repo_id = "logasanjeev/emotions-analyzer-bert"
@@ -25,10 +26,36 @@ _, _ = predict_emotions("dummy text")
 emotion_labels = inference_module.EMOTION_LABELS
 default_thresholds = inference_module.THRESHOLDS
 
-# Prediction function with grouped bar chart
-def predict_emotions_with_details(text, confidence_threshold=0.0):
+# Function to merge subwords and their scores
+def merge_subwords(tokens, scores):
+    words = []
+    word_scores = []
+    current_word = ""
+    current_score = 0.0
+    count = 0
+    for t, s in zip(tokens, scores):
+        if t in ['[CLS]', '[SEP]', '[PAD]']:
+            continue
+        if t.startswith('##'):
+            current_word += t[2:]
+            current_score += s
+            count += 1
+        else:
+            if current_word:
+                words.append(current_word)
+                word_scores.append(current_score / count)
+            current_word = t
+            current_score = s
+            count = 1
+    if current_word:
+        words.append(current_word)
+        word_scores.append(current_score / count)
+    return words, word_scores
+
+# Prediction function with grouped bar chart and optional heatmap
+def predict_emotions_with_details(text, confidence_threshold=0.0, show_heatmap=False):
     if not text.strip():
-        return "Please enter some text.", "", "", None
+        return "Please enter some text.", "", "", None, ""
     
     predictions_str, processed_text = predict_emotions(text)
     
@@ -51,7 +78,11 @@ def predict_emotions_with_details(text, confidence_threshold=0.0):
     attention_mask = encodings['attention_mask'].to(inference_module.DEVICE)
     
     with torch.no_grad():
-        outputs = inference_module.MODEL(input_ids, attention_mask=attention_mask)
+        outputs = inference_module.MODEL(
+            input_ids, 
+            attention_mask=attention_mask, 
+            output_attentions=show_heatmap
+        )
         logits = torch.sigmoid(outputs.logits).cpu().numpy()[0]
     
     # All emotions for Top 5
@@ -120,7 +151,39 @@ def predict_emotions_with_details(text, confidence_threshold=0.0):
             font=dict(color="#e5e7eb")
         )
     
-    return processed_text, thresholded_output, top_5_output, fig
+    # Generate heatmap if enabled
+    heatmap_html = ""
+    if show_heatmap:
+        attentions = outputs.attentions[-1]  # Last layer attention [batch, heads, seq, seq]
+        cls_att = attentions[0, :, 0, :].mean(dim=0).cpu().numpy()  # Average over heads, from CLS
+        seq_len = int(attention_mask[0].sum())
+        att_scores = cls_att[:seq_len]
+        input_tokens = inference_module.TOKENIZER.convert_ids_to_tokens(input_ids[0][:seq_len])
+        
+        words, word_scores = merge_subwords(input_tokens, att_scores)
+        
+        # Normalize scores to 0-1
+        if max(word_scores) > min(word_scores):
+            word_scores = [(s - min(word_scores)) / (max(word_scores) - min(word_scores)) for s in word_scores]
+        else:
+            word_scores = [0.0] * len(words)
+        
+        # Generate HTML with colored spans
+        html = ""
+        for word, score in zip(words, word_scores):
+            alpha = score
+            color = f"rgba(255, 100, 100, {alpha:.2f})"  # Gradient from transparent to red
+            html += f'<span style="background-color: {color}; padding: 2px 4px; margin: 0 2px; border-radius: 4px; color: black;">{word}</span>'
+        
+        heatmap_html = f"""
+        <div style="padding: 16px; background: rgba(55, 65, 81, 0.7); border-radius: 12px; border: 1px solid rgba(255, 255, 255, 0.1); margin-top: 24px;">
+            <h4 style="color: #e5e7eb; margin-bottom: 12px;">Attention Heatmap (Focus Areas)</h4>
+            <div style="overflow-x: auto; white-space: nowrap;">{html}</div>
+            <p style="color: #d1d5db; font-size: 0.875rem; margin-top: 8px;">Red intensity indicates model's focus (based on CLS token attention in last layer).</p>
+        </div>
+        """
+    
+    return processed_text, thresholded_output, top_5_output, fig, heatmap_html
 
 # Enhanced CSS with modern design
 custom_css = """
@@ -343,6 +406,11 @@ with gr.Blocks(css=custom_css) as demo:
                     info="Filter emotions below this confidence level",
                     elem_classes=["input-slider"]
                 )
+                show_heatmap = gr.Checkbox(
+                    label="Show Attention Heatmap",
+                    value=False,
+                    info="Visualize model focus on text (explainability)"
+                )
                 submit_btn = gr.Button("Analyze Emotions", variant="primary")
     
     # Output Section
@@ -372,6 +440,10 @@ with gr.Blocks(css=custom_css) as demo:
                     label="Emotion Confidence Visualization",
                     elem_classes=["output-plot"]
                 )
+                heatmap_output = gr.HTML(
+                    label="Attention Heatmap",
+                    visible=True
+                )
     
     # Example carousel
     with gr.Group():
@@ -405,8 +477,8 @@ with gr.Blocks(css=custom_css) as demo:
     # Bind predictions
     submit_btn.click(
         fn=predict_emotions_with_details,
-        inputs=[text_input, confidence_slider],
-        outputs=[processed_text_output, thresholded_output, top_5_output, output_plot]
+        inputs=[text_input, confidence_slider, show_heatmap],
+        outputs=[processed_text_output, thresholded_output, top_5_output, output_plot, heatmap_output]
     )
 
 # Launch