benchmark_check.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""`aquin benchmark --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 = "benchmark-check.json"
_PNG_NAME = "benchmark-check.png"
 
_METRICS = (
    ("InterpScore", "score"),
    ("Purity", "purity_score"),
    ("MUI", "mui_score"),
)
 
 
def _normalize_benchmark_result(result: dict[str, Any]) -> dict[str, Any]:
    if not isinstance(result, dict):
        return {"error": "Invalid benchmark result"}
 
    data = result
    if data.get("type") in ("benchmark", "benchmarkSuite") and isinstance(data.get("data"), dict):
        data = data["data"]
    if isinstance(data.get("content"), dict):
        data = data["content"]
        if data.get("type") in ("benchmark", "benchmarkSuite") and isinstance(data.get("data"), dict):
            data = data["data"]
 
    return dict(data)
 
 
def has_benchmark_payload(result: dict[str, Any]) -> bool:
    data = _normalize_benchmark_result(result)
    if any(data.get(k) is not None for k in ("score", "purity_score", "mui_score")):
        return True
    return data.get("feature_idx") is not None
 
 
def write_benchmark_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_benchmark_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_benchmark_check(flat, png_path)
    return json_path, png_path
 
 
def _score_color(val: float) -> str:
    if val >= 0.75:
        return "#34d399"
    if val >= 0.5:
        return "#fbbf24"
    return "#f87171"
 
 
def _plot_benchmark_check(result: dict[str, Any], png_path: Path) -> None:
    import matplotlib
 
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
 
    err = result.get("error")
    has_scores = any(result.get(k) is not None for k in ("score", "purity_score", "mui_score"))
 
    if err and not has_scores:
        fig, ax = plt.subplots(figsize=(6, 2), facecolor="#0f1117")
        ax.axis("off")
        ax.text(0.5, 0.5, f"Error: {err}", ha="center", va="center", wrap=True, color="#e5e7eb")
        fig.savefig(png_path, dpi=140, bbox_inches="tight", facecolor="#0f1117")
        plt.close(fig)
        return
 
    ref = result.get("feature_ref") or result.get("feature_idx", "?")
    model = str(result.get("model_id") or "")
    label = str(result.get("label") or "")
 
    title = f"Feature benchmark — {ref}"
    if model:
        title += f"  ·  {model}"
    if label:
        title += f"\n{label[:48]}"
 
    fig, axes = plt.subplots(1, 2, figsize=(11, 4.5), facecolor="#0f1117")
    fig.suptitle(title, color="#e5e7eb", fontsize=11, y=1.04)
 
    _plot_metric_panel(axes[0], result)
    _plot_activation_panel(axes[1], result)
 
    if err:
        fig.text(0.5, 0.01, f"Note: {err}", ha="center", fontsize=7, color="#fbbf24", wrap=True)
 
    fig.tight_layout(rect=(0, 0.04 if err else 0, 1, 0.96))
    fig.savefig(png_path, dpi=140, bbox_inches="tight", facecolor="#0f1117")
    plt.close(fig)
 
 
def _plot_metric_panel(ax: Any, result: dict[str, Any]) -> None:
    ax.set_facecolor("#0f1117")
    ax.set_title("InterpScore · Purity · MUI", color="#e5e7eb", fontsize=10)
 
    rows = [(name, result.get(key)) for name, key in _METRICS if result.get(key) is not None]
    if not rows:
        ax.axis("off")
        ax.text(0.5, 0.5, "No scores", ha="center", va="center", color="#9ca3af")
        return
 
    labels = [r[0] for r in rows]
    vals = [float(r[1]) for r in rows]
    colors = [_score_color(v) for v in vals]
    y_pos = list(range(len(rows)))
 
    ax.barh(y_pos, vals, color=colors, height=0.65, alpha=0.92)
    ax.set_yticks(y_pos)
    ax.set_yticklabels(labels, fontsize=9, color="#d1d5db")
    ax.invert_yaxis()
    ax.set_xlim(0, 1.05)
    ax.set_xlabel("Score (0–1)", color="#9ca3af", fontsize=8)
    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")
    for i, v in enumerate(vals):
        ax.text(min(v + 0.02, 0.98), i, f"{v:.3f}", va="center", fontsize=8, color="#9ca3af")
 
 
def _plot_activation_panel(ax: Any, result: dict[str, Any]) -> None:
    ax.set_facecolor("#0f1117")
    ax.set_title("Activation separation", color="#e5e7eb", fontsize=10)
 
    pos_mean = result.get("positive_mean")
    neg_mean = result.get("negative_mean")
    if pos_mean is None and neg_mean is None:
        pos_ex = result.get("positive_examples") or []
        if pos_ex:
            _plot_examples_panel(ax, pos_ex[:5])
            return
        ax.axis("off")
        ax.text(0.5, 0.5, "No activation stats", ha="center", va="center", color="#9ca3af")
        return
 
    labels = ["fires on", "silent on"]
    vals = [float(pos_mean or 0), float(neg_mean or 0)]
    colors = ["#34d399", "#78716c"]
    y_pos = [0, 1]
    ax.barh(y_pos, vals, color=colors, height=0.55, alpha=0.9)
    ax.set_yticks(y_pos)
    ax.set_yticklabels(labels, fontsize=9, color="#d1d5db")
    ax.invert_yaxis()
    ax.set_xlabel("Mean activation", color="#9ca3af", fontsize=8)
    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")
    for i, v in enumerate(vals):
        ax.text(v + max(vals) * 0.02 + 0.01, i, f"{v:.3f}", va="center", fontsize=8, color="#9ca3af")
 
 
def _plot_examples_panel(ax: Any, examples: list[dict[str, Any]]) -> None:
    ax.set_title("Top positive examples", color="#e5e7eb", fontsize=10)
    labels = [_trunc(str(ex.get("sentence") or ""), 36) for ex in examples]
    vals = [float(ex.get("activation") or 0) for ex in examples]
    y_pos = list(range(len(examples)))
    ax.barh(y_pos, vals, color="#6366f1", height=0.65, alpha=0.9)
    ax.set_yticks(y_pos)
    ax.set_yticklabels(labels, fontsize=7, color="#d1d5db")
    ax.invert_yaxis()
    ax.set_xlabel("Activation", 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 _trunc(text: str, limit: int) -> str:
    t = text.strip()
    if len(t) <= limit:
        return t or "·"
    return t[: limit - 1] + "…"