Create app.py

app.py

@@ -0,0 +1,231 @@
+import gradio as gr
+import torch
+import numpy as np
+import pandas as pd
+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
+
+# Initialize BERT model and tokenizer
+tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+model = BertModel.from_pretrained('bert-base-uncased')
+model.eval()
+
+# Alternative MLP model (uncomment to use instead of BERT)
+"""
+# import torch.nn as nn
+# class SimpleMLP(nn.Module):
+#     def __init__(self, input_size=10, hidden_sizes=[64, 32], output_size=2):
+#         super(SimpleMLP, self).__init__()
+#         layers = []
+#         prev_size = input_size
+#         for hidden_size in hidden_sizes:
+#             layers.append(nn.Linear(prev_size, hidden_size))
+#             layers.append(nn.ReLU())
+#             prev_size = hidden_size
+#         layers.append(nn.Linear(prev_size, output_size))
+#         self.network = nn.Sequential(*layers)
+#     def forward(self, x):
+#         return self.network(x)
+# model = SimpleMLP()
+# model.eval()
+"""
+
+# Store intermediate activations
+activations = {}
+def hook_fn(module, input, output, name):
+    activations[name] = output
+
+# Register hooks for BERT layers (or MLP layers)
+for name, layer in model.named_modules():
+    if 'layer' in name or 'embeddings' in name:  # Focus on transformer layers
+        layer.register_forward_hook(lambda m, i, o, n=name: hook_fn(m, i, o, n))
+    # For MLP, replace with:
+    # if isinstance(layer, nn.Linear) or isinstance(layer, nn.ReLU):
+    #     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):
+    """
+    Process input text, compute embeddings, activations, and visualizations.
+    Parameters:
+    - input_text: User-provided text input
+    - layer_name: Selected layer for visualization
+    - visualize_option: 'Embeddings', 'Attention', or 'Activations'
+    - attribution_target: Target class for attribution (0 or 1 for binary classification)
+    Returns:
+    - Dictionary with plots and dataframes
+    """
+    global activations
+    activations = {}  # Reset activations
+
+    # Tokenize input
+    inputs = tokenizer(input_text, return_tensors='pt', padding=True, truncation=True, max_length=512)
+    input_ids = inputs['input_ids']
+    attention_mask = inputs['attention_mask']
+
+    # Forward pass
+    with torch.no_grad():
+        outputs = model(input_ids, attention_mask=attention_mask, output_attentions=True, output_hidden_states=True)
+        embeddings = outputs.last_hidden_state  # [batch, seq_len, hidden_size]
+        attentions = outputs.attentions  # List of attention weights
+        hidden_states = outputs.hidden_states  # List of hidden states
+
+    # Convert token IDs to tokens
+    tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
+
+    # Initialize output dictionary
+    results = {
+        "plots": [],
+        "dataframes": [],
+        "text": []
+    }
+
+    # Visualization: Embeddings (t-SNE)
+    if visualize_option == "Embeddings":
+        emb = embeddings[0].detach().numpy()  # [seq_len, hidden_size]
+        if emb.shape[0] > 1:  # Need at least 2 points for t-SNE
+            tsne = TSNE(n_components=2, random_state=42, perplexity=min(5, emb.shape[0]-1))
+            reduced = tsne.fit_transform(emb)
+            fig, ax = plt.subplots()
+            scatter = 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")
+            # Convert plot to base64 for Gradio
+            buf = io.BytesIO()
+            plt.savefig(buf, format='png')
+            buf.seek(0)
+            img = Image.open(buf)
+            img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
+            results["plots"].append(f"data:image/png;base64,{img_base64}")
+            plt.close()
+
+    # Visualization: Attention Weights
+    if visualize_option == "Attention":
+        if attentions:
+            attn = attentions[-1][0, 0].detach().numpy()  # Last layer, first head
+            fig, ax = plt.subplots()
+            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)
+            # Convert plot to base64
+            buf = io.BytesIO()
+            plt.savefig(buf, format='png')
+            buf.seek(0)
+            img = Image.open(buf)
+            img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
+            results["plots"].append(f"data:image/png;base64,{img_base64}")
+            plt.close()
+
+    # Visualization: Activations
+    if visualize_option == "Activations":
+        if layer_name in activations:
+            act = activations[layer_name]
+            if isinstance(act, tuple):  # Handle attention outputs
+                act = act[0]
+            act = act[0].detach().numpy()  # [seq_len, hidden_size]
+            df = pd.DataFrame(act, index=tokens)
+            results["dataframes"].append(df)
+            # Plot mean activation per token
+            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}")
+            # Convert plot to base64
+            buf = io.BytesIO()
+            plt.savefig(buf, format='png')
+            buf.seek(0)
+            img = Image.open(buf)
+            img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
+            results["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[:, attribution_target]
+    
+    ig = IntegratedGradients(forward_func)
+    attributions, delta = ig.attribute(
+        inputs=input_ids,
+        additional_forward_args=(attention_mask,),
+        target=attribution_target,
+        return_convergence_delta=True
+    )
+    attr = attributions[0].detach().numpy()
+    attr_df = pd.DataFrame({"Token": tokens, "Attribution": attr.sum(axis=1)})
+    results["dataframes"].append(attr_df)
+    
+    # Plot attributions
+    fig, ax = plt.subplots()
+    ax.bar(range(len(attr_df)), attr_df["Attribution"])
+    ax.set_xticks(range(len(attr_df)))
+    ax.set_xticklabels(tokens, rotation=45)
+    ax.set_title("Integrated Gradients Attribution")
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    img = Image.open(buf)
+    img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
+    results["plots"].append(f"data:image/png;base64,{img_base64}")
+    plt.close()
+
+    return (
+        results["plots"] if results["plots"] else None,
+        results["dataframes"] if results["dataframes"] else None,
+        "\n".join(results["text"]) if results["text"] else "Processing complete."
+    )
+
+# 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 the paths of a BERT model from input to output. Enter text, select a layer, and choose a visualization option.")
+        
+        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=[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", "Activations"],
+                    value="Embeddings"
+                )
+                attribution_target = gr.Slider(
+                    label="Attribution Target Class (0 or 1 for binary classification)",
+                    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")
+                text_output = gr.Textbox(label="Messages")
+        
+        submit_btn.click(
+            fn=process_input,
+            inputs=[input_text, layer_name, visualize_option, attribution_target],
+            outputs=[plot_output, dataframe_output, text_output]
+        )
+    
+    return demo
+
+# Launch the demo
+if __name__ == "__main__":
+    demo = create_gradio_interface()
+    demo.launch()