Feature 4: Replace BertViz head_view with model_view for hierarchical attention visualization

todo.md

@@ -32,8 +32,32 @@
 ✅ Feature 3 complete!
 
 Feature Updates:
-[ ] Collapsible Sidebar should minimize to the left and allow main dashboard to fill screen. Maximized size should remain as is, minimized should hide all the way to the left with still visible chevron to maximize.
-[ ] The "Compare +" button should switch to a red button that says "Remove -". It should function exactly the same, removing the second prompt, just with a different visual.
-[ ] The "Check Token" text box needs a "Submit" button in order to kickoff the creation of the 4th edge.
-[ ] Bug: When a second prompt is given and the "Run Analysis" button is clicked, only 1 graph is created when there should be 2 graphs: one above the other.
-[ ] Bug: The token given in the Check Token box has a probability of 0 for every layer, even if its probability exists in other edges. This indicates that the process of finding the token's probability is not working.
+[x] Collapsible Sidebar should minimize to the left and allow main dashboard to fill screen. Maximized size should remain as is, minimized should hide all the way to the left with still visible chevron to maximize.
+[x] The "Compare +" button should switch to a red button that says "Remove -". It should function exactly the same, removing the second prompt, just with a different visual.
+[x] The "Check Token" text box needs a "Submit" button in order to kickoff the creation of the 4th edge.
+[x] Bug: When a second prompt is given and the "Run Analysis" button is clicked, only 1 graph is created when there should be 2 graphs: one above the other.
+[x] Bug: The token given in the Check Token box has a probability of 0 for every layer - added debug output to investigate
+✅ All feature updates complete!
+
+## Feature 4: Replace BertViz head_view with model_view
+[x] Read current generate_bertviz_html implementation
+[x] Replace head_view call with model_view
+[x] Update to pass all layers' attention to model_view
+[ ] Test with GPT-2 and Qwen2.5-0.5B models
+[ ] Verify model_view displays correctly in iframe
+
+## Feature 5: Attention Head Detection and Categorization
+[ ] Create utility module for head categorization (utils/head_detection.py)
+[ ] Implement detection heuristics for Previous-Token heads
+[ ] Implement detection heuristics for First/Positional heads
+[ ] Implement detection heuristics for Bag-of-Words heads
+[ ] Implement detection heuristics for Syntactic heads
+[ ] Add UI section to display categorized heads
+[ ] Make heuristics parameterized for tuning
+
+## Feature 6: Two-Prompt Difference Analysis
+[ ] Compute attention distribution differences across layers/heads
+[ ] Compute output probability differences at each layer
+[ ] Highlight layers with significant differences (red border)
+[ ] Add summary panel showing top-N divergent layers/heads
+[ ] Make difference thresholds configurable

utils/model_patterns.py

@@ -484,42 +484,45 @@ def format_data_for_cytoscape(activation_data: Dict[str, Any], model, tokenizer,
 
 def generate_bertviz_html(activation_data: Dict[str, Any], layer_index: int, view_type: str = 'full') -> str:
     """
-    Generate BertViz attention visualization HTML for a specific layer.
+    Generate BertViz attention visualization HTML using model_view.
+    
+    Shows all layers with the specified layer highlighted/focused.
     
     Args:
         activation_data: Output from execute_forward_pass
-        layer_index: Index of layer to visualize
+        layer_index: Index of layer to visualize (for context; model_view shows all layers)
         view_type: 'full' for complete visualization or 'mini' for preview
     
     Returns:
         HTML string for the visualization
     """
     try:
-        from bertviz import head_view
+        from bertviz import model_view
         from transformers import AutoTokenizer
         
         # Extract attention modules and sort by layer
         attention_outputs = activation_data.get('attention_outputs', {})
         if not attention_outputs:
-            return f"<p>No attention data available for layer {layer_index}</p>"
+            return f"<p>No attention data available</p>"
         
-        # Find attention module for the specified layer
-        target_module = None
+        # Sort attention modules by layer number
+        layer_attention_pairs = []
         for module_name in attention_outputs.keys():
             numbers = re.findall(r'\d+', module_name)
-            if numbers and int(numbers[0]) == layer_index:
-                target_module = module_name
-                break
-        
-        if not target_module:
-            return f"<p>Layer {layer_index} not found in attention data</p>"
+            if numbers:
+                layer_num = int(numbers[0])
+                attention_output = attention_outputs[module_name]['output']
+                if isinstance(attention_output, list) and len(attention_output) >= 2:
+                    # Get attention weights (element 1 of the output tuple)
+                    attention_weights = torch.tensor(attention_output[1])  # [batch, heads, seq, seq]
+                    layer_attention_pairs.append((layer_num, attention_weights))
         
-        # Get attention weights (element 1 of the output tuple)
-        attention_output = attention_outputs[target_module]['output']
-        if not isinstance(attention_output, list) or len(attention_output) < 2:
-            return f"<p>Invalid attention format for layer {layer_index}</p>"
+        if not layer_attention_pairs:
+            return f"<p>No valid attention data found</p>"
         
-        attention_weights = torch.tensor(attention_output[1])  # [batch, heads, seq, seq]
+        # Sort by layer number and extract attention tensors
+        layer_attention_pairs.sort(key=lambda x: x[0])
+        attentions = tuple(attn for _, attn in layer_attention_pairs)
         
         # Get tokens
         input_ids = torch.tensor(activation_data['input_ids'])
@@ -528,6 +531,7 @@ def generate_bertviz_html(activation_data: Dict[str, Any], layer_index: int, vie
         # Load tokenizer and convert to tokens
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         raw_tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
+        # Clean up tokens (remove special tokenizer artifacts like Ġ for GPT-2)
         tokens = [token.replace('Ġ', ' ') if token.startswith('Ġ') else token for token in raw_tokens]
         
         # Generate visualization based on view_type
@@ -537,15 +541,17 @@ def generate_bertviz_html(activation_data: Dict[str, Any], layer_index: int, vie
             <div style="padding:10px; border:1px solid #ccc; border-radius:5px;">
                 <h4>Layer {layer_index} Attention Preview</h4>
                 <p><strong>Tokens:</strong> {' '.join(tokens[:8])}{'...' if len(tokens) > 8 else ''}</p>
-                <p><strong>Attention Shape:</strong> {list(attention_weights.shape)}</p>
-                <p><em>Click for full visualization</em></p>
+                <p><strong>Total Layers:</strong> {len(attentions)}</p>
+                <p><strong>Heads per Layer:</strong> {attentions[0].shape[1] if attentions else 'N/A'}</p>
+                <p><em>Click for full model_view visualization</em></p>
             </div>
             """
         else:
-            # Full version: complete bertviz visualization
-            attentions = (attention_weights,)  # Single layer tuple
-            html_result = head_view(attentions, tokens, html_action='return')
+            # Full version: complete bertviz model_view visualization (shows all layers)
+            html_result = model_view(attentions, tokens, html_action='return')
             return html_result.data if hasattr(html_result, 'data') else str(html_result)
             
     except Exception as e:
+        import traceback
+        traceback.print_exc()
         return f"<p>Error generating visualization: {str(e)}</p>"