feature_logits_check.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""`aquin feature logit --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 = "feature-logits-check.json"
_PNG_NAME = "feature-logits-check.png"
 
 
def _flatten_feature_logits_result(result: dict[str, Any]) -> dict[str, Any]:
    if result.get("error"):
        return dict(result)
 
    data = result.get("content") or result
    if not isinstance(data, dict):
        return {"error": "Invalid feature-logits result"}
 
    boosts = data.get("boosts") or data.get("top") or []
    suppresses = data.get("suppresses") or data.get("bottom") or []
    return {
        **data,
        "boosts": boosts,
        "suppresses": suppresses,
        "top": boosts,
        "bottom": suppresses,
    }
 
 
def write_feature_logits_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 = _flatten_feature_logits_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_feature_logits_check(flat, png_path)
    return json_path, png_path
 
 
def _plot_feature_logits_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
 
    ref = result.get("feature_ref") or result.get("feature_idx", "?")
    layer = result.get("layer", "?")
    label = str(result.get("label") or result.get("causal_label") or "")
    boosts = result.get("boosts") or []
    suppresses = result.get("suppresses") or []
 
    title = f"Feature logits — {ref}  ·  layer {layer}"
    if label:
        title += f"  ·  {label[:32]}"
 
    fig, axes = plt.subplots(1, 2, figsize=(11, max(4, min(12, max(len(boosts), len(suppresses)) * 0.32 + 1.8))), facecolor="#0f1117")
    fig.suptitle(title, color="#e5e7eb", fontsize=11, y=1.02)
 
    _plot_token_panel(axes[0], boosts, "Boosted tokens", "#34d399", positive=True)
    _plot_token_panel(axes[1], suppresses, "Suppressed tokens", "#f87171", positive=False)
 
    fig.tight_layout()
    fig.savefig(png_path, dpi=140, bbox_inches="tight", facecolor="#0f1117")
    plt.close(fig)
 
 
def _plot_token_panel(
    ax: Any,
    rows: list[dict[str, Any]],
    title: str,
    color: str,
    *,
    positive: bool,
) -> None:
    ax.set_facecolor("#0f1117")
    ax.set_title(title, color="#e5e7eb", fontsize=10)
 
    if not rows:
        ax.axis("off")
        ax.text(0.5, 0.5, "No data", ha="center", va="center", color="#9ca3af")
        return
 
    labels = [_short_token(str(r.get("token") or "")) for r in rows]
    vals = [float(r.get("logit") or 0) for r in rows]
    scale = max((abs(v) for v in vals), default=1.0) or 1.0
    widths = [abs(v) / scale for v in vals]
 
    y_pos = list(range(len(rows)))
    ax.barh(y_pos, widths, color=color, height=0.72, alpha=0.9)
    ax.set_yticks(y_pos)
    ax.set_yticklabels(labels, fontsize=8, color="#d1d5db")
    ax.invert_yaxis()
    ax.set_xlabel("Relative |logit|", 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")
 
    for i, (v, w) in enumerate(zip(vals, widths)):
        sign = "+" if positive else ""
        ax.text(w + 0.02, i, f"{sign}{v:.3f}", va="center", fontsize=7, color="#9ca3af")
 
 
def _short_token(tok: str, limit: int = 16) -> str:
    t = tok.replace("Ġ", " ").replace("##", "").strip()
    if len(t) <= limit:
        return t or "·"
    return t[: limit - 1] + "…"