interpkit 0.1.0__py3-none-any.whl

interpkit/ops/ablate.py

@@ -0,0 +1,90 @@
+"""ablate — zero or mean ablate a module and measure the effect on output."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from interpkit.ops.patch import _get_module
+
+if TYPE_CHECKING:
+    from interpkit.core.model import Model
+
+
+def run_ablate(
+    model: "Model",
+    input_data: Any,
+    *,
+    at: str,
+    method: str = "zero",
+) -> dict[str, Any]:
+    """Ablate module *at* and measure the effect on output logits.
+
+    Parameters
+    ----------
+    method:
+        ``"zero"`` replaces the module output with zeros.
+        ``"mean"`` replaces it with the mean activation across the sequence dimension.
+    """
+    from interpkit.core.render import render_ablate
+
+    model_input = model._prepare(input_data)
+    target_mod = _get_module(model._model, at)
+
+    # 1. Clean forward — get baseline logits
+    clean_logits = model._forward(model_input)
+
+    # 2. Ablated forward
+    def _ablate_hook(_mod: torch.nn.Module, _inp: Any, output: Any) -> Any:
+        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 None:
+            return output
+
+        if method == "zero":
+            replacement = torch.zeros_like(t)
+        elif method == "mean":
+            if t.dim() >= 3:
+                replacement = t.mean(dim=-2, keepdim=True).expand_as(t)
+            else:
+                replacement = t.mean(dim=-1, keepdim=True).expand_as(t)
+        else:
+            raise ValueError(f"Unknown ablation method: {method!r}. Use 'zero' or 'mean'.")
+
+        if isinstance(output, torch.Tensor):
+            return replacement
+        return (replacement,) + tuple(output[1:])
+
+    handle = target_mod.register_forward_hook(_ablate_hook)
+    ablated_logits = model._forward(model_input)
+    handle.remove()
+
+    effect = _compute_ablation_effect(clean_logits, ablated_logits)
+
+    result = {
+        "module": at,
+        "method": method,
+        "effect": effect,
+        "clean_logits": clean_logits,
+        "ablated_logits": ablated_logits,
+    }
+    render_ablate(result)
+    return result
+
+
+def _compute_ablation_effect(clean: torch.Tensor, ablated: torch.Tensor) -> float:
+    """Measure how much ablation changed the output (0 = no change, 1 = max change)."""
+    clean_flat = clean.view(-1, clean.shape[-1]).float()
+    ablated_flat = ablated.view(-1, ablated.shape[-1]).float()
+
+    if clean_flat.shape[0] > 1:
+        clean_flat = clean_flat[-1:]
+        ablated_flat = ablated_flat[-1:]
+
+    clean_probs = torch.softmax(clean_flat, dim=-1)
+    ablated_probs = torch.softmax(ablated_flat, dim=-1)
+
+    cosine_sim = torch.nn.functional.cosine_similarity(clean_probs, ablated_probs, dim=-1)
+    return (1.0 - cosine_sim.item())

interpkit/ops/activations.py

@@ -0,0 +1,67 @@
+"""activations — extract raw activation tensors at any named module."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from interpkit.ops.patch import _get_module
+
+if TYPE_CHECKING:
+    from interpkit.core.model import Model
+
+
+def run_activations(
+    model: "Model",
+    input_data: Any,
+    *,
+    at: str | list[str],
+    print_stats: bool = True,
+) -> dict[str, torch.Tensor] | torch.Tensor:
+    """Extract activations at one or more named modules.
+
+    Returns a single tensor if *at* is a string, or a dict if *at* is a list.
+    """
+    model_input = model._prepare(input_data)
+    single = isinstance(at, str)
+    module_names = [at] if single else list(at)
+
+    # Check activation cache first (pass prepared input to avoid re-tokenizing)
+    cached = model._get_cached(input_data, module_names, _prepared_input=model_input)
+    if cached is not None:
+        cache = cached
+    else:
+        cache = {}
+
+        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:
+                    cache[name] = t.detach().clone()
+            return hook_fn
+
+        hooks = []
+        for name in module_names:
+            mod = _get_module(model._model, name)
+            hooks.append(mod.register_forward_hook(_make_hook(name)))
+
+        with torch.no_grad():
+            model._forward(model_input)
+
+        for h in hooks:
+            h.remove()
+
+    if print_stats:
+        from interpkit.core.render import render_activations
+
+        render_activations(cache)
+
+    if single:
+        if at not in cache:
+            raise RuntimeError(f"Module '{at}' produced no tensor output.")
+        return cache[at]
+
+    return cache

