backend/test_ablation.py

"""
Unit tests for ablation functionality
Tests that hooks are correctly applied and model components are properly disabled
"""

import torch
import numpy as np
from transformers import AutoModelForCausalLM, AutoTokenizer
import pytest
import logging
from typing import Dict, Set, Any, List
import json

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class AblationTester:
    """Test suite for ablation functionality"""
    
    def __init__(self):
        self.model = None
        self.tokenizer = None
        self.device = torch.device("cpu")
        
    def setup(self):
        """Load model for testing"""
        logger.info("Loading model for ablation tests...")
        self.model = AutoModelForCausalLM.from_pretrained(
            "Salesforce/codegen-350M-mono",
            torch_dtype=torch.float32,
            low_cpu_mem_usage=True
        ).to(self.device)
        self.tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen-350M-mono")
        self.tokenizer.pad_token = self.tokenizer.eos_token
        logger.info("Model loaded successfully")
        
    def test_model_architecture(self):
        """Test 1: Verify model architecture matches expectations"""
        logger.info("\n=== Test 1: Model Architecture ===")
        
        # Check number of layers
        assert self.model.config.n_layer == 20, f"Expected 20 layers, got {self.model.config.n_layer}"
        logger.info(f"✓ Model has {self.model.config.n_layer} layers")
        
        # Check number of attention heads
        assert self.model.config.n_head == 16, f"Expected 16 heads, got {self.model.config.n_head}"
        logger.info(f"✓ Model has {self.model.config.n_head} attention heads per layer")
        
        # Check layer structure
        for i in range(self.model.config.n_layer):
            layer = self.model.transformer.h[i]
            assert hasattr(layer, 'attn'), f"Layer {i} missing attention module"
            assert hasattr(layer, 'mlp'), f"Layer {i} missing MLP/FFN module"
            assert hasattr(layer, 'ln_1'), f"Layer {i} missing layer norm 1"
            assert hasattr(layer, 'ln_2'), f"Layer {i} missing layer norm 2"
        logger.info("✓ All layers have correct structure (attn, mlp, ln_1, ln_2)")
        
        return True
    
    def test_attention_hook_attachment(self):
        """Test 2: Verify attention hooks can be attached and work"""
        logger.info("\n=== Test 2: Attention Hook Attachment ===")
        
        # Create a hook that counts calls
        hook_calls = {'count': 0, 'output_shape': None}
        
        def test_hook(module, input, output):
            hook_calls['count'] += 1
            if isinstance(output, tuple):
                hook_calls['output_shape'] = output[0].shape
            else:
                hook_calls['output_shape'] = output.shape
            return output
        
        # Attach hook to first layer attention
        handle = self.model.transformer.h[0].attn.register_forward_hook(test_hook)
        
        # Run a forward pass
        inputs = self.tokenizer("test", return_tensors="pt").to(self.device)
        with torch.no_grad():
            outputs = self.model(**inputs)
        
        # Verify hook was called
        assert hook_calls['count'] > 0, "Hook was not called"
        logger.info(f"✓ Hook called {hook_calls['count']} times")
        logger.info(f"✓ Attention output shape: {hook_calls['output_shape']}")
        
        # Clean up
        handle.remove()
        return True
    
    def test_attention_zeroing(self):
        """Test 3: Verify attention can be zeroed out"""
        logger.info("\n=== Test 3: Attention Zeroing ===")
        
        # Get baseline output
        inputs = self.tokenizer("def test():", return_tensors="pt").to(self.device)
        with torch.no_grad():
            baseline_output = self.model(**inputs)
            baseline_logits = baseline_output.logits[0, -1, :].cpu().numpy()
        
        # Create hook that zeros attention
        def zero_attention_hook(module, input, output):
            if isinstance(output, tuple):
                return (torch.zeros_like(output[0]),) + output[1:]
            return torch.zeros_like(output)
        
        # Apply hook to all attention layers
        handles = []
        for i in range(self.model.config.n_layer):
            handle = self.model.transformer.h[i].attn.register_forward_hook(zero_attention_hook)
            handles.append(handle)
        
        # Get ablated output
        with torch.no_grad():
            ablated_output = self.model(**inputs)
            ablated_logits = ablated_output.logits[0, -1, :].cpu().numpy()
        
        # Clean up hooks
        for handle in handles:
            handle.remove()
        
        # Verify outputs are different
        difference = np.mean(np.abs(baseline_logits - ablated_logits))
        assert difference > 0.1, f"Outputs too similar (diff={difference}), ablation may not be working"
        logger.info(f"✓ Ablated output differs from baseline (mean diff: {difference:.4f})")
        
        # Check that ablated output has lower confidence (higher entropy)
        baseline_probs = torch.softmax(torch.tensor(baseline_logits), dim=0)
        ablated_probs = torch.softmax(torch.tensor(ablated_logits), dim=0)
        
        baseline_entropy = -torch.sum(baseline_probs * torch.log(baseline_probs + 1e-10))
        ablated_entropy = -torch.sum(ablated_probs * torch.log(ablated_probs + 1e-10))
        
        logger.info(f"  Baseline entropy: {baseline_entropy:.4f}")
        logger.info(f"  Ablated entropy: {ablated_entropy:.4f}")
        
        return True
    
    def test_ffn_ablation(self):
        """Test 4: Verify FFN can be disabled"""
        logger.info("\n=== Test 4: FFN Ablation ===")
        
        # Get baseline
        inputs = self.tokenizer("def test():", return_tensors="pt").to(self.device)
        with torch.no_grad():
            baseline_output = self.model(**inputs)
            baseline_logits = baseline_output.logits[0, -1, :].cpu().numpy()
        
        # Hook to disable FFN
        def zero_ffn_hook(module, input, output):
            return torch.zeros_like(output)
        
        # Apply to all FFN layers
        handles = []
        for i in range(self.model.config.n_layer):
            handle = self.model.transformer.h[i].mlp.register_forward_hook(zero_ffn_hook)
            handles.append(handle)
        
        # Get ablated output
        with torch.no_grad():
            ablated_output = self.model(**inputs)
            ablated_logits = ablated_output.logits[0, -1, :].cpu().numpy()
        
        # Clean up
        for handle in handles:
            handle.remove()
        
        # Verify difference
        difference = np.mean(np.abs(baseline_logits - ablated_logits))
        assert difference > 0.1, f"FFN ablation not working (diff={difference})"
        logger.info(f"✓ FFN ablation changes output (mean diff: {difference:.4f})")
        
        return True
    
    def test_partial_attention_ablation(self):
        """Test 5: Verify partial attention head disabling"""
        logger.info("\n=== Test 5: Partial Attention Ablation ===")
        
        # Get baseline
        inputs = self.tokenizer("def test():", return_tensors="pt").to(self.device)
        with torch.no_grad():
            baseline_output = self.model(**inputs)
            baseline_logits = baseline_output.logits[0, -1, :].cpu().numpy()
        
        # Hook to scale attention (simulating partial disable)
        def scale_attention_hook(module, input, output):
            scale = 0.5  # Disable half the heads (simplified)
            if isinstance(output, tuple):
                return (output[0] * scale,) + output[1:]
            return output * scale
        
        # Apply to layer 0
        handle = self.model.transformer.h[0].attn.register_forward_hook(scale_attention_hook)
        
        # Get partially ablated output
        with torch.no_grad():
            ablated_output = self.model(**inputs)
            ablated_logits = ablated_output.logits[0, -1, :].cpu().numpy()
        
        # Clean up
        handle.remove()
        
        # Verify outputs are different but not as different as full ablation
        difference = np.mean(np.abs(baseline_logits - ablated_logits))
        assert 0.01 < difference < 0.5, f"Partial ablation unexpected difference: {difference}"
        logger.info(f"✓ Partial ablation works (mean diff: {difference:.4f})")
        
        return True
    
    def test_data_format_conversion(self):
        """Test 6: Verify frontend data format is correctly parsed"""
        logger.info("\n=== Test 6: Data Format Conversion ===")
        
        # Simulate frontend data (JSON with string keys)
        frontend_data = {
            "layers": [0, 1, 2],
            "attention_heads": {
                "0": [0, 1, 2, 3],
                "1": [4, 5, 6, 7],
                "2": list(range(16))  # All heads
            },
            "ffn_layers": [3, 4],
            "embeddings": False,
            "layer_norm": []
        }
        
        # Parse as backend would
        disabled_layers = set(frontend_data.get('layers', []))
        disabled_attention_raw = frontend_data.get('attention_heads', {})
        disabled_attention = {int(k) if isinstance(k, str) else k: v 
                             for k, v in disabled_attention_raw.items()}
        disabled_ffn = set(frontend_data.get('ffn_layers', []))
        
        # Verify parsing
        assert disabled_layers == {0, 1, 2}, f"Layers parsed incorrectly: {disabled_layers}"
        assert 0 in disabled_attention, "String key '0' not converted to int 0"
        assert disabled_attention[0] == [0, 1, 2, 3], f"Attention heads parsed incorrectly"
        assert len(disabled_attention[2]) == 16, "Full layer disable not parsed"
        assert disabled_ffn == {3, 4}, f"FFN layers parsed incorrectly: {disabled_ffn}"
        
        logger.info("✓ Frontend data format correctly parsed")
        logger.info(f"  Disabled layers: {disabled_layers}")
        logger.info(f"  Disabled attention heads: {list(disabled_attention.keys())}")
        logger.info(f"  Disabled FFN: {disabled_ffn}")
        
        return True
    
    def test_generation_with_ablation(self):
        """Test 7: Full generation test with various ablations"""
        logger.info("\n=== Test 7: Generation with Ablation ===")
        
        prompt = "def fibonacci(n):"
        
        # Test configurations
        configs = [
            {"name": "No ablation", "components": {}},
            {"name": "All attention", "components": {
                "attention_heads": {str(i): list(range(16)) for i in range(20)}
            }},
            {"name": "All FFN", "components": {
                "ffn_layers": list(range(20))
            }},
            {"name": "Layers 0-9", "components": {
                "layers": list(range(10))
            }}
        ]
        
        results = []
        for config in configs:
            logger.info(f"\n  Testing: {config['name']}")
            
            # Apply ablation
            disabled_components = config['components']
            
            # Parse components
            disabled_layers = set(disabled_components.get('layers', []))
            disabled_attention_raw = disabled_components.get('attention_heads', {})
            disabled_attention = {int(k) if isinstance(k, str) else k: v 
                                 for k, v in disabled_attention_raw.items()}
            disabled_ffn = set(disabled_components.get('ffn_layers', []))
            
            # Apply hooks
            handles = []
            for layer_idx in range(self.model.config.n_layer):
                if layer_idx in disabled_layers:
                    def layer_hook(module, input, output):
                        if isinstance(output, tuple):
                            return (input[0],) + output[1:]
                        return input[0]
                    handle = self.model.transformer.h[layer_idx].register_forward_hook(layer_hook)
                    handles.append(handle)
                else:
                    if layer_idx in disabled_attention:
                        heads = disabled_attention[layer_idx]
                        if len(heads) == 16:
                            def attention_hook(module, input, output):
                                if isinstance(output, tuple):
                                    return (torch.zeros_like(output[0]),) + output[1:]
                                return torch.zeros_like(output)
                            handle = self.model.transformer.h[layer_idx].attn.register_forward_hook(attention_hook)
                            handles.append(handle)
                    
                    if layer_idx in disabled_ffn:
                        def ffn_hook(module, input, output):
                            return torch.zeros_like(output)
                        handle = self.model.transformer.h[layer_idx].mlp.register_forward_hook(ffn_hook)
                        handles.append(handle)
            
            # Generate
            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
            with torch.no_grad():
                output_ids = self.model.generate(
                    **inputs,
                    max_new_tokens=20,
                    temperature=0.7,
                    do_sample=True
                )
            
            generated_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
            
            # Clean up hooks
            for handle in handles:
                handle.remove()
            
            results.append({
                "config": config['name'],
                "output": generated_text
            })
            logger.info(f"    Output: {generated_text[:50]}...")
        
        # Verify all outputs are different (except baseline)
        outputs = [r['output'] for r in results]
        unique_outputs = len(set(outputs))
        logger.info(f"\n✓ Generated {unique_outputs} unique outputs from {len(configs)} configs")
        
        for result in results:
            logger.info(f"  {result['config']}: {result['output'][:80]}...")
        
        return True
    
    def run_all_tests(self):
        """Run all ablation tests"""
        logger.info("=" * 60)
        logger.info("ABLATION FUNCTIONALITY TEST SUITE")
        logger.info("=" * 60)
        
        self.setup()
        
        tests = [
            self.test_model_architecture,
            self.test_attention_hook_attachment,
            self.test_attention_zeroing,
            self.test_ffn_ablation,
            self.test_partial_attention_ablation,
            self.test_data_format_conversion,
            self.test_generation_with_ablation
        ]
        
        passed = 0
        failed = 0
        
        for test in tests:
            try:
                if test():
                    passed += 1
                    logger.info(f"  ✅ {test.__name__} PASSED\n")
            except Exception as e:
                failed += 1
                logger.error(f"  ❌ {test.__name__} FAILED: {e}\n")
        
        logger.info("=" * 60)
        logger.info(f"TEST RESULTS: {passed} passed, {failed} failed")
        logger.info("=" * 60)
        
        return failed == 0

if __name__ == "__main__":
    tester = AblationTester()
    success = tester.run_all_tests()
    exit(0 if success else 1)