Add simplified attention view with top-3 attended tokens and BertViz toggle button

app.py

@@ -609,6 +609,37 @@ def create_layer_accordions(activation_data, model_name):
             else:
                 content_items.append(html.P("No predictions available"))
             
+            # Add attention view (top-3 attended tokens)
+            top_attended = layer.get('top_attended_tokens', [])
+            if top_attended:
+                attention_items = []
+                for token, weight in top_attended:
+                    attention_items.append(html.Div([
+                        html.Span(f"'{token}'", style={'fontWeight': 'bold', 'marginRight': '8px'}),
+                        html.Span(f"{weight:.3f}", style={'color': '#6c757d', 'fontSize': '13px'})
+                    ], style={'marginBottom': '4px'}))
+                
+                content_items.append(html.Div([
+                    html.Hr(style={'margin': '15px 0'}),
+                    html.H6("Attention (current position)", style={'marginBottom': '8px', 'fontSize': '14px', 'color': '#495057'}),
+                    html.Div(attention_items, style={'marginBottom': '10px'}),
+                    html.Button(
+                        [html.I(className="fas fa-expand", style={'marginRight': '5px'}), "Open full interactive view"],
+                        id={'type': 'bertviz-btn', 'layer': layer_num},
+                        className='bertviz-button',
+                        style={
+                            'padding': '6px 12px',
+                            'fontSize': '12px',
+                            'backgroundColor': '#667eea',
+                            'color': 'white',
+                            'border': 'none',
+                            'borderRadius': '4px',
+                            'cursor': 'pointer'
+                        }
+                    ),
+                    html.Div(id={'type': 'bertviz-container', 'layer': layer_num}, style={'marginTop': '10px'})
+                ]))
+            
             panel = html.Details([
                 html.Summary(summary_text, className="layer-summary"),
                 html.Div(content_items, className="layer-content")
@@ -693,6 +724,38 @@ def toggle_comparison_mode(n_clicks, is_comparing):
     
     return new_comparing, style, style, button_content, button_class
 
+# Callback to show BertViz visualization when button is clicked
+@app.callback(
+    Output({'type': 'bertviz-container', 'layer': dash.dependencies.MATCH}, 'children'),
+    [Input({'type': 'bertviz-btn', 'layer': dash.dependencies.MATCH}, 'n_clicks')],
+    [State('session-activation-store', 'data'),
+     State({'type': 'bertviz-btn', 'layer': dash.dependencies.MATCH}, 'id')],
+    prevent_initial_call=True
+)
+def toggle_bertviz_view(n_clicks, activation_data, button_id):
+    """Show/hide BertViz full model view when button is clicked."""
+    if not n_clicks or not activation_data:
+        return None
+    
+    try:
+        from utils import generate_bertviz_html
+        layer_num = button_id['layer']
+        
+        # Toggle: even clicks hide, odd clicks show
+        if n_clicks % 2 == 0:
+            return None
+        
+        # Generate BertViz HTML for this layer
+        bertviz_html = generate_bertviz_html(activation_data, layer_num, 'full')
+        
+        return html.Iframe(
+            srcDoc=bertviz_html,
+            style={'width': '100%', 'height': '500px', 'border': '1px solid #ddd', 'borderRadius': '4px', 'marginTop': '10px'}
+        )
+        
+    except Exception as e:
+        return html.P(f"Error loading BertViz: {str(e)}", style={'color': '#dc3545', 'fontSize': '13px'})
+
 # Node click callback for analysis results  
 @app.callback(
     Output('results-container', 'children'),

todo.md

@@ -24,10 +24,10 @@ Note: Minimal-change approach. Reuse existing files (`app.py`, `components/main_
 - [x] Add a spinning "Loading visuals..." after loading data until all the visualizations are loaded
 
 ## Feature: Simplified attention view + open full interactive view
-- [ ] From `activation_data['attention_outputs']`, compute top-3 attended input tokens for current position (per layer)
-- [ ] Render a simple list: token text + attention weight (rounded), most-to-least
-- [ ] Add button/link: "Open full interactive view" → shows existing BertViz `model_view` for that layer
-- [ ] Keep current BertViz plumbing; reuse callbacks to load the selected layer
+- [x] From `activation_data['attention_outputs']`, compute top-3 attended input tokens for current position (per layer)
+- [x] Render a simple list: token text + attention weight (rounded), most-to-least
+- [x] Add button/link: "Open full interactive view" → shows existing BertViz `model_view` for that layer
+- [x] Keep current BertViz plumbing; reuse callbacks to load the selected layer
 
 ## Feature: Tokenization example of initial prompt (top of page)
 - [ ] Add a new section above panels: "How the tokenizer splits your prompt"

utils/model_patterns.py

@@ -420,12 +420,75 @@ def _compute_certainty(probs: List[float]) -> float:
     return max(0.0, min(1.0, certainty))  # Clamp to [0, 1]
 
 
+def _get_top_attended_tokens(activation_data: Dict[str, Any], layer_num: int, tokenizer, top_k: int = 3) -> Optional[List[Tuple[str, float]]]:
+    """
+    Get top-K attended input tokens for the current position (last token) in a layer.
+    Averages attention across all heads.
+    
+    Args:
+        activation_data: Output from execute_forward_pass
+        layer_num: Layer number to analyze
+        tokenizer: Tokenizer for decoding tokens
+        top_k: Number of top attended tokens to return
+    
+    Returns:
+        List of (token_string, attention_weight) tuples, sorted by weight (highest first)
+    """
+    try:
+        attention_outputs = activation_data.get('attention_outputs', {})
+        input_ids = activation_data.get('input_ids', [])
+        
+        if not attention_outputs or not input_ids:
+            return None
+        
+        # Find attention output for this layer
+        target_module = None
+        for module_name in attention_outputs.keys():
+            numbers = re.findall(r'\d+', module_name)
+            if numbers and int(numbers[0]) == layer_num:
+                target_module = module_name
+                break
+        
+        if not target_module:
+            return None
+        
+        attention_output = attention_outputs[target_module]['output']
+        if not isinstance(attention_output, list) or len(attention_output) < 2:
+            return None
+        
+        # Get attention weights: [batch, heads, seq_len, seq_len]
+        attention_weights = torch.tensor(attention_output[1])
+        
+        # Average across heads: [seq_len, seq_len]
+        avg_attention = attention_weights[0].mean(dim=0)
+        
+        # Get attention from last position to all positions
+        last_pos_attention = avg_attention[-1, :]  # [seq_len]
+        
+        # Get top-K attended positions
+        top_values, top_indices = torch.topk(last_pos_attention, min(top_k, len(last_pos_attention)))
+        
+        # Convert to tokens
+        input_ids_tensor = torch.tensor(input_ids[0]) if isinstance(input_ids[0], list) else torch.tensor(input_ids)
+        result = []
+        for idx, weight in zip(top_indices, top_values):
+            token_id = input_ids_tensor[idx].item()
+            token_str = tokenizer.decode([token_id], skip_special_tokens=False)
+            result.append((token_str, weight.item()))
+        
+        return result
+        
+    except Exception as e:
+        print(f"Warning: Could not compute attended tokens for layer {layer_num}: {e}")
+        return None
+
+
 def extract_layer_data(activation_data: Dict[str, Any], model, tokenizer) -> List[Dict[str, Any]]:
     """
-    Extract layer-by-layer data for accordion display with top-5, deltas, and certainty.
+    Extract layer-by-layer data for accordion display with top-5, deltas, certainty, and attention.
     
     Returns:
-        List of dicts with: layer_num, top_token, top_prob, top_5_tokens, deltas, certainty
+        List of dicts with: layer_num, top_token, top_prob, top_5_tokens, deltas, certainty, top_attended_tokens
     """
     layer_modules = activation_data.get('block_modules', [])
     if not layer_modules:
@@ -444,6 +507,9 @@ def extract_layer_data(activation_data: Dict[str, Any], model, tokenizer) -> Lis
     for layer_num, module_name in layer_info:
         top_tokens = _get_top_tokens(activation_data, module_name, model, tokenizer, top_k=5) if logit_lens_enabled else None
         
+        # Get top-3 attended tokens for this layer
+        top_attended = _get_top_attended_tokens(activation_data, layer_num, tokenizer, top_k=3)
+        
         if top_tokens:
             top_token, top_prob = top_tokens[0]
             
@@ -465,7 +531,8 @@ def extract_layer_data(activation_data: Dict[str, Any], model, tokenizer) -> Lis
                 'top_3_tokens': top_tokens[:3],  # Keep for backward compatibility
                 'top_5_tokens': top_tokens[:5],  # New: top-5 for bar chart
                 'deltas': deltas,
-                'certainty': certainty
+                'certainty': certainty,
+                'top_attended_tokens': top_attended  # New: attention view
             })
             
             # Update previous layer probabilities
@@ -479,7 +546,8 @@ def extract_layer_data(activation_data: Dict[str, Any], model, tokenizer) -> Lis
                 'top_3_tokens': [],
                 'top_5_tokens': [],
                 'deltas': {},
-                'certainty': 0.0
+                'certainty': 0.0,
+                'top_attended_tokens': top_attended
             })
     
     return layer_data