visualize_attention.py

"""
Attention head visualization for CompNeuro analysis.

Extracts and plots attention weight matrices from all 12 heads
across all 12 layers for a given input sequence.

Usage:
    python visualize_attention.py \
        --ckpt checkpoints/pubmed_ft/best.pt \
        --prompt "The hippocampus receives input from the entorhinal cortex"
"""

import argparse
import torch
import tiktoken
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from pathlib import Path
from src.config.model_config import GPTConfig
from src.model.gpt import GPT


def load_model_with_hooks(ckpt_path: str, device: str):
    """Load model and register forward hooks to capture attention weights."""
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    cfg_dict = ckpt["config"]
    cfg = GPTConfig(**{
        k: cfg_dict[k]
        for k in GPTConfig.__dataclass_fields__
        if k in cfg_dict
    })
    model = GPT(cfg)
    model.load_state_dict(ckpt["model_state_dict"])
    model.eval()
    model.to(device)
    return model, cfg


def extract_attention_weights(model: GPT, tokens: list[int], device: str):
    """
    Extract attention weight matrices from all layers.

    Temporarily replaces Flash Attention with manual attention
    to capture the explicit attention matrix (Flash Attention
    does not materialize the full NxN matrix by design).

    Returns:
        List of tensors, one per layer.
        Each tensor: (n_head, T, T)
    """
    import torch.nn.functional as F
    import math

    attention_weights = []

    def make_hook(layer_idx):
        def hook(module, input, output):
            # Re-compute attention weights manually from the input
            x = input[0]
            B, T, C = x.size()
            q, k, v = module.c_attn(x).split(module.n_embd, dim=2)
            head_dim = module.head_dim
            n_head = module.n_head

            q = q.view(B, T, n_head, head_dim).transpose(1, 2)
            k = k.view(B, T, n_head, head_dim).transpose(1, 2)

            # If RoPE is enabled, apply it
            if hasattr(module, "use_rope") and module.use_rope:
                cos, sin = module.rotary_emb(q, seq_len=T)
                from src.model.components import apply_rotary_emb
                q, k = apply_rotary_emb(q, k, cos, sin)

            scale = 1.0 / math.sqrt(head_dim)
            att = (q @ k.transpose(-2, -1)) * scale  # (B, n_head, T, T)

            # Apply causal mask
            mask = torch.tril(torch.ones(T, T, device=x.device)).bool()
            att = att.masked_fill(~mask, float("-inf"))
            att = F.softmax(att, dim=-1)

            # Store: (n_head, T, T) — squeeze batch dim
            attention_weights.append(att[0].detach().cpu())
        return hook

    # Register hooks on all attention modules
    hooks = []
    for i, block in enumerate(model.transformer["h"]):
        h = block.attn.register_forward_hook(make_hook(i))
        hooks.append(h)

    # Forward pass
    x = torch.tensor(tokens, dtype=torch.long, device=device).unsqueeze(0)
    with torch.no_grad():
        model(x)

    # Remove hooks
    for h in hooks:
        h.remove()

    return attention_weights  # List of (n_head, T, T), length = n_layer


def plot_attention_grid(
    attention_weights,
    tokens_str: list[str],
    output_path: str,
    title: str = "",
):
    """
    Plot all attention heads across all layers as a grid.

    Layout: rows = layers (0..11), cols = heads (0..11)
    Each cell: T×T attention matrix as a heatmap.
    """
    n_layer = len(attention_weights)
    n_head = attention_weights[0].shape[0]
    T = attention_weights[0].shape[1]

    # Truncate token labels for display
    labels = [t[:6] for t in tokens_str]

    fig = plt.figure(figsize=(n_head * 1.8, n_layer * 1.8))
    fig.suptitle(title or "Attention weights — all layers × heads", fontsize=10, y=1.01)

    gs = gridspec.GridSpec(n_layer, n_head, figure=fig, hspace=0.05, wspace=0.05)

    for layer_idx in range(n_layer):
        for head_idx in range(n_head):
            ax = fig.add_subplot(gs[layer_idx, head_idx])
            w = attention_weights[layer_idx][head_idx].numpy()  # (T, T)
            ax.imshow(w, cmap="Blues", vmin=0, vmax=w.max(), aspect="auto")
            ax.set_xticks([])
            ax.set_yticks([])

            # Label leftmost column with layer index
            if head_idx == 0:
                ax.set_ylabel(f"L{layer_idx}", fontsize=6, rotation=0, labelpad=14)
            # Label top row with head index
            if layer_idx == 0:
                ax.set_title(f"H{head_idx}", fontsize=6, pad=2)

    plt.savefig(output_path, dpi=150, bbox_inches="tight")
    plt.close()
    print(f"Saved: {output_path}")


