conductor setup

.cursor/rules/minimal_changes.mdc

@@ -20,6 +20,7 @@ alwaysApply: true
 - Plan first:
   - Update `todo.md` with the smallest next actions tied to `plans.md`.
   - Keep tasks atomic and check them off as you go.
+  - Use the `conductor` folder to learn about the project. Maintain this folder after every change to the code in order to keep running memory (only make changes if necessary).
 
 - Keep edits minimal:
   - Prefer small, surgical changes over refactors.

README.md

@@ -1,210 +1,105 @@
-# Transformer Activation Capture and Visualization
+# Transformer Explanation Dashboard
 
-This project provides tools for capturing activations from transformer models and visualizing attention patterns using bertviz and an interactive Dash web application.
+A comprehensive, interactive tool for capturing, visualizing, and experimenting with Transformer-based Large Language Models (LLMs). This project demystifies the inner workings of models by transforming abstract architectural concepts into tangible, observable phenomena.
 
-## Overview
+## Vision
 
-The project consists of multiple components:
-1. **Interactive Dashboard** (`app.py`) - Web-based visualization with automatic model family detection
-2. **Model Configuration** (`utils/model_config.py`) - Hard-coded mappings for common model families
-3. **Activation Capture** (`utils/model_patterns.py`) - PyVene-based activation capture utilities
-4. **Legacy Tools** (`agnostic_capture.py`, `bertviz_head_model_view.py`) - Command-line tools
+To foster a deep, intuitive understanding of how powerful models process information by combining interactive visualizations with hands-on experimentation capabilities.
 
-## New Feature: Automatic Model Family Detection
+## Key Features
 
-The dashboard now automatically detects model families and pre-fills dropdown selections with appropriate modules and parameters. This eliminates manual selection for common architectures.
+### 🔍 Interactive Pipeline Visualization
+Follow the data flow step-by-step through the model's architecture:
+1.  **Tokenization**: See how text is split and assigned IDs.
+2.  **Embedding**: Visualize the look-up of semantic vectors.
+3.  **Attention**: Explore head-level attention patterns using **BertViz**.
+4.  **MLP (Feed-Forward)**: Understand where factual knowledge is stored.
+5.  **Output Selection**: View probability distributions and top predictions.
 
-### Supported Model Families
+### 🧪 Experiments & Investigation
+Go beyond static observation with interactive experiments:
+*   **Ablation Studies**: Selectively disable specific attention heads across different layers to observe their impact on generation and probability.
+*   **Token Attribution**: Use **Integrated Gradients** to see which input tokens contributed most to a specific prediction.
+*   **Beam Search Analysis**: Visualize how multiple generation choices are explored.
 
-- **LLaMA-like**: LLaMA 2/3, Mistral, Mixtral, Qwen2/2.5
-- **GPT-2**: GPT-2, GPT-2 Medium/Large/XL
-- **OPT**: Facebook OPT models (125M - 30B)
-- **GPT-NeoX**: EleutherAI Pythia, GPT-NeoX-20B
-- **BLOOM**: BigScience BLOOM models
-- **Falcon**: TII Falcon models
-- **MPT**: MosaicML MPT models
+### 🤖 Broad Model Support
+The dashboard features **Automatic Model Family Detection**, supporting a wide range of architectures without manual configuration:
+*   **LLaMA-like**: LLaMA 2/3, Mistral, Mixtral, Qwen2/2.5
+*   **GPT-2**: GPT-2 (Small/Medium/Large/XL)
+*   **OPT**: Facebook OPT models
+*   **GPT-NeoX**: Pythia, GPT-NeoX
+*   **BLOOM**: BigScience BLOOM
+*   **Falcon**: TII Falcon
+*   **MPT**: MosaicML MPT
 
-### How It Works
+## Getting Started
 
-1. Select a model from the dropdown
-2. The app detects the model family (e.g., "gpt2", "llama_like")
-3. Dropdowns auto-fill with family-specific patterns:
-   - **Attention modules**: e.g., `transformer.h.{N}.attn` for GPT-2
-   - **MLP modules**: e.g., `model.layers.{N}.mlp` for LLaMA
-   - **Normalization parameters**: e.g., `model.norm.weight` for LLaMA
-   - **Logit lens parameter**: e.g., `lm_head.weight`
-4. You can still manually adjust selections if needed
+### Prerequisites
+*   Python 3.11+ recommended
+*   PyTorch
 
-### Adding New Models
+### Installation
 
