Update app.py

app.py

@@ -6,10 +6,8 @@ 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
-from PIL import Image
 import logging
 
 # Set up logging
@@ -35,13 +33,16 @@ 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, layer_name, visualize_option, attribution_target=0):
+def process_input(input_text, visualize_option):
     """
-    Process input text, compute embeddings, activations, and visualizations.
+    Process input text and generate visualizations for BERT embeddings or attention.
+    Parameters:
+    - input_text: User-provided text
+    - visualize_option: 'Embeddings' or 'Attention'
     Returns:
-    - List of base64-encoded plot images
-    - List of dictionaries for dataframe display
-    - Status message
+    - Base64-encoded plot image (str)
+    - Dataframe dictionary (dict)
+    - Status message (str)
     """
     global activations
     activations = {}  # Reset activations
@@ -49,7 +50,7 @@ def process_input(input_text, layer_name, visualize_option, attribution_target=0
     try:
         # Validate input
         if not input_text.strip():
-            return [], [], "Error: Input text cannot be empty."
+            return None, {"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)
@@ -66,141 +67,96 @@ def process_input(input_text, layer_name, visualize_option, attribution_target=0
         tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
 
         # Initialize outputs
-        plots = []
-        dataframes = []
+        plot_data = None
+        dataframe = None
 
         # Visualization: Embeddings (t-SNE)
         if visualize_option == "Embeddings":
-            emb = embeddings[0].detach().numpy()
+            emb = embeddings[0].detach().numpy()  # [seq_len, hidden_size]
             if emb.shape[0] > 1:
                 try:
                     tsne = TSNE(n_components=2, random_state=42, perplexity=min(5, emb.shape[0]-1))
                     reduced = tsne.fit_transform(emb)
-                    fig, ax = plt.subplots()
+                    fig, ax = plt.subplots(figsize=(8, 6))
                     ax.scatter(reduced[:, 0], reduced[:, 1], c='blue')
                     for i, token in enumerate(tokens):
                         ax.annotate(token, (reduced[i, 0], reduced[i, 1]))
                     ax.set_title("t-SNE of Token Embeddings")
                     buf = io.BytesIO()
-                    plt.savefig(buf, format='png')
+                    plt.savefig(buf, format='png', bbox_inches='tight')
                     buf.seek(0)
                     img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
-                    plots.append(f"data:image/png;base64,{img_base64}")
+                    plot_data = f"data:image/png;base64,{img_base64}"
                     plt.close()
+                    # Dataframe for coordinates
+                    dataframe = pd.DataFrame({
+                        "Token": tokens,
+                        "t-SNE X": reduced[:, 0],
+                        "t-SNE Y": reduced[:, 1]
+                    }).to_dict()
                 except Exception as e:
                     logger.warning(f"t-SNE failed: {e}")
-                    dataframes.append({"Error": ["t-SNE could not be computed."]})
+                    dataframe = {"Error": [str(e)]}
+                    return None, dataframe, f"Error: t-SNE computation failed: {e}"
+            else:
+                dataframe = {"Error": ["Too few tokens for t-SNE."]}
+                return None, dataframe, "Error: Too few tokens for t-SNE."
 
         # Visualization: Attention Weights
-        if visualize_option == "Attention":
+        elif visualize_option == "Attention":
             if attentions:
-                attn = attentions[-1][0, 0].detach().numpy()
-                fig, ax = plt.subplots()
+                attn = attentions[-1][0, 0].detach().numpy()  # Last layer, first head
+                fig, ax = plt.subplots(figsize=(8, 6))
                 sns.heatmap(attn, xticklabels=tokens, yticklabels=tokens, cmap='viridis', ax=ax)
                 ax.set_title("Attention Weights (Last Layer, Head 0)")
                 plt.xticks(rotation=45)
                 plt.yticks(rotation=0)
                 buf = io.BytesIO()
-                plt.savefig(buf, format='png')
+                plt.savefig(buf, format='png', bbox_inches='tight')
                 buf.seek(0)
                 img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
-                plots.append(f"data:image/png;base64,{img_base64}")
+                plot_data = f"data:image/png;base64,{img_base64}"
                 plt.close()
+                # Dataframe for attention weights
+                dataframe = pd.DataFrame(attn, index=tokens, columns=tokens).to_dict()
+            else:
+                dataframe = {"Error": ["No attention weights available."]}
+                return None, dataframe, "Error: No attention weights available."
 
-        # Visualization: Activations
-        if visualize_option == "Activations" and layer_name in activations:
-            act = activations[layer_name]
-            if isinstance(act, tuple):
-                act = act[0]
-            act = act[0].detach().numpy()
-            df = pd.DataFrame(act, index=tokens)
-            dataframes.append(df.to_dict())  # Convert to dict for serialization
-            fig, ax = plt.subplots()
-            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')
-            buf.seek(0)
-            img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
-            plots.append(f"data:image/png;base64,{img_base64}")
-            plt.close()
-
-        # 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})
-            dataframes.append(attr_df.to_dict())
-            fig, ax = plt.subplots()
-            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')
-            buf.seek(0)
-            img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
-            plots.append(f"data:image/png;base64,{img_base64}")
-            plt.close()
-        except Exception as e:
-            logger.warning(f"Integrated Gradients failed: {e}")
-            dataframes.append({"Error": ["Attribution could not be computed."]})
-
-        return plots, dataframes, "Processing complete."
+        return plot_data, dataframe, "Processing complete."
 
     except Exception as e:
         logger.error(f"Processing failed: {e}")
-        return [], [{"Error": [str(e)]}], f"Error: {e}"
+        return None, {"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("Analyze BERT's neural network paths. Enter text, select a layer, and choose a visualization.")
+        gr.Markdown("Visualize BERT embeddings or attention weights. Enter text and select a visualization type.")
         
         with gr.Row():
             with gr.Column():
-                input_text = gr.Textbox(label="Input Text", value="The quick brown fox jumps over the lazy dog.")
-                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"
+                input_text = gr.Textbox(
+                    label="Input Text",
+                    value="The quick brown fox jumps over the lazy dog.",
+                    placeholder="Enter text here..."
                 )
                 visualize_option = gr.Radio(
                     label="Visualization Type",
-                    choices=["Embeddings", "Attention", "Activations"],
+                    choices=["Embeddings", "Attention"],
                     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.Gallery(label="Visualizations")
-                dataframe_output = gr.Dataframe(label="Data Outputs")
+                plot_output = gr.Image(label="Visualization", type="pil")
+                dataframe_output = gr.Dataframe(label="Data Output")
                 text_output = gr.Textbox(label="Messages")
         
         submit_btn.click(
             fn=process_input,
-            inputs=[input_text, layer_name, visualize_option, attribution_target],
+            inputs=[input_text, visualize_option],
             outputs=[plot_output, dataframe_output, text_output]
         )
     
@@ -213,4 +169,4 @@ if __name__ == "__main__":
         demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
     except Exception as e:
         logger.error(f"Failed to launch Gradio demo: {e}")
-        print(f"Error launching demo: {e}. Try running locally without share=True.")
+        print(f"Error launching demo: {e}. Try running locally with a different port or without share=True.")