perturbation_check.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""`aquin check perturbation --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 = "perturbation-check.json"
_PNG_NAME = "perturbation-check.png"
 
 
def write_perturbation_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_perturbation_check(result, tool_name, png_path)
    return json_path, png_path
 
 
 
def _plot_perturbation_check(
    result: dict[str, Any],
    tool_name: str | None,
    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
 
    _ = tool_name
    _plot_llm_perturbation(result, png_path, plt)
 
 
def _plot_llm_perturbation(result: dict[str, Any], png_path: Path, plt: Any) -> None:
    raw = result.get("channels") or result.get("perturbation_results") or []
    rows = sorted(
        [
            {
                "channel": int(r.get("channel", i)),
                "kl": float(r.get("kl_mean", r.get("kl_divergence")) or 0),
            }
            for i, r in enumerate(raw)
        ],
        key=lambda x: x["kl"],
        reverse=True,
    )[:24]
 
    layer = result.get("layer", "—")
    prompt = str(result.get("prompt") or "")[:56]
    mean_kl = float(result.get("global_mean_kl") or 0)
    max_kl = float(result.get("global_max_kl") or 0)
 
    fig, ax = plt.subplots(figsize=(9, max(4, len(rows) * 0.28 + 1.5)), facecolor="#0f1117")
    fig.suptitle(
        f"Perturbation sensitivity (LLM) — layer {layer}",
        color="#e5e7eb",
        fontsize=11,
        y=0.98,
    )
    ax.set_title(f'"{prompt}"  ·  mean KL {mean_kl:.4f}  ·  max {max_kl:.4f}', color="#9ca3af", fontsize=9, pad=8)
 
    if not rows:
        ax.axis("off")
        ax.text(0.5, 0.5, "No channel results.", ha="center", va="center", color="#9ca3af")
    else:
        labels = [f"ch {r['channel']}" for r in rows]
        vals = [r["kl"] for r in rows]
        colors = ["#f87171" if v >= max_kl * 0.85 else "#6366f1" if v >= mean_kl else "#78716c" for v in vals]
        y_pos = list(range(len(rows)))
        ax.barh(y_pos, vals, color=colors, height=0.72)
        ax.set_yticks(y_pos)
        ax.set_yticklabels(labels, fontsize=8, color="#d1d5db")
        ax.invert_yaxis()
        ax.set_xlabel("KL divergence", color="#9ca3af", fontsize=9)
        ax.tick_params(axis="x", colors="#6b7280", labelsize=8)
        ax.grid(axis="x", alpha=0.2, color="#4b5563")
        for spine in ax.spines.values():
            spine.set_color("#374151")
 
    fig.tight_layout()
    fig.savefig(png_path, dpi=140, bbox_inches="tight", facecolor="#0f1117")
    plt.close(fig)
 
 
 
def _short_token(tok: str, limit: int = 14) -> str:
    t = tok.replace("Ġ", " ").replace("##", "").strip()
    if len(t) <= limit:
        return t or "·"
    return t[: limit - 1] + "…"