Ablation updated for full sequence, needs refactor for front-end and workflow

app.py

@@ -11,7 +11,8 @@ import json
 import torch
 from utils import (load_model_and_get_patterns, execute_forward_pass, extract_layer_data,
                    categorize_single_layer_heads, format_categorization_summary,
-                   compute_layer_wise_summaries, perform_beam_search, compute_sequence_trajectory)
+                   compute_layer_wise_summaries, perform_beam_search, compute_sequence_trajectory,
+                   execute_forward_pass_with_head_ablation, evaluate_sequence_ablation, score_sequence)
 from utils.model_config import get_auto_selections, get_model_family
 
 # Import modular components
@@ -1729,10 +1730,11 @@ def handle_head_selection(n_clicks_list, selected_heads):
     [State({'type': 'selected-heads-store', 'layer': ALL}, 'data'),
      State('session-activation-store', 'data'),
      State('model-dropdown', 'value'),
-     State('prompt-input', 'value')],
+     State('prompt-input', 'value'),
+     State('generation-results-store', 'data')],
     prevent_initial_call=True
 )
-def run_head_ablation(n_clicks_list, selected_heads_list, activation_data, model_name, prompt):
+def run_head_ablation(n_clicks_list, selected_heads_list, activation_data, model_name, prompt_input, generation_results):
     """Run forward pass with selected heads ablated."""
     # Identify which button was clicked
     ctx = dash.callback_context
