feat(ablation): Unify ablation execution into execute_forward_pass

tests/test_unified_ablation.py

@@ -0,0 +1,56 @@
+
+import sys
+import os
+import torch
+import pytest
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Add project root to path
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+from utils.model_patterns import execute_forward_pass
+
+def test_unified_ablation():
+    """
+    Verify that execute_forward_pass can handle ablation configuration.
+    """
+    model_name = "gpt2"
+    prompt = "The quick brown fox jumps over the"
+    
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    model.eval()
+
+    config = {
+        "attention_modules": ["transformer.h.0.attn"],
+        "block_modules": ["transformer.h.0"],
+        "norm_parameters": [],
+        "logit_lens_parameter": "transformer.ln_f.weight" 
+    }
+
+    # 1. Baseline
+    baseline = execute_forward_pass(model, tokenizer, prompt, config)
+    baseline_prob = baseline['actual_output']['probability']
+    
+    # 2. Ablated via execute_forward_pass (New API we want to support)
+    heads_to_ablate = {0: [0]} # Layer 0, Head 0
+    
+    # We expect this to fail currently as the argument doesn't exist
+    try:
+        ablated = execute_forward_pass(
+            model, tokenizer, prompt, config, 
+            ablation_config=heads_to_ablate
+        )
+        ablated_prob = ablated['actual_output']['probability']
+        
+        print(f"Baseline: {baseline_prob}, Ablated: {ablated_prob}")
+        
+        # Assert change
+        assert abs(baseline_prob - ablated_prob) > 1e-6
+        assert ablated.get('ablated_heads_by_layer') == heads_to_ablate
+        
+    except TypeError:
+        pytest.fail("execute_forward_pass does not accept ablation_config argument yet")
+
+if __name__ == "__main__":
+    test_unified_ablation()

utils/model_patterns.py

@@ -148,7 +148,7 @@ def get_actual_model_output(model_output, tokenizer) -> Tuple[str, float]:
         return token_str, top_prob.item()
 
 
-def execute_forward_pass(model, tokenizer, prompt: str, config: Dict[str, Any]) -> Dict[str, Any]:
+def execute_forward_pass(model, tokenizer, prompt: str, config: Dict[str, Any], ablation_config: Optional[Dict[int, List[int]]] = None) -> Dict[str, Any]:
     """
     Execute forward pass with PyVene IntervenableModel to capture activations from specified modules.
     
@@ -157,10 +157,14 @@ def execute_forward_pass(model, tokenizer, prompt: str, config: Dict[str, Any])
         tokenizer: Loaded tokenizer
         prompt: Input text prompt
         config: Dict with module lists like {"attention_modules": [...], "block_modules": [...], ...}
+        ablation_config: Optional dict mapping layer numbers to list of head indices to ablate.
     
     Returns:
         JSON-serializable dict with captured activations and metadata
     """
+    if ablation_config:
+        return execute_forward_pass_with_multi_layer_head_ablation(model, tokenizer, prompt, config, ablation_config)
+
     print(f"Executing forward pass with prompt: '{prompt}'")
     
     # Extract module lists from config