interpkit 0.1.0__py3-none-any.whl

interpkit/core/render.py

@@ -0,0 +1,465 @@
+"""Terminal rendering — rich tables, trees, unicode bar charts, heatmap export."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+from rich.console import Console
+from rich.table import Table
+from rich.tree import Tree
+
+if TYPE_CHECKING:
+    from interpkit.core.discovery import ModelArchInfo, ModuleInfo
+
+console = Console()
+
+_ROLE_TAGS = {
+    "attention": "[attn]",
+    "mlp": "[mlp]",
+    "head": "[head]",
+    "norm": "[norm]",
+    "embed": "[embed]",
+}
+
+
+# ------------------------------------------------------------------
+# Inspect rendering
+# ------------------------------------------------------------------
+
+
+def render_inspect(arch_info: "ModelArchInfo") -> None:
+    """Print a module tree with types, param counts, and detected roles."""
+    header_parts = []
+    if arch_info.arch_family:
+        header_parts.append(arch_info.arch_family)
+    if arch_info.num_layers is not None:
+        header_parts.append(f"{arch_info.num_layers} layers")
+    if arch_info.hidden_size is not None:
+        header_parts.append(f"hidden={arch_info.hidden_size}")
+    if arch_info.vocab_size is not None:
+        header_parts.append(f"vocab={arch_info.vocab_size}")
+
+    total_params = sum(m.param_count for m in arch_info.modules)
+    header_parts.append(f"{_format_params(total_params)} params total")
+
+    console.print(f"\n[bold]{' | '.join(header_parts)}[/bold]")
+
+    table = Table(show_header=True, header_style="bold", show_lines=False, pad_edge=False)
+    table.add_column("Module", style="cyan", no_wrap=True)
+    table.add_column("Type", style="dim")
+    table.add_column("Params", justify="right")
+    table.add_column("Output Shape", style="dim")
+    table.add_column("Role", style="bold yellow")
+
+    for m in arch_info.modules:
+        role_tag = _ROLE_TAGS.get(m.role or "", "")
+        shape_str = str(m.output_shape) if m.output_shape else ""
+        param_str = _format_params(m.param_count) if m.param_count > 0 else ""
+        table.add_row(m.name, m.type_name, param_str, shape_str, role_tag)
+
+    console.print(table)
+    console.print()
+
+
+# ------------------------------------------------------------------
+# Causal trace rendering
+# ------------------------------------------------------------------
+
+
+def render_trace(
+    results: list[dict[str, Any]],
+    model_name: str,
+    total_modules: int,
+    top_k: int | None,
+) -> None:
+    """Print a ranked bar chart of causal tracing results."""
+    scanned = len(results)
+    if top_k is not None:
+        title = f"Causal Trace: {model_name}  (top {scanned} of {total_modules} modules)"
+    else:
+        title = f"Causal Trace: {model_name}  ({total_modules} modules)"
+
+    console.print(f"\n[bold]{title}[/bold]")
+
+    if not results:
+        console.print("  No significant causal effects found.")
+        return
+
+    max_effect = max(r["effect"] for r in results) if results else 1.0
+    bar_width = 30
+
+    table = Table(show_header=False, show_lines=False, pad_edge=False, box=None)
+    table.add_column("Module", style="cyan", no_wrap=True, min_width=35)
+    table.add_column("Role", style="yellow", min_width=8)
+    table.add_column("Bar", no_wrap=True)
+    table.add_column("Effect", justify="right", style="bold")
+
+    for r in results:
+        fill = int(bar_width * r["effect"] / max_effect) if max_effect > 0 else 0
+        bar = "█" * fill
+        role = _ROLE_TAGS.get(r.get("role") or "", "")
+        table.add_row(r["module"], role, f"[green]{bar}[/green]", f"{r['effect']:.3f}")
+
+    console.print(table)
+
+    if results:
+        best = results[0]
+        console.print(f"\n  Top component: [bold cyan]{best['module']}[/bold cyan] (effect: {best['effect']:.3f})")
+
+    if top_k is not None and scanned < total_modules:
+        console.print(
+            f"  Run with --top-k 0 to scan all {total_modules} modules.\n"
+        )
+    else:
+        console.print()
+
+
+# ------------------------------------------------------------------
+# Logit lens rendering
+# ------------------------------------------------------------------
+
+
+def render_lens(
+    predictions: list[dict[str, Any]],
+    model_name: str,
+) -> None:
+    """Print logit lens top predictions per layer."""
+    console.print(f"\n[bold]Logit Lens: {model_name}[/bold]")
+
+    table = Table(show_header=True, header_style="bold", show_lines=False)
+    table.add_column("Layer", style="cyan")
+    table.add_column("Top-1 Token", style="bold")
+    table.add_column("Prob", justify="right")
+    table.add_column("Top-5 Tokens", style="dim")
+
+    for pred in predictions:
+        top5_str = ", ".join(
+            f"{tok} ({prob:.2f})" for tok, prob in zip(pred["top5_tokens"], pred["top5_probs"])
+        )
+        table.add_row(
+            pred["layer_name"],
+            pred["top1_token"],
+            f"{pred['top1_prob']:.3f}",
+            top5_str,
+        )
+
+    console.print(table)
+    console.print()
+
+
+# ------------------------------------------------------------------
+# Attribution rendering
+# ------------------------------------------------------------------
+
+
+def render_attribution_tokens(
+    tokens: list[str],
+    scores: list[float],
+) -> None:
+    """Print tokens coloured by attribution score (terminal)."""
+    console.print("\n[bold]Attribution (gradient saliency)[/bold]")
+
+    if not scores:
+        console.print("  No attribution scores computed.")
+        return
+
+    max_score = max(abs(s) for s in scores) if scores else 1.0
+
+    parts: list[str] = []
+    for tok, score in zip(tokens, scores):
+        intensity = abs(score) / max_score if max_score > 0 else 0
+        if intensity > 0.7:
+            parts.append(f"[bold red]{tok}[/bold red]")
+        elif intensity > 0.4:
+            parts.append(f"[yellow]{tok}[/yellow]")
+        elif intensity > 0.15:
+            parts.append(f"[dim]{tok}[/dim]")
+        else:
+            parts.append(tok)
+
+    console.print("  " + "".join(parts))
+
+    # Also show ranked list
+    ranked = sorted(zip(tokens, scores), key=lambda x: abs(x[1]), reverse=True)
+    console.print()
+    for tok, score in ranked[:10]:
+        bar_len = int(20 * abs(score) / max_score) if max_score > 0 else 0
+        bar = "█" * bar_len
+        console.print(f"  {tok:>15s}  [green]{bar}[/green]  {score:.4f}")
+
+    console.print()
+
+
+def render_attribution_heatmap(
+    attribution: torch.Tensor,
+    output_path: str = "attribution_heatmap.png",
+) -> None:
+    """Save a vision attribution heatmap to a file."""
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    attr_np = attribution.detach().cpu().numpy()
+
+    # Collapse channel dim if present
+    if attr_np.ndim == 3:
+        attr_np = attr_np.mean(axis=0)
+    elif attr_np.ndim == 4:
+        attr_np = attr_np[0].mean(axis=0)
+
+    fig, ax = plt.subplots(1, 1, figsize=(6, 6))
+    im = ax.imshow(np.abs(attr_np), cmap="hot", interpolation="bilinear")
+    ax.set_title("Gradient Attribution")
+    ax.axis("off")
+    fig.colorbar(im, ax=ax, fraction=0.046)
+    fig.savefig(output_path, bbox_inches="tight", dpi=150)
+    plt.close(fig)
+
+    console.print(f"\n  Attribution heatmap saved to [bold]{output_path}[/bold]\n")
+
+
+# ------------------------------------------------------------------
+# Patch rendering
+# ------------------------------------------------------------------
+
+
+def render_patch(result: dict[str, Any]) -> None:
+    """Print the result of a single activation patch."""
+    console.print(f"\n[bold]Activation Patch at: {result['module']}[/bold]")
+    console.print(f"  Normalised effect: [bold]{result['effect']:.4f}[/bold]")
+    console.print()
+
+
+# ------------------------------------------------------------------
+# Activations rendering
+# ------------------------------------------------------------------
+
+
+def render_activations(cache: dict[str, torch.Tensor]) -> None:
+    """Print a summary table of extracted activations."""
+    console.print("\n[bold]Activations[/bold]")
+
+    table = Table(show_header=True, header_style="bold", show_lines=False)
+    table.add_column("Module", style="cyan", no_wrap=True)
+    table.add_column("Shape", style="dim")
+    table.add_column("Norm", justify="right")
+    table.add_column("Mean", justify="right")
+    table.add_column("Std", justify="right")
+    table.add_column("Min", justify="right")
+    table.add_column("Max", justify="right")
+
+    for name, tensor in cache.items():
+        t = tensor.float()
+        table.add_row(
+            name,
+            str(tuple(tensor.shape)),
+            f"{t.norm():.3f}",
+            f"{t.mean():.4f}",
+            f"{t.std():.4f}",
+            f"{t.min():.4f}",
+            f"{t.max():.4f}",
+        )
+
+    console.print(table)
+    console.print()
+
+
+# ------------------------------------------------------------------
+# Ablation rendering
+# ------------------------------------------------------------------
+
+
+def render_ablate(result: dict[str, Any]) -> None:
+    """Print the result of an ablation."""
+    method = result.get("method", "zero")
+    console.print(f"\n[bold]Ablation ({method}) at: {result['module']}[/bold]")
+    console.print(f"  Effect on output: [bold]{result['effect']:.4f}[/bold]")
+    console.print()
+
+
+# ------------------------------------------------------------------
+# Attention rendering
+# ------------------------------------------------------------------
+
+
+def render_attention(
+    attention_data: list[dict[str, Any]],
+    tokens: list[str] | None,
+    model_name: str,
+) -> None:
+    """Print attention summary per layer/head."""
+    console.print(f"\n[bold]Attention Patterns: {model_name}[/bold]")
+
+    if not attention_data:
+        console.print("  No attention data captured.")
+        return
+
+    table = Table(show_header=True, header_style="bold", show_lines=False)
+    table.add_column("Layer", style="cyan")
+    table.add_column("Head", style="cyan", justify="right")
+    table.add_column("Top Attention", style="dim")
+    table.add_column("Entropy", justify="right")
+
+    for entry in attention_data:
+        top_attn_parts = []
+        for src, tgt, score in entry.get("top_pairs", [])[:3]:
+            src_tok = tokens[src] if tokens and src < len(tokens) else str(src)
+            tgt_tok = tokens[tgt] if tokens and tgt < len(tokens) else str(tgt)
+            top_attn_parts.append(f"{src_tok}->{tgt_tok} ({score:.2f})")
+        table.add_row(
+            str(entry["layer"]),
+            str(entry["head"]),
+            ", ".join(top_attn_parts),
+            f"{entry.get('entropy', 0.0):.2f}",
+        )
+
+    console.print(table)
+    console.print()
+
+
+# ------------------------------------------------------------------
+# Steering rendering
+# ------------------------------------------------------------------
+
+
+def render_steer(
+    original_tokens: list[tuple[str, float]],
+    steered_tokens: list[tuple[str, float]],
+    module_name: str,
+    scale: float,
+) -> None:
+    """Print side-by-side comparison of top tokens with and without steering."""
+    console.print(f"\n[bold]Steering at: {module_name} (scale={scale})[/bold]")
+
+    table = Table(show_header=True, header_style="bold", show_lines=False)
+    table.add_column("Rank", justify="right", style="dim")
+    table.add_column("Original Token", style="cyan")
+    table.add_column("Prob", justify="right")
+    table.add_column("Steered Token", style="green")
+    table.add_column("Prob", justify="right")
+
+    n = max(len(original_tokens), len(steered_tokens))
+    for i in range(min(n, 10)):
+        orig_tok = original_tokens[i][0] if i < len(original_tokens) else ""
+        orig_prob = f"{original_tokens[i][1]:.3f}" if i < len(original_tokens) else ""
+        steer_tok = steered_tokens[i][0] if i < len(steered_tokens) else ""
+        steer_prob = f"{steered_tokens[i][1]:.3f}" if i < len(steered_tokens) else ""
+        table.add_row(str(i + 1), orig_tok, orig_prob, steer_tok, steer_prob)
+
+    console.print(table)
+    console.print()
+
+
+# ------------------------------------------------------------------
+# Probe rendering
+# ------------------------------------------------------------------
+
+
+def render_probe(result: dict[str, Any]) -> None:
+    """Print probe results — accuracy and top features."""
+    console.print(f"\n[bold]Linear Probe at: {result['module']}[/bold]")
+    console.print(f"  Accuracy: [bold]{result['accuracy']:.3f}[/bold]")
+
+    if result.get("train_accuracy") is not None:
+        console.print(f"  Train accuracy: {result['train_accuracy']:.3f}")
+
+    if result.get("top_features"):
+        console.print("\n  Top features by weight magnitude:")
+        for idx, weight in result["top_features"][:10]:
+            bar_len = int(20 * abs(weight) / abs(result["top_features"][0][1])) if result["top_features"][0][1] != 0 else 0
+            bar = "█" * bar_len
+            console.print(f"    dim {idx:>5d}  [green]{bar}[/green]  {weight:.4f}")
+
+    console.print()
+
+
+# ------------------------------------------------------------------
+# Diff rendering
+# ------------------------------------------------------------------
+
+
+def render_diff(
+    results: list[dict[str, Any]],
+    model_a_name: str,
+    model_b_name: str,
+) -> None:
+    """Print per-module activation distance between two models."""
+    console.print(f"\n[bold]Model Diff: {model_a_name} vs {model_b_name}[/bold]")
+
+    if not results:
+        console.print("  No differences computed.")
+        return
+
+    max_dist = max(r["distance"] for r in results) if results else 1.0
+    bar_width = 30
+
+    table = Table(show_header=False, show_lines=False, pad_edge=False, box=None)
+    table.add_column("Module", style="cyan", no_wrap=True, min_width=35)
+    table.add_column("Bar", no_wrap=True)
+    table.add_column("Cosine Dist", justify="right", style="bold")
+
+    for r in results:
+        fill = int(bar_width * r["distance"] / max_dist) if max_dist > 0 else 0
+        bar = "█" * fill
+        table.add_row(r["module"], f"[green]{bar}[/green]", f"{r['distance']:.4f}")
+
+    console.print(table)
+
+    if results:
+        best = results[0]
+        console.print(f"\n  Most changed: [bold cyan]{best['module']}[/bold cyan] (distance: {best['distance']:.4f})")
+
+    console.print()
+
+
+# ------------------------------------------------------------------
+# SAE features rendering
+# ------------------------------------------------------------------
+
+
+def render_features(result: dict[str, Any]) -> None:
+    """Print SAE feature decomposition results."""
+    console.print(f"\n[bold]SAE Features at: {result['module']}[/bold]")
+    console.print(
+        f"  Active features: [bold]{result['num_active_features']}[/bold] / {result['total_features']}  "
+        f"| Sparsity: {result['sparsity']:.2%}  "
+        f"| Reconstruction error: {result['reconstruction_error']:.4f}"
+    )
+
+    top = result.get("top_features", [])
+    if not top:
+        console.print("  No active features found.")
+        console.print()
+        return
+
+    max_val = max(abs(v) for _, v in top) if top else 1.0
+
+    table = Table(show_header=True, header_style="bold", show_lines=False)
+    table.add_column("Rank", justify="right", style="dim")
+    table.add_column("Feature", style="cyan", justify="right")
+    table.add_column("Activation", justify="right")
+    table.add_column("Bar", no_wrap=True)
+
+    bar_width = 25
+    for rank, (idx, val) in enumerate(top, 1):
+        fill = int(bar_width * abs(val) / max_val) if max_val > 0 else 0
+        bar = "█" * fill
+        table.add_row(str(rank), str(idx), f"{val:.4f}", f"[green]{bar}[/green]")
+
+    console.print(table)
+    console.print()
+
+
+# ------------------------------------------------------------------
+# Helpers
+# ------------------------------------------------------------------
+
+
+def _format_params(n: int) -> str:
+    if n >= 1_000_000_000:
+        return f"{n / 1_000_000_000:.1f}B"
+    if n >= 1_000_000:
+        return f"{n / 1_000_000:.1f}M"
+    if n >= 1_000:
+        return f"{n / 1_000:.1f}K"
+    return str(n)

interpkit/core/tl_compat.py

@@ -0,0 +1,174 @@
+"""TransformerLens interop — bidirectional name translation between native and TL hook names."""
+
+from __future__ import annotations
+
+import re
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from interpkit.core.discovery import ModelArchInfo
+
+# ── TL canonical hook names ──────────────────────────────────────
+#
+# TL standardizes all transformer architectures to:
+#   blocks.{N}.hook_resid_pre
+#   blocks.{N}.hook_resid_post
+#   blocks.{N}.attn.hook_q / hook_k / hook_v / hook_z / hook_result
+#   blocks.{N}.attn.hook_pattern
+#   blocks.{N}.hook_attn_out
+#   blocks.{N}.hook_mlp_out
+#   blocks.{N}.mlp.hook_pre / hook_post
+#   blocks.{N}.ln1 / ln2
+#
+# These map to native names like:
+#   transformer.h.{N}           -> blocks.{N}
+#   transformer.h.{N}.attn      -> blocks.{N}.attn
+#   transformer.h.{N}.mlp       -> blocks.{N}.mlp
+#   model.layers.{N}.self_attn  -> blocks.{N}.attn
+#   model.layers.{N}.mlp        -> blocks.{N}.mlp
+
+# Patterns for extracting layer index and component from native names
+_NATIVE_LAYER_RE = re.compile(
+    r"^(?P<prefix>.+?)[.\[](?P<idx>\d+)[.\]]*(?P<suffix>.*)$"
+)
+
+# Common native -> TL component mappings
+_COMPONENT_TO_TL: list[tuple[re.Pattern, str]] = [
+    (re.compile(r"\.self_attn$|\.attn$|\.attention$", re.I), ".attn"),
+    (re.compile(r"\.mlp$|\.ffn$|\.feed_forward$", re.I), ".mlp"),
+    (re.compile(r"\.ln_?1$|\.input_layernorm$", re.I), ".ln1"),
+    (re.compile(r"\.ln_?2$|\.post_attention_layernorm$", re.I), ".ln2"),
+    (re.compile(r"\.self_attn\.q_proj$|\.attn\.q_proj$", re.I), ".attn.hook_q"),
+    (re.compile(r"\.self_attn\.k_proj$|\.attn\.k_proj$", re.I), ".attn.hook_k"),
+    (re.compile(r"\.self_attn\.v_proj$|\.attn\.v_proj$", re.I), ".attn.hook_v"),
+    (re.compile(r"\.self_attn\.o_proj$|\.attn\.c_proj$|\.attn\.out_proj$", re.I), ".attn.hook_result"),
+]
+
+# TL hook -> native component suffix patterns
+_TL_TO_COMPONENT: list[tuple[str, list[str]]] = [
+    (".attn", ["attn", "self_attn", "attention"]),
+    (".mlp", ["mlp", "ffn", "feed_forward"]),
+    (".ln1", ["ln_1", "ln1", "input_layernorm"]),
+    (".ln2", ["ln_2", "ln2", "post_attention_layernorm"]),
+    (".attn.hook_q", ["attn.q_proj", "self_attn.q_proj"]),
+    (".attn.hook_k", ["attn.k_proj", "self_attn.k_proj"]),
+    (".attn.hook_v", ["attn.v_proj", "self_attn.v_proj"]),
+    (".attn.hook_result", ["attn.c_proj", "attn.out_proj", "self_attn.o_proj"]),
+]
+
+
+def to_tl_name(native_name: str, arch_info: "ModelArchInfo | None" = None) -> str:
+    """Translate a native PyTorch module name to the corresponding TL hook name.
+
+    Examples::
+
+        to_tl_name("transformer.h.8.mlp")   -> "blocks.8.mlp"
+        to_tl_name("transformer.h.8.attn")  -> "blocks.8.attn"
+        to_tl_name("model.layers.3.self_attn.q_proj") -> "blocks.3.attn.hook_q"
+    """
+    m = _NATIVE_LAYER_RE.match(native_name)
+    if m is None:
+        return native_name
+
+    idx = m.group("idx")
+    raw_suffix = m.group("suffix")
+
+    # Normalize suffix to always start with "." for pattern matching
+    clean_suffix = raw_suffix.lstrip(".")
+    dotted_suffix = f".{clean_suffix}" if clean_suffix else ""
+
+    # Try specific component mappings first
+    for pattern, tl_suffix in _COMPONENT_TO_TL:
+        if dotted_suffix and pattern.search(dotted_suffix):
+            return f"blocks.{idx}{tl_suffix}"
+
+    # Bare layer reference (e.g. "transformer.h.8")
+    if not clean_suffix:
+        return f"blocks.{idx}"
+
+    # Fallback: preserve suffix as-is under blocks.{N}
+    return f"blocks.{idx}.{clean_suffix}"
+
+
+def to_native_name(
+    tl_name: str,
+    arch_info: "ModelArchInfo | None" = None,
+) -> str:
+    """Translate a TL hook name back to the most likely native module name.
+
+    Requires ``arch_info`` from a loaded model to resolve the native layer prefix
+    (e.g. ``transformer.h`` vs ``model.layers``). Without it, returns a best-guess.
+
+    Examples::
+
+        to_native_name("blocks.8.mlp", arch_info) -> "transformer.h.8.mlp"
+        to_native_name("blocks.3.attn.hook_q", arch_info) -> "transformer.h.3.attn.q_proj"
+    """
+    # Parse TL name: blocks.{N}.{rest}
+    tl_match = re.match(r"^blocks\.(\d+)(?:\.(.+))?$", tl_name)
+    if tl_match is None:
+        return tl_name
+
+    idx = tl_match.group(1)
+    tl_suffix = tl_match.group(2) or ""
+
+    # Determine native layer prefix from arch_info
+    prefix = _infer_native_prefix(arch_info)
+
+    if not tl_suffix:
+        return f"{prefix}.{idx}"
+
+    # Try specific TL -> native mappings
+    for tl_component, native_candidates in _TL_TO_COMPONENT:
+        tl_component_clean = tl_component.lstrip(".")
+        if tl_suffix == tl_component_clean:
+            # Pick the first candidate that exists in the module tree, or fall back to first
+            if arch_info is not None:
+                module_names = {m.name for m in arch_info.modules}
+                for candidate in native_candidates:
+                    full = f"{prefix}.{idx}.{candidate}"
+                    if full in module_names:
+                        return full
+            return f"{prefix}.{idx}.{native_candidates[0]}"
+
+    # Strip TL-specific "hook_" prefixes for unknown suffixes
+    clean = re.sub(r"hook_", "", tl_suffix)
+    return f"{prefix}.{idx}.{clean}"
+
+
+def list_tl_hooks(model: Any) -> list[str]:
+    """List all TL hook point names on a HookedTransformer.
+
+    Returns an empty list if the model is not a HookedTransformer.
+    """
+    hook_dict = getattr(model, "hook_dict", None)
+    if hook_dict is not None:
+        return sorted(hook_dict.keys())
+
+    # Fallback: look for HookPoint modules
+    hooks = []
+    for name, mod in model.named_modules():
+        if type(mod).__name__ == "HookPoint":
+            hooks.append(name)
+    return sorted(hooks)
+
+
+def _infer_native_prefix(arch_info: "ModelArchInfo | None") -> str:
+    """Infer the native layer name prefix (e.g. 'transformer.h', 'model.layers')."""
+    if arch_info is None:
+        return "blocks"
+
+    if arch_info.layer_names:
+        first = arch_info.layer_names[0]
+        # Strip trailing .{digit} to get prefix
+        m = re.match(r"^(.+?)\.\d+$", first)
+        if m:
+            return m.group(1)
+
+    # Scan modules for repeating indexed patterns
+    for mod in arch_info.modules:
+        m = re.match(r"^(.+?)\.\d+$", mod.name)
+        if m:
+            return m.group(1)
+
+    return "blocks"