@@ -1745,7 +1747,6 @@ def run_head_ablation(n_clicks_list, selected_heads_list, activation_data, model
         return no_update, no_update, no_update
     
     # Find the index in the states_list that corresponds to this layer
-    # ctx.states_list contains the State values in order
     button_index = None
     if hasattr(ctx, 'states_list') and ctx.states_list:
         # states_list[0] corresponds to selected-heads-store
@@ -1756,7 +1757,6 @@ def run_head_ablation(n_clicks_list, selected_heads_list, activation_data, model
     
     # Fallback: if states_list doesn't work, try matching by iterating
     if button_index is None:
-        # This shouldn't happen, but as a fallback, just return error
         return no_update, no_update, html.Div([
             html.I(className="fas fa-exclamation-circle", style={'marginRight': '8px', 'color': '#dc3545'}),
             f"Could not determine button index for layer {layer_num}"
@@ -1772,12 +1772,14 @@ def run_head_ablation(n_clicks_list, selected_heads_list, activation_data, model
     
     try:
         from transformers import AutoModelForCausalLM, AutoTokenizer
-        from utils import execute_forward_pass_with_head_ablation
         
         # Save original activation data before ablation
         import copy
         original_data = copy.deepcopy(activation_data)
         
+        # Determine the sequence to analyze (prefer activation data prompt over input box)
+        sequence_text = activation_data.get('prompt', prompt_input)
+        
         # Load model and tokenizer
         model = AutoModelForCausalLM.from_pretrained(model_name, attn_implementation='eager')
         tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -1789,18 +1791,86 @@ def run_head_ablation(n_clicks_list, selected_heads_list, activation_data, model
             'norm_parameters': activation_data.get('norm_parameters', [])
         }
         
-        # Run ablation
+        # 1. Run Standard Ablation (Forward Pass)
         ablated_data = execute_forward_pass_with_head_ablation(
-            model, tokenizer, prompt, config, layer_num, selected_heads
+            model, tokenizer, sequence_text, config, layer_num, selected_heads
+        )
+        
+        # 2. Compute Full Sequence Metrics (KL Divergence, Delta Probs)
+        # This requires re-running passes (Original & Ablated) on the full sequence
+        # We use a helper that handles the ablation hooking internally for the metric pass
+        seq_metrics = evaluate_sequence_ablation(
+            model, tokenizer, sequence_text, config, 
+            ablation_type='head', ablation_target=(layer_num, selected_heads)
         )
+        ablated_data['sequence_metrics'] = seq_metrics
+        
+        # 3. Re-score Top Generated Sequences (if available)
+        if generation_results:
+            top_sequences_comparison = []
+            
+            # Helper to run ablation for scoring (we need to apply hook again)
+            # Since we can't easily pass 'ablated model' around, we re-apply hooks
+            # Simplification: We already have 'evaluate_sequence_ablation'.
+            # But that compares Ref vs Abl.
+            # Here we just want Ablated Score.
+            # Actually, `score_sequence` runs valid forward pass. 
+            # We need to apply ablation hooks to validly score.
+            
+            # Define localized hook manager for scoring
+            def get_ablated_score(seq_text):
+                # Apply hook
+                hooks = []
+                def head_ablation_hook(module, input, output):
+                   # Similar to evaluate_sequence_ablation hook
+                    if isinstance(output, tuple): h = output[0]
+                    else: h = output
+                    if not isinstance(h, torch.Tensor): h = torch.tensor(h)
+                    
+                    num_heads = model.config.num_attention_heads
+                    head_dim = h.shape[-1] // num_heads
+                    new_shape = h.shape[:-1] + (num_heads, head_dim)
+                    reshaped = h.view(new_shape).clone()
+                    for h_idx in selected_heads: reshaped[..., h_idx, :] = 0
+                    ablated = reshaped.view(h.shape)
+                    return (ablated,) + output[1:] if isinstance(output, tuple) else ablated
+
+                # Register
+                target_module = None
+                for name, mod in model.named_modules():
+                    if f"layers.{layer_num}.self_attn" in name or f"h.{layer_num}.attn" in name:
+                        if "k_proj" not in name: 
+                            target_module = mod; break
+                
+                if target_module:
+                    hooks.append(target_module.register_forward_hook(head_ablation_hook))
+                
+                try:
+                    score = score_sequence(model, tokenizer, seq_text)
+                finally:
+                    for hook in hooks: hook.remove()
+                return score
+
+            for res in generation_results:
+                txt = res['text']
+                orig_score = res['score']
+                new_score = get_ablated_score(txt)
+                top_sequences_comparison.append({
+                    'text': txt,
+                    'original_score': orig_score,
+                    'ablated_score': new_score,
+                    'delta': new_score - orig_score
+                })
+            
+            ablated_data['top_sequences_comparison'] = top_sequences_comparison
+
         
         # Update activation data with ablated results
-        # Mark as ablated for visual indication
         ablated_data['ablated'] = True
         ablated_data['ablated_layer'] = layer_num
         ablated_data['ablated_heads'] = selected_heads
         
-        # Preserve input_ids from original data if not present (prompt is unchanged)
+        # Preserve input_ids if needed
         if 'input_ids' not in ablated_data and 'input_ids' in activation_data:
             ablated_data['input_ids'] = activation_data['input_ids']
         
@@ -1808,7 +1878,7 @@ def run_head_ablation(n_clicks_list, selected_heads_list, activation_data, model
         heads_str = ', '.join([f"H{h}" for h in sorted(selected_heads)])
         success_message = html.Div([
             html.I(className="fas fa-check-circle", style={'marginRight': '8px', 'color': '#28a745'}),
-            f"Ablation complete: Layer {layer_num}, Heads {heads_str} removed"
+            f"Ablation complete: Layer {layer_num}, Heads {heads_str} removed. Scroll down for sequence analysis."
         ], className="status-success")
         
         return ablated_data, original_data, success_message
@@ -1870,5 +1940,133 @@ def reset_ablation(n_clicks, original_data):
     return original_data, {}, success_message
 
 
+
+# Callback to update sequence ablation analysis view
+@app.callback(
+    [Output('sequence-ablation-results-container', 'children'),
+     Output('sequence-ablation-results-container', 'style')],
+    Input('session-activation-store', 'data'),
+    prevent_initial_call=False
+)
+def update_sequence_ablation_view(activation_data):
+    """Update the sequence ablation results view (KL Divergence, Sequence Comparison)."""
+    if not activation_data or not activation_data.get('ablated', False):
+        return [], {'display': 'none'}
+    
+    try:
+        import plotly.graph_objs as go
+        from dash import html, dcc
+        
+        children = []
+        
+        # 1. Header
+        children.append(html.H3("Full Sequence Ablation Analysis", style={'marginTop': '0', 'marginBottom': '20px', 'color': '#2d3748'}))
+        
+        # 2. Top-5 Sequence Comparison Table
+        top_seqs = activation_data.get('top_sequences_comparison', [])
+        if top_seqs:
+            rows = []
+            for i, seq in enumerate(top_seqs):
+                delta = seq['delta']
+                delta_color = '#28a745' if delta > 0 else '#dc3545' if delta < 0 else '#6c757d'
+                
+                rows.append(html.Tr([
+                    html.Td(f"#{i+1}", style={'fontWeight': 'bold'}),
+                    html.Td(seq['text'], style={'fontFamily': 'monospace', 'maxWidth': '400px', 'overflow': 'hidden', 'textOverflow': 'ellipsis', 'whiteSpace': 'nowrap'}),
+                    html.Td(f"{seq['original_score']:.4f}"),
+                    html.Td(f"{seq['ablated_score']:.4f}"),
+                    html.Td(f"{delta:+.4f}", style={'color': delta_color, 'fontWeight': 'bold'})
+                ]))
+            
+            table_header = html.Thead(html.Tr([
+                html.Th("Rank"), html.Th("Sequence"), html.Th("Original Score"), html.Th("Ablated Score"), html.Th("Delta")
+            ]))
+            table_body = html.Tbody(rows)
+            
+            children.append(html.Div([
+                html.H5("Top Sequences Impact", style={'marginBottom': '10px'}),
+                html.Table([table_header, table_body], className="table table-striped table-bordered")
+            ], style={'marginBottom': '30px', 'padding': '15px', 'backgroundColor': '#fff', 'borderRadius': '8px', 'boxShadow': '0 2px 4px rgba(0,0,0,0.05)'}))
+        
+        # 3. KL Divergence Chart
+        seq_metrics = activation_data.get('sequence_metrics', {})
+        kl_divs = seq_metrics.get('kl_divergence', [])
+        tokens = seq_metrics.get('tokens', [])
+        
+        if kl_divs:
+            # KL Chart
+            fig_kl = go.Figure()
+            fig_kl.add_trace(go.Scatter(
+                x=list(range(len(kl_divs))),
+                y=kl_divs,
+                mode='lines+markers',
+                name='KL Divergence',
+                line=dict(color='#6610f2', width=2),
+                hovertext=[f"Token: {t}<br>KL: {v:.4f}" for t, v in zip(tokens, kl_divs)],
+                hoverinfo='text'
+            ))
+            
+            fig_kl.update_layout(
+                title="KL Divergence per Position (Distribution Shift)",
+                xaxis_title="Position / Token",
+                yaxis_title="KL Divergence (nats)",
+                margin=dict(l=20, r=20, t=40, b=20),
+                height=300,
+                xaxis=dict(
+                    tickmode='array',
+                    tickvals=list(range(len(tokens))),
+                    ticktext=tokens
+                )
+            )
+            
+            children.append(html.Div([
+                dcc.Graph(figure=fig_kl, config={'displayModeBar': False})
+            ], style={'marginBottom': '20px', 'padding': '15px', 'backgroundColor': '#fff', 'borderRadius': '8px', 'boxShadow': '0 2px 4px rgba(0,0,0,0.05)'}))
+        
+        # 4. Target Probability Deltas Chart
+        prob_deltas = seq_metrics.get('probability_deltas', [])
+        if prob_deltas:
+            # Shift tokens for x-axis (deltas are for prediction of next token)
+            # Input: T0, T1, T2
+            # Delta 0: Change in P(T1|T0)
+            # So x-axis should be T1, T2...
+            target_tokens = tokens[1:] if len(tokens) > 1 else []
+            
+            fig_delta = go.Figure()
+            fig_delta.add_trace(go.Bar(
+                x=list(range(len(prob_deltas))),
+                y=prob_deltas,
+                name='Prob Delta',
+                marker_color=['#28a745' if v >= 0 else '#dc3545' for v in prob_deltas],
+                hovertext=[f"Target: {t}<br>Change: {v:+.4f}" for t, v in zip(target_tokens, prob_deltas)],
+                hoverinfo='text'
+            ))
+            
+            fig_delta.update_layout(
+                title="Target Probability Change per Position",
+                xaxis_title="Target Token",
+                yaxis_title="Probability Delta",
+                margin=dict(l=20, r=20, t=40, b=20),
+                height=300,
+                xaxis=dict(
+                    tickmode='array',
+                    tickvals=list(range(len(target_tokens))),
+                    ticktext=target_tokens
+                )
+            )
+            
+            children.append(html.Div([
+                dcc.Graph(figure=fig_delta, config={'displayModeBar': False})
+            ], style={'marginBottom': '20px', 'padding': '15px', 'backgroundColor': '#fff', 'borderRadius': '8px', 'boxShadow': '0 2px 4px rgba(0,0,0,0.05)'}))
+        
+        return children, {'display': 'block', 'marginTop': '30px', 'paddingTop': '30px', 'borderTop': '1px solid #dee2e6'}
+        
+    except Exception as e:
+        print(f"Error in ablation view: {e}")
+        import traceback
+        traceback.print_exc()
+        return html.Div(f"Error loading visualization: {str(e)}"), {'display': 'block'}
+
+
 if __name__ == '__main__':
     app.run(debug=True, port=8050)

components/main_panel.py

@@ -105,7 +105,10 @@ def create_main_panel():
                             html.I(className="fas fa-spinner fa-spin", style={'fontSize': '24px', 'color': '#667eea', 'marginRight': '10px'}),
                             html.Span("Loading visuals...", style={'fontSize': '16px', 'color': '#495057'})
                         ], style={'display': 'flex', 'alignItems': 'center', 'justifyContent': 'center', 'padding': '2rem'})
-                    )
+                    ),
+                    
+                    # Sequence Ablation Results (New)
+                    html.Div(id='sequence-ablation-results-container', style={'marginTop': '30px', 'display': 'none'})
                 ], className="visualization-section")
             ])
         ], id="analysis-view-container", style={'display': 'none'}) # Hidden by default

