sae_stats.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
"""Multi-layer SAE statistics export for LLMs."""
 
from __future__ import annotations
 
import json
import os
from pathlib import Path
from typing import Any
 
import torch
 
from aquin.compute.device import resolve_compute_device
 
DEVICE = resolve_compute_device()
 
_PROBE_TEXT_KEYS = ("text", "prompt", "input", "query", "instruction")
_META_KEYS = ("id", "stressor", "lang", "quant_run_id", "condition", "label", "group")
 
 
def _sae_cache_dir(model_id: str) -> Path:
    return Path.home() / ".aquin" / "sae" / model_id
 
 
def _parse_layers_wanted(layers: str | None) -> list[int] | None:
    if not layers or layers.strip().lower() == "all":
        return None
    wanted: list[int] = []
    for part in layers.split(","):
        part = part.strip()
        if not part:
            continue
        try:
            wanted.append(int(part))
        except ValueError as exc:
            raise ValueError(f"Invalid --layers value {part!r} (expected integers like 8 or 8,9,10,11)") from exc
    if not wanted:
        raise ValueError("--layers must list at least one layer index, e.g. --layers 11 or --layers 8,9,10,11")
    return wanted
 
 
def _format_load_sae_commands(model_id: str, layers: list[int]) -> str:
    return "\n".join(f"  aquin load sae {model_id}-l{layer}" for layer in layers)
 
 
def _missing_sae_error(
    model_id: str,
    *,
    mode: str,
    available: list[int],
    layers_arg: str | None,
) -> str:
    """Actionable error when SAE checkpoint files are missing on disk."""
    sae_dir = _sae_cache_dir(model_id)
    wanted = _parse_layers_wanted(layers_arg)
    lines: list[str] = []
 
    if not available:
        lines.append(f"No LLM SAE checkpoints found for {model_id}.")
        lines.append(f"Expected directory: {sae_dir}/")
        lines.append("Expected files like: sae_layer<N>.pt")
    else:
        lines.append(f"Some requested LLM SAE layers are not downloaded for {model_id}.")
        lines.append(f"Downloaded layers: {available}")
 
    if wanted:
        missing = [layer for layer in wanted if layer not in available]
        if missing:
            lines.append(f"Missing layer(s): {missing}")
    elif not available:
        lines.append("No --layers specified; need at least one downloaded checkpoint.")
 
    lines.append("")
    lines.append("Pull checkpoints first (or use aquin load sae --path):")
    pull_layers = (
        [layer for layer in (wanted or []) if layer not in available]
        if wanted
        else [8]
    )
    if not pull_layers and wanted:
        pull_layers = wanted
    lines.append(_format_load_sae_commands(model_id, pull_layers))
 
    lines.append("")
    lines.append("Load an SAE: aquin load sae <model-l{n}> or --path <file.pt>")
    return "\n".join(lines)
 
 
def _expand(path: str) -> Path:
    return Path(os.path.expanduser(path)).resolve()
 
 
def _rows_to_probes(rows: list[Any]) -> list[dict[str, Any]]:
    probes: list[dict[str, Any]] = []
    for i, row in enumerate(rows):
        if isinstance(row, str):
            probes.append({"id": f"probe_{i}", "text": row})
            continue
        if not isinstance(row, dict):
            raise ValueError(f"Probe row {i} must be string or object")
        text = next((str(row[k]) for k in _PROBE_TEXT_KEYS if row.get(k)), None)
        if not text:
            raise ValueError(f"Probe row {i} missing text field ({', '.join(_PROBE_TEXT_KEYS)})")
        out: dict[str, Any] = {"id": row.get("id") or f"probe_{i}", "text": text}
        for key in _META_KEYS:
            if row.get(key) is not None:
                out[key] = row[key]
        probes.append(out)
    if not probes:
        raise ValueError("No probes found")
    return probes
 
 