interpkit/ops/attention.py

@@ -0,0 +1,234 @@
+"""attention — capture and display attention patterns for transformer models."""
+
+from __future__ import annotations
+
+import re
+from typing import TYPE_CHECKING, Any
+
+import torch
+import torch.nn.functional as F
+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_attention(
+    model: "Model",
+    input_data: Any,
+    *,
+    layer: int | None = None,
+    head: int | None = None,
+    save: str | None = None,
+    html: str | None = None,
+) -> list[dict[str, Any]] | None:
+    """Capture attention weights and display a summary.
+
+    Computes attention weights manually from Q/K projections via hooks,
+    since modern transformer implementations (SDPA, FlashAttention) don't
+    return attention weights.
+    """
+    from interpkit.core.render import render_attention
+
+    arch = model.arch_info
+    attn_modules = [m for m in arch.modules if m.role == "attention"]
+
+    if not attn_modules:
+        console.print(
+            "\n  [yellow]attention not available:[/yellow] no attention modules detected"
+            f" for {arch.arch_family or 'this model'}.\n"
+        )
+        return None
+
+    model_input = model._prepare(input_data)
+
+    tokens = None
+    if model._tokenizer is not None and isinstance(input_data, str):
+        encoded = model._tokenizer(input_data, return_tensors="pt")
+        token_ids = encoded["input_ids"][0].tolist()
+        tokens = model._tokenizer.convert_ids_to_tokens(token_ids)
+
+    # Capture Q and K projections to compute attention weights manually
+    qk_cache: dict[str, dict[str, torch.Tensor]] = {}
+    hooks = []
+
+    for mod_info in attn_modules:
+        if layer is not None:
+            layer_match = re.search(r"\.(\d+)\.", mod_info.name)
+            if layer_match and int(layer_match.group(1)) != layer:
+                continue
+
+        attn_mod = _get_module(model._model, mod_info.name)
+
+        # Find Q/K projection submodules
+        for child_name, child_mod in attn_mod.named_modules():
+            full_name = f"{mod_info.name}.{child_name}" if child_name else mod_info.name
+            is_qkv = any(p in child_name.lower() for p in ("c_attn", "qkv", "q_proj", "query"))
+            is_k = any(p in child_name.lower() for p in ("k_proj", "key"))
+
+            if is_qkv or is_k:
+                def _make_qk_hook(name: str, attn_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)) else None
+                        )
+                        if t is not None:
+                            qk_cache.setdefault(attn_name, {})[name] = t.detach()
+                    return hook_fn
+
+                hooks.append(child_mod.register_forward_hook(_make_qk_hook(child_name, mod_info.name)))
+
+    with torch.no_grad():
+        model._forward(model_input)
+
+    for h in hooks:
+        h.remove()
+
+    # Compute attention weights from cached Q/K
+    results: list[dict[str, Any]] = []
+
+    for attn_name, projections in qk_cache.items():
+        layer_match = re.search(r"\.(\d+)\.", attn_name)
+        layer_idx = int(layer_match.group(1)) if layer_match else 0
+
+        attn_weights = _compute_attention_from_projections(
+            projections, arch.num_attention_heads or 12
+        )
+
+        if attn_weights is None:
+            continue
+
+        num_heads = attn_weights.shape[0]
+        for head_idx in range(num_heads):
+            if head is not None and head_idx != head:
+                continue
+
+            head_attn = attn_weights[head_idx]
+            top_pairs = _get_top_pairs(head_attn, k=5)
+            entropy = _attention_entropy(head_attn)
+
+            results.append({
+                "layer": layer_idx,
+                "head": head_idx,
+                "top_pairs": top_pairs,
+                "entropy": entropy,
+                "weights": head_attn,
+            })
+
+    if not results:
+        console.print(
+            "\n  [yellow]attention:[/yellow] could not compute attention weights.\n"
+        )
+        return None
+
+    model_name = arch.arch_family or "model"
+    render_attention(results, tokens, model_name)
+
+    if save is not None:
+        from interpkit.core.plot import plot_attention, plot_attention_multi
+
+        if layer is not None and head is not None and len(results) == 1:
+            plot_attention(results[0]["weights"], tokens, layer=results[0]["layer"],
+                           head=results[0]["head"], save_path=save)
+        else:
+            plot_attention_multi(results, tokens, save_path=save)
+
+    if html is not None:
+        from interpkit.core.html import html_attention as gen_html_attention
+        from interpkit.core.html import save_html
+
+        serializable = []
+        for r in results:
+            entry = {**r}
+            w = r.get("weights")
+            if isinstance(w, torch.Tensor):
+                entry["weights"] = w.tolist()
+            serializable.append(entry)
+        save_html(gen_html_attention(serializable, tokens), html)
+
+    return results
+
+
+def _compute_attention_from_projections(
+    projections: dict[str, torch.Tensor],
+    num_heads: int,
+) -> torch.Tensor | None:
+    """Compute attention weights from captured QKV or Q/K projections."""
+    # GPT-2 style: c_attn produces [Q, K, V] concatenated
+    for key, tensor in projections.items():
+        if "c_attn" in key or "qkv" in key:
+            # tensor shape: (batch, seq, 3 * hidden) or (seq, 3 * hidden)
+            if tensor.dim() == 3:
+                tensor = tensor[0]  # drop batch
+            hidden = tensor.shape[-1] // 3
+            q, k, _v = tensor.split(hidden, dim=-1)
+            return _qk_to_attention(q, k, num_heads)
+
+    # Separate Q and K projections
+    q_tensor = None
+    k_tensor = None
+    for key, tensor in projections.items():
+        if "q_proj" in key or "query" in key:
+            q_tensor = tensor
+        elif "k_proj" in key or "key" in key:
+            k_tensor = tensor
+
+    if q_tensor is not None and k_tensor is not None:
+        if q_tensor.dim() == 3:
+            q_tensor = q_tensor[0]
+        if k_tensor.dim() == 3:
+            k_tensor = k_tensor[0]
+        return _qk_to_attention(q_tensor, k_tensor, num_heads)
+
+    return None
+
+
+def _qk_to_attention(
+    q: torch.Tensor, k: torch.Tensor, num_heads: int,
+) -> torch.Tensor:
+    """Compute attention weights from Q and K tensors.
+
+    q, k: (seq_len, hidden_size)
+    Returns: (num_heads, seq_len, seq_len)
+    """
+    seq_len, hidden = q.shape
+    head_dim = hidden // num_heads
+
+    q = q.view(seq_len, num_heads, head_dim).transpose(0, 1)  # (heads, seq, head_dim)
+    k = k.view(seq_len, num_heads, head_dim).transpose(0, 1)
+
+    scale = head_dim ** 0.5
+    scores = torch.matmul(q, k.transpose(-2, -1)) / scale  # (heads, seq, seq)
+
+    # Apply causal mask
+    causal_mask = torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=q.device), diagonal=1)
+    scores.masked_fill_(causal_mask.unsqueeze(0), float("-inf"))
+
+    return F.softmax(scores, dim=-1)
+
+
+def _get_top_pairs(
+    attn: torch.Tensor, k: int = 5,
+) -> list[tuple[int, int, float]]:
+    """Find top-k (source_pos, target_pos, score) pairs in an attention matrix."""
+    flat = attn.view(-1)
+    topk_vals, topk_idxs = flat.topk(min(k, flat.numel()))
+    seq_len = attn.shape[-1]
+    pairs = []
+    for val, idx in zip(topk_vals.tolist(), topk_idxs.tolist()):
+        src = idx // seq_len
+        tgt = idx % seq_len
+        pairs.append((src, tgt, val))
+    return pairs
+
+
+def _attention_entropy(attn: torch.Tensor) -> float:
+    """Mean entropy of attention distributions across query positions."""
+    eps = 1e-10
+    log_attn = torch.log(attn + eps)
+    entropy_per_query = -(attn * log_attn).sum(dim=-1)
+    return entropy_per_query.mean().item()