utils/__init__.py

@@ -1,14 +1,17 @@
-from .model_patterns import load_model_and_get_patterns, execute_forward_pass, logit_lens_transformation, extract_layer_data, generate_bertviz_html, generate_category_bertviz_html, get_check_token_probabilities, execute_forward_pass_with_layer_ablation, execute_forward_pass_with_head_ablation, merge_token_probabilities, compute_global_top5_tokens, detect_significant_probability_increases, compute_layer_wise_summaries
+from .model_patterns import load_model_and_get_patterns, execute_forward_pass, logit_lens_transformation, extract_layer_data, generate_bertviz_html, generate_category_bertviz_html, get_check_token_probabilities, execute_forward_pass_with_layer_ablation, execute_forward_pass_with_head_ablation, merge_token_probabilities, compute_global_top5_tokens, detect_significant_probability_increases, compute_layer_wise_summaries, evaluate_sequence_ablation
 from .model_config import get_model_family, get_family_config, get_auto_selections, MODEL_TO_FAMILY, MODEL_FAMILIES
 from .head_detection import categorize_all_heads, categorize_single_layer_heads, format_categorization_summary, HeadCategorizationConfig
 from .prompt_comparison import compare_attention_layers, compare_output_probabilities, format_comparison_summary, ComparisonConfig
 from .beam_search import perform_beam_search, compute_sequence_trajectory