-Edit `utils/model_config.py` and add entries to `MODEL_TO_FAMILY`:
+1.  Clone the repository:
+    ```bash
+    git clone https://github.com/yourusername/transformer-dashboard.git
+    cd transformer-dashboard
+    ```
 
-```python
-MODEL_TO_FAMILY = {
-    "your-org/your-model": "llama_like",  # or gpt2, opt, etc.
-    # ...
-}
-```
-
-No code changes needed if the model follows an existing family's architecture!
+2.  Install dependencies:
+    ```bash
+    pip install -r requirements.txt
+    ```
 
-## Files
+### Running the Dashboard
 
-### `agnostic_capture.py`
-A model-agnostic activation capture tool that hooks into transformer modules and saves their outputs.
+Launch the application:
 
-**Key Features:**
-- Automatically categorizes modules into attention, MLP, and other types
-- Interactive module selection by pattern
-- Supports both PyTorch hooks and PyVene integration
-- Saves data in organized JSON structure for easy retrieval
-
-**Usage:**
 ```bash
-# Basic usage
-python agnostic_capture.py --model "Qwen/Qwen2.5-0.5B" --prompt "Once upon a time"
-
-# Capture attention weights for bertviz
-python agnostic_capture.py --model "Qwen/Qwen2.5-0.5B" --prompt "Once upon a time" --output my_activations.json
-
-# Auto-select patterns (attention:0, mlp:0 selects first pattern from each)
-python agnostic_capture.py --auto-select "attn:0;mlp:0;other:" --model "Qwen/Qwen2.5-0.5B"
+python app.py
 ```
 
-**Interactive Selection:**
-When run without `--auto-select`, the script will:
-1. Show available module patterns grouped by type
-2. Allow you to select patterns by index, name, or suffix
-3. Selected patterns apply to ALL layers that contain them
-
-### `bertviz_head_model_view.py`
-Creates interactive HTML visualizations of attention patterns using the bertviz library.
-
-**Features:**
-- Generates head view (attention patterns per head)
-- Generates model view (attention patterns across layers)
-- Automatically extracts attention weights from captured data
-- Saves HTML files that can be opened in any browser
+Open your browser and navigate to `http://127.0.0.1:8050/`.
 
-**Usage:**
-```bash
-python bertviz_head_model_view.py
-```
+## Usage Guide
 
-**Output:**
-- `bertviz/attention_head_view_{model_name}.html` - Head-level attention patterns
-- `bertviz/attention_model_view_{model_name}.html` - Model-level attention patterns
-
-## Data Structure
-
-The captured data is organized in the following JSON structure:
-
-```json
-{
-  "model": "model_name",
-  "prompt": "input_text",
-  "input_ids": [[token_ids]],
-  "selected_patterns": {
-    "attention": ["pattern1", "pattern2"],
-    "mlp": ["pattern1"],
-    "other": []
-  },
-  "selected_modules": {
-    "attention": ["model.layers.0.self_attn", "model.layers.1.self_attn", ...],
-    "mlp": ["model.layers.0.mlp", "model.layers.1.mlp", ...],
-    "other": []
-  },
-  "captured": {
-    "attention_outputs": {
-      "model.layers.0.self_attn": {
-        "output": [
-          [[...]], // Attention output (processed values)
-          [[...]]  // Attention weights (used by bertviz)
-        ]
-      }
-    },
-    "mlp_outputs": { ... },
-    "other_outputs": { ... }
-  }
-}
-```
+1.  **Select a Model**: Choose from the predefined list or enter a HuggingFace model ID. The system will auto-detect the architecture.
+2.  **Enter a Prompt**: Type a sentence to analyze.
+3.  **Configure Generation**: Adjust "Number of New Tokens" and "Number of Generation Choices" (Beam Width).
+4.  **Run Analysis**: Click "Analyze" to run the forward pass.
+5.  **Explore the Pipeline**: Click on the pipeline stages (Tokenization, Attention, etc.) to expand detailed views.
+6.  **Run Experiments**:
+    *   Use the **Investigation Panel** at the bottom to switch between Ablation and Attribution tabs.
+    *   In **Ablation**, select layers and heads to disable, then click "Run Ablation Experiment".
+    *   In **Attribution**, select a target token and method to visualize feature importance.
 
-## Workflow
+## Project Structure
 