def load_probes(path: str | list[Any]) -> list[dict[str, Any]]:
    """Load probe rows from a path, inline JSON, or an in-memory list."""
    if isinstance(path, list):
        return _rows_to_probes(path)
    raw = str(path).strip()
    if raw.startswith("[") or raw.startswith("{"):
        data = json.loads(raw)
        if isinstance(data, list):
            return _rows_to_probes(data)
        if isinstance(data, dict) and isinstance(data.get("probes"), list):
            return _rows_to_probes(data["probes"])
        raise ValueError("Expected a JSON array of probe strings")
 
    from aquin.compute.activation_capture import resolve_prompts_path
 
    resolved = resolve_prompts_path(_expand(raw))
    if resolved is None:
        if "/" not in raw and "\\" not in raw and not raw.lower().endswith((".json", ".jsonl")):
            return _rows_to_probes([raw])
        raise FileNotFoundError(
            f"Probes file not found: {path} "
            f"(tried cwd and repo parent)"
        )
    p = resolved
 
    if p.suffix.lower() == ".jsonl":
        rows: list[Any] = []
        with p.open(encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    rows.append(json.loads(line))
    else:
        with p.open(encoding="utf-8") as f:
            loaded = json.load(f)
        if isinstance(loaded, list):
            rows = loaded
        elif isinstance(loaded, dict) and isinstance(loaded.get("probes"), list):
            rows = loaded["probes"]
        else:
            raise ValueError(f"Expected JSON array or {{probes: [...]}} in {p}")
    return _rows_to_probes(rows)
 
 
def parse_layers_arg(
    layers: str | None,
    available: list[int],
    *,
    model_id: str | None = None,
    mode: str = "llm",
) -> list[int]:
    if layers is not None and not isinstance(layers, str):
        layers = str(layers)
    if not layers or layers.strip().lower() == "all":
        if not available:
            raise ValueError(
                _missing_sae_error(model_id or "model", mode=mode, available=[], layers_arg=layers)
                if model_id
                else "No SAE layers available on disk."
            )
        return list(available)
    wanted = _parse_layers_wanted(layers)
    if wanted is None:
        return list(available)
    missing = [layer for layer in wanted if layer not in available]
    if missing:
        if model_id:
            raise ValueError(
                _missing_sae_error(model_id, mode=mode, available=available, layers_arg=layers)
            )
        raise ValueError(
            f"No SAE checkpoint for layer(s) {missing}. Downloaded: {available or 'none'}"
        )
    return wanted
 
 
def list_llm_sae_layers(model_id: str) -> list[int]:
    from aquin.compute.model_loader import get_available_sae_layers, resolve_model_id
 
    return get_available_sae_layers(resolve_model_id(model_id))
 
 
def _feature_stats(mean_acts: torch.Tensor, *, top_k: int) -> dict[str, Any]:
    """mean_acts: (n_features,)"""
    mean_acts = torch.nan_to_num(mean_acts, nan=0.0, posinf=0.0, neginf=0.0)
    n_features = int(mean_acts.shape[0])
    active = mean_acts > 1e-6
    l0 = float(active.sum().item())
    sparsity = round(1.0 - l0 / max(n_features, 1), 6)
    top = mean_acts.topk(min(top_k, n_features))
    top_features = [
        {"feature_idx": int(i), "mean_activation": round(float(v), 6)}
        for i, v in zip(top.indices.tolist(), top.values.tolist())
    ]
    return {
        "n_features": n_features,
        "mean_l0": round(l0, 4),
        "sparsity": sparsity,
        "mean_activation": round(float(mean_acts.mean().item()), 6),
        "max_activation": round(float(mean_acts.max().item()), 6),
        "top_features": top_features,
    }
 
 
def _llm_probe_acts(
    model: Any,
    sae: Any,
    model_id: str,
    layer: int,
    text: str,
) -> torch.Tensor:
    from aquin.compute.feature_analysis import normalize
 
    tokens = model.to_tokens(text[:512])
    hook = f"blocks.{layer}.hook_resid_post"
    with torch.no_grad():
        _, cache = model.run_with_cache(tokens, names_filter=hook, return_type=None)
        resid = cache[hook][0]
        encoded = sae.encode(normalize(resid, model_id, layer))
    return torch.nan_to_num(encoded.mean(dim=0).cpu(), nan=0.0, posinf=0.0, neginf=0.0)
 
 
def run_llm_sae_stats(
    model: Any,
    model_id: str,
    probes: list[dict[str, Any]],
    *,
    layers: list[int],
    top_k: int = 10,
) -> dict[str, Any]:
    from aquin.compute.feature_analysis import load_sae
 
    layer_stats: list[dict[str, Any]] = []
    heatmap_rows: list[str] = []
    heatmap_cols = [f"L{layer}" for layer in layers]
    heatmap_values: list[list[float]] = []
 
    for layer in layers:
        try:
            sae = load_sae(model_id, layer)
        except Exception as exc:
            layer_stats.append({
                "layer": layer,
                "sae_available": False,
                "error": str(exc),
            })
            continue
 
        per_probe: list[dict[str, Any]] = []
        accum: list[torch.Tensor] = []
        for probe in probes:
            acts = _llm_probe_acts(model, sae, model_id, layer, probe["text"])
            accum.append(acts)
            l0 = float((acts > 1e-6).sum().item())
            per_probe.append({
                "probe_id": probe["id"],
                **{k: probe[k] for k in _META_KEYS if k in probe},
                "mean_l0": round(l0, 4),
                "mean_activation": round(float(acts.mean().item()), 6),
            })
 
        mean_acts = torch.stack(accum).mean(dim=0)
        layer_stats.append({
            "layer": layer,
            "sae_available": True,
            **_feature_stats(mean_acts, top_k=top_k),
            "per_probe": per_probe,
        })
 
    for probe in probes:
        heatmap_rows.append(str(probe["id"]))
        row_vals: list[float] = []
        for layer in layers:
            layer_entry = next((x for x in layer_stats if x.get("layer") == layer), None)
            if not layer_entry or not layer_entry.get("sae_available"):
                row_vals.append(0.0)
                continue
            match = next(
                (p for p in layer_entry.get("per_probe", []) if p.get("probe_id") == probe["id"]),
                None,
            )
            row_vals.append(float(match["mean_l0"]) if match else 0.0)
        heatmap_values.append(row_vals)
 
    return _build_payload(
        model_id=model_id,
        mode="llm",
        probes=probes,
        layers=layers,
        layer_stats=layer_stats,
        heatmap_rows=heatmap_rows,
        heatmap_cols=heatmap_cols,
        heatmap_values=heatmap_values,
        top_k=top_k,
    )
 
 
def _build_payload(
    *,
    model_id: str,
    mode: str,
    probes: list[dict[str, Any]],
    layers: list[int],
    layer_stats: list[dict[str, Any]],
    heatmap_rows: list[str],
    heatmap_cols: list[str],
    heatmap_values: list[list[float]],
    top_k: int,
) -> dict[str, Any]:
    profile = [
        {
            "layer": entry["layer"],
            "mean_l0": entry.get("mean_l0"),
            "sparsity": entry.get("sparsity"),
            "mean_activation": entry.get("mean_activation"),
            "sae_available": entry.get("sae_available", False),
        }
        for entry in layer_stats
        if entry.get("sae_available")
    ]
 
    return {
        "schema_version": 1,
        "model_id": model_id,
        "mode": mode,
        "n_probes": len(probes),
        "layers_requested": layers,
        "top_k": top_k,
        "probes": probes,
        "layer_stats": layer_stats,
        "layer_profile": profile,
        "heatmap": {
            "metric": "mean_l0",
            "rows": heatmap_rows,
            "cols": heatmap_cols,
            "values": heatmap_values,
        },
    }
 
 
def run_sae_stats(args: dict[str, Any]) -> dict[str, Any]:
    from aquin.compute.model_loader import (
        get_loaded_model,
        load_model,
        resolve_model_id,
    )
 
    model_id = args.get("model_id") or "llama-3.2-1b"
    try:
        model_id = resolve_model_id(model_id)
    except ValueError as exc:
        return {"error": str(exc)}
 
    prompts_path = args.get("prompts") or args.get("prompt")
    if not prompts_path:
        probes = [{"id": "probe_0", "text": "The capital of France is"}]
    else:
        try:
            probes = load_probes(prompts_path)
        except (OSError, ValueError, json.JSONDecodeError, TypeError) as exc:
            return {"error": str(exc)}
 
    top_k = int(args.get("top_k") or 10)
    layers_arg = args.get("layers")
 
    model = get_loaded_model()
    if model is None:
        try:
            model = load_model(model_id)
        except Exception as exc:
            return {"error": f"Could not load model: {exc}"}
 
    available = list_llm_sae_layers(model_id)
    if not available:
        return {
            "error": _missing_sae_error(
                model_id, mode="llm", available=[], layers_arg=layers_arg,
            ),
        }
    try:
        layers = parse_layers_arg(layers_arg, available, model_id=model_id, mode="llm")
    except ValueError as exc:
        return {"error": str(exc)}
 
    return run_llm_sae_stats(model, model_id, probes, layers=layers, top_k=top_k)