+from .ablation_metrics import compute_kl_divergence, score_sequence, get_token_probability_deltas
+
 
 __all__ = [
     'load_model_and_get_patterns', 
     'execute_forward_pass',
     'execute_forward_pass_with_layer_ablation',
     'execute_forward_pass_with_head_ablation',
+    'evaluate_sequence_ablation',
     'logit_lens_transformation',
     'extract_layer_data',
     'generate_bertviz_html',
@@ -32,5 +35,9 @@ __all__ = [
     'format_comparison_summary',
     'ComparisonConfig',
     'perform_beam_search',
-    'compute_sequence_trajectory'
+    'perform_beam_search',
+    'compute_sequence_trajectory',
+    'compute_kl_divergence',
+    'score_sequence',
+    'get_token_probability_deltas'
 ]

utils/ablation_metrics.py

@@ -0,0 +1,105 @@
+
+import torch
+import torch.nn.functional as F
+from typing import List, Dict, Any, Tuple, Optional
+
+def compute_kl_divergence(logits_p: torch.Tensor, logits_q: torch.Tensor) -> List[float]:
+    """
+    Compute KL Divergence KL(P || Q) for each position in the sequence.
+    P is the reference distribution (logits_p), Q is the ablated distribution (logits_q).
+    
+    Args:
+        logits_p: Reference logits [batch, seq_len, vocab_size]
+        logits_q: Ablated logits [batch, seq_len, vocab_size]
+        
+    Returns:
+        List of KL divergence values for each position.
+    """
+    with torch.no_grad():
+        # Ensure batch size 1 or handle appropriately
+        if logits_p.dim() == 3:
+            logits_p = logits_p.squeeze(0)
+        if logits_q.dim() == 3:
+            logits_q = logits_q.squeeze(0)
+            
+        # P = softmax(logits_p)
+        # Q = softmax(logits_q)
+        # KL(P||Q) = sum(P * (log P - log Q))
+        
+        # Use log_softmax for stability
+        log_probs_p = F.log_softmax(logits_p, dim=-1)
+        log_probs_q = F.log_softmax(logits_q, dim=-1)
+        probs_p = torch.exp(log_probs_p)
+        
+        # Element-wise KL
+        kl_divs = torch.sum(probs_p * (log_probs_p - log_probs_q), dim=-1)
+        
+        return kl_divs.tolist()
+
+def score_sequence(model, tokenizer, text: str) -> float:
+    """
+    Compute the total log probability (score) of a text sequence.
+    
+    Args:
+        model: HuggingFace model
+        tokenizer: Tokenizer
+        text: The sequence to score
+        
+    Returns:
+        Total log probability.
+    """
+    inputs = tokenizer(text, return_tensors="pt")
+    input_ids = inputs["input_ids"].to(model.device)
+    
+    with torch.no_grad():
+        outputs = model(input_ids)
+        logits = outputs.logits # [1, seq_len, vocab_size]
+        
+        # We want P(token_i | tokens_<i)
+        # The logits at position i-1 predict position i
+        
+        # Shift logits and labels
+        shift_logits = logits[0, :-1, :].contiguous()
+        shift_labels = input_ids[0, 1:].contiguous()
+        
+        # Helper to pick specific token probabilities
+        # log_softmax
+        log_probs_all = F.log_softmax(shift_logits, dim=-1)
+        
+        # Gather only the target label log probs
+        # gather needs index column vector
+        target_log_probs = log_probs_all.gather(1, shift_labels.unsqueeze(1)).squeeze(1)
+        
+        total_score = target_log_probs.sum().item()
+        
+        return total_score
+
+def get_token_probability_deltas(logits_ref: torch.Tensor, logits_abl: torch.Tensor, input_ids: torch.Tensor) -> List[float]:
+    """
+    Compute the change in probability (Prob_abl - Prob_ref) for the actual target tokens.
+    
+    Args:
+        logits_ref: Reference logits
+        logits_abl: Ablated logits
+        input_ids: The sequence token IDs [1, seq_len]
+        
+    Returns:
+        List of probability deltas for each position (starting from first prediction).
+    """
+    with torch.no_grad():
+        if logits_ref.dim() == 3: logits_ref = logits_ref.squeeze(0)
+        if logits_abl.dim() == 3: logits_abl = logits_abl.squeeze(0)
+        
+        target_ids = input_ids[0, 1:] # Targets are from index 1 onwards
+        
+        # Probabilities
+        probs_ref = F.softmax(logits_ref[:-1], dim=-1) # Predicts 1..N
+        probs_abl = F.softmax(logits_abl[:-1], dim=-1)
+        
+        # Gather target probs
+        ref_target_probs = probs_ref.gather(1, target_ids.unsqueeze(1)).squeeze(1)
+        abl_target_probs = probs_abl.gather(1, target_ids.unsqueeze(1)).squeeze(1)
+        
+        deltas = (abl_target_probs - ref_target_probs).tolist()
+        
+        return deltas

utils/model_patterns.py

@@ -614,6 +614,160 @@ def execute_forward_pass_with_layer_ablation(model, tokenizer, prompt: str, conf
     return result
 
 
+def evaluate_sequence_ablation(model, tokenizer, sequence_text: str, config: Dict[str, Any],
+                             ablation_type: str, ablation_target: Any) -> Dict[str, Any]:
+    """
+    Evaluate the impact of ablation on a full sequence.
+    
+    This runs TWO forward passes on the FULL sequence:
+    1. Reference pass (original model) -> Capture logits/probs
+    2. Ablated pass (modified model) -> Capture logits/probs
+    
+    Then computes metrics: KL Divergence, Target Prob Changes.
+    
+    Args:
+        model: Loaded transformer model
+        tokenizer: Tokenizer
+        sequence_text: The full text sequence to evaluate
+        config: Module configuration (needed for ablation setup)
+        ablation_type: 'head' or 'layer'
+        ablation_target: tuple (layer, head_indices) or int (layer_num)
+        
+    Returns:
+        Dict with evaluation metrics.
+    """
+    from .ablation_metrics import compute_kl_divergence, get_token_probability_deltas
+    
+    print(f"Evaluating sequence ablation: Type={ablation_type}, Target={ablation_target}")
+    
+    inputs = tokenizer(sequence_text, return_tensors="pt")
+    input_ids = inputs["input_ids"].to(model.device)
+    
+    # --- 1. Reference Pass ---
+    with torch.no_grad():
+        outputs_ref = model(input_ids)
+        logits_ref = outputs_ref.logits # [1, seq_len, vocab_size]
+        
+    # --- 2. Ablated Pass ---
+    # Setup ablation based on type
+    
+    # We need to wrap the model using PyVene logic or custom hooks just for this pass
+    # Since we already have logic in execute_forward_pass_with_..._ablation, we can reuse the Hook logic
+    # But we want the full logits, not just captured activations.
+    
+    # Let's manually register hooks here for simplicity and control
+    hooks = []
+    
+    def head_ablation_hook_factory(layer_idx, head_indices):
+        def hook(module, input, output):
+            # output is (hidden_states, ...) or hidden_states
+            if isinstance(output, tuple):
+                hidden_states = output[0]
+            else:
+                hidden_states = output
+                
+            # Assume hidden_states is [batch, seq, hidden]
+            # Reshape, zero out heads, Reshape back
+            if not isinstance(hidden_states, torch.Tensor):
+                 if isinstance(hidden_states, list): hidden_states = torch.tensor(hidden_states)
+            
+            # Move to device if needed? They should be on device.
+            
+            num_heads = model.config.num_attention_heads
+            head_dim = hidden_states.shape[-1] // num_heads
+            
+            # view: [batch, seq, heads, dim]
+            new_shape = hidden_states.shape[:-1] + (num_heads, head_dim)
+            reshaped = hidden_states.view(new_shape)
+            
+            # Create mask or just zero out
+            # We can't modify in place securely with autograd usually, but here no_grad is on.
+            # Clone to be safe
+            reshaped = reshaped.clone()
+            
+            for h_idx in head_indices:
+                reshaped[..., h_idx, :] = 0
+                
+            ablated_hidden = reshaped.view(hidden_states.shape)
+            
+            if isinstance(output, tuple):
+                return (ablated_hidden,) + output[1:]
+            return ablated_hidden
+        return hook
+
+    # Hook for Layer Ablation (Identity/Skip or Zero)
+    # We'll use Identity (Skip Layer) as a simpler approximation of "removing logic" 
+    # OR Mean Ablation if we had the mean. 
+    # For now, let's just do nothing for layer ablation or return error, 
+    # as the user primarily asks for "ablation experiment updates" which often means Heads.
+    # But to be safe, let's implement the same Mean Ablation if possible, or Identity.
+    # Identity (Skip) is easier:
+    def identity_hook(module, input, output):
+        # input is tuple (hidden_states, ...)
+        return input if isinstance(input, tuple) else (input,)
+
+    try:
+        if ablation_type == 'head':
+            layer_num, head_indices = ablation_target
+            # Find module
+            # Standard transformers: model.layers[i].self_attn
+            # We need the exact module name map standard to HuggingFace
+            # Or use the config's mapping if available.
+            # Let's rely on standard naming or search
+            
+            # Simple heuristic: find 'layers.X.self_attn' or 'h.X.attn'
+            target_module = None
+            for name, mod in model.named_modules():
+                # Check for standard patterns
+                # layer_num is int
+                if f"layers.{layer_num}.self_attn" in name or f"h.{layer_num}.attn" in name or f"blocks.{layer_num}.attn" in name:
+                     if "k_proj" not in name and "v_proj" not in name and "q_proj" not in name: # avoid submodules
+                         target_module = mod
+                         break
+            
+            if target_module:
+                hooks.append(target_module.register_forward_hook(head_ablation_hook_factory(layer_num, head_indices)))
+            else:
+                print(f"Warning: Could not find attention module for layer {layer_num}")
+
+        elif ablation_type == 'layer':
+            layer_num = ablation_target
+            target_module = None
+            for name, mod in model.named_modules():
+                # Layers are usually 'model.layers.X' or 'transformer.h.X'
+                # We want the module that corresponds to the layer block
+                # Be careful not to pick 'layers.X.mlp'
+                if (f"layers.{layer_num}" in name or f"h.{layer_num}" in name) and name.count('.') <= 2: # heuristic for top-level layer
+                     target_module = mod
+                     break
+            
+            if target_module:
+                 # Skip layer (Identity)
+                 hooks.append(target_module.register_forward_hook(lambda m, i, o: i[0] if isinstance(i, tuple) else i))
+
+        # Run Ablated Pass
+        with torch.no_grad():
+            outputs_abl = model(input_ids)
+            logits_abl = outputs_abl.logits
+
+    finally:
+        for hook in hooks:
+            hook.remove()
+            
+    # --- 3. Compute Metrics ---
+    # KL Divergence [seq_len]
+    kl_div = compute_kl_divergence(logits_ref, logits_abl)
+    
+    # Prob Deltas for actual tokens [seq_len-1] (shifted)
+    prob_deltas = get_token_probability_deltas(logits_ref, logits_abl, input_ids)
+    
+    return {
+        "kl_divergence": kl_div,
+        "probability_deltas": prob_deltas,
+        "tokens": [tokenizer.decode([tid]) for tid in input_ids[0].tolist()]
+    }
+
+
 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.