-1. **Capture Activations:**
-   ```bash
-   python agnostic_capture.py --model "Qwen/Qwen2.5-0.5B" --prompt "Your text here"
-   ```
-   - Select attention patterns (e.g., `model.{layer}.self_attn`)
-   - This creates `agnostic_activations.json`
+*   `app.py`: Main application entry point and layout orchestration.
+*   `components/`: Modular UI components.
+    *   `pipeline.py`: The core 5-stage visualization.
+    *   `investigation_panel.py`: Ablation and attribution interfaces.
+    *   `ablation_panel.py`: Specific logic for head ablation UI.
+    *   `tokenization_panel.py`: Token visualization.
+*   `utils/`: Backend logic and helper functions.
+    *   `model_patterns.py`: Activation capture and hooking logic.
+    *   `model_config.py`: Model family definitions and auto-detection.
+    *   `head_detection.py`: Attention head categorization logic.
+    *   `beam_search.py`: Beam search implementation.
+*   `tests/`: Comprehensive test suite ensuring stability.
+*   `conductor/`: Detailed project documentation and product guidelines.
 
-2. **Generate Visualizations:**
-   ```bash
-   python bertviz_head_model_view.py
-   ```
-   - Reads from `agnostic_activations.json`
-   - Creates HTML visualization files in `bertviz/` directory
+## Documentation
 
-3. **View Results:**
-   - Open the generated HTML files in your browser
-   - Explore attention patterns across heads and layers
+For more detailed information on the project's background and technical details, check the `conductor/` directory:
+*   [Product Definition](conductor/product.md)
+*   [Tech Stack](conductor/tech-stack.md)
+*   [Workflow](conductor/workflow.md)
 
-## Requirements
+## Contributing
 
-```bash
-pip install torch transformers bertviz
-```
+Contributions are welcome! Please ensure that any new features include appropriate tests in the `tests/` directory. Run the test suite before submitting:
 
-Optional:
 ```bash
-pip install pyvene  # For enhanced hooking capabilities
+pytest tests/
 ```
-
-## Important Notes
-
-### For Attention Visualization:
-- **Must capture `self_attn` modules** (not `self_attn.o_proj`) for bertviz to work
-- Attention modules return tuples: `(output, attention_weights)`
-- bertviz uses the attention weights (element 1) for visualization
-
-### Module Selection:
-- Patterns use `{layer}` placeholder (e.g., `model.{layer}.self_attn`)
-- Selected patterns apply to ALL layers automatically
-- Use indices, exact names, or unique suffixes for selection
-
-### File Outputs:
-- `agnostic_activations.json` - Captured activation data
-- `bertviz/attention_head_view_{model}.html` - Per-head attention visualization
-- `bertviz/attention_model_view_{model}.html` - Cross-layer attention visualization
-
-## Troubleshooting
-
-**"Attention tensor does not have correct dimensions"**
-- Ensure you captured `self_attn` modules, not output projections
-- Check that attention weights have shape `(batch, heads, seq_len, seq_len)`
-
-**"Module not found"**
-- Verify module patterns match your model architecture
-- Use the interactive selection to see available patterns
-
-**"No data captured"**
-- Check hook registration succeeded
-- Ensure selected modules exist in the model
-- Verify the model actually runs forward pass
-
-## Example Session
-
-```bash
-# 1. Capture attention data
-python agnostic_capture.py --model "Qwen/Qwen2.5-0.5B" --prompt "The cat sat on the mat"
-# Select attention patterns: 0 (for model.{layer}.self_attn)
-# Select MLP patterns: (press enter to skip)
-# Select other patterns: (press enter to skip)
-
-# 2. Generate visualizations
-python bertviz_head_model_view.py
-
-# 3. Open bertviz/attention_head_view_Qwen_Qwen2.5-0.5B.html in browser
-```
-
-This will show you how the model attends to different tokens when processing "The cat sat on the mat".

app.py

@@ -639,26 +639,6 @@ def switch_investigation_tab(abl_clicks, attr_clicks, current_tab):
 # CALLBACKS: Investigation Panel - Ablation (Updated for New UI)
 # ============================================================================
 
