attention_check.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""`aquin check attention --check` — save JSON + diagram to the current working directory."""
 
from __future__ import annotations
 
import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
 
_JSON_NAME = "attention-check.json"
_PNG_NAME = "attention-check.png"
 
 
def write_attention_check(
    result: dict[str, Any],
    *,
    tool_name: str | None,
    cwd: str | Path,
) -> tuple[Path, Path]:
    cwd = Path(cwd)
    json_path = cwd / _JSON_NAME
    png_path = cwd / _PNG_NAME
 
    payload = {
        "saved_at": datetime.now(timezone.utc).isoformat(),
        "tool": tool_name,
        **result,
    }
    json_path.write_text(json.dumps(payload, indent=2, default=str), encoding="utf-8")
    _plot_attention_check(result, tool_name, png_path)
    return json_path, png_path
 
 
 
def _plot_attention_check(
    result: dict[str, Any],
    tool_name: str | None,
    png_path: Path,
) -> None:
    import matplotlib
 
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    import numpy as np
 
    if result.get("error"):
        fig, ax = plt.subplots(figsize=(6, 2))
        ax.axis("off")
        ax.text(0.5, 0.5, f"Error: {result['error']}", ha="center", va="center", wrap=True)
        fig.savefig(png_path, dpi=140, bbox_inches="tight", facecolor="#0f1117")
        plt.close(fig)
        return
 
    _ = tool_name
    _plot_llm_attention(result, png_path, plt, np)
 
 
def _plot_llm_attention(result: dict[str, Any], png_path: Path, plt: Any, np: Any) -> None:
    n_layers = int(result.get("n_layers") or 0)
    n_heads = int(result.get("n_heads") or 0)
    heads = result.get("heads") or []
 
    sink = np.zeros((max(n_layers, 1), max(n_heads, 1)))
    induction = np.zeros_like(sink)
    for row in heads:
        layer = int(row.get("layer", 0))
        head = int(row.get("head", 0))
        if layer < sink.shape[0] and head < sink.shape[1]:
            sink[layer, head] = float(row.get("sink_score") or 0)
            induction[layer, head] = float(row.get("induction_score") or 0)
 
    prompt = str(result.get("prompt") or "")[:72]
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.5), facecolor="#0f1117")
    fig.suptitle(f"Attention routing — {prompt}", color="#e5e7eb", fontsize=11, y=1.02)
 
    for ax, data, title in (
        (axes[0], sink, "Sink score"),
        (axes[1], induction, "Induction score"),
    ):
        im = ax.imshow(data, aspect="auto", cmap="viridis", origin="lower")
        ax.set_xlabel("Head", color="#9ca3af", fontsize=9)
        ax.set_ylabel("Layer", color="#9ca3af", fontsize=9)
        ax.set_title(title, color="#e5e7eb", fontsize=10)
        ax.tick_params(colors="#6b7280", labelsize=8)
        fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
 
    fig.tight_layout()
    fig.savefig(png_path, dpi=140, bbox_inches="tight", facecolor="#0f1117")
    plt.close(fig)
 
 
 
def _short_token(tok: str, limit: int = 10) -> str:
    t = tok.replace("Ġ", " ").replace("##", "")
    if len(t) <= limit:
        return t
    return t[: limit - 1] + "…"