Update app.py

app.py

@@ -6,6 +6,7 @@ import matplotlib.pyplot as plt
 from transformers import BertTokenizer, BertModel
 from sklearn.manifold import TSNE
 import seaborn as sns
+from captum.attr import IntegratedGradients
 import io
 import base64
 import logging
@@ -33,16 +34,18 @@ for name, layer in model.named_modules():
     if 'layer' in name or 'embeddings' in name:
         layer.register_forward_hook(lambda m, i, o, n=name: hook_fn(m, i, o, n))
 
-def process_input(input_text, visualize_option):
+def process_input(input_text, layer_name, visualize_option, attribution_target=0):
     """
-    Process input text and generate visualizations for BERT embeddings or attention.
+    Process input text, compute embeddings, activations, attention, and attribution.
     Parameters:
     - input_text: User-provided text
-    - visualize_option: 'Embeddings' or 'Attention'
+    - layer_name: Selected layer for activation visualization
+    - visualize_option: 'Embeddings', 'Attention', or 'Activations'
+    - attribution_target: Target class for attribution (0 or 1)
     Returns:
-    - Base64-encoded plot image (str)
-    - Dataframe dictionary (dict)
-    - Status message (str)
+    - HTML string with base64-encoded image(s)
+    - List of dataframe dictionaries
+    - Status message
     """
     global activations
     activations = {}  # Reset activations
@@ -50,7 +53,7 @@ def process_input(input_text, visualize_option):
     try:
         # Validate input
         if not input_text.strip():
-            return None, {"Error": ["Input text cannot be empty."]}, "Error: Input text cannot be empty."
+            return "<p>Error: Input text cannot be empty.</p>", [{"Error": ["Input text cannot be empty."]}], "Error: Input text cannot be empty."
 
         # Tokenize input
         inputs = tokenizer(input_text, return_tensors='pt', padding=True, truncation=True, max_length=512)
@@ -67,8 +70,8 @@ def process_input(input_text, visualize_option):
         tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
 
         # Initialize outputs
-        plot_data = None
-        dataframe = None
+        html_plots = []
+        dataframes = []
 
         # Visualization: Embeddings (t-SNE)
         if visualize_option == "Embeddings":
@@ -86,7 +89,7 @@ def process_input(input_text, visualize_option):
                     plt.savefig(buf, format='png', bbox_inches='tight')
                     buf.seek(0)
                     img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
-                    plot_data = f"data:image/png;base64,{img_base64}"
+                    html_plots.append(f'<img src="data:image/png;base64,{img_base64}" alt="t-SNE Plot" style="max-width:100%;"/>')
                     plt.close()
                     # Dataframe for coordinates
                     dataframe = pd.DataFrame({
@@ -94,13 +97,14 @@ def process_input(input_text, visualize_option):
                         "t-SNE X": reduced[:, 0],
                         "t-SNE Y": reduced[:, 1]
                     }).to_dict()
+                    dataframes.append(dataframe)
                 except Exception as e:
                     logger.warning(f"t-SNE failed: {e}")
-                    dataframe = {"Error": [str(e)]}
-                    return None, dataframe, f"Error: t-SNE computation failed: {e}"
+                    dataframes.append({"Error": [str(e)]})
+                    html_plots.append("<p>Error: t-SNE computation failed.</p>")
             else:
-                dataframe = {"Error": ["Too few tokens for t-SNE."]}
-                return None, dataframe, "Error: Too few tokens for t-SNE."
+                dataframes.append({"Error": ["Too few tokens for t-SNE."]})
+                html_plots.append("<p>Error: Too few tokens for t-SNE.</p>")
 
         # Visualization: Attention Weights
         elif visualize_option == "Attention":
@@ -115,25 +119,88 @@ def process_input(input_text, visualize_option):
                 plt.savefig(buf, format='png', bbox_inches='tight')
                 buf.seek(0)
                 img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
-                plot_data = f"data:image/png;base64,{img_base64}"
+                html_plots.append(f'<img src="data:image/png;base64,{img_base64}" alt="Attention Heatmap" style="max-width:100%;"/>')
                 plt.close()
                 # Dataframe for attention weights
                 dataframe = pd.DataFrame(attn, index=tokens, columns=tokens).to_dict()
+                dataframes.append(dataframe)
             else:
-                dataframe = {"Error": ["No attention weights available."]}
-                return None, dataframe, "Error: No attention weights available."
-
-        return plot_data, dataframe, "Processing complete."
+                dataframes.append({"Error": ["No attention weights available."]})
+                html_plots.append("<p>Error: No attention weights available.</p>")
+
+        # Visualization: Activations
+        elif visualize_option == "Activations":
+            if layer_name in activations:
+                act = activations[layer_name]
+                if isinstance(act, tuple):
+                    act = act[0]
+                act = act[0].detach().numpy()  # [seq_len, hidden_size]
+                dataframe = pd.DataFrame(act, index=tokens).to_dict()
+                dataframes.append(dataframe)
+                # Plot mean activation per token
+                fig, ax = plt.subplots(figsize=(8, 6))
+                mean_act = np.mean(act, axis=1)
+                ax.bar(range(len(mean_act)), mean_act)
+                ax.set_xticks(range(len(mean_act)))
+                ax.set_xticklabels(tokens, rotation=45)
+                ax.set_title(f"Mean Activations in {layer_name}")
+                buf = io.BytesIO()
+                plt.savefig(buf, format='png', bbox_inches='tight')
+                buf.seek(0)
+                img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
+                html_plots.append(f'<img src="data:image/png;base64,{img_base64}" alt="Activations Plot" style="max-width:100%;"/>')
+                plt.close()
+            else:
+                dataframes.append({"Error": [f"Layer {layer_name} not found."]})
+                html_plots.append(f"<p>Error: Layer {layer_name} not found.</p>")
+
+        # Attribution: Integrated Gradients
+        def forward_func(inputs, attention_mask=None):
+            outputs = model(inputs, attention_mask=attention_mask)
+            return outputs.pooler_output[:, int(attribution_target)]
+
+        ig = IntegratedGradients(forward_func)
+        try:
+            attributions, _ = ig.attribute(
+                inputs=input_ids,
+                additional_forward_args=(attention_mask,),
+                target=int(attribution_target),
+                return_convergence_delta=True
+            )
+            attr = attributions[0].detach().numpy().sum(axis=1)
+            attr_df = pd.DataFrame({"Token": tokens, "Attribution": attr}).to_dict()
+            dataframes.append(attr_df)
+            # Plot attributions
+            fig, ax = plt.subplots(figsize=(8, 6))
+            ax.bar(range(len(attr)), attr)
+            ax.set_xticks(range(len(attr)))
+            ax.set_xticklabels(tokens, rotation=45)
+            ax.set_title("Integrated Gradients Attribution")
+            buf = io.BytesIO()
+            plt.savefig(buf, format='png', bbox_inches='tight')
+            buf.seek(0)
+            img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
+            html_plots.append(f'<img src="data:image/png;base64,{img_base64}" alt="Attribution Plot" style="max-width:100%;"/>')
+            plt.close()
+        except Exception as e:
+            logger.warning(f"Integrated Gradients failed: {e}")
+            dataframes.append({"Error": [str(e)]})
+            html_plots.append("<p>Error: Attribution computation failed.</p>")
+
+        # Combine HTML plots
+        html_output = "<div>" + "".join(html_plots) + "</div>"
+
+        return html_output, dataframes, "Processing complete."
 
     except Exception as e:
         logger.error(f"Processing failed: {e}")
-        return None, {"Error": [str(e)]}, f"Error: {e}"
+        return f"<p>Error: {e}</p>", [{"Error": [str(e)]}], f"Error: {e}"
 
 # Gradio Interface
 def create_gradio_interface():
     with gr.Blocks(title="Neural Network Visualization Demo") as demo:
         gr.Markdown("# Neural Network Visualization Demo")
-        gr.Markdown("Visualize BERT embeddings or attention weights. Enter text and select a visualization type.")
+        gr.Markdown("Analyze BERT's neural network paths. Enter text, select a layer, and choose a visualization.")
         
         with gr.Row():
             with gr.Column():
@@ -142,21 +209,33 @@ def create_gradio_interface():
                     value="The quick brown fox jumps over the lazy dog.",
                     placeholder="Enter text here..."
                 )
+                layer_name = gr.Dropdown(
+                    label="Select Layer",
+                    choices=[str(name) for name, _ in model.named_modules() if 'layer' in name or 'embeddings' in name],
+                    value="embeddings"
+                )
                 visualize_option = gr.Radio(
                     label="Visualization Type",
-                    choices=["Embeddings", "Attention"],
+                    choices=["Embeddings", "Attention", "Activations"],
                     value="Embeddings"
                 )
+                attribution_target = gr.Slider(
+                    label="Attribution Target Class (0 or 1)",
+                    minimum=0,
+                    maximum=1,
+                    step=1,
+                    value=0
+                )
                 submit_btn = gr.Button("Analyze")
             
             with gr.Column():
-                plot_output = gr.Image(label="Visualization", type="pil")
-                dataframe_output = gr.Dataframe(label="Data Output")
+                plot_output = gr.HTML(label="Visualizations")
+                dataframe_output = gr.Dataframe(label="Data Outputs")
                 text_output = gr.Textbox(label="Messages")
         
         submit_btn.click(
             fn=process_input,
-            inputs=[input_text, visualize_option],
+            inputs=[input_text, layer_name, visualize_option, attribution_target],
             outputs=[plot_output, dataframe_output, text_output]
         )