sae_catalog_metrics.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Pre-publish metrics for public SAE catalog rows (UMAP + layer InterpScore)."""
 
from __future__ import annotations
 
import json
from pathlib import Path
from typing import Any
 
import torch
 
from aquin.compute.sae import SparseAutoencoder
 
CATALOG_PROBE_PROMPTS = [
    "The capital of France is",
    "In quantum mechanics, particles can",
    "The recipe calls for two cups of",
    "She opened the letter and read",
    "function merge_sort(arr):",
]
 
 
def compute_umap_points_from_sae(sae_path: str | Path) -> list[dict[str, Any]]:
    """3D UMAP of decoder rows (same recipe as bridge._compute_umap_points)."""
    sae = SparseAutoencoder.load(str(sae_path), device="cpu")
    w = sae.W_dec.float().cpu().detach()
    try:
        import umap
 
        reducer = umap.UMAP(n_components=3, random_state=42, n_neighbors=15)
        embedding = reducer.fit_transform(w.numpy())
        return [
            {"feature_idx": i, "x": float(r[0]), "y": float(r[1]), "z": float(r[2])}
            for i, r in enumerate(embedding)
        ]
    except ImportError:
        w_centered = w - w.mean(0)
        _, _, vt = torch.linalg.svd(w_centered, full_matrices=False)
        proj = (w_centered @ vt[:3].T).numpy()
        return [
            {"feature_idx": i, "x": float(r[0]), "y": float(r[1]), "z": float(r[2])}
            for i, r in enumerate(proj)
        ]
 
 
def _collect_top_features(
    *,
    model_id: str,
    layer: int,
    prompts: list[str],
    n_per_prompt: int,
    max_features: int,
) -> list[tuple[int, float, str]]:
    from aquin.compute.causal_trace import run_chat
    from aquin.compute.feature_analysis import run_feature_analysis_unlabeled
    from aquin.compute.model_loader import get_loaded_model, load_model
 
    model = get_loaded_model()
    if model is None:
        model = load_model(model_id)
 
    best: dict[int, tuple[float, str]] = {}
    for prompt in prompts:
        try:
            response = run_chat(prompt, model_id=model_id, max_new_tokens=120, temperature=0.7)
        except Exception as exc:
            print(f"[catalog-metrics] generation failed ({prompt[:40]!r}): {exc}", flush=True)
            continue
        try:
            feat = run_feature_analysis_unlabeled(
                prompt, response, model, model_id=model_id, layer=layer,
            )
        except Exception as exc:
            print(f"[catalog-metrics] feature pass failed ({prompt[:40]!r}): {exc}", flush=True)
            continue
        for row in (feat.get("top_response_features") or [])[:n_per_prompt]:
            idx = int(row["feature_idx"])
            act = float(row.get("activation") or 0.0)
            prev = best.get(idx)
            if prev is None or act > prev[0]:
                best[idx] = (act, prompt)
 
    ranked = sorted(
        ((idx, act, prompt) for idx, (act, prompt) in best.items()),
        key=lambda t: t[1],
        reverse=True,
    )
    return ranked[:max_features]
 
 
def compute_layer_interp_score(
    *,
    model_id: str,
    layer: int,
    sae_path: str | Path,
    prompts: list[str] | None = None,
    n_features: int = 8,
    n_per_prompt: int = 5,
    n_samples: int = 8,
    client: Any | None = None,
) -> dict[str, Any]:
    """Mean InterpScore over top response features from catalog probe prompts."""
    from aquin.compute.interp_score import run_interp_score
    from aquin.compute.model_loader import get_loaded_model, load_model
    from aquin.compute.openai_client import get_openai_client
    from aquin.compute.user_sae import activate_user_sae
 
    sae_path = Path(sae_path).resolve()
    if not sae_path.is_file():
        raise FileNotFoundError(f"SAE not found: {sae_path}")
 
    activate_user_sae(model_id=model_id, layer=layer, path=sae_path, name=sae_path.parent.name)
 
    if client is None:
        client = get_openai_client({})
    if client is None:
        return {
            "interp_score": None,
            "feature_scores": [],
            "n_benchmarked": 0,
            "error": "OpenAI not configured. Set OPENAI_API_KEY.",
        }
 
    probe_prompts = list(prompts or CATALOG_PROBE_PROMPTS)
    candidates = _collect_top_features(
        model_id=model_id,
        layer=layer,
        prompts=probe_prompts,
        n_per_prompt=n_per_prompt,
        max_features=n_features,
    )
    if not candidates:
        return {
            "interp_score": None,
            "feature_scores": [],
            "n_benchmarked": 0,
            "error": "No active features found on probe prompts.",
        }
 
    model = get_loaded_model()
    if model is None:
        model = load_model(model_id)
 
    from aquin.compute.feature_analysis import load_sae
 
    sae = load_sae(model_id, layer)
    feature_scores: list[dict[str, Any]] = []
    scores: list[float] = []
 
    for feature_idx, activation, prompt in candidates:
        print(f"[catalog-metrics] InterpScore f{feature_idx} (act={activation:.3f})", flush=True)
        interp = run_interp_score(
            feature_idx,
            prompt,
            model,
            sae,
            client,
            n_samples=n_samples,
            model_id=model_id,
            layer=layer,
        )
        score = interp.get("score")
        row = {
            "feature_idx": feature_idx,
            "activation": round(activation, 4),
            "prompt": prompt,
            "score": score,
            "purity_score": interp.get("purity_score"),
            "label": interp.get("label"),
            "error": interp.get("error"),
        }
        feature_scores.append(row)
        if score is not None:
            scores.append(float(score))
 
    mean_score = round(sum(scores) / len(scores), 4) if scores else None
    return {
        "interp_score": mean_score,
        "feature_scores": feature_scores,
        "n_benchmarked": len(scores),
        "n_candidates": len(candidates),
        "prompts": probe_prompts,
    }
 
 
def write_catalog_sidecars(
    *,
    sae_path: str | Path,
    umap_points: list[dict[str, Any]],
    metrics: dict[str, Any],
) -> tuple[Path, Path]:
    sae_path = Path(sae_path).resolve()
    umap_path = sae_path.with_suffix(".umap.json")
    metrics_path = sae_path.with_suffix(".metrics.json")
    umap_path.write_text(json.dumps({"points": umap_points}, indent=2), encoding="utf-8")
    metrics_path.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    return umap_path, metrics_path
 
 
def compute_and_write_catalog_sidecars(
    *,
    model_id: str,
    layer: int,
    sae_path: str | Path,
    skip_interp: bool = False,
    **interp_kwargs: Any,
) -> dict[str, Any]:
    sae_path = Path(sae_path).resolve()
    print(f"[catalog-metrics] UMAP on {sae_path.name}…", flush=True)
    umap_points = compute_umap_points_from_sae(sae_path)
    print(f"[catalog-metrics] UMAP done ({len(umap_points)} points)", flush=True)
 
    metrics: dict[str, Any] = {"layer": layer, "model_id": model_id, "n_umap_points": len(umap_points)}
    if not skip_interp:
        from aquin.compute.loader_shim import apply as shim_apply
 
        shim_apply()
        print(f"[catalog-metrics] InterpScore layer {layer}…", flush=True)
        interp = compute_layer_interp_score(
            model_id=model_id,
            layer=layer,
            sae_path=sae_path,
            **interp_kwargs,
        )
        metrics.update(interp)
        if interp.get("interp_score") is not None:
            print(f"[catalog-metrics] layer {layer} interp_score={interp['interp_score']}", flush=True)
        elif interp.get("error"):
            print(f"[catalog-metrics] interp skipped: {interp['error']}", flush=True)
 
    umap_path, metrics_path = write_catalog_sidecars(
        sae_path=sae_path, umap_points=umap_points, metrics=metrics,
    )
    print(f"[catalog-metrics] wrote {umap_path.name}, {metrics_path.name}", flush=True)
    return {
        "umap_path": str(umap_path),
        "metrics_path": str(metrics_path),
        "interp_score": metrics.get("interp_score"),
        "n_umap_points": len(umap_points),
    }