interpkit/ops/attribute.py

@@ -0,0 +1,206 @@
+"""attribute — gradient saliency over input tokens or pixels."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+if TYPE_CHECKING:
+    from interpkit.core.model import Model
+
+
+def run_attribute(
+    model: "Model",
+    input_data: Any,
+    *,
+    target: int | None = None,
+    save: str | None = None,
+    html: str | None = None,
+) -> None:
+    """Compute gradient-based saliency and render results.
+
+    For text inputs: shows coloured tokens by importance.
+    For image/tensor inputs: saves a heatmap to disk.
+    """
+    is_text = isinstance(input_data, str) and not _is_image_path(input_data)
+    is_image = isinstance(input_data, str) and _is_image_path(input_data)
+
+    if is_text:
+        _attribute_text(model, input_data, target=target, save=save, html=html)
+    elif is_image:
+        _attribute_image(model, input_data, target=target, save=save)
+    else:
+        _attribute_tensor(model, input_data, target=target)
+
+
+def _attribute_text(model: "Model", text: str, *, target: int | None, save: str | None = None, html: str | None = None) -> None:
+    from interpkit.core.render import render_attribution_tokens
+
+    if model._tokenizer is None:
+        raise ValueError("No tokenizer available for text attribution.")
+
+    encoded = model._tokenizer(text, return_tensors="pt")
+    input_ids = encoded["input_ids"].to(model._device)
+
+    # Get embedding layer
+    embed_layer = _find_embedding(model._model)
+    if embed_layer is None:
+        raise RuntimeError("Could not find embedding layer for gradient attribution.")
+
+    embeddings = embed_layer(input_ids)
+    embeddings = embeddings.detach().requires_grad_(True)
+
+    # Replace embedding output with our gradient-tracked version
+    original_forward = embed_layer.forward
+
+    def _patched_forward(*args: Any, **kwargs: Any) -> torch.Tensor:
+        return embeddings
+
+    embed_layer.forward = _patched_forward  # type: ignore[assignment]
+
+    try:
+        model_kwargs = {k: v.to(model._device) for k, v in encoded.items() if k != "input_ids"}
+        out = model._model(input_ids, **model_kwargs)
+
+        logits = out.logits if hasattr(out, "logits") else (out[0] if isinstance(out, (tuple, list)) else out)
+
+        # Pick target: last-position argmax if not specified
+        if logits.dim() == 3:
+            logits_last = logits[0, -1, :]
+        else:
+            logits_last = logits[0]
+
+        if target is None:
+            target = logits_last.argmax().item()
+
+        score = logits_last[target]
+        score.backward()
+    finally:
+        embed_layer.forward = original_forward  # type: ignore[assignment]
+
+    if embeddings.grad is None:
+        raise RuntimeError("Gradient computation failed — no gradients on embeddings.")
+
+    # Per-token importance: L2 norm of gradient over the embedding dimension
+    token_grads = embeddings.grad[0]  # (seq_len, hidden)
+    token_scores = token_grads.norm(dim=-1).tolist()  # (seq_len,)
+
+    tokens = model._tokenizer.convert_ids_to_tokens(input_ids[0].tolist())
+    render_attribution_tokens(tokens, token_scores)
+
+    if save is not None:
+        from interpkit.core.plot import plot_attribution
+
+        plot_attribution(tokens, token_scores, save_path=save)
+
+    if html is not None:
+        from interpkit.core.html import html_attribution as gen_html_attribution
+        from interpkit.core.html import save_html
+
+        save_html(gen_html_attribution(tokens, token_scores), html)
+
+
+def _attribute_image(model: "Model", image_path: str, *, target: int | None, save: str | None = None) -> None:
+    from interpkit.core.inputs import _load_image
+    from interpkit.core.render import render_attribution_heatmap
+
+    processed = _load_image(
+        image_path,
+        image_processor=model._image_processor,
+        device=model._device,
+    )
+
+    if isinstance(processed, dict):
+        pixel_key = "pixel_values" if "pixel_values" in processed else list(processed.keys())[0]
+        pixel_values = processed[pixel_key].requires_grad_(True)
+        model_input = {**processed, pixel_key: pixel_values}
+        out = model._model(**model_input)
+    else:
+        pixel_values = processed.requires_grad_(True)
+        out = model._model(pixel_values)
+
+    logits = out.logits if hasattr(out, "logits") else (out[0] if isinstance(out, (tuple, list)) else out)
+
+    if logits.dim() > 1:
+        logits_flat = logits[0]
+    else:
+        logits_flat = logits
+
+    if target is None:
+        target = logits_flat.argmax().item()
+
+    score = logits_flat[target]
+    score.backward()
+
+    if pixel_values.grad is None:
+        raise RuntimeError("Gradient computation failed — no gradients on pixel values.")
+
+    out_path = save or "attribution_heatmap.png"
+    render_attribution_heatmap(pixel_values.grad[0], output_path=out_path)
+
+
+def _attribute_tensor(model: "Model", tensor_input: Any, *, target: int | None) -> None:
+    from interpkit.core.render import render_attribution_tokens
+
+    inp = model._prepare(tensor_input)
+
+    if isinstance(inp, dict):
+        # Pick the first tensor-valued entry
+        for k, v in inp.items():
+            if isinstance(v, torch.Tensor) and v.is_floating_point():
+                inp[k] = v.requires_grad_(True)
+                grad_tensor = v
+                break
+        else:
+            raise ValueError("No floating-point tensor found in input dict.")
+        out = model._model(**inp)
+    else:
+        inp = inp.requires_grad_(True)
+        grad_tensor = inp
+        out = model._model(inp)
+
+    logits = out.logits if hasattr(out, "logits") else (out[0] if isinstance(out, (tuple, list)) else out)
+    logits_flat = logits.view(-1)
+
+    if target is None:
+        target = logits_flat.argmax().item()
+
+    score = logits_flat[target]
+    score.backward()
+
+    if grad_tensor.grad is None:
+        raise RuntimeError("Gradient computation failed.")
+
+    # Flatten to feature importance
+    grad = grad_tensor.grad.detach().float()
+    if grad.dim() > 1:
+        feature_scores = grad.view(grad.shape[0], -1).norm(dim=0).tolist()
+    else:
+        feature_scores = grad.abs().tolist()
+
+    labels = [f"feat_{i}" for i in range(len(feature_scores))]
+    render_attribution_tokens(labels, feature_scores)
+
+
+def _find_embedding(model: torch.nn.Module) -> torch.nn.Module | None:
+    """Find the token embedding layer."""
+    for name, mod in model.named_modules():
+        if isinstance(mod, torch.nn.Embedding) and "token" not in name.lower().replace("token", ""):
+            pass
+        if isinstance(mod, torch.nn.Embedding):
+            # Pick the largest embedding (usually token embeddings, not position)
+            if mod.num_embeddings > 1000:
+                return mod
+
+    # Fallback: first embedding
+    for _name, mod in model.named_modules():
+        if isinstance(mod, torch.nn.Embedding):
+            return mod
+
+    return None
+
+
+def _is_image_path(s: str) -> bool:
+    import os
+    return os.path.splitext(s)[1].lower() in {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".webp"}

interpkit/ops/diff.py

@@ -0,0 +1,79 @@
+"""diff — compare activations between two models on the same input."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+if TYPE_CHECKING:
+    from interpkit.core.model import Model
+
+
+def run_diff(
+    model_a: "Model",
+    model_b: "Model",
+    input_data: Any,
+    *,
+    save: str | None = None,
+) -> list[dict[str, Any]]:
+    """Compare activations between two models at all discovered layers.
+
+    Returns a list of dicts sorted by cosine distance (highest change first).
+    """
+    from interpkit.core.render import render_diff
+    from interpkit.ops.activations import run_activations
+
+    # Find shared layer-like modules
+    layers_a = set(model_a.arch_info.layer_names or [])
+    layers_b = set(model_b.arch_info.layer_names or [])
+
+    # If no layers detected, fall back to all named modules
+    if not layers_a:
+        layers_a = {m.name for m in model_a.arch_info.modules if m.param_count > 0}
+    if not layers_b:
+        layers_b = {m.name for m in model_b.arch_info.modules if m.param_count > 0}
+
+    shared_layers = sorted(layers_a & layers_b)
+
+    if not shared_layers:
+        from rich.console import Console
+        Console().print("\n  [yellow]diff:[/yellow] no shared modules found between the two models.\n")
+        return []
+
+    acts_a = run_activations(model_a, input_data, at=shared_layers, print_stats=False)
+    acts_b = run_activations(model_b, input_data, at=shared_layers, print_stats=False)
+
+    results: list[dict[str, Any]] = []
+    for name in shared_layers:
+        if name not in acts_a or name not in acts_b:
+            continue
+
+        a = acts_a[name].float().view(-1)
+        b = acts_b[name].float().view(-1)
+
+        if a.shape != b.shape:
+            min_size = min(a.numel(), b.numel())
+            a = a[:min_size]
+            b = b[:min_size]
+
+        cosine_sim = torch.nn.functional.cosine_similarity(a.unsqueeze(0), b.unsqueeze(0), dim=-1)
+        distance = (1.0 - cosine_sim.item())
+
+        results.append({
+            "module": name,
+            "distance": distance,
+        })
+
+    results.sort(key=lambda r: r["distance"], reverse=True)
+
+    model_a_name = model_a.arch_info.arch_family or "model_a"
+    model_b_name = model_b.arch_info.arch_family or "model_b"
+    render_diff(results, model_a_name, model_b_name)
+
+    if save is not None:
+        from interpkit.core.plot import plot_diff
+
+        plot_diff(results, model_a_name=model_a_name, model_b_name=model_b_name, save_path=save)
+
+    return results

interpkit/ops/inspect.py

@@ -0,0 +1,14 @@
+"""inspect — module tree with types, param counts, output shapes, and detected roles."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from interpkit.core.model import Model
+
+
+def run_inspect(model: "Model") -> None:
+    from interpkit.core.render import render_inspect
+
+    render_inspect(model.arch_info)