check_weights_check.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""`aquin check-weights --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 = "check-weights-check.json"
_PNG_NAME = "check-weights-check.png"
 
_STATUS_COLORS = {
    "high_risk": "#f87171",
    "suspicious": "#fbbf24",
    "clean": "#78716c",
}
 
 
def write_check_weights_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_check_weights_check(result, png_path)
    return json_path, png_path
 
 
def _plot_check_weights_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 "")
    verdict = str(result.get("verdict") or "—")
    composite = float(result.get("composite_risk") or 0)
    rank = result.get("rank") if isinstance(result.get("rank"), dict) else {}
 
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.8), facecolor="#0f1117")
    fig.suptitle(
        f"Weight health — {model}  ·  {verdict}  ·  risk {composite:.1%}",
        color="#e5e7eb",
        fontsize=11,
        y=1.02,
    )
 
    _plot_trojan_panel(axes[0], result)
    _plot_rank_panel(axes[1], rank)
 
    fig.tight_layout()
    fig.savefig(png_path, dpi=140, bbox_inches="tight", facecolor="#0f1117")
    plt.close(fig)
 
 
def _plot_trojan_panel(ax: Any, result: dict[str, Any]) -> None:
    scored = result.get("scored_tensors") or []
    top = sorted(scored, key=lambda t: float(t.get("risk_score") or 0), reverse=True)[:18]
 
    ax.set_title("Trojan scan (top tensors)", color="#e5e7eb", fontsize=10)
    if not top:
        ax.axis("off")
        ax.text(0.5, 0.5, "No tensor scores.", ha="center", va="center", color="#9ca3af")
        return
 
    labels = [
        f"L{t.get('layer_idx', '?')} {str(t.get('name', '')).split('.')[-1][:12]}"
        for t in top
    ]
    vals = [float(t.get("risk_score") or 0) for t in top]
    colors = [_STATUS_COLORS.get(str(t.get("status") or ""), "#78716c") for t in top]
    y_pos = list(range(len(top)))
 
    ax.barh(y_pos, vals, color=colors, height=0.72)
    ax.set_yticks(y_pos)
    ax.set_yticklabels(labels, fontsize=7, color="#d1d5db")
    ax.invert_yaxis()
    ax.set_xlim(0, max(max(vals) * 1.1, 0.1))
    ax.set_xlabel("Risk score", color="#9ca3af", fontsize=8)
    ax.tick_params(axis="x", colors="#6b7280", labelsize=7)
    ax.grid(axis="x", alpha=0.2, color="#4b5563")
    for spine in ax.spines.values():
        spine.set_color("#374151")
 
 
def _plot_rank_panel(ax: Any, rank: dict[str, Any]) -> None:
    matrices = rank.get("matrices") or []
    threshold = float(rank.get("collapse_threshold") or 0.1)
 
    ax.set_title("Stable rank by layer (min per layer)", color="#e5e7eb", fontsize=10)
    if not matrices:
        ax.axis("off")
        ax.text(0.5, 0.5, "No rank data.", ha="center", va="center", color="#9ca3af")
        return
 
    by_layer: dict[int, float] = {}
    for row in matrices:
        layer = int(row.get("layer", 0))
        stable = float(row.get("stable_rank") or 0)
        by_layer[layer] = min(by_layer.get(layer, stable), stable)
 
    layers = sorted(by_layer)
    vals = [by_layer[li] for li in layers]
 
    ax.plot(layers, vals, color="#6366f1", linewidth=2, marker="o", markersize=3)
    ax.axhline(threshold, color="#f87171", linestyle="--", linewidth=1, alpha=0.7, label=f"collapse < {threshold}")
    for li, v in zip(layers, vals):
        if v < threshold:
            ax.scatter([li], [v], color="#f87171", s=28, zorder=3)
 
    ax.set_xlabel("Layer", color="#9ca3af", fontsize=8)
    ax.set_ylabel("Min stable rank", color="#9ca3af", fontsize=8)
    ax.tick_params(colors="#6b7280", labelsize=7)
    ax.grid(alpha=0.2, color="#4b5563")
    ax.legend(fontsize=7, facecolor="#1f2937", edgecolor="#374151", labelcolor="#d1d5db")
    for spine in ax.spines.values():
        spine.set_color("#374151")