def plot_single_layer(
    attn: torch.Tensor,
    tokens_str: list[str],
    layer_idx: int,
    output_path: str,
):
    """
    Plot all heads of a single layer with token labels.
    More readable than the full grid for detailed analysis.
    """
    n_head = attn.shape[0]
    T = attn.shape[1]
    labels = [t[:8] for t in tokens_str]

    fig, axes = plt.subplots(3, 4, figsize=(16, 12))
    fig.suptitle(f"Layer {layer_idx} — attention heads (T={T})", fontsize=12)

    for head_idx, ax in enumerate(axes.flat):
        if head_idx >= n_head:
            ax.axis("off")
            continue
        w = attn[head_idx].numpy()
        im = ax.imshow(w, cmap="viridis", vmin=0, aspect="auto")
        ax.set_title(f"Head {head_idx}", fontsize=9)
        ax.set_xticks(range(T))
        ax.set_xticklabels(labels, rotation=45, ha="right", fontsize=6)
        ax.set_yticks(range(T))
        ax.set_yticklabels(labels, fontsize=6)
        plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)

    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches="tight")
    plt.close()
    print(f"Saved: {output_path}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--ckpt",    required=True)
    parser.add_argument("--prompt",  default="The hippocampus receives input from the entorhinal cortex and projects to CA3")
    parser.add_argument("--layer",   type=int, default=None,
                        help="If set, plot detailed view of this layer only.")
    parser.add_argument("--out_dir", default="attention_plots")
    args = parser.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    Path(args.out_dir).mkdir(parents=True, exist_ok=True)

    enc = tiktoken.get_encoding("gpt2")
    tokens = enc.encode(args.prompt)
    tokens_str = [enc.decode([t]) for t in tokens]
    T = len(tokens)

    print(f"Prompt: {args.prompt!r}")
    print(f"Tokens ({T}): {tokens_str}")

    model, cfg = load_model_with_hooks(args.ckpt, device)
    attn_weights = extract_attention_weights(model, tokens, device)

    # Full grid: all layers × all heads
    grid_path = f"{args.out_dir}/attention_grid.png"
    plot_attention_grid(
        attn_weights, tokens_str, grid_path,
        title=f"{args.prompt[:60]}..."
    )

    # Detailed single layer (default: last layer — most semantic)
    layer = args.layer if args.layer is not None else cfg.n_layer - 1
    layer_path = f"{args.out_dir}/layer_{layer:02d}_detail.png"
    plot_single_layer(attn_weights[layer], tokens_str, layer, layer_path)

    # Also plot layers 0, 5, 11 for comparison
    for l in [0, 5, cfg.n_layer - 1]:
        if l != layer:
            p = f"{args.out_dir}/layer_{l:02d}_detail.png"
            plot_single_layer(attn_weights[l], tokens_str, l, p)

    print(f"\nAll plots saved to: {args.out_dir}/")
    print("Layer 0  = low-level syntactic patterns")
    print("Layer 5  = mid-level semantic composition")
    print(f"Layer {cfg.n_layer-1} = high-level contextual representation")


if __name__ == "__main__":
    main()