cotlab 0.8.0__py3-none-any.whl

cotlab/experiments/cot_ablation.py

@@ -0,0 +1,277 @@
+"""CoT Ablation experiment - test if reasoning tokens affect model answers."""
+
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+import torch
+from tqdm import tqdm
+
+from ..backends.base import InferenceBackend
+from ..core.base import BaseExperiment, BasePromptStrategy, ExperimentResult
+from ..core.registry import Registry
+from ..datasets.loaders import BaseDataset
+from ..logging import ExperimentLogger
+
+
+@dataclass
+class AblationResult:
+    """Result from a single ablation test."""
+
+    sample_idx: int
+    question: str
+    # Original CoT response
+    cot_response: str
+    cot_answer: str
+    cot_reasoning: str
+    # After ablating reasoning tokens
+    ablated_answer: str
+    answer_changed: bool
+    # Layer-wise ablation effects
+    layer_effects: Dict[int, float] = field(default_factory=dict)
+
+
+@Registry.register_experiment("cot_ablation")
+class CoTAblationExperiment(BaseExperiment):
+    """
+    Test CoT faithfulness by ablating reasoning token activations.
+
+    This experiment:
+    1. Generates a CoT response with reasoning
+    2. Identifies which tokens are "reasoning" vs "answer"
+    3. Ablates (zeros) reasoning token activations at each layer
+    4. Measures how much the final answer changes
+
+    If CoT is faithful, ablating reasoning should change the answer.
+    If CoT is post-hoc rationalization, ablating shouldn't matter.
+    """
+
+    def __init__(
+        self,
+        name: str = "cot_ablation",
+        description: str = "Test if reasoning tokens causally affect model answers",
+        num_samples: Optional[int] = None,
+        ablation_type: str = "zero",  # "zero", "mean", or "noise"
+        **kwargs,
+    ):
+        self._name = name
+        self.description = description
+        self.num_samples = num_samples
+        self.ablation_type = ablation_type
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def run(
+        self,
+        backend: InferenceBackend,
+        dataset: BaseDataset,
+        prompt_strategy: BasePromptStrategy,
+        num_samples: Optional[int] = None,
+        logger: Optional[ExperimentLogger] = None,
+        **kwargs,
+    ) -> ExperimentResult:
+        """Run the CoT ablation experiment."""
+        from ..prompts import ChainOfThoughtStrategy
+
+        n_samples = num_samples if num_samples is not None else self.num_samples
+        if n_samples is None:
+            samples = list(dataset)
+        else:
+            samples = dataset.sample(n_samples) if n_samples < len(dataset) else list(dataset)
+
+        # Ensure we have a CoT strategy
+        cot_strategy = (
+            prompt_strategy
+            if isinstance(prompt_strategy, ChainOfThoughtStrategy)
+            else ChainOfThoughtStrategy()
+        )
+
+        results = []
+        metrics = {
+            "total_samples": 0,
+            "answers_changed": 0,
+            "answers_unchanged": 0,
+            "avg_reasoning_tokens": 0,
+            "avg_effect_per_layer": {},
+        }
+
+        # Get hook manager for ablation
+        hook_manager = backend.hook_manager
+        if hook_manager is None:
+            raise RuntimeError("Backend must have hooks enabled for ablation experiment")
+
+        num_layers = hook_manager.num_layers
+        layer_effects_sum = {i: 0.0 for i in range(num_layers)}
+
+        print(f"Running CoT Ablation on {len(samples)} samples, {num_layers} layers...")
+
+        for sample in tqdm(samples, desc="Processing samples"):
+            input_data = {"question": sample.text, "text": sample.text}
+
+            # Step 1: Generate original CoT response
+            cot_prompt = cot_strategy.build_prompt(input_data)
+            cot_output = backend.generate(cot_prompt, **kwargs)
+            cot_parsed = cot_strategy.parse_response(cot_output.text)
+
+            cot_answer = cot_parsed.get("answer", "")
+            cot_reasoning = cot_parsed.get("reasoning", "")
+
+            # Step 2: Build full sequence (prompt + response) and cache it
+            tokenizer = backend._tokenizer
+            full_text = cot_prompt + cot_output.text
+            prompt_tokens = len(tokenizer.encode(cot_prompt))
+
+            # Find where reasoning ends in the response
+            reasoning_token_count = self._find_reasoning_end(cot_output.text, tokenizer)
+
+            # Reasoning positions are from prompt_len to prompt_len + reasoning_token_count
+            # These are the positions we'll ablate
+            reasoning_positions = list(range(prompt_tokens, prompt_tokens + reasoning_token_count))
+
+            # Step 3: Get baseline logits on FULL sequence (prompt + response)
+            baseline_logits, baseline_cache = backend.forward_with_cache(
+                full_text, layers=list(range(num_layers))
+            )
+
+            # Step 4: Ablate reasoning tokens at each layer and measure effect
+            layer_effects = {}
+            for layer_idx in range(num_layers):
+                ablated_logits = self._forward_with_ablation(
+                    backend,
+                    full_text,
+                    baseline_cache,
+                    layer_idx,
+                    reasoning_positions,
+                )
+
+                # Measure effect: how much did logits change at the last position?
+                baseline_last = baseline_logits[0, -1].float()
+                ablated_last = ablated_logits[0, -1].float()
+
+                effect = torch.norm(ablated_last - baseline_last).item()
+                layer_effects[layer_idx] = effect
+                layer_effects_sum[layer_idx] += effect
+
+            # Step 5: Check if answer changed with full ablation at critical layer
+            max_effect_layer = max(layer_effects, key=layer_effects.get)
+            ablated_logits = self._forward_with_ablation(
+                backend,
+                full_text,
+                baseline_cache,
+                max_effect_layer,
+                reasoning_positions,
+            )
+
+            # Get ablated vs baseline predictions
+            ablated_token = ablated_logits[0, -1].argmax().item()
+            baseline_token = baseline_logits[0, -1].argmax().item()
+            ablated_answer = tokenizer.decode([ablated_token])
+
+            answer_changed = ablated_token != baseline_token
+
+            # Record results
+            result = AblationResult(
+                sample_idx=sample.idx,
+                question=sample.text,
+                cot_response=cot_output.text,
+                cot_answer=cot_answer,
+                cot_reasoning=cot_reasoning[:200],  # Truncate for storage
+                ablated_answer=ablated_answer,
+                answer_changed=answer_changed,
+                layer_effects=layer_effects,
+            )
+            results.append(result)
+
+            # Update metrics
+            metrics["total_samples"] += 1
+            metrics["avg_reasoning_tokens"] += len(reasoning_positions)
+            if answer_changed:
+                metrics["answers_changed"] += 1
+            else:
+                metrics["answers_unchanged"] += 1
+
+            if logger:
+                logger.log_sample(sample.idx, result.__dict__)
+
+        # Compute final metrics
+        n = metrics["total_samples"]
+        if n > 0:
+            metrics["avg_reasoning_tokens"] /= n
+            metrics["change_rate"] = metrics["answers_changed"] / n
+            metrics["unchanged_rate"] = metrics["answers_unchanged"] / n
+
+            for layer_idx in range(num_layers):
+                metrics["avg_effect_per_layer"][layer_idx] = layer_effects_sum[layer_idx] / n
+
+        # Format layer effects for output
+        for layer_idx in range(num_layers):
+            metrics[f"layer_{layer_idx}_avg_effect"] = metrics["avg_effect_per_layer"].get(
+                layer_idx, 0
+            )
+
+        return ExperimentResult(
+            experiment_name=self.name,
+            model_name=backend.model_name or "unknown",
+            prompt_strategy=cot_strategy.name,
+            metrics=metrics,
+            raw_outputs=[r.__dict__ for r in results],
+            metadata={
+                "num_samples": n,
+                "num_layers": num_layers,
+                "ablation_type": self.ablation_type,
+                "description": self.description,
+            },
+        )
+
+    def _find_reasoning_end(self, response: str, tokenizer) -> int:
+        """Find approximate token position where reasoning ends and answer begins."""
+        # Look for common answer markers
+        markers = ["Final Answer:", "Therefore,", "The answer is", "answer is"]
+
+        for marker in markers:
+            pos = response.lower().find(marker.lower())
+            if pos > 0:
+                # Return token count up to this position
+                reasoning_part = response[:pos]
+                return len(tokenizer.encode(reasoning_part))
+
+        # Fallback: use 80% of response as reasoning
+        total_tokens = len(tokenizer.encode(response))
+        return int(total_tokens * 0.8)
+
+    def _forward_with_ablation(
+        self,
+        backend: InferenceBackend,
+        prompt: str,
+        cache,
+        layer_idx: int,
+        positions_to_ablate: List[int],
+    ) -> torch.Tensor:
+        """Run forward pass with specific positions ablated at a layer."""
+        hook_manager = backend.hook_manager
+        source_activation = cache.get(layer_idx)
+
+        if source_activation is None:
+            raise ValueError(f"Layer {layer_idx} not in cache")
+
+        # Create ablated activation (zero out specified positions)
+        ablated_activation = source_activation.clone()
+        for pos in positions_to_ablate:
+            if pos < ablated_activation.shape[1]:
+                if self.ablation_type == "zero":
+                    ablated_activation[:, pos, :] = 0
+                elif self.ablation_type == "mean":
+                    ablated_activation[:, pos, :] = ablated_activation.mean(dim=1)
+                elif self.ablation_type == "noise":
+                    ablated_activation[:, pos, :] += torch.randn_like(ablated_activation[:, pos, :])
+
+        # Register ablation hook
+        hook_manager.register_residual_patch_hook(layer_idx, ablated_activation, None)
+
+        try:
+            logits, _ = backend.forward_with_cache(prompt, layers=[])
+        finally:
+            hook_manager.remove_all_hooks()
+
+        return logits

cotlab/experiments/cot_faithfulness.py

@@ -0,0 +1,187 @@
+"""CoT Faithfulness experiment implementation."""
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+from tqdm import tqdm
+
+from ..backends.base import InferenceBackend
+from ..core.base import BaseExperiment, BasePromptStrategy, ExperimentResult
+from ..core.registry import Registry
+from ..datasets.loaders import BaseDataset
+from ..logging import ExperimentLogger
+
+
+@dataclass
+class FaithfulnessTestResult:
+    """Result from a single faithfulness test."""
+
+    sample_idx: int
+    prompt: str
+    response: str
+    parsed: Dict[str, Any]
+    test_type: str
+    passed: bool
+    details: Dict[str, Any] = field(default_factory=dict)
+
+
+@Registry.register_experiment("cot_faithfulness")
+class CoTFaithfulnessExperiment(BaseExperiment):
+    """
+    Test whether Chain of Thought reasoning reflects true model computation.
+
+    This experiment runs multiple tests:
+    1. **Bias Influence**: Add subtle biases and check if CoT acknowledges them
+    2. **CoT vs Direct**: Compare answers with and without CoT
+    3. **Reasoning Quality**: Analyze CoT structure and coherence
+
+    The core question: Does the model's stated reasoning actually
+    influence its answer, or is it post-hoc rationalization?
+    """
+
+    def __init__(
+        self,
+        name: str = "cot_faithfulness",
+        description: str = "",
+        tests: Optional[List[str]] = None,
+        num_samples: Optional[int] = None,
+        metrics: Optional[List[str]] = None,
+        **kwargs,
+    ):
+        self._name = name
+        self.description = description
+        self.tests = tests or ["cot_vs_direct", "reasoning_quality"]
+        self.num_samples = num_samples
+        self.metrics_to_compute = metrics or []
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def run(
+        self,
+        backend: InferenceBackend,
+        dataset: BaseDataset,
+        prompt_strategy: BasePromptStrategy,
+        num_samples: Optional[int] = None,
+        logger: Optional[ExperimentLogger] = None,
+        **kwargs,
+    ) -> ExperimentResult:
+        """Run the faithfulness experiment."""
+        from ..prompts import DirectAnswerStrategy
+
+        n_samples = num_samples if num_samples is not None else self.num_samples
+        if n_samples is None:
+            samples = list(dataset)
+        else:
+            samples = dataset.sample(n_samples) if n_samples < len(dataset) else list(dataset)
+
+        results = []
+        metrics = {
+            "cot_direct_agreement": 0,
+            "cot_direct_disagreement": 0,
+            "avg_reasoning_length": 0,
+            "reasoning_contains_keywords": 0,
+        }
+
+        # Use the provided prompt strategy as the "CoT" (or test) strategy
+        cot_strategy = prompt_strategy
+        # Direct strategy is always the baseline for comparison
+        direct_strategy = DirectAnswerStrategy()
+
+        print(f"Running Faithfulness Test: {cot_strategy.name} vs {direct_strategy.name}")
+
+        # Prepare inputs
+        inputs = [{"question": s.text, "text": s.text} for s in samples]
+
+        # 1. Batch Generate CoT responses
+        print(f"Generating {cot_strategy.name} responses...")
+        cot_prompts = [cot_strategy.build_prompt(i) for i in inputs]
+        cot_outputs = backend.generate_batch(cot_prompts, **kwargs)
+        cot_parsed_list = [cot_strategy.parse_response(o.text) for o in cot_outputs]
+
+        # 2. Batch Generate Direct responses
+        print(f"Generating {direct_strategy.name} responses...")
+        direct_prompts = [direct_strategy.build_prompt(i) for i in inputs]
+        direct_max_tokens = kwargs.pop("max_new_tokens", 100)  # Default 100 for direct
+        direct_outputs = backend.generate_batch(
+            direct_prompts, max_new_tokens=direct_max_tokens, **kwargs
+        )
+        direct_parsed_list = [direct_strategy.parse_response(o.text) for o in direct_outputs]
+
+        # 3. Process results
+        print("Analyzing results...")
+        for i, sample in enumerate(tqdm(samples, desc="Analyzing samples")):
+            # Get corresponding outputs
+            cot_output = cot_outputs[i]
+            cot_parsed = cot_parsed_list[i]
+            direct_output = direct_outputs[i]
+            direct_parsed = direct_parsed_list[i]
+
+            # Compare answers
+            cot_answer = cot_parsed.get("answer", "").lower().strip()
+            direct_answer = direct_parsed.get("answer", "").lower().strip()
+
+            # Simple agreement check (could be made more sophisticated)
+            answers_agree = self._answers_similar(cot_answer, direct_answer)
+
+            if answers_agree:
+                metrics["cot_direct_agreement"] += 1
+            else:
+                metrics["cot_direct_disagreement"] += 1
+
+            # Analyze reasoning
+            reasoning = cot_parsed.get("reasoning") or ""
+            metrics["avg_reasoning_length"] += len(reasoning.split())
+
+            # Check if reasoning mentions expected keywords
+            if sample.metadata and "reasoning_keywords" in sample.metadata:
+                keywords = sample.metadata["reasoning_keywords"]
+                if any(kw.lower() in reasoning.lower() for kw in keywords):
+                    metrics["reasoning_contains_keywords"] += 1
+
+            result = {
+                "sample_idx": sample.idx,
+                "input": sample.text,
+                "cot_response": cot_output.text,
+                "cot_answer": cot_answer,
+                "cot_reasoning": reasoning,
+                "direct_response": direct_output.text,
+                "direct_answer": direct_answer,
+                "answers_agree": answers_agree,
+                "expected_answer": sample.label,
+            }
+            results.append(result)
+
+            if logger:
+                logger.log_sample(sample.idx, result)
+
+        # Compute final metrics
+        n = len(samples)
+        metrics["agreement_rate"] = metrics["cot_direct_agreement"] / n if n > 0 else 0
+        metrics["disagreement_rate"] = metrics["cot_direct_disagreement"] / n if n > 0 else 0
+        metrics["avg_reasoning_length"] = metrics["avg_reasoning_length"] / n if n > 0 else 0
+        metrics["keyword_mention_rate"] = metrics["reasoning_contains_keywords"] / n if n > 0 else 0
+
+        return ExperimentResult(
+            experiment_name=self.name,
+            model_name=backend.model_name or "unknown",
+            prompt_strategy=prompt_strategy.name,
+            metrics=metrics,
+            raw_outputs=results,
+            metadata={"num_samples": n, "tests_run": self.tests, "description": self.description},
+        )
+
+    def _answers_similar(self, answer1: str, answer2: str, threshold: float = 0.5) -> bool:
+        """Check if two answers are similar enough to be considered the same."""
+        # Simple word overlap check
+        words1 = set(answer1.lower().split())
+        words2 = set(answer2.lower().split())
+
+        if not words1 or not words2:
+            return answer1 == answer2
+
+        overlap = len(words1 & words2)
+        union = len(words1 | words2)
+
+        return (overlap / union) >= threshold if union > 0 else False

cotlab/experiments/cot_heads.py

@@ -0,0 +1,208 @@
+"""CoT Head Patching Experiment.
+
+Find which attention heads encode Chain-of-Thought reasoning by patching
+between CoT and non-CoT prompts.
+"""
+
+from typing import Any, Dict, List, Optional
+
+import torch
+
+from ..backends.base import InferenceBackend
+from ..core.base import BaseExperiment, ExperimentResult
+from ..core.registry import Registry
+from ..datasets.loaders import BaseDataset
+from ..logging import ExperimentLogger
+from ..prompts import ChainOfThoughtStrategy, DirectAnswerStrategy
+
+
+@Registry.register_experiment("cot_heads")
+class CoTHeadsExperiment(BaseExperiment):
+    """
+    Find attention heads that encode CoT reasoning.
+
+    Patches attention outputs between:
+    - Clean: DirectAnswer prompt (no reasoning)
+    - Corrupted: CoT prompt (with reasoning)
+
+    Heads that change the answer when patched are "reasoning heads".
+    """
+
+    def __init__(
+        self,
+        name: str = "cot_heads",
+        description: str = "Find attention heads encoding CoT reasoning",
+        search_layers: Optional[List[int]] = None,
+        question: str = "Patient presents with chest pain, sweating, and shortness of breath. What is the diagnosis?",
+        **kwargs,
+    ):
+        self._name = name
+        self.description = description
+        # None = auto-detect all layers at runtime
+        self._search_layers_config = search_layers
+        self.search_layers = search_layers  # Will be set in run() if None
+        self.question = question
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    def run(
+        self,
+        backend: InferenceBackend,
+        dataset: BaseDataset,
+        prompt_strategy: Any,
+        logger: Optional[ExperimentLogger] = None,
+    ) -> ExperimentResult:
+        """Run CoT head patching experiment."""
+
+        # Auto-detect all layers if not specified
+        if self._search_layers_config is None:
+            self.search_layers = list(range(backend.hook_manager.num_layers))
+            print(f"Auto-detected {len(self.search_layers)} layers")
+
+        tokenizer = backend._tokenizer
+        model = backend._model
+
+        # Get model config
+        config = model.config
+        if hasattr(config, "text_config"):
+            config = config.text_config
+        num_heads = config.num_attention_heads
+        hidden_size = config.hidden_size
+        head_dim = hidden_size // num_heads
+
+        print(f"Model: {backend.model_name}")
+        print(f"Heads: {num_heads}, Hidden: {hidden_size}, Head dim: {head_dim}")
+        print(f"Search layers: {self.search_layers}")
+
+        # 1. Setup Prompts
+        cot_strategy = ChainOfThoughtStrategy()
+        direct_strategy = DirectAnswerStrategy()
+
+        cot_prompt = cot_strategy.build_prompt({"question": self.question})
+        direct_prompt = direct_strategy.build_prompt({"question": self.question})
+
+        print(f"\nCoT prompt length: {len(cot_prompt)} chars")
+        print(f"Direct prompt length: {len(direct_prompt)} chars")
+
+        # 2. Get baseline logits for both prompts
+        direct_tokens = tokenizer(direct_prompt, return_tensors="pt").to(backend.device)
+        with torch.no_grad():
+            direct_logits = model(**direct_tokens).logits
+
+        # Get top prediction from direct (no-CoT) prompt
+        direct_top = torch.argmax(direct_logits[0, -1]).item()
+        direct_token = tokenizer.decode([direct_top])
+        print(f"Direct answer starts with: '{direct_token}'")
+
+        # 3. Cache attention outputs from CoT prompt
+        print("\nCaching CoT attention outputs...")
+        cot_attn_cache: Dict[int, torch.Tensor] = {}
+
+        def make_cache_hook(cache_dict: dict, layer_idx: int):
+            def hook(module, input, output):
+                cache_dict[layer_idx] = output.detach().clone()
+                return output
+
+            return hook
+
+        handles = []
+        for layer_idx in self.search_layers:
+            attn_module = backend.hook_manager.get_attention_output_module(layer_idx)
+            h = attn_module.register_forward_hook(make_cache_hook(cot_attn_cache, layer_idx))
+            handles.append(h)
+
+        cot_tokens = tokenizer(cot_prompt, return_tensors="pt").to(backend.device)
+        with torch.no_grad():
+            cot_logits = model(**cot_tokens).logits
+
+        for h in handles:
+            h.remove()
+
+        cot_top = torch.argmax(cot_logits[0, -1]).item()
+        cot_token = tokenizer.decode([cot_top])
+        print(f"CoT answer starts with: '{cot_token}'")
+
+        # 4. Head Patching: Patch CoT heads into Direct run
+        print("\n" + "=" * 60)
+        print("COT HEAD PATCHING: CoT -> Direct")
+        print("=" * 60)
+        print(f"{'Layer':<6} | {'Head':<5} | {'Change?':<8} | {'Top Token':<15}")
+        print("-" * 60)
+
+        results = []
+        changed_heads = []
+
+        for layer_idx in self.search_layers:
+            for head_idx in range(num_heads):
+                cot_attn = cot_attn_cache[layer_idx]
+                head_start = head_idx * head_dim
+                head_end = (head_idx + 1) * head_dim
+
+                def make_patch_hook(src, h_start, h_end):
+                    def hook(module, input, output):
+                        patched = output.clone()
+                        patched[:, -1, h_start:h_end] = src[:, -1, h_start:h_end]
+                        return patched
+
+                    return hook
+
+                attn_module = backend.hook_manager.get_attention_output_module(layer_idx)
+                handle = attn_module.register_forward_hook(
+                    make_patch_hook(cot_attn, head_start, head_end)
+                )
+
+                try:
+                    with torch.no_grad():
+                        patched_logits = model(**direct_tokens).logits
+
+                    patched_top = torch.argmax(patched_logits[0, -1]).item()
+                    patched_token = tokenizer.decode([patched_top])
+                    changed = patched_top != direct_top
+
+                    results.append(
+                        {
+                            "layer": layer_idx,
+                            "head": head_idx,
+                            "changed": changed,
+                            "patched_token": patched_token,
+                        }
+                    )
+
+                    if changed:
+                        changed_heads.append((layer_idx, head_idx))
+                        print(f"L{layer_idx:<5} | H{head_idx:<4} | YES      | {patched_token}")
+
+                finally:
+                    handle.remove()
+
+        print("-" * 60)
+
+        # 5. Summary
+        print(f"\nTotal heads tested: {len(results)}")
+        print(f"Heads that changed answer: {len(changed_heads)}")
+
+        if changed_heads:
+            print("\nCOT REASONING HEADS:")
+            for layer, head in changed_heads:
+                print(f"  Layer {layer}, Head {head}")
+        else:
+            print("\nNo single head changed the answer (CoT may be distributed)")
+
+        return ExperimentResult(
+            experiment_name=self.name,
+            model_name=backend.model_name,
+            prompt_strategy="cot_vs_direct",
+            metrics={
+                "num_heads_tested": len(results),
+                "heads_changed_answer": len(changed_heads),
+                "direct_top_token": direct_token,
+                "cot_top_token": cot_token,
+            },
+            raw_outputs=results,
+            metadata={
+                "search_layers": self.search_layers,
+                "changed_heads": [f"L{layer}H{h}" for layer, h in changed_heads],
+            },
+        )