interpkit 0.1.0__py3-none-any.whl

interpkit/ops/lens.py

@@ -0,0 +1,151 @@
+"""lens — logit lens: project each layer's output to vocabulary space."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+from rich.console import Console
+
+from interpkit.ops.patch import _get_module
+
+if TYPE_CHECKING:
+    from interpkit.core.model import Model
+
+console = Console()
+
+
+def run_lens(model: "Model", text: Any, *, save: str | None = None) -> list[dict[str, Any]] | None:
+    """Project each layer's hidden state through the unembedding matrix.
+
+    Only works for language models with a detectable output head.
+    """
+    from interpkit.core.render import render_lens
+
+    arch = model.arch_info
+
+    if not arch.is_language_model or arch.unembedding_name is None:
+        console.print(
+            f"\n  [yellow]lens not available:[/yellow] no unembedding matrix detected"
+            f" for {arch.arch_family or 'this model'}.\n"
+        )
+        return None
+
+    if not arch.layer_names:
+        console.print(
+            "\n  [yellow]lens not available:[/yellow] no layer structure detected.\n"
+        )
+        return None
+
+    if model._tokenizer is None:
+        console.print(
+            "\n  [yellow]lens not available:[/yellow] no tokenizer loaded.\n"
+        )
+        return None
+
+    text_input = model._prepare(text)
+
+    # Get the unembedding weight matrix
+    unembed_mod = _get_module(model._model, arch.unembedding_name)
+    unembed_weight = unembed_mod.weight  # shape: (vocab_size, hidden_size)
+
+    # Capture hidden states at the output of each layer
+    layer_outputs: dict[str, torch.Tensor] = {}
+
+    def _make_hook(name: str):
+        def hook_fn(_mod: torch.nn.Module, _inp: Any, output: Any) -> None:
+            t = output if isinstance(output, torch.Tensor) else (
+                output[0] if isinstance(output, (tuple, list)) and isinstance(output[0], torch.Tensor) else None
+            )
+            if t is not None:
+                layer_outputs[name] = t.detach()
+        return hook_fn
+
+    hooks = []
+    for layer_name in arch.layer_names:
+        try:
+            mod = _get_module(model._model, layer_name)
+            hooks.append(mod.register_forward_hook(_make_hook(layer_name)))
+        except AttributeError:
+            continue
+
+    with torch.no_grad():
+        model._forward(text_input)
+
+    for h in hooks:
+        h.remove()
+
+    if not layer_outputs:
+        console.print("\n  [yellow]lens:[/yellow] no layer outputs captured.\n")
+        return None
+
+    # Apply layer norm if the model has a final layer norm before the head
+    # (common pattern: ln_f, model.norm, etc.)
+    final_norm = _find_final_norm(model._model, arch)
+
+    predictions: list[dict[str, Any]] = []
+
+    for layer_name in arch.layer_names:
+        if layer_name not in layer_outputs:
+            continue
+
+        hidden = layer_outputs[layer_name].float()
+
+        # Take the last token position
+        if hidden.dim() == 3:
+            hidden = hidden[:, -1, :]  # (batch, hidden)
+        elif hidden.dim() == 2:
+            hidden = hidden[-1:, :]
+
+        # Apply final norm if found
+        if final_norm is not None:
+            hidden = final_norm(hidden)
+
+        # Project through unembedding: logits = hidden @ W^T
+        logits = hidden @ unembed_weight.float().T  # (batch, vocab)
+        probs = torch.softmax(logits, dim=-1)
+
+        top5_probs, top5_ids = probs[0].topk(5)
+        top5_tokens = [model._tokenizer.decode([tid]) for tid in top5_ids.tolist()]
+        top5_probs_list = top5_probs.tolist()
+
+        predictions.append({
+            "layer_name": layer_name,
+            "top1_token": top5_tokens[0],
+            "top1_prob": top5_probs_list[0],
+            "top5_tokens": top5_tokens,
+            "top5_probs": top5_probs_list,
+        })
+
+    model_name = arch.arch_family or "model"
+    render_lens(predictions, model_name)
+
+    if save is not None:
+        from interpkit.core.plot import plot_lens
+
+        plot_lens(predictions, save_path=save)
+
+    return predictions
+
+
+def _find_final_norm(model: torch.nn.Module, arch: Any) -> torch.nn.Module | None:
+    """Try to find the final layer norm applied before the LM head."""
+    import re
+
+    norm_pattern = re.compile(
+        r"^(model\.norm|transformer\.ln_f|gpt_neox\.final_layer_norm|"
+        r"model\.final_layernorm|backbone\.norm_f)$",
+        re.IGNORECASE,
+    )
+    for name, mod in model.named_modules():
+        if norm_pattern.match(name):
+            return mod
+
+    # Generic fallback: look for a top-level norm module
+    for name, mod in model.named_modules():
+        if name.count(".") <= 1 and isinstance(mod, (torch.nn.LayerNorm,)):
+            type_name = type(mod).__name__
+            if "norm" in name.lower() or "Norm" in type_name:
+                return mod
+
+    return None

interpkit/ops/patch.py

@@ -0,0 +1,112 @@
+"""patch — activation patching at a named module between clean and corrupted inputs."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+if TYPE_CHECKING:
+    from interpkit.core.model import Model
+
+
+def _get_module(model: torch.nn.Module, name: str) -> torch.nn.Module:
+    parts = name.split(".")
+    mod = model
+    for part in parts:
+        mod = getattr(mod, part)
+    return mod
+
+
+def run_patch(
+    model: "Model",
+    clean: Any,
+    corrupted: Any,
+    *,
+    at: str,
+) -> dict[str, Any]:
+    """Patch the output of module *at* from the clean run into the corrupted run.
+
+    Returns a dict with ``effect`` — a normalised scalar in [0, 1] measuring how
+    much the patched corrupted run's output shifted toward the clean output.
+    """
+    from interpkit.core.render import render_patch
+
+    clean_input, corrupted_input = model._prepare_pair(clean, corrupted)
+
+    # 1. Clean forward — cache the target module's output
+    cached_activation: list[torch.Tensor] = []
+
+    target_mod = _get_module(model._model, at)
+
+    def _cache_hook(_mod: torch.nn.Module, _inp: Any, output: Any) -> None:
+        if isinstance(output, torch.Tensor):
+            cached_activation.append(output.detach().clone())
+        elif isinstance(output, (tuple, list)):
+            cached_activation.append(output[0].detach().clone())
+
+    handle = target_mod.register_forward_hook(_cache_hook)
+    clean_logits = model._forward(clean_input)
+    handle.remove()
+
+    if not cached_activation:
+        raise RuntimeError(f"Module '{at}' produced no tensor output during clean forward pass.")
+
+    # 2. Corrupted forward (baseline)
+    corrupted_logits = model._forward(corrupted_input)
+
+    # 3. Patched forward — replace target module's output with cached clean activation
+    def _patch_hook(_mod: torch.nn.Module, _inp: Any, output: Any) -> Any:
+        if isinstance(output, torch.Tensor):
+            return cached_activation[0]
+        elif isinstance(output, (tuple, list)):
+            return (cached_activation[0],) + tuple(output[1:])
+        return output
+
+    handle = target_mod.register_forward_hook(_patch_hook)
+    patched_logits = model._forward(corrupted_input)
+    handle.remove()
+
+    # 4. Compute normalised effect
+    effect = _compute_effect(clean_logits, corrupted_logits, patched_logits)
+
+    result = {
+        "module": at,
+        "effect": effect,
+        "clean_logits": clean_logits,
+        "corrupted_logits": corrupted_logits,
+        "patched_logits": patched_logits,
+    }
+    render_patch(result)
+    return result
+
+
+def _compute_effect(
+    clean: torch.Tensor,
+    corrupted: torch.Tensor,
+    patched: torch.Tensor,
+) -> float:
+    """Normalised patching effect: 0 = patched == corrupted, 1 = patched == clean."""
+    # Use KL divergence on the last-token logits as the distance metric
+    clean_flat = clean.view(-1, clean.shape[-1]).float()
+    corrupted_flat = corrupted.view(-1, corrupted.shape[-1]).float()
+    patched_flat = patched.view(-1, patched.shape[-1]).float()
+
+    # Take last position for sequence models
+    if clean_flat.shape[0] > 1:
+        clean_flat = clean_flat[-1:]
+        corrupted_flat = corrupted_flat[-1:]
+        patched_flat = patched_flat[-1:]
+
+    clean_probs = torch.softmax(clean_flat, dim=-1)
+    corrupted_probs = torch.softmax(corrupted_flat, dim=-1)
+    patched_probs = torch.softmax(patched_flat, dim=-1)
+
+    dist_corrupted_clean = torch.norm(corrupted_probs - clean_probs).item()
+    dist_patched_clean = torch.norm(patched_probs - clean_probs).item()
+
+    if dist_corrupted_clean < 1e-8:
+        return 0.0
+
+    effect = 1.0 - (dist_patched_clean / dist_corrupted_clean)
+    return max(0.0, min(1.0, effect))

