layer_analysis_check.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""`aquin check layer --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 = "layer-analysis-check.json"
_PNG_NAME = "layer-analysis-check.png"
 
 
def write_layer_analysis_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_layer_analysis_check(result, tool_name, png_path)
    return json_path, png_path
 
 
 
def _plot_layer_analysis_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), 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_layer_analysis(result, png_path, plt, np)
 
 
def _plot_llm_layer_analysis(result: dict[str, Any], png_path: Path, plt: Any, np: Any) -> None:
    stability = result.get("stability") or {}
    ood = result.get("ood") or {}
    localize = result.get("localize") if isinstance(result.get("localize"), dict) else None
    stab_layers = stability.get("layers") or []
    ood_layers = ood.get("layers") or []
    loc_layers = (localize or {}).get("layers") or []
 
    n_cols = 3 if loc_layers else 2
    fig, axes = plt.subplots(1, n_cols, figsize=(5.5 * n_cols, 4.5), facecolor="#0f1117")
    if n_cols == 1:
        axes = [axes]
    model = str(result.get("model_id") or "")
    title = f"Layer analysis (LLM) — {model}"
    if localize and localize.get("collapse_layer") is not None:
        title += f"  ·  collapse L{localize.get('collapse_layer')}"
    fig.suptitle(title, color="#e5e7eb", fontsize=11, y=1.02)
 
    ax0 = axes[0]
    if stab_layers:
        xs = [int(r["layer"]) for r in stab_layers]
        pc1 = [float(r.get("top1_variance_ratio") or 0) for r in stab_layers]
        colors = []
        for r in stab_layers:
            status = str(r.get("status") or "ok")
            colors.append({"dead": "#f87171", "collapsed": "#fbbf24", "ok": "#34d399"}.get(status, "#9ca3af"))
        ax0.bar(xs, pc1, color=colors, width=0.8)
        ax0.axhline(0.85, color="#fbbf24", linestyle="--", linewidth=0.8, alpha=0.6, label="collapse")
        ax0.set_xlabel("Layer", color="#9ca3af", fontsize=9)
        ax0.set_ylabel("PC1 variance ratio", color="#9ca3af", fontsize=9)
        ax0.set_title("Activation stability", color="#e5e7eb", fontsize=10)
        ax0.tick_params(colors="#6b7280", labelsize=8)
    else:
        ax0.axis("off")
        ax0.text(0.5, 0.5, "No stability data", ha="center", va="center", color="#9ca3af")
 
    ax1 = axes[1]
    if ood_layers:
        xs = [int(r["layer"]) for r in ood_layers]
        sep = [float(r.get("separation") or 0) for r in ood_layers]
        ax1.plot(xs, sep, color="#34d399", marker="o", markersize=4, linewidth=1.5)
        peak = ood.get("peak_layer")
        if peak is not None:
            ax1.axvline(int(peak), color="#facc15", linestyle="--", linewidth=0.8, alpha=0.7)
        ax1.set_xlabel("Layer", color="#9ca3af", fontsize=9)
        ax1.set_ylabel("OOD separation", color="#9ca3af", fontsize=9)
        ax1.set_title("In-domain vs OOD", color="#e5e7eb", fontsize=10)
        ax1.tick_params(colors="#6b7280", labelsize=8)
    else:
        ax1.axis("off")
        ax1.text(0.5, 0.5, "No OOD data", ha="center", va="center", color="#9ca3af")
 
    if loc_layers:
        ax2 = axes[2]
        xs = [int(r["layer"]) for r in loc_layers]
        sig = [float(r.get("signal") or 0) for r in loc_layers]
        ax2.plot(xs, sig, color="#a78bfa", marker="o", markersize=4, linewidth=1.5)
        peak = localize.get("peak_layer") if localize else None
        collapse = localize.get("collapse_layer") if localize else None
        if peak is not None:
            ax2.axvline(int(peak), color="#34d399", linestyle="--", linewidth=0.8, alpha=0.8, label="peak")
        if collapse is not None:
            ax2.axvline(int(collapse), color="#f87171", linestyle="--", linewidth=0.8, alpha=0.8, label="collapse")
        ax2.set_xlabel("Layer", color="#9ca3af", fontsize=9)
        ax2.set_ylabel("Deception signal", color="#9ca3af", fontsize=9)
        ax2.set_title("Localize collapse", color="#e5e7eb", fontsize=10)
        ax2.tick_params(colors="#6b7280", labelsize=8)
        ax2.legend(fontsize=7, labelcolor="#9ca3af", frameon=False)
 
    fig.tight_layout()
    fig.savefig(png_path, dpi=140, bbox_inches="tight", facecolor="#0f1117")
    plt.close(fig)