sae_stats_check.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""`aquin sae-stats --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 = "sae-stats-check.json"
_PNG_NAME = "sae-stats-check.png"
 
 
def _normalize_sae_stats_result(result: dict[str, Any]) -> dict[str, Any]:
    if not isinstance(result, dict):
        return {"error": "Invalid sae-stats result"}
 
    data = result
    card = result.get("card")
    if isinstance(card, dict) and isinstance(card.get("data"), dict):
        card_data = card["data"]
        data = {
            "schema_version": result.get("schema_version", 1),
            "model_id": card_data.get("modelId") or result.get("model_id"),
            "mode": card_data.get("mode") or result.get("mode"),
            "n_probes": card_data.get("nProbes") or result.get("n_probes"),
            "layers_requested": card_data.get("layersRequested") or result.get("layers_requested"),
            "top_k": card_data.get("topK") or result.get("top_k"),
            "probes": card_data.get("probes") or result.get("probes", []),
            "layer_stats": card_data.get("layerStats") or result.get("layer_stats", []),
            "layer_profile": card_data.get("layerProfile") or result.get("layer_profile", []),
            "heatmap": card_data.get("heatmap") or result.get("heatmap", {}),
            "saved_to": card_data.get("savedTo") or result.get("saved_to"),
        }
    elif isinstance(result.get("content"), dict):
        data = result["content"]
 
    if result.get("error"):
        data = {**data, "error": result["error"]}
    return data
 
 
def write_sae_stats_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
 
    flat = _normalize_sae_stats_result(result)
    payload = {
        "saved_at": datetime.now(timezone.utc).isoformat(),
        "tool": tool_name,
        **flat,
    }
    json_path.write_text(json.dumps(payload, indent=2, default=str), encoding="utf-8")
    _plot_sae_stats_check(flat, png_path)
    return json_path, png_path
 
 
def _plot_sae_stats_check(result: dict[str, Any], png_path: Path) -> None:
    import matplotlib
 
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
 
    if result.get("error"):
        fig, ax = plt.subplots(figsize=(6, 2), facecolor="#0f1117")
        ax.axis("off")
        ax.text(0.5, 0.5, f"Error: {result['error']}", ha="center", va="center", wrap=True, color="#e5e7eb")
        fig.savefig(png_path, dpi=140, bbox_inches="tight", facecolor="#0f1117")
        plt.close(fig)
        return
 
    model = str(result.get("model_id") or "")
    mode = str(result.get("mode") or "")
    n_probes = int(result.get("n_probes") or 0)
    profile = [p for p in (result.get("layer_profile") or []) if p.get("sae_available", True)]
    heatmap = result.get("heatmap") or {}
 
    title = f"SAE stats — {model}  ·  {mode}  ·  {n_probes} probes"
    fig = plt.figure(figsize=(12, 7.5), facecolor="#0f1117")
    gs = fig.add_gridspec(2, 2, height_ratios=[1, 1.35], hspace=0.38, wspace=0.28)
    ax_profile = fig.add_subplot(gs[0, 0])
    ax_sparse = fig.add_subplot(gs[0, 1])
    ax_heat = fig.add_subplot(gs[1, :])
 
    fig.suptitle(title, color="#e5e7eb", fontsize=11, y=0.98)
 
    _plot_layer_profile(ax_profile, profile, metric="mean_l0", title="Mean L0 per layer", color="#6366f1")
    _plot_layer_profile(ax_sparse, profile, metric="sparsity", title="Sparsity per layer", color="#34d399", pct=True)
    _plot_heatmap(ax_heat, heatmap, result.get("probes") or [])
 
    for ax in (ax_profile, ax_sparse, ax_heat):
        ax.set_facecolor("#0f1117")
        for spine in ax.spines.values():
            spine.set_color("#374151")
 
    fig.savefig(png_path, dpi=140, bbox_inches="tight", facecolor="#0f1117")
    plt.close(fig)
 
 
def _plot_layer_profile(
    ax: Any,
    profile: list[dict[str, Any]],
    *,
    metric: str,
    title: str,
    color: str,
    pct: bool = False,
) -> None:
    ax.set_title(title, color="#e5e7eb", fontsize=10)
 
    if not profile:
        ax.axis("off")
        ax.text(0.5, 0.5, "No layer profile", ha="center", va="center", color="#9ca3af")
        return
 
    layers = [int(p.get("layer", i)) for i, p in enumerate(profile)]
    vals = [float(p.get(metric) or 0) for p in profile]
    ax.bar(layers, vals, color=color, width=0.75, alpha=0.9)
    ax.set_xlabel("Layer", color="#9ca3af", fontsize=8)
    ylab = "Sparsity" if pct else "Mean L0"
    ax.set_ylabel(ylab, color="#9ca3af", fontsize=8)
    ax.tick_params(colors="#6b7280", labelsize=8)
    ax.grid(axis="y", alpha=0.2, color="#4b5563")
    if pct:
        ax.set_ylim(0, min(1.05, max(vals) * 1.15 + 0.05))
 
 
def _plot_heatmap(ax: Any, heatmap: dict[str, Any], probes: list[dict[str, Any]]) -> None:
    rows = heatmap.get("rows") or []
    cols = heatmap.get("cols") or []
    values = heatmap.get("values") or []
    metric = str(heatmap.get("metric") or "mean_l0")
 
    ax.set_title(f"Probe × layer heatmap ({metric})", color="#e5e7eb", fontsize=10, pad=10)
 
    if not rows or not cols or not values:
        ax.axis("off")
        ax.text(0.5, 0.5, "No heatmap data", ha="center", va="center", color="#9ca3af")
        return
 
    import numpy as np
 
    data = np.array(values, dtype=float)
    if data.ndim != 2:
        ax.axis("off")
        ax.text(0.5, 0.5, "Invalid heatmap shape", ha="center", va="center", color="#9ca3af")
        return
 
    im = ax.imshow(data, aspect="auto", cmap="viridis", origin="upper")
    ax.set_xticks(range(len(cols)))
    ax.set_xticklabels(cols, fontsize=7, color="#d1d5db", rotation=45, ha="right")
    ax.set_yticks(range(len(rows)))
    ylabels = [_probe_label(row_id, probes) for row_id in rows]
    ax.set_yticklabels(ylabels, fontsize=7, color="#d1d5db")
    ax.tick_params(axis="x", colors="#6b7280")
    ax.tick_params(axis="y", colors="#6b7280")
    cbar = ax.figure.colorbar(im, ax=ax, fraction=0.025, pad=0.02)
    cbar.ax.tick_params(colors="#9ca3af", labelsize=7)
    cbar.set_label(metric, color="#9ca3af", fontsize=8)
 
 
def _probe_label(row_id: str, probes: list[dict[str, Any]]) -> str:
    for p in probes:
        if str(p.get("id")) == str(row_id):
            stressor = p.get("stressor")
            if stressor:
                return f"{row_id} ({stressor})"[:24]
            text = str(p.get("text") or p.get("prompt") or "")
            if text:
                return _trunc(text, 22)
            break
    return str(row_id)[:22]
 
 
def _trunc(text: str, limit: int) -> str:
    t = text.strip()
    if len(t) <= limit:
        return t or "·"
    return t[: limit - 1] + "…"