sae_stats_cli.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""aquin sae-stats — multi-layer SAE statistics export."""
 
import sys
 
from aquin.cli_flags import reject_legacy_output_flags
from pathlib import Path
from typing import Any
 
 
def _parse_flag(args: list[str], name: str) -> str | None:
    for i, a in enumerate(args):
        if a == name and i + 1 < len(args):
            return args[i + 1]
    return None
 
 
def _has_flag(args: list[str], name: str) -> bool:
    return name in args
 
 
def _parse_int_flag(args: list[str], name: str, default: int) -> int:
    raw = _parse_flag(args, name)
    if raw is None:
        return default
    try:
        return max(1, int(raw))
    except ValueError:
        print(f"Error: {name} must be an integer")
        sys.exit(1)
 
 
def _parse_topk_flag(args: list[str], default: int) -> int:
    for name in ("--topk", "--top-k", "--top_k"):
        raw = _parse_flag(args, name)
        if raw is not None:
            try:
                return max(1, int(raw))
            except ValueError:
                print(f"Error: {name} must be an integer")
                sys.exit(1)
    return default
 
 
def _ensure_compute_env() -> None:
    from aquin.compute.loader_shim import apply as _shim_apply
    from aquin.engine.local_server import start as _start_local_server
 
    _shim_apply()
    _start_local_server()
 
 
def _require_loaded_model_id() -> str:
    from aquin.compute.model_loader import get_active_model_id, resolve_model_id
 
    active = (get_active_model_id() or "").strip()
    if not active:
        print("Error: no model loaded. Run: aquin load --model <id>")
        sys.exit(1)
    try:
        return resolve_model_id(active)
    except ValueError as e:
        print(f"Error: {e}")
        sys.exit(1)
 
 
def _print_help() -> None:
    print("Export multi-layer SAE statistics over a probe dataset (LLM).")
    print("")
    print("Prerequisite: aquin load --model <id>")
    print("  aquin load sae <model-l{n}>")
    print("")
    print("Usage: aquin sae-stats --prompts <json|jsonl>")
    print("       [--layers <all|9|0,9,15>]")
    print("       [--topk <n>]")
    print("       [--save <path>]")
    print("       [--check] [--umap]")
    print("")
    print("  --prompts   Probe file. Each row: text (or prompt) + optional id, stressor, lang, quant_run_id.")
    print("  --layers    all (default) or comma-separated layer indices with SAE checkpoints on disk.")
    print("  --topk      Top firing features per layer (default: 10).")
    print("  --save      Write schema_version=1 JSON export to this path.")
    print("  --check     Save sae-stats-check.json and sae-stats-check.png in the current directory.")
    print("  --umap      Load SAE UMAP projection after the export (web explorer).")
    print("")
    print("Example:")
    print("  aquin sae-stats --prompts probes.jsonl --layers 8 --topk 10 --check")
    print("  aquin sae-stats --prompts probes.jsonl --layers 8 --umap")
    print("")
    print("Docs: https://aquin.app/docs/inspection-sae")
 
 
def cmd_sae_stats(args: list[str]) -> None:
    if _has_flag(args, "--help") or _has_flag(args, "-h"):
        _print_help()
        return
 
    reject_legacy_output_flags(args)
 
    prompts = _parse_flag(args, "--prompts")
    if not prompts:
        print("Error: --prompts is required.")
        print("")
        _print_help()
        sys.exit(1)
 
    from aquin.compute.activation_capture import resolve_prompts_path
 
    resolved = resolve_prompts_path(Path(prompts).expanduser())
    if resolved is None:
        print(f"Error: probe file not found: {prompts}")
        print("       --prompts must be a path to a .json or .jsonl file (inline JSON is not supported).")
        print("")
        print("Quick fix:")
        print('  printf \'%s\\n\' \'{"id":"p1","text":"The cat sat on the mat"}\' \\')
        print('    \'{"id":"p2","text":"Paris is the capital of France"}\' > /tmp/probes.jsonl')
        print("  aquin sae-stats --prompts /tmp/probes.jsonl --layers 11")
        sys.exit(1)
    prompts = str(resolved)
 
    layers = _parse_flag(args, "--layers")
    top_k = _parse_topk_flag(args, 10)
    save_path = _parse_flag(args, "--save")
    do_check = _has_flag(args, "--check")
    want_umap = _has_flag(args, "--umap")
 
    if _parse_flag(args, "--model") is not None:
        print("Error: sae-stats uses the loaded session model only.")
        print("       Run: aquin load --model <id>")
        sys.exit(1)
 
    _ensure_compute_env()
 
    from aquin.cli import _build_tool_ctx
    from aquin.engine.sync_dispatch import dispatch_with_sync, require_active_session
 
    mid = _require_loaded_model_id()
    ctx = _build_tool_ctx(model_id=mid)
    require_active_session(ctx, label="aquin sae-stats")
 
    tool_args: dict[str, Any] = {
        "model_id": mid,
        "prompts": prompts,
        "top_k": top_k,
    }
    if layers:
        tool_args["layers"] = layers
    if save_path:
        tool_args["save"] = save_path
 
    try:
        print(f"[sae-stats] model={mid}  probes={prompts}  layers={layers or 'all'}  top-k={top_k}")
        result = dispatch_with_sync("run_sae_stats", tool_args, ctx)
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)
 
    from aquin.cli_output import print_tool_result
 
    print_tool_result("sae-stats", result)
 
    if isinstance(result, dict) and result.get("error"):
        sys.exit(1)
 
    if do_check:
        if isinstance(result, dict) and result.get("error"):
            print(f"[sae-stats --check] skipped save: {result['error']}", file=sys.stderr)
        else:
            import os
            from aquin.sae_stats_check import write_sae_stats_check
 
            try:
                json_path, png_path = write_sae_stats_check(
                    result, tool_name="run_sae_stats", cwd=os.getcwd(),
                )
                print(f"Saved {json_path}")
                print(f"Saved {png_path}")
            except Exception as exc:
                print(f"[sae-stats --check] failed to save files: {exc}", file=sys.stderr)
                import traceback
                traceback.print_exc()
                sys.exit(1)
 
    if want_umap:
        from aquin.cli import _run_umap_followup
 
        _run_umap_followup(
            ctx,
            result=result,
            tool_args=tool_args,
            ensure_model=mid,
        )