cotlab 0.8.0__py3-none-any.whl

cotlab/patching/patcher.py

@@ -0,0 +1,439 @@
+"""Activation patcher for causal intervention experiments."""
+
+from dataclasses import dataclass
+from typing import Dict, Iterable, List, Optional, Tuple
+
+import torch
+
+from .cache import ActivationCache
+from .hooks import HookManager
+
+
+@dataclass
+class PatchingResult:
+    """Result from a patching operation."""
+
+    output_text: str
+    patched_layers: List[int]
+    patched_positions: Optional[List[int]]
+    original_answer: Optional[str] = None
+    patched_answer: Optional[str] = None
+
+    @property
+    def answer_changed(self) -> bool:
+        """Check if patching changed the answer."""
+        if self.original_answer is None or self.patched_answer is None:
+            return False
+        return self.original_answer.strip() != self.patched_answer.strip()
+
+
+@dataclass
+class ForwardPatchResult:
+    """Result from forward-only patching (no generation)."""
+
+    layer_idx: int
+    clean_logits: torch.Tensor
+    corrupted_logits: torch.Tensor
+    patched_logits: torch.Tensor
+    effect_score: float  # How much patching recovered clean behavior
+
+    @property
+    def kl_from_clean(self) -> float:
+        """KL divergence of patched from clean distribution."""
+        clean_probs = torch.softmax(self.clean_logits[:, -1, :], dim=-1)
+        patched_probs = torch.softmax(self.patched_logits[:, -1, :], dim=-1)
+        return torch.nn.functional.kl_div(
+            patched_probs.log(), clean_probs, reduction="batchmean"
+        ).item()
+
+    @property
+    def kl_from_corrupted(self) -> float:
+        """KL divergence of patched from corrupted distribution."""
+        corrupted_probs = torch.softmax(self.corrupted_logits[:, -1, :], dim=-1)
+        patched_probs = torch.softmax(self.patched_logits[:, -1, :], dim=-1)
+        return torch.nn.functional.kl_div(
+            patched_probs.log(), corrupted_probs, reduction="batchmean"
+        ).item()
+
+
+class ActivationPatcher:
+    """
+    Perform activation patching interventions for causal analysis.
+
+    Activation patching is a technique to test causal importance of
+    specific model components by replacing activations from one run
+    with activations from another run.
+
+    Example:
+        >>> patcher = ActivationPatcher(backend)
+        >>> clean_out, clean_cache = backend.generate_with_cache(clean_prompt)
+        >>> results = patcher.patch_run(
+        ...     corrupted_prompt,
+        ...     source_cache=clean_cache,
+        ...     target_layers=[5, 10, 15]
+        ... )
+    """
+
+    def __init__(self, backend: "TransformersBackend"):
+        """
+        Initialize patcher with a backend.
+
+        Args:
+            backend: TransformersBackend with hook support
+        """
+        if not backend.supports_activations:
+            raise ValueError("Backend does not support activation access")
+
+        self.backend = backend
+
+    @property
+    def hook_manager(self) -> HookManager:
+        return self.backend.hook_manager
+
+    def patch_run(
+        self,
+        prompt: str,
+        source_cache: ActivationCache,
+        target_layers: List[int],
+        token_positions: Optional[List[int]] = None,
+        **gen_kwargs,
+    ) -> PatchingResult:
+        """
+        Run generation while patching activations from source_cache.
+
+        This is the core patching operation: run the model on `prompt`
+        but replace activations at `target_layers` with values from
+        `source_cache` (typically from a "clean" run).
+
+        Args:
+            prompt: Input prompt (typically the "corrupted" version)
+            source_cache: Activations from another run to patch in
+            target_layers: Which layers to patch
+            token_positions: Which token positions to patch (None = all)
+            **gen_kwargs: Additional generation arguments
+
+        Returns:
+            PatchingResult with output and metadata
+        """
+        # Register patching hooks
+        for layer_idx in target_layers:
+            source_activation = source_cache.get(layer_idx)
+            if source_activation is None:
+                raise ValueError(f"Layer {layer_idx} not in source cache")
+
+            self.hook_manager.register_patch_hook(layer_idx, source_activation, token_positions)
+
+        try:
+            output = self.backend.generate(prompt, **gen_kwargs)
+        finally:
+            self.hook_manager.remove_all_hooks()
+
+        return PatchingResult(
+            output_text=output.text, patched_layers=target_layers, patched_positions=token_positions
+        )
+
+    def sweep_layers(
+        self,
+        clean_prompt: str,
+        corrupted_prompt: str,
+        layers: Optional[List[int]] = None,
+        **gen_kwargs,
+    ) -> Dict[int, PatchingResult]:
+        """
+        Sweep patching across layers to find causal importance.
+
+        For each layer, patches that layer's activations from the clean
+        run into the corrupted run and measures the effect.
+
+        Args:
+            clean_prompt: Prompt that elicits desired behavior
+            corrupted_prompt: Prompt that elicits different behavior
+            layers: Which layers to sweep (None = all)
+            **gen_kwargs: Generation arguments
+
+        Returns:
+            Dict mapping layer_idx -> ForwardPatchResult
+        """
+        # Get clean logits and activations
+        clean_logits, clean_cache = self.backend.forward_with_cache(clean_prompt, layers=layers)
+
+        # Get corrupted logits (no patching)
+        corrupted_logits, _ = self.backend.forward_with_cache(corrupted_prompt, layers=layers)
+
+        target_layers = layers if layers is not None else clean_cache.layers
+        results = {}
+
+        for layer_idx in target_layers:
+            # Run forward pass with patching on this layer
+            patched_logits = self._forward_with_patch(corrupted_prompt, clean_cache, layer_idx)
+
+            # Compute effect score (how much patching moves toward clean)
+            effect = self._compute_logit_effect(clean_logits, corrupted_logits, patched_logits)
+
+            results[layer_idx] = ForwardPatchResult(
+                layer_idx=layer_idx,
+                clean_logits=clean_logits.cpu(),
+                corrupted_logits=corrupted_logits.cpu(),
+                patched_logits=patched_logits.cpu(),
+                effect_score=effect,
+            )
+
+        return results
+
+    def sweep_heads(
+        self,
+        clean_prompt: str,
+        corrupted_prompt: str,
+        layers: Optional[List[int]] = None,
+        head_indices: Optional[Iterable[int]] = None,
+        target_heads: Optional[Dict[int, Iterable[int]]] = None,
+    ) -> Dict[Tuple[int, int], ForwardPatchResult]:
+        """
+        Sweep patching across attention heads to find causal importance.
+
+        Args:
+            clean_prompt: Prompt that elicits desired behavior
+            corrupted_prompt: Prompt that elicits different behavior
+            layers: Which layers to sweep (None = all)
+            head_indices: Heads to patch for all layers (e.g., [0,1,2])
+            target_heads: Dict mapping layer -> list of head indices
+
+        Returns:
+            Dict mapping (layer_idx, head_idx) -> ForwardPatchResult
+        """
+        if head_indices is None and not target_heads:
+            raise ValueError("Provide head_indices or target_heads for head patching.")
+        if head_indices is not None and target_heads:
+            raise ValueError("Use either head_indices or target_heads, not both.")
+
+        layers_to_sweep = layers
+        if target_heads:
+            layers_to_sweep = list(target_heads.keys())
+
+        # Get clean logits and attention outputs
+        clean_logits, clean_cache = self.backend.forward_with_attention_cache(
+            clean_prompt, layers=layers_to_sweep
+        )
+
+        # Get corrupted logits (no patching)
+        corrupted_logits, _ = self.backend.forward_with_cache(corrupted_prompt, layers=[])
+
+        num_heads, head_dim = self._get_head_info()
+
+        results: Dict[Tuple[int, int], ForwardPatchResult] = {}
+        layers_to_sweep = layers_to_sweep if layers_to_sweep is not None else clean_cache.layers
+
+        for layer_idx in layers_to_sweep:
+            heads = list(target_heads.get(layer_idx, [])) if target_heads else list(head_indices)
+            for head_idx in heads:
+                if head_idx < 0 or head_idx >= num_heads:
+                    raise ValueError(f"Head index {head_idx} out of range (0..{num_heads - 1})")
+
+                patched_logits = self._forward_with_head_patch(
+                    corrupted_prompt, clean_cache, layer_idx, [head_idx], head_dim
+                )
+                effect = self._compute_logit_effect(clean_logits, corrupted_logits, patched_logits)
+
+                results[(layer_idx, head_idx)] = ForwardPatchResult(
+                    layer_idx=layer_idx,
+                    clean_logits=clean_logits.cpu(),
+                    corrupted_logits=corrupted_logits.cpu(),
+                    patched_logits=patched_logits.cpu(),
+                    effect_score=effect,
+                )
+
+        return results
+
+    def _forward_with_patch(
+        self,
+        prompt: str,
+        source_cache: ActivationCache,
+        layer_idx: int,
+    ) -> torch.Tensor:
+        """Run a single forward pass with patching at specified layer's residual stream."""
+        source_activation = source_cache.get(layer_idx)
+        if source_activation is None:
+            raise ValueError(f"Layer {layer_idx} not in source cache")
+
+        # Use residual patch hook for safer patching
+        self.hook_manager.register_residual_patch_hook(layer_idx, source_activation, None)
+
+        try:
+            logits, _ = self.backend.forward_with_cache(prompt, layers=[])
+        finally:
+            self.hook_manager.remove_all_hooks()
+
+        return logits
+
+    def _forward_with_head_patch(
+        self,
+        prompt: str,
+        source_cache: ActivationCache,
+        layer_idx: int,
+        head_indices: List[int],
+        head_dim: int,
+    ) -> torch.Tensor:
+        """Run forward pass with attention head patching at a specific layer."""
+        source_activation = source_cache.get(layer_idx)
+        if source_activation is None:
+            raise ValueError(f"Layer {layer_idx} not in source cache")
+
+        self.hook_manager.register_multi_head_patch_hook(
+            layer_idx, head_indices, source_activation, head_dim
+        )
+
+        try:
+            logits, _ = self.backend.forward_with_cache(prompt, layers=[])
+        finally:
+            self.hook_manager.remove_all_hooks()
+
+        return logits
+
+    def _get_head_info(self) -> Tuple[int, int]:
+        """Return (num_heads, head_dim) from the model config."""
+        cfg = getattr(self.backend, "model", None).config
+        # Handle multimodal models with nested text_config
+        if hasattr(cfg, "text_config"):
+            cfg = cfg.text_config
+        num_heads = (
+            getattr(cfg, "num_attention_heads", None)
+            or getattr(cfg, "n_head", None)
+            or getattr(cfg, "num_heads", None)
+        )
+        head_dim = getattr(cfg, "head_dim", None)
+        hidden_size = (
+            getattr(cfg, "hidden_size", None)
+            or getattr(cfg, "hidden_dim", None)
+            or getattr(cfg, "d_model", None)
+        )
+
+        if head_dim is None and num_heads and hidden_size:
+            head_dim = hidden_size // num_heads
+
+        if not num_heads or not head_dim:
+            raise ValueError(
+                "Could not determine head information from model config "
+                f"(num_heads={num_heads}, head_dim={head_dim})."
+            )
+
+        return int(num_heads), int(head_dim)
+
+    def _compute_logit_effect(
+        self,
+        clean_logits: torch.Tensor,
+        corrupted_logits: torch.Tensor,
+        patched_logits: torch.Tensor,
+    ) -> float:
+        """
+        Compute how much patching moves logits from corrupted toward clean.
+
+        Returns value from 0 (no effect) to 1 (full recovery).
+        """
+        # Use last token logits
+        clean_last = clean_logits[:, -1, :]
+        corrupted_last = corrupted_logits[:, -1, :]
+        patched_last = patched_logits[:, -1, :]
+
+        # Compute distances
+        clean_to_corrupted = torch.norm(clean_last - corrupted_last).item()
+        patched_to_corrupted = torch.norm(patched_last - corrupted_last).item()
+
+        if clean_to_corrupted == 0:
+            return 0.0
+
+        # Effect = how much we moved toward clean (normalized)
+        effect = 1.0 - (patched_to_corrupted / clean_to_corrupted)
+        return max(0.0, min(1.0, effect))
+
+    def sweep_positions(
+        self,
+        clean_prompt: str,
+        corrupted_prompt: str,
+        layer_idx: int,
+        positions: Optional[List[int]] = None,
+        **gen_kwargs,
+    ) -> Dict[int, PatchingResult]:
+        """
+        Sweep patching across token positions at a fixed layer.
+
+        Args:
+            clean_prompt: Prompt that elicits desired behavior
+            corrupted_prompt: Prompt that elicits different behavior
+            layer_idx: Which layer to patch
+            positions: Which positions to sweep (None = all)
+            **gen_kwargs: Generation arguments
+
+        Returns:
+            Dict mapping position -> PatchingResult
+        """
+        # Get clean activations
+        _, clean_cache = self.backend.generate_with_cache(
+            clean_prompt, layers=[layer_idx], **gen_kwargs
+        )
+
+        # Get corrupted baseline
+        corrupted_output = self.backend.generate(corrupted_prompt, **gen_kwargs)
+
+        # Determine positions to sweep
+        clean_act = clean_cache.get(layer_idx)
+        target_positions = positions if positions is not None else list(range(clean_act.shape[1]))
+
+        results = {}
+
+        for pos in target_positions:
+            result = self.patch_run(
+                corrupted_prompt,
+                clean_cache,
+                target_layers=[layer_idx],
+                token_positions=[pos],
+                **gen_kwargs,
+            )
+            result.original_answer = corrupted_output.text
+            result.patched_answer = result.output_text
+            results[pos] = result
+
+        return results
+
+    def ablate_layer(
+        self, prompt: str, layer_idx: int, ablation_type: str = "zero", **gen_kwargs
+    ) -> PatchingResult:
+        """
+        Ablate a layer by zeroing or replacing with mean.
+
+        Args:
+            prompt: Input prompt
+            layer_idx: Which layer to ablate
+            ablation_type: "zero" or "mean"
+            **gen_kwargs: Generation arguments
+
+        Returns:
+            PatchingResult from ablated run
+        """
+
+        def zero_hook(module, input, output):
+            if isinstance(output, tuple):
+                hidden_states = output[0]
+                rest = output[1:]
+                return (torch.zeros_like(hidden_states),) + rest
+            return torch.zeros_like(output)
+
+        def mean_hook(module, input, output):
+            if isinstance(output, tuple):
+                hidden_states = output[0]
+                rest = output[1:]
+                mean_val = hidden_states.mean(dim=(0, 1), keepdim=True)
+                return (mean_val.expand_as(hidden_states),) + rest
+            mean_val = output.mean(dim=(0, 1), keepdim=True)
+            return mean_val.expand_as(output)
+
+        hook = zero_hook if ablation_type == "zero" else mean_hook
+        self.hook_manager.register_forward_hook(layer_idx, hook)
+
+        try:
+            output = self.backend.generate(prompt, **gen_kwargs)
+        finally:
+            self.hook_manager.remove_all_hooks()
+
+        return PatchingResult(
+            output_text=output.text, patched_layers=[layer_idx], patched_positions=None
+        )

cotlab/patching/sae.py

@@ -0,0 +1,181 @@
+"""GemmaScope-2 JumpReLU Sparse Autoencoder loader.
+
+Loads SAE weights directly from HuggingFace (no SAELens / TransformerLens
+dependency) and exposes a minimal encode() forward pass compatible with the
+CoTLab TransformersBackend hook infrastructure.
+
+Architecture (JumpReLU SAE, Lieberum et al. 2024)
+--------------------------------------------------
+    h        = x - b_dec                  # centre around decoder bias
+    pre_act  = h @ w_enc + b_enc          # linear projection  [..., d_sae]
+    features = pre_act * (pre_act > θ)    # JumpReLU gate
+
+Weight layout in params.safetensors (GemmaScope-2 convention)
+--------------------------------------------------------------
+    w_enc     float32  [d_model, d_sae]
+    b_enc     float32  [d_sae]
+    w_dec     float32  [d_sae, d_model]
+    b_dec     float32  [d_model]
+    threshold float32  [d_sae]
+
+HF path pattern
+---------------
+    {site}/layer_{N}_width_{width}_l0_{l0_label}/params.safetensors
+    e.g. resid_post_all/layer_9_width_16k_l0_small/params.safetensors
+"""
+
+import re
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+
+class GemmaScopeLayer(nn.Module):
+    """JumpReLU SAE for a single residual-stream layer."""
+
+    def __init__(
+        self,
+        w_enc: torch.Tensor,
+        b_enc: torch.Tensor,
+        w_dec: torch.Tensor,
+        b_dec: torch.Tensor,
+        threshold: torch.Tensor,
+        layer: int,
+        repo_id: str,
+        source_path: str,
+    ):
+        super().__init__()
+        self.register_buffer("w_enc", w_enc)  # [d_model, d_sae]
+        self.register_buffer("b_enc", b_enc)  # [d_sae]
+        self.register_buffer("w_dec", w_dec)  # [d_sae, d_model]
+        self.register_buffer("b_dec", b_dec)  # [d_model]
+        self.register_buffer("threshold", threshold)  # [d_sae]
+        self.layer = layer
+        self.repo_id = repo_id
+        self.source_path = source_path
+
+    @property
+    def d_model(self) -> int:
+        return self.w_enc.shape[0]
+
+    @property
+    def d_sae(self) -> int:
+        return self.w_enc.shape[1]
+
+    # ------------------------------------------------------------------
+    # Factory
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        repo_id: str,
+        layer: int,
+        site: str = "resid_post_all",
+        width: str = "16k",
+        l0_label: str = "small",
+        token: Optional[str] = None,
+    ) -> "GemmaScopeLayer":
+        """Download and load a GemmaScope-2 SAE from HuggingFace.
+
+        Args:
+            repo_id:   HF repo, e.g. ``"google/gemma-scope-2-270m-it"``
+            layer:     Residual stream layer index (0-based).
+            site:      SAE training site.  ``"resid_post_all"`` covers every
+                       layer; other options are ``"resid_post"`` (4 depths only),
+                       ``"attn_out_all"``, ``"mlp_out_all"``.
+            width:     Feature dictionary width: ``"16k"`` or ``"262k"``.
+            l0_label:  Sparsity label used in the directory name: ``"small"``,
+                       ``"medium"``, or ``"big"``.  If not found, falls back to
+                       any available file for that layer/width combination.
+            token:     HuggingFace API token (optional; reads HF_TOKEN env var
+                       automatically via huggingface_hub).
+        """
+        from huggingface_hub import hf_hub_download, list_repo_files  # noqa: PLC0415
+        from safetensors import safe_open  # noqa: PLC0415
+
+        # Try the canonical direct path first (avoids full repo listing).
+        direct = f"{site}/layer_{layer}_width_{width}_l0_{l0_label}/params.safetensors"
+        print(f"  [SAE] Fetching layer={layer} ({direct}) …")
+
+        try:
+            local_path = hf_hub_download(repo_id=repo_id, filename=direct, token=token)
+            chosen = direct
+        except Exception:
+            print("  [SAE] Direct path not found — scanning repo …")
+            all_files = list(list_repo_files(repo_id, token=token))
+            layer_tag = f"layer_{layer}_"
+            width_tag = f"width_{width}_"
+            candidates = [
+                f
+                for f in all_files
+                if site in f
+                and layer_tag in f
+                and width_tag in f
+                and f.endswith("params.safetensors")
+            ]
+            if not candidates:
+                available = sorted(
+                    {
+                        re.search(r"layer_(\d+)_", f).group(1)
+                        for f in all_files
+                        if "params.safetensors" in f and re.search(r"layer_(\d+)_", f)
+                    },
+                    key=int,
+                )
+                raise FileNotFoundError(
+                    f"No SAE found for site={site!r}, layer={layer}, width={width!r} "
+                    f"in {repo_id}.\nAvailable layers (any site/width): {available}"
+                )
+            preferred = [f for f in candidates if f"l0_{l0_label}" in f]
+            chosen = preferred[0] if preferred else candidates[0]
+            local_path = hf_hub_download(repo_id=repo_id, filename=chosen, token=token)
+
+        with safe_open(local_path, framework="pt") as f:
+            w_enc = f.get_tensor("w_enc").float()  # [d_model, d_sae]
+            b_enc = f.get_tensor("b_enc").float()  # [d_sae]
+            w_dec = f.get_tensor("w_dec").float()  # [d_sae, d_model]
+            b_dec = f.get_tensor("b_dec").float()  # [d_model]
+            threshold = f.get_tensor("threshold").float()  # [d_sae]
+
+        print(f"  [SAE] Loaded  layer={layer}  d_model={w_enc.shape[0]}  d_sae={w_enc.shape[1]}")
+        return cls(
+            w_enc=w_enc,
+            b_enc=b_enc,
+            w_dec=w_dec,
+            b_dec=b_dec,
+            threshold=threshold,
+            layer=layer,
+            repo_id=repo_id,
+            source_path=chosen,
+        )
+
+    # ------------------------------------------------------------------
+    # Forward
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """JumpReLU encode.
+
+        Args:
+            x: Residual stream activations, shape ``[..., d_model]``.
+               Must be on the same device as the SAE buffers.
+
+        Returns:
+            Sparse feature activations, shape ``[..., d_sae]``.
+        """
+        orig_dtype = x.dtype
+        x = x.float()
+        h = x - self.b_dec
+        pre = h @ self.w_enc + self.b_enc
+        features = pre * (pre > self.threshold).float()
+        return features.to(orig_dtype)
+
+    def __repr__(self) -> str:
+        return (
+            f"GemmaScopeLayer(layer={self.layer}, "
+            f"d_model={self.d_model}, d_sae={self.d_sae}, "
+            f"repo={self.repo_id})"
+        )

cotlab/prompts/__init__.py

@@ -0,0 +1,43 @@
+"""Prompt strategies module."""
+
+from .cardiology import CardiologyPromptStrategy
+from .histopathology import HistopathologyPromptStrategy
+from .length_matched_strategies import (
+    ChainOfThoughtMatchedStrategy,
+    ContrarianMatchedStrategy,
+    DirectAnswerMatchedStrategy,
+)
+from .mcq import MCQPromptStrategy
+from .neurology import NeurologyPromptStrategy
+from .oncology import OncologyPromptStrategy
+from .plab import PLABPromptStrategy
+from .pubhealthbench import PubHealthBenchMCQPromptStrategy
+from .radiology import RadiologyPromptStrategy
+from .strategies import (
+    ArroganceStrategy,
+    ChainOfThoughtStrategy,
+    DirectAnswerStrategy,
+    NoInstructionStrategy,
+    SimplePromptStrategy,
+    create_prompt_strategy,
+)
+
+__all__ = [
+    "SimplePromptStrategy",
+    "ChainOfThoughtStrategy",
+    "DirectAnswerStrategy",
+    "ArroganceStrategy",
+    "NoInstructionStrategy",
+    "CardiologyPromptStrategy",
+    "HistopathologyPromptStrategy",
+    "MCQPromptStrategy",
+    "NeurologyPromptStrategy",
+    "OncologyPromptStrategy",
+    "PLABPromptStrategy",
+    "PubHealthBenchMCQPromptStrategy",
+    "RadiologyPromptStrategy",
+    "create_prompt_strategy",
+    "ContrarianMatchedStrategy",
+    "ChainOfThoughtMatchedStrategy",
+    "DirectAnswerMatchedStrategy",
+]