-@app.callback(
-    Output('ablation-model-view-container', 'children'),
-    [Input('session-activation-store', 'data')]
-)
-def update_ablation_model_view(activation_data):
-    """Update BertViz model view when new analysis is run."""
-    if not activation_data:
-        return html.Div("Run analysis to see attention visualization.", 
-                       style={'padding': '20px', 'color': '#6c757d', 'textAlign': 'center'})
-    
-    try:
-        html_content = generate_bertviz_model_view_html(activation_data)
-        return html.Iframe(
-            srcDoc=html_content,
-            style={'width': '100%', 'height': '100%', 'border': 'none'}
-        )
-    except Exception as e:
-        return html.Div(f"Error generating visualization: {str(e)}", style={'color': 'red', 'padding': '20px'})
-
-
 @app.callback(
     [Output('ablation-layer-select', 'options'),
      Output('ablation-head-select', 'options')],
@@ -768,11 +748,13 @@ def manage_ablation_heads(add_clicks, clear_clicks, remove_clicks,
      State('session-activation-store', 'data'),
      State('model-dropdown', 'value'),
      State('prompt-input', 'value'),
-     State('session-selected-beam-store', 'data')],
+     State('session-selected-beam-store', 'data'),
+     State('max-new-tokens-slider', 'value'),
+     State('beam-width-slider', 'value')],
     prevent_initial_call=True
 )
-def run_ablation_experiment(n_clicks, selected_heads, activation_data, model_name, prompt, selected_beam):
-    """Run ablation on ORIGINAL PROMPT and compare results."""
+def run_ablation_experiment(n_clicks, selected_heads, activation_data, model_name, prompt, selected_beam, max_new_tokens, beam_width):
+    """Run ablation on ORIGINAL PROMPT and compare results, including beam generation."""
     if not n_clicks or not selected_heads or not activation_data:
         return no_update, no_update
     
@@ -809,7 +791,7 @@ def run_ablation_experiment(n_clicks, selected_heads, activation_data, model_nam
         if not heads_by_layer:
             return html.Div("No valid heads selected.", style={'color': '#dc3545'}), no_update
         
-        # Run ablation
+        # Run ablation for analysis (single pass)
         ablated_data = execute_forward_pass_with_multi_layer_head_ablation(
             model, tokenizer, sequence_text, config, heads_by_layer
         )
@@ -821,9 +803,25 @@ def run_ablation_experiment(n_clicks, selected_heads, activation_data, model_nam
         ablated_token = ablated_output.get('token', '')
         ablated_prob = ablated_output.get('probability', 0)
         
+        # Run ablation for generation
+        ablated_beam = None
+        try:
+            # Always perform beam search during ablation to show comparison
+            beam_results = perform_beam_search(
+                model, tokenizer, sequence_text, 
+                beam_width=beam_width, 
+                max_new_tokens=max_new_tokens,
+                ablation_config=heads_by_layer
+            )
+            if beam_results:
+                # Select the top beam
+                ablated_beam = {'text': beam_results[0]['text'], 'score': beam_results[0].get('score', 0)}
+        except Exception as e:
+            print(f"Error during ablated generation: {e}")
+        
         results_display = create_ablation_results_display(
             original_token, ablated_token, original_prob, ablated_prob,
-            selected_heads, selected_beam
+            selected_heads, selected_beam, ablated_beam
         )
         
         return results_display, ablated_data
@@ -832,11 +830,6 @@ def run_ablation_experiment(n_clicks, selected_heads, activation_data, model_nam
         import traceback
         traceback.print_exc()
         return html.Div(f"Ablation error: {str(e)}", style={'color': '#dc3545'}), no_update
-        
-    except Exception as e:
-        import traceback
-        traceback.print_exc()
-        return html.Div(f"Ablation error: {str(e)}", style={'color': '#dc3545'})
 
 
 # ============================================================================

components/ablation_panel.py

@@ -22,17 +22,6 @@ def create_ablation_panel():
             ], style={'color': '#6c757d', 'fontSize': '14px', 'marginBottom': '16px'})
         ]),
         