interpkit/ops/probe.py

@@ -0,0 +1,128 @@
+"""probe — train a linear probe on activations to test linear separability."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+from rich.console import Console
+
+if TYPE_CHECKING:
+    from interpkit.core.model import Model
+
+console = Console()
+
+
+def run_probe(
+    model: "Model",
+    texts: list[str],
+    labels: list[int],
+    *,
+    at: str,
+) -> dict[str, Any]:
+    """Train a linear probe on activations at module *at*.
+
+    Uses LogisticRegression from scikit-learn. Falls back to a simple
+    torch-based probe if sklearn is not installed.
+    """
+    from interpkit.core.render import render_probe
+    from interpkit.ops.activations import run_activations
+
+    if len(texts) != len(labels):
+        raise ValueError(f"texts ({len(texts)}) and labels ({len(labels)}) must have the same length.")
+
+    # Extract activations for all texts
+    features = []
+    for text in texts:
+        act = run_activations(model, text, at=at, print_stats=False)
+        # Take last-token hidden state for sequence models
+        if act.dim() == 3:
+            vec = act[0, -1, :]  # (hidden,)
+        elif act.dim() == 2:
+            vec = act[-1, :]
+        else:
+            vec = act.view(-1)
+        features.append(vec.cpu().float().numpy())
+
+    import numpy as np
+
+    X = np.stack(features)
+    y = np.array(labels)
+
+    try:
+        result = _probe_sklearn(X, y)
+    except ImportError:
+        result = _probe_torch(X, y)
+
+    result["module"] = at
+    render_probe(result)
+    return result
+
+
+def _probe_sklearn(X: Any, y: Any) -> dict[str, Any]:
+    from sklearn.linear_model import LogisticRegression
+    from sklearn.model_selection import cross_val_score
+
+    n_samples = len(y)
+
+    if n_samples >= 10:
+        cv_folds = min(5, n_samples)
+        clf = LogisticRegression(max_iter=1000, solver="lbfgs")
+        scores = cross_val_score(clf, X, y, cv=cv_folds, scoring="accuracy")
+        accuracy = float(scores.mean())
+    else:
+        accuracy = None
+
+    # Train on full data for feature analysis
+    clf = LogisticRegression(max_iter=1000, solver="lbfgs")
+    clf.fit(X, y)
+    train_accuracy = float(clf.score(X, y))
+
+    # Top features by weight magnitude
+    weights = clf.coef_[0] if clf.coef_.ndim == 2 else clf.coef_
+    top_indices = list(reversed(sorted(range(len(weights)), key=lambda i: abs(weights[i]))))[:20]
+    top_features = [(int(i), float(weights[i])) for i in top_indices]
+
+    return {
+        "accuracy": accuracy if accuracy is not None else train_accuracy,
+        "train_accuracy": train_accuracy,
+        "top_features": top_features,
+    }
+
+
+def _probe_torch(X: Any, y: Any) -> dict[str, Any]:
+    """Fallback probe using pure PyTorch when sklearn is not available."""
+    import numpy as np
+
+    X_t = torch.tensor(X, dtype=torch.float32)
+    y_t = torch.tensor(y, dtype=torch.long)
+
+    n_features = X_t.shape[1]
+    n_classes = len(set(y))
+
+    linear = torch.nn.Linear(n_features, n_classes)
+    optimizer = torch.optim.Adam(linear.parameters(), lr=0.01)
+    criterion = torch.nn.CrossEntropyLoss()
+
+    linear.train()
+    for _ in range(500):
+        optimizer.zero_grad()
+        logits = linear(X_t)
+        loss = criterion(logits, y_t)
+        loss.backward()
+        optimizer.step()
+
+    linear.eval()
+    with torch.no_grad():
+        preds = linear(X_t).argmax(dim=-1)
+        train_accuracy = float((preds == y_t).float().mean().item())
+
+    weights = linear.weight.detach()[0].numpy() if n_classes == 2 else linear.weight.detach().mean(dim=0).numpy()
+    top_indices = list(reversed(sorted(range(len(weights)), key=lambda i: abs(weights[i]))))[:20]
+    top_features = [(int(i), float(weights[i])) for i in top_indices]
+
+    return {
+        "accuracy": train_accuracy,
+        "train_accuracy": train_accuracy,
+        "top_features": top_features,
+    }

interpkit/ops/sae.py

@@ -0,0 +1,212 @@
+"""sae — load pre-trained Sparse Autoencoders and decompose activations into features."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from interpkit.ops.patch import _get_module
+
+if TYPE_CHECKING:
+    from interpkit.core.model import Model
+
+
+@dataclass
+class SAE:
+    """A loaded Sparse Autoencoder with weights ready for inference."""
+
+    W_enc: torch.Tensor
+    W_dec: torch.Tensor
+    b_enc: torch.Tensor
+    b_dec: torch.Tensor
+    d_in: int = 0
+    d_sae: int = 0
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.relu((x - self.b_dec) @ self.W_enc + self.b_enc)
+
+    def decode(self, features: torch.Tensor) -> torch.Tensor:
+        return features @ self.W_dec + self.b_dec
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return (features, reconstruction)."""
+        features = self.encode(x)
+        x_hat = self.decode(features)
+        return features, x_hat
+
+
+def load_sae(hf_id: str, *, device: str | torch.device = "cpu") -> SAE:
+    """Download and load a Sparse Autoencoder from HuggingFace.
+
+    Expects the repo to contain ``sae_weights.safetensors`` (or ``.pt``)
+    and optionally ``cfg.json`` with metadata.
+    """
+    from huggingface_hub import hf_hub_download
+
+    device = torch.device(device)
+
+    # Try safetensors first, fall back to .pt
+    weights = _download_weights(hf_id)
+
+    required_keys = {"W_enc", "W_dec", "b_enc", "b_dec"}
+    missing = required_keys - set(weights.keys())
+    if missing:
+        raise KeyError(
+            f"SAE weights from {hf_id!r} are missing keys: {missing}. "
+            f"Found keys: {list(weights.keys())}. "
+            f"interpkit expects the SAELens format (W_enc, W_dec, b_enc, b_dec)."
+        )
+
+    W_enc = weights["W_enc"].to(device).float()
+    W_dec = weights["W_dec"].to(device).float()
+    b_enc = weights["b_enc"].to(device).float()
+    b_dec = weights["b_dec"].to(device).float()
+
+    metadata = _download_config(hf_id)
+
+    d_in = W_enc.shape[0]
+    d_sae = W_enc.shape[1]
+
+    return SAE(
+        W_enc=W_enc,
+        W_dec=W_dec,
+        b_enc=b_enc,
+        b_dec=b_dec,
+        d_in=d_in,
+        d_sae=d_sae,
+        metadata=metadata,
+    )
+
+
+def load_sae_from_tensors(
+    W_enc: torch.Tensor,
+    W_dec: torch.Tensor,
+    b_enc: torch.Tensor,
+    b_dec: torch.Tensor,
+    *,
+    metadata: dict[str, Any] | None = None,
+) -> SAE:
+    """Create an SAE from raw weight tensors (useful for testing)."""
+    return SAE(
+        W_enc=W_enc,
+        W_dec=W_dec,
+        b_enc=b_enc,
+        b_dec=b_dec,
+        d_in=W_enc.shape[0],
+        d_sae=W_enc.shape[1],
+        metadata=metadata or {},
+    )
+
+
+def run_features(
+    model: "Model",
+    input_data: Any,
+    *,
+    at: str,
+    sae: SAE,
+    top_k: int = 20,
+    print_results: bool = True,
+) -> dict[str, Any]:
+    """Decompose activations at *at* through the SAE and return top features.
+
+    Returns a dict with ``top_features``, ``reconstruction_error``, ``sparsity``.
+    """
+    from interpkit.ops.activations import run_activations
+
+    act = run_activations(model, input_data, at=at, print_stats=False)
+    if not isinstance(act, torch.Tensor):
+        raise TypeError(f"Expected tensor from activations, got {type(act).__name__}")
+
+    # Flatten to 2D: (batch * seq, d_model)
+    if act.dim() == 1:
+        flat = act.unsqueeze(0).float()
+    else:
+        flat = act.view(-1, act.shape[-1]).float()
+
+    if flat.shape[-1] != sae.d_in:
+        raise ValueError(
+            f"Activation dim ({flat.shape[-1]}) does not match SAE input dim ({sae.d_in}). "
+            f"Make sure the SAE was trained on the same layer."
+        )
+
+    features, x_hat = sae.forward(flat)
+
+    # Reconstruction error (mean L2 across positions)
+    recon_error = (flat - x_hat).norm(dim=-1).mean().item()
+
+    # Sparsity: fraction of features that are zero
+    sparsity = (features == 0).float().mean().item()
+
+    # Top-K features by mean activation (across all positions)
+    mean_activations = features.mean(dim=0)
+    topk_vals, topk_idxs = mean_activations.topk(min(top_k, sae.d_sae))
+
+    top_features = [
+        (idx.item(), val.item())
+        for idx, val in zip(topk_idxs, topk_vals)
+    ]
+
+    result = {
+        "module": at,
+        "top_features": top_features,
+        "reconstruction_error": recon_error,
+        "sparsity": sparsity,
+        "num_active_features": int((mean_activations > 0).sum().item()),
+        "total_features": sae.d_sae,
+        "feature_activations": features.detach(),
+    }
+
+    if print_results:
+        from interpkit.core.render import render_features
+
+        render_features(result)
+
+    return result
+
+
+def _download_weights(hf_id: str) -> dict[str, torch.Tensor]:
+    """Download SAE weights from HuggingFace."""
+    from huggingface_hub import hf_hub_download
+    from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError
+
+    # Try safetensors first
+    try:
+        path = hf_hub_download(hf_id, filename="sae_weights.safetensors")
+        from safetensors.torch import load_file
+
+        return load_file(path)
+    except (EntryNotFoundError, FileNotFoundError):
+        pass
+    except RepositoryNotFoundError:
+        raise FileNotFoundError(
+            f"HuggingFace repository {hf_id!r} not found. "
+            f"Check the repo ID and your network/auth settings."
+        )
+
+    # Fall back to .pt
+    try:
+        path = hf_hub_download(hf_id, filename="sae_weights.pt")
+        return torch.load(path, map_location="cpu", weights_only=True)
+    except (EntryNotFoundError, FileNotFoundError):
+        pass
+
+    raise FileNotFoundError(
+        f"Could not find sae_weights.safetensors or sae_weights.pt in {hf_id!r}. "
+        f"The HF repo should contain one of these files."
+    )
+
+
+def _download_config(hf_id: str) -> dict[str, Any]:
+    """Download SAE config from HuggingFace (optional)."""
+    from huggingface_hub import hf_hub_download
+
+    try:
+        path = hf_hub_download(hf_id, filename="cfg.json")
+        return json.loads(Path(path).read_text())
+    except Exception:
+        return {}