-        # BertViz Model View Visualization
-        html.Div([
-            html.H6("Attention Visualization (BertViz Model View)", style={'color': '#495057', 'marginBottom': '8px'}),
-            html.P("Use this view to identify interesting heads to ablate.", style={'color': '#6c757d', 'fontSize': '12px'}),
-            dcc.Loading(
-                id="loading-ablation-viz",
-                type="default",
-                children=html.Div(id='ablation-model-view-container', style={'height': '500px', 'width': '100%', 'overflow': 'auto', 'border': '1px solid #e2e8f0', 'borderRadius': '8px'})
-            )
-        ], style={'marginBottom': '20px'}),
-        
         # Head Selector Interface
         html.Div([
             html.Label("Add Head to Ablation List:", className="input-label", style={'marginBottom': '8px', 'display': 'block'}),
@@ -170,21 +159,10 @@ def create_selected_heads_display(selected_heads):
 
 
 def create_ablation_results_display(original_token, ablated_token, original_prob, ablated_prob,
-                                    selected_heads, selected_beam=None):
+                                    selected_heads, selected_beam=None, ablated_beam=None):
     """
-    Create the ablation results display.
-    
-    Args:
-        original_token: Original predicted token
-        ablated_token: Predicted token after ablation
-        original_prob: Original prediction probability
-        ablated_prob: Ablated prediction probability
-        selected_heads: List of {layer: N, head: M} dicts
-        selected_beam: Optional data for original beam comparison
+    Create the ablation results display focusing on full generation comparison.
     """
-    output_changed = original_token != ablated_token
-    prob_delta = ablated_prob - original_prob
-    
     # Format selected heads for display
     all_heads_formatted = [f"L{item['layer']}-H{item['head']}" for item in selected_heads if isinstance(item, dict)]
 
@@ -200,87 +178,76 @@ def create_ablation_results_display(original_token, ablated_token, original_prob
         ], style={'marginBottom': '16px'})
     ]))
     
-    # Before/After comparison
-    results.append(html.Div([
-        # Original
-        html.Div([
-            html.Div("Original", style={'fontSize': '12px', 'color': '#6c757d', 'marginBottom': '6px'}),
-            html.Div(original_token, style={
-                'padding': '10px 16px',
-                'backgroundColor': '#e8f5e9',
-                'border': '2px solid #4caf50',
-                'borderRadius': '6px',
-                'fontFamily': 'monospace',
-                'fontWeight': '600',
-                'textAlign': 'center'
-            }),
-            html.Div(f"{original_prob:.1%}", style={
-                'fontSize': '12px',
-                'color': '#6c757d',
-                'marginTop': '4px',
-                'textAlign': 'center'
-            })
-        ], style={'flex': '1'}),
-        
-        # Arrow
-        html.Div([
-            html.I(className='fas fa-arrow-right', style={
-                'fontSize': '24px',
-                'color': '#dc3545' if output_changed else '#6c757d'
-            })
-        ], style={'display': 'flex', 'alignItems': 'center', 'padding': '0 20px'}),
-        
-        # Ablated
-        html.Div([
-            html.Div("After Ablation", style={'fontSize': '12px', 'color': '#6c757d', 'marginBottom': '6px'}),
-            html.Div(ablated_token, style={
-                'padding': '10px 16px',
-                'backgroundColor': '#ffebee' if output_changed else '#e8f5e9',
-                'border': f'2px solid {"#dc3545" if output_changed else "#4caf50"}',
-                'borderRadius': '6px',
-                'fontFamily': 'monospace',
-                'fontWeight': '600',
-                'textAlign': 'center'
-            }),
-            html.Div(f"{ablated_prob:.1%} ({prob_delta:+.1%})", style={
-                'fontSize': '12px',
-                'color': '#dc3545' if prob_delta < 0 else '#4caf50' if prob_delta > 0 else '#6c757d',
-                'marginTop': '4px',
-                'textAlign': 'center'
-            })
-        ], style={'flex': '1'})
+    # Generation Comparison (Main Display)
+    if selected_beam and selected_beam.get('text') and ablated_beam and ablated_beam.get('text'):
+        gen_changed = selected_beam['text'] != ablated_beam['text']
         
-    ], style={
-        'display': 'flex',
-        'alignItems': 'stretch',
-        'padding': '16px',
-        'backgroundColor': 'white',
-        'borderRadius': '8px',
-        'border': '1px solid #e2e8f0'
-    }))
-    
-    # Interpretation
-    results.append(html.Div([
-        html.I(className='fas fa-lightbulb', style={'color': '#ffc107', 'marginRight': '8px'}),
-        html.Span(
-            "The prediction changed! These heads are important for this input."
-            if output_changed else
-            "The prediction stayed the same. These heads may not be critical for this specific input.",
-            style={'color': '#6c757d', 'fontSize': '13px'}
-        )
-    ], style={'marginTop': '16px', 'padding': '12px', 'backgroundColor': '#fff8e1', 'borderRadius': '6px'}))
-    
-    # Selected beam comparison context
-    if selected_beam and selected_beam.get('text'):
         results.append(html.Div([
-            html.Hr(style={'margin': '15px 0', 'borderColor': '#dee2e6'}),
+            html.H6("Full Generation Comparison", style={'color': '#495057', 'marginBottom': '15px'}),
+            
+            # Comparison Grid
+            html.Div([
+                # Original Generation
+                html.Div([
+                    html.Div("Original Generation", style={'fontSize': '12px', 'fontWeight': 'bold', 'color': '#28a745', 'marginBottom': '8px'}),
+                    html.Div(selected_beam['text'], style={
+                        'fontFamily': 'monospace',
+                        'fontSize': '13px',
+                        'padding': '12px',
+                        'backgroundColor': '#f8f9fa',
+                        'border': '1px solid #dee2e6', 
+                        'borderRadius': '6px',
+                        'whiteSpace': 'pre-wrap',
+                        'height': '100%'
+                    })
+                ], style={'flex': '1', 'marginRight': '10px'}),
+                
+                # Ablated Generation
+                html.Div([
+                    html.Div("Ablated Generation", style={'fontSize': '12px', 'fontWeight': 'bold', 'color': '#dc3545', 'marginBottom': '8px'}),
+                    html.Div(ablated_beam['text'], style={
+                        'fontFamily': 'monospace',
+                        'fontSize': '13px',
+                        'padding': '12px',
+                        'backgroundColor': '#fff5f5' if gen_changed else '#f8f9fa',
+                        'border': '1px solid #dee2e6',
+                        'borderRadius': '6px',
+                        'whiteSpace': 'pre-wrap',
+                        'height': '100%'
+                    })
+                ], style={'flex': '1', 'marginLeft': '10px'})
+            ], style={'display': 'flex', 'alignItems': 'stretch'}),
+            
+            # Generation Change Indicator
             html.Div([
-                html.I(className='fas fa-info-circle', style={'color': '#6c757d', 'marginRight': '8px'}),
-                html.Span("Original generation for reference: ", style={'fontWeight': '500', 'color': '#495057'}),
-                html.Span(selected_beam['text'], style={'fontFamily': 'monospace', 'backgroundColor': '#f8f9fa', 
-                                                       'padding': '4px 8px', 'borderRadius': '4px'})
-            ], style={'fontSize': '13px'})
+                html.I(className=f"fas {'fa-exclamation-circle' if gen_changed else 'fa-check-circle'}", 
+                       style={'color': '#dc3545' if gen_changed else '#28a745', 'marginRight': '8px'}),
+                html.Span(
+                    "The generated sequence changed significantly after ablation."
+                    if gen_changed else
+                    "The generated sequence remained identical.",
+                    style={'fontWeight': '500', 'color': '#495057'}
+                )
+            ], style={'marginTop': '15px', 'fontSize': '13px'}),
+
+            # Probability info as secondary context
+            html.Div([
+                html.Hr(style={'margin': '15px 0', 'borderTop': '1px dotted #dee2e6'}),
+                html.Span("Immediate next-token probability: ", style={'color': '#6c757d', 'fontSize': '12px'}),
+                html.Span(f"{original_prob:.1%} → {ablated_prob:.1%} ", 
+                         style={'fontSize': '12px', 'fontWeight': 'bold', 'color': '#495057'}),
+                html.Span(f"({ablated_prob - original_prob:+.1%})", 
+                         style={'fontSize': '12px', 'color': '#dc3545' if ablated_prob < original_prob else '#28a745'})
+            ], style={'marginTop': '10px'})
+            
         ]))
+    else:
+        # Fallback if beam data is missing
+        results.append(html.Div([
+            html.I(className='fas fa-info-circle', style={'color': '#6c757d', 'marginRight': '8px'}),
+            html.Span("Run a full analysis first to select a generation for comparison.", 
+                     style={'color': '#6c757d', 'fontSize': '14px'})
+        ], style={'padding': '20px', 'backgroundColor': '#f8f9fa', 'borderRadius': '8px', 'border': '1px solid #dee2e6'}))
         
     return html.Div(results, style={
         'padding': '20px',

components/pipeline.py

@@ -792,6 +792,21 @@ def create_output_content(top_tokens=None, predicted_token=None, predicted_prob=
             ], style={'backgroundColor': 'white', 'borderRadius': '8px', 'border': '1px solid #e2e8f0'})
         )
     
+    # Disclaimer about token selection drivers
+    content_items.append(
+        html.Div([
+            html.I(className='fas fa-info-circle', style={'color': '#6c757d', 'marginRight': '8px'}),
+            html.Span([
+                html.Strong("Note on Token Selection: "),
+                "While the probabilities above show the model's raw preference for the immediate next token, the final choice ",
+                "can be influenced by other factors. Techniques like ", html.Strong("Beam Search"), 
+                " look ahead at multiple possible sequences to find the best overall result, rather than just the single most likely token at each step. ",
+                "Additionally, architectures like ", html.Strong("Mixture of Experts (MoE)"),
+                " might route processing through different specialized internal networks which can impact the final output distribution."
+            ], style={'color': '#6c757d', 'fontSize': '13px'})
+        ], style={'marginTop': '16px', 'padding': '12px', 'backgroundColor': '#f8f9fa', 'borderRadius': '6px', 'border': '1px solid #dee2e6'})
+    )
+    
     return html.Div(content_items)
 
 

conductor/tracks/ablation_20260129/index.md

@@ -1,5 +0,0 @@
-# Track ablation_20260129 Context
-
-- [Specification](./spec.md)
-- [Implementation Plan](./plan.md)
-- [Metadata](./metadata.json)

conductor/tracks/ablation_20260129/metadata.json

@@ -1,8 +0,0 @@
-{
-  "track_id": "ablation_20260129",
-  "type": "feature",
-  "status": "new",
-  "created_at": "2026-01-29T12:40:00Z",
-  "updated_at": "2026-01-29T12:40:00Z",
-  "description": "Implement interactive ablation studies for attention heads in the Dash dashboard."
-}

conductor/tracks/ablation_20260129/plan.md

@@ -1,31 +0,0 @@
-# Implementation Plan - Interactive Attention Head Ablation
-
-## Phase 1: Backend Support for Ablation [checkpoint: 5fc7374]
-- [x] Task: Create a reproduction script to test manual PyVene interventions for head ablation. [43cf4ff]
-    - [x] Sub-task: Write a standalone script that loads a small model (e.g., GPT-2) and uses PyVene to zero out a specific head (e.g., L0H0).
-    - [x] Sub-task: Verify that the output logits change compared to the baseline run.
-- [x] Task: Extend `utils/model_patterns.py` (or creating `utils/ablation.py`) to support dynamic head masking. [890f413]
-    - [x] Sub-task: Write tests for the new ablation utility function.
-    - [x] Sub-task: Implement a function `apply_ablation_mask(model, heads_to_ablate)` that registers the necessary PyVene hooks.
-- [x] Task: Update the main inference pipeline to accept an ablation configuration. [d2ea949]
-    - [ ] Sub-task: Modify the capture logic to check for an "ablation list" in the request.
-    - [ ] Sub-task: Ensure the pipeline correctly applies the mask before running the forward pass.
-- [ ] Task: Conductor - User Manual Verification 'Backend Support for Ablation' (Protocol in workflow.md)
-
-## Phase 2: Frontend Control Panel [checkpoint: 24bd049]
-- [x] Task: Create an `AblationPanel` component in `components/`. [e2e8f0]
-    - [x] Sub-task: Integrate BertViz Model View to visualize attention heads (replacing custom heatmap/grid).
-    - [x] Sub-task: Implement the callback to handle clicks on the grid and update a `dcc.Store` with the list of disabled heads.
-- [x] Task: Integrate the `AblationPanel` into `app.py`. [e2e8f0]
-    - [ ] Sub-task: Add the panel to the main layout (likely in a new "Experiments" tab or collapsible sidebar).
-    - [ ] Sub-task: Connect the global "Run" or "Update" callback to include the ablation state from the store.
-- [ ] Task: Conductor - User Manual Verification 'Frontend Control Panel' (Protocol in workflow.md)
-
-## Phase 3: Visualization & Feedback Loop
-- [~] Task: Connect the Frontend Ablation State to the Backend Inference.
-    - [ ] Sub-task: Update the main `app.py` callback to pass the `disabled_heads` list to the backend capture function.
-    - [ ] Sub-task: Verify that toggling a head in the UI updates the Logit Lens/Output display.
-- [ ] Task: Visual Polish for Ablated State.
-    - [ ] Sub-task: Ensure the Attention Map visualization shows disabled heads as blank or "inactive".
-    - [ ] Sub-task: Add a "Reset Ablations" button to quickly restore the original model state.
-- [ ] Task: Conductor - User Manual Verification 'Visualization & Feedback Loop' (Protocol in workflow.md)

conductor/tracks/ablation_20260129/spec.md

@@ -1,36 +0,0 @@
-# Track Specification: Interactive Attention Head Ablation
-
-## Overview
-This track introduces interactive ablation capabilities to the Dash dashboard. Users will be able to selectively disable (zero out) specific attention heads in the transformer model and observe the resulting changes in model output (logits/probabilities) and attention patterns. This directly supports the "Interactive Experimentation" core value proposition.
-
-## Goals
-- Enable users to toggle specific attention heads on/off via the UI.
-- Update the model's forward pass to respect these ablation masks.
-- Visualize the "ablated" state compared to the "original" state (if feasible) or simply show the new state.
-- Provide immediate feedback on how head removal affects token prediction.
-
-## User Stories
-- As a student, I want to turn off a specific attention head to see if it is responsible for a particular grammatical dependency (e.g., matching plural subjects to verbs).
-- As a researcher, I want to ablate a group of heads to test a hypothesis about distributed representations.
-- As a user, I want clear visual indicators of which heads are currently active or disabled.
-
-## Requirements
-
-### Frontend (Dash)
-- **Ablation Control Panel:** A UI component (e.g., a grid of toggles or a heatmap with clickable cells) representing all attention heads in the model (Layers x Heads).
-- **State Management:** Store the set of "disabled heads" in the Dash app state (`dcc.Store`).
-- **Visual Feedback:** 
-    - Disabled heads should be visually distinct (e.g., grayed out) in the visualization.
-    - The output (Logit Lens or Top-K tokens) must update dynamically when heads are toggled.
-
-### Backend (Model Logic)
-- **Intervention Mechanism:** Modify the `model_patterns.py` or `agnostic_capture.py` logic to accept an "ablation mask".
-- **PyVene Integration:** Use PyVene's intervention capabilities to zero out the activations of specific heads during the forward pass.
-    - *Technical Note:* This might require defining a specific intervention function that takes the head output and multiplies it by 0 if the index matches the ablated head.
-
-### Visualization
-- Update the attention map visualization to reflect that the ablated head is contributing nothing (blank map or "Disabled" overlay).
-
-## Non-Functional Requirements
-- **Latency:** The update loop (Toggle -> Inference -> Update UI) should be fast enough for interactive exploration (< 2-3 seconds for small/medium models).
-- **Clarity:** It must be obvious to the user that they have modified the model. A "Reset All" button is essential.

utils/model_patterns.py

@@ -979,6 +979,18 @@ def evaluate_sequence_ablation(model, tokenizer, sequence_text: str, config: Dic
     }
 
 
+def _prepare_hidden_state(layer_output: Any) -> torch.Tensor:
+    """Helper to convert layer output to tensor, handling tuple outputs."""
+    # Handle PyVene captured tuple outputs where 2nd element is None (e.g. use_cache=False)
+    if isinstance(layer_output, (list, tuple)) and len(layer_output) > 1 and layer_output[1] is None:
+        layer_output = layer_output[0]
+        
+    hidden = torch.tensor(layer_output) if not isinstance(layer_output, torch.Tensor) else layer_output
+    if hidden.dim() == 4:
+        hidden = hidden.squeeze(0)
+    return hidden
+
+
 def logit_lens_transformation(layer_output: Any, norm_data: List[Any], model, tokenizer, norm_parameter: Optional[str] = None, top_k: int = 5) -> List[Tuple[str, float]]:
     """
     Transform layer output to top K token probabilities using logit lens.
@@ -1003,9 +1015,7 @@ def logit_lens_transformation(layer_output: Any, norm_data: List[Any], model, to
     """
     with torch.no_grad():
         # Convert to tensor and ensure proper shape [batch, seq_len, hidden_dim]
-        hidden = torch.tensor(layer_output) if not isinstance(layer_output, torch.Tensor) else layer_output
-        if hidden.dim() == 4:
-            hidden = hidden.squeeze(0)
+        hidden = _prepare_hidden_state(layer_output)
         
         # Step 1: Apply final layer normalization (critical for intermediate layers)
         final_norm = get_norm_layer_from_parameter(model, norm_parameter)
@@ -1096,9 +1106,7 @@ def _get_token_probabilities_for_layer(activation_data: Dict[str, Any], module_n
         lm_head = model.get_output_embeddings()
         
         with torch.no_grad():
-            hidden = torch.tensor(layer_output) if not isinstance(layer_output, torch.Tensor) else layer_output
-            if hidden.dim() == 4:
-                hidden = hidden.squeeze(0)
+            hidden = _prepare_hidden_state(layer_output)
             
             if final_norm is not None:
                 hidden = final_norm(hidden)