sae_cli.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""aquin diff sae — base vs checkpoint SAE activation diff (via aquin diff sae)."""
 
import json
import sys
from pathlib import Path
 
from aquin.cli_flags import reject_legacy_output_flags
 
 
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 _resolve_sae_model(model_id: str) -> str:
    """Return canonical model id."""
    from aquin.compute.model_loader import resolve_model_id
 
    return resolve_model_id(model_id)
 
 
def _ensure_compute_env() -> None:
    """Match cmd_tool: shared model + local HTTP server for web popovers."""
    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 _build_sae_card(tool_name: str, payload: dict, sync_args: dict | None = None) -> dict | None:
    if tool_name == "run_sae_diff":
        return {"type": "saeDiff", "data": payload}
    if tool_name == "run_sae_train":
        args = sync_args or {}
        return {
            "type": "saeTrain",
            "data": {
                "modelId": args.get("model_id") or payload.get("model_id"),
                "layer": args.get("layer") if args.get("layer") is not None else payload.get("layer"),
                "quick": bool(args.get("quick", payload.get("quick"))),
                "name": args.get("name") or payload.get("name"),
                "outputPath": payload.get("output_path") or args.get("save"),
                "checkpoint": args.get("checkpoint") or payload.get("checkpoint"),
                "status": payload.get("status", "done"),
            },
        }
    if tool_name == "run_sae_align":
        pairs = payload.get("pairs") or []
        sorted_pairs = sorted(pairs, key=lambda x: float(x.get("cosine", 0)))
        args = sync_args or {}
        return {
            "type": "saeAlign",
            "data": {
                "meanCosine": payload.get("mean_cosine"),
                "nPairs": payload.get("n_pairs"),
                "saeA": args.get("sae_a") or payload.get("sae_a"),
                "saeB": args.get("sae_b") or payload.get("sae_b"),
                "lowestPairs": sorted_pairs[:8],
                "highestPairs": list(reversed(sorted_pairs[-8:])),
            },
        }
    return None
 
 
def _sync_sae_result(
    ctx: dict,
    tool_name: str,
    args: dict,
    payload: dict,
) -> None:
    from aquin.engine.sync_dispatch import sync_cli_result
 
    card = _build_sae_card(tool_name, payload, args)
    sync_cli_result(ctx, tool_name, args, payload, card=card)
 
 
def _print_sae_diff(payload: dict) -> None:
    print(f"model          : {payload.get('baseModelId')}")
    print(f"checkpoint     : {payload.get('ftCheckpointName')}")
    print(f"layer          : {payload.get('layer')}")
    print(f"changed        : {payload.get('nChanged')}/{payload.get('nFeatures')}")
    print(f"mean |delta|   : {payload.get('meanAbsDelta')}")
    print(f"max  |delta|   : {payload.get('maxAbsDelta')}")
    print("\ntop feature deltas:")
    for row in (payload.get("featureDeltas") or [])[:15]:
        print(
            f"  f{row['feature_idx']:>5}  base={row['base_act']:.4f}  "
            f"ft={row['ft_act']:.4f}  Δ={row['delta']:+.4f}"
        )
 
 
def _require_session_model(model_id: str | None) -> str:
    from aquin.compute.model_loader import get_active_model_id, resolve_model_id
 
    locked = (get_active_model_id() or "").strip()
    if not locked:
        print("Error: no model loaded.")
        print("       Run: aquin load --model <id>")
        sys.exit(1)
    if model_id and model_id != locked:
        try:
            if resolve_model_id(model_id) != resolve_model_id(locked):
                print(f"Error: loaded model is '{locked}'.")
                print("       --model must match the loaded model or be omitted.")
                sys.exit(1)
        except ValueError:
            pass
    try:
        return resolve_model_id(locked)
    except ValueError:
        return locked
 
 
def _resolve_train_model_id(model_id: str | None, activations: str | None) -> str:
    """Session model, explicit --model, or capture-dir manifest (activations-only)."""
    from aquin.compute.activation_store import read_manifest
    from aquin.compute.model_loader import get_active_model_id, resolve_model_id
 
    locked = (get_active_model_id() or "").strip()
    if locked:
        return _require_session_model(model_id)
 
    if model_id:
        try:
            return resolve_model_id(model_id)
        except ValueError:
            raise
 
    if activations:
        root = Path(activations).expanduser()
        manifest = read_manifest(root)
        if not manifest:
            # capture dirs nest training acts under _acts_layerN after resolve
            for child in root.glob("**/manifest.json"):
                manifest = read_manifest(child.parent)
                if manifest and manifest.get("model_id"):
                    break
        mid = (manifest or {}).get("model_id") if manifest else None
        if isinstance(mid, str) and mid.strip():
            try:
                return resolve_model_id(mid.strip())
            except ValueError:
                pass
 
    print("Error: no model loaded.")
    print("       Run: aquin load --model <id>")
    print("       Or pass --model <id> / train from --activations with a capture manifest.")
    sys.exit(1)
 
 
def cmd_sae_diff(args: list[str]) -> None:
    reject_legacy_output_flags(args)
    model_id = _parse_flag(args, "--model")
    checkpoint = _parse_flag(args, "--checkpoint")
    prompts_path = _parse_flag(args, "--prompts")
    layer_s = _parse_flag(args, "--layer")
    sae_path = _parse_flag(args, "--sae")
    name = _parse_flag(args, "--name") or Path(checkpoint or "checkpoint").stem
    save_path = _parse_flag(args, "--save")
 
    if not checkpoint:
        print("Usage: aquin diff sae [--model <id>] --checkpoint <path> [--prompts <json|jsonl>]")
        print("       [--layer N] [--sae <path>] [--name <label>] [--save <json>]")
        sys.exit(1)
 
    model_id = _require_session_model(model_id)
 
    _ensure_compute_env()
 
    from aquin.cli import _build_tool_ctx
    from aquin.compute.sae_diff import load_prompts, run_sae_diff
 
    try:
        mid = _resolve_sae_model(model_id)
    except ValueError as e:
        print(f"Error: {e}")
        sys.exit(1)
 
    ckpt = Path(checkpoint)
    if not ckpt.exists():
        print(f"Checkpoint not found: {ckpt}")
        sys.exit(1)
 
    layer = int(layer_s) if layer_s else None
    prompts = load_prompts(prompts_path)
    ctx = _build_tool_ctx(model_id=mid)
    from aquin.engine.sync_dispatch import require_active_session
 
    require_active_session(ctx, label="aquin diff sae")
    sync_args = {
        "model_id": mid,
        "checkpoint": str(ckpt),
        "name": name,
        "prompts": prompts_path or "",
        "model_mode": "llm",
    }
 
    try:
        print(f"[diff sae] llm base={mid}  target={ckpt.name}  prompts={len(prompts)}")
        payload = run_sae_diff(
            mid,
            prompts,
            target_checkpoint=str(ckpt),
            checkpoint_name=name,
            layer=layer,
            sae_path=sae_path,
        )
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)
 
    _print_sae_diff(payload)
    if save_path:
        Path(save_path).write_text(json.dumps(payload, indent=2), encoding="utf-8")
        print(f"\n[diff sae] wrote {save_path}")
    _sync_sae_result(ctx, "run_sae_diff", sync_args, payload)
 
 
def _parse_train_layers(spec: str, n_layers: int) -> list[int]:
    raw = (spec or "").strip().lower()
    if raw in ("all", "*"):
        return list(range(n_layers))
    out: list[int] = []
    for part in raw.split(","):
        part = part.strip()
        if not part:
            continue
        layer = int(part)
        if layer < 0 or layer >= n_layers:
            raise ValueError(f"Layer {layer} out of range for model (0..{n_layers - 1})")
        out.append(layer)
    if not out:
        raise ValueError("--layers must be 'all' or comma-separated indices, e.g. --layers 0,2,8,15")
    return out
 
 
def _resolve_train_name(
    name: str | None,
    layer: int,
    *,
    multi: bool,
    checkpoint: str | None = None,
) -> str:
    if name:
        if "{layer}" in name:
            return name.replace("{layer}", str(layer))
        if multi:
            return f"{name}-l{layer}"
        return name
    if checkpoint:
        stem = Path(checkpoint).stem
        return f"{stem}-l{layer}" if multi else stem
    return f"base-l{layer}" if multi else "base"
 
 
def _train_one_layer(
    *,
    mid: str,
    layer: int,
    tag: str,
    ctx: dict,
    checkpoint: str | None,
    corpus: str | None,
    activations: str | None,
    quick: bool,
    balance: bool,
    balance_group: str | None,
    save_path: str | None,
    max_steps: int | None = None,
    max_epochs: int | None = None,
) -> None:
    from aquin.compute.sae_train import default_user_sae_path, train_sae
 
    out = Path(save_path) if save_path else default_user_sae_path(mid, tag, layer)
    sync_args = {
        "model_id": mid,
        "layer": layer,
        "checkpoint": checkpoint or "",
        "quick": quick,
        "name": tag,
        "output": str(out),
        "activations": activations or "",
        "balance": balance,
        "group": balance_group or "",
        "model_mode": "llm",
    }
    acts_note = f" activations={activations}" if activations else ""
    bal_note = f" balance=True group={balance_group}" if balance and balance_group else (" balance=True" if balance else "")
    print(f"[sae train] llm model={mid} layer={layer} quick={quick}{acts_note}{bal_note} -> {out}")
 
    train_sae(
        mid,
        layer,
        out,
        checkpoint_path=checkpoint,
        corpus_path=corpus,
        quick=quick,
        activations_dir=activations,
        balance=balance,
        balance_group=balance_group,
        max_steps=max_steps,
        max_epochs=max_epochs,
    )
 
    payload = {**sync_args, "status": "done", "output_path": str(out)}
    print(f"[sae train] done: {out}")
    print(f"[sae train] load for inspect: aquin load sae --user {tag} --layer {layer}")
    print(f"[sae train] diff vs public: aquin diff sae --checkpoint <path>  |  aquin sae align --sae-a <public> --sae-b {out}")
    _sync_sae_result(ctx, "run_sae_train", sync_args, payload)
 
 
def cmd_sae_train(args: list[str]) -> None:
    reject_legacy_output_flags(args)
    model_id = _parse_flag(args, "--model")
    checkpoint = _parse_flag(args, "--checkpoint")
    layer_s = _parse_flag(args, "--layer")
    layers_s = _parse_flag(args, "--layers")
    corpus = _parse_flag(args, "--corpus")
    activations = _parse_flag(args, "--activations")
    name = _parse_flag(args, "--name")
    save_path = _parse_flag(args, "--save")
    quick = _has_flag(args, "--quick")
    balance = _has_flag(args, "--balance")
    balance_group = _parse_flag(args, "--group")
    max_steps_s = _parse_flag(args, "--max-steps")
    max_epochs_s = _parse_flag(args, "--max-epochs")
    max_steps = int(max_steps_s) if max_steps_s else None
    max_epochs = int(max_epochs_s) if max_epochs_s else None
 
    if not layer_s and not layers_s:
        print("Usage: aquin sae train [--model <id>] --layer <n> | --layers <all|0,1,2,...>")
        print("       [--checkpoint <path>] [--corpus <json|jsonl>] [--activations <dir>]")
        print("       [--quick] [--max-steps N] [--max-epochs N] [--balance] [--group <field>]")
        print("       [--name <tag>] [--save <path>]")
        print("")
        print("  Full train (default): 2M tokens, up to 50k steps / 200 epochs per layer")
        print("  --name 'lfm-instruct-l{layer}-v1'  use {layer} when training multiple layers")
        print("")
        print("  aquin sae train --layers all --name 'lfm-instruct-l{layer}-v1'")
        sys.exit(1)
 
    if layer_s and layers_s:
        print("Error: pass --layer or --layers, not both.")
        sys.exit(1)
 
    if save_path and layers_s:
        print("Error: --save is only valid with a single --layer.")
        sys.exit(1)
 
    model_id = _resolve_train_model_id(model_id, activations)
 
    _ensure_compute_env()
 
    from aquin.cli import _build_tool_ctx
    from aquin.compute.model_loader import get_active_model_id, get_config
 
    try:
        mid = _resolve_sae_model(model_id)
    except ValueError as e:
        print(f"Error: {e}")
        sys.exit(1)
 
    n_layers = int(get_config(mid)["n_layers"])
    if layers_s:
        try:
            layers = _parse_train_layers(layers_s, n_layers)
        except ValueError as e:
            print(f"Error: {e}")
            sys.exit(1)
    else:
        layers = [int(layer_s)]
 
    ctx = _build_tool_ctx(model_id=mid)
    from aquin.engine.sync_dispatch import require_active_session
 
    # Activations-only train never needs a resident GPU model.
    if not activations or (get_active_model_id() or "").strip():
        require_active_session(ctx, label="aquin sae train")
 
    if checkpoint and not Path(checkpoint).exists():
        print(f"Checkpoint not found: {checkpoint}")
        sys.exit(1)
 
    if activations and not Path(activations).expanduser().exists():
        print(f"Activations directory not found: {activations}")
        sys.exit(1)
 
    if activations and corpus:
        print("Note: --corpus is ignored when --activations is set (no forward passes).")
    if balance_group and not balance:
        print("Note: --group is ignored unless --balance is set.")
    if balance and not activations:
        print("Note: --balance currently applies when training from metadata-tagged capture dirs via --activations.")
 
    multi = len(layers) > 1
    if multi:
        print(f"[sae train] training {len(layers)} layers: {layers}")
 
    for i, layer in enumerate(layers):
        if multi:
            print(f"\n[sae train] layer {layer} ({i + 1}/{len(layers)})")
        tag = _resolve_train_name(name, layer, multi=multi, checkpoint=checkpoint)
        try:
            _train_one_layer(
                mid=mid,
                layer=layer,
                tag=tag,
                ctx=ctx,
                checkpoint=checkpoint,
                corpus=corpus,
                activations=activations,
                quick=quick,
                balance=balance,
                balance_group=balance_group,
                save_path=save_path,
                max_steps=max_steps,
                max_epochs=max_epochs,
            )
        except Exception as e:
            print(f"Error on layer {layer}: {e}")
            sys.exit(1)
 
    if multi:
        print(f"\n[sae train] all {len(layers)} layers complete.")
 
 
def _resolve_user_sae_path(model_id: str, run_name: str, layer: int) -> Path:
    from aquin.compute.model_loader import resolve_model_id
    from aquin.compute.user_sae import USER_SAE_ROOT
 
    slug = resolve_model_id(model_id)
    path = USER_SAE_ROOT / slug / run_name.format(layer=layer) / f"sae_layer{layer}.pt"
    if not path.is_file():
        raise FileNotFoundError(str(path))
    return path
 
 
def cmd_sae_catalog_metrics(args: list[str]) -> None:
    model_id = _parse_flag(args, "--model")
    sae_path_s = _parse_flag(args, "--sae")
    layer_s = _parse_flag(args, "--layer")
    user_run = _parse_flag(args, "--user-run")
    layers_s = _parse_flag(args, "--layers")
    skip_interp = _has_flag(args, "--skip-interp")
    n_features_s = _parse_flag(args, "--n-features")
 
    if not sae_path_s and not user_run:
        print("Usage: aquin sae catalog-metrics [--model <id>]")
        print("       --sae <path> --layer <n>")
        print("       | --user-run <name> --layers <all|0,1,2,...>")
        print("       [--skip-interp] [--n-features N]")
        print("")
        print("  Writes sae_layerN.umap.json + sae_layerN.metrics.json next to each checkpoint.")
        print("  Then publish the SAE files from this directory.")
        print("")
        print("  aquin sae catalog-metrics --user-run 'lfm-230m-l{layer}-v1' --layers all")
        sys.exit(1)
 
    if user_run and not layers_s:
        print("Error: --user-run requires --layers (e.g. --layers all)")
        sys.exit(1)
    if sae_path_s and not layer_s and not user_run:
        print("Error: --sae requires --layer")
        sys.exit(1)
    if user_run and sae_path_s:
        print("Error: pass --sae + --layer, or --user-run + --layers, not both.")
        sys.exit(1)
 
    model_id = _require_session_model(model_id)
    _ensure_compute_env()
 
    from aquin.compute.model_loader import get_config
    from aquin.compute.sae_catalog_metrics import compute_and_write_catalog_sidecars
 
    try:
        from aquin.compute.model_loader import resolve_model_id
 
        mid = resolve_model_id(model_id)
    except ValueError as e:
        print(f"Error: {e}")
        sys.exit(1)
 
    n_layers = int(get_config(mid)["n_layers"])
    n_features = int(n_features_s) if n_features_s else 8
 
    if user_run:
        try:
            layers = _parse_train_layers(layers_s or "", n_layers)
        except ValueError as e:
            print(f"Error: {e}")
            sys.exit(1)
        jobs = [(layer, _resolve_user_sae_path(mid, user_run, layer)) for layer in layers]
    else:
        jobs = [(int(layer_s), Path(sae_path_s).expanduser().resolve())]
 
    for i, (layer, sae_path) in enumerate(jobs):
        if len(jobs) > 1:
            print(f"\n[sae catalog-metrics] layer {layer} ({i + 1}/{len(jobs)})")
        try:
            compute_and_write_catalog_sidecars(
                model_id=mid,
                layer=layer,
                sae_path=sae_path,
                skip_interp=skip_interp,
                n_features=n_features,
            )
        except Exception as e:
            print(f"Error on layer {layer}: {e}")
            sys.exit(1)
 
    if len(jobs) > 1:
        print(f"\n[sae catalog-metrics] all {len(jobs)} layers complete.")
 
 
def _default_user_sae_checkpoint() -> Path | None:
    from aquin.compute.user_sae import list_user_saes
 
    rows = list_user_saes()
    if not rows:
        return None
    smoke = [r for r in rows if str(r.get("name") or "") == "chat-smoke"]
    pick = smoke[0] if smoke else rows[0]
    path = pick.get("path")
    return Path(str(path)).expanduser() if path else None
 
 
def _is_dictionary_sae_path(path: Path) -> bool:
    sae_root = (Path.home() / ".aquin" / "sae").resolve()
    try:
        resolved = path.expanduser().resolve()
    except OSError:
        return False
    if not resolved.is_file():
        return False
    try:
        resolved.relative_to(sae_root)
    except ValueError:
        return False
    return resolved.name.startswith("sae_layer") and resolved.suffix == ".pt"
 
 
def _resolve_align_sae_path(raw: str | None) -> Path | None:
    default = _default_user_sae_checkpoint()
    if not raw:
        return default
    path = Path(raw).expanduser()
    if _is_dictionary_sae_path(path):
        return path.resolve()
    if default is not None:
        print(
            f"[sae align] {path} is not an SAE dictionary under ~/.aquin/sae; using {default}",
            flush=True,
        )
        return default
    return path if path.is_file() else None
 
 
def cmd_sae_align(args: list[str]) -> None:
    reject_legacy_output_flags(args)
    sae_a = _parse_flag(args, "--sae-a") or _parse_flag(args, "--sae_a")
    sae_b = _parse_flag(args, "--sae-b") or _parse_flag(args, "--sae_b")
    save_path = _parse_flag(args, "--save")
    max_f = _parse_flag(args, "--max-features") or _parse_flag(args, "--max_features")
 
    path_a = _resolve_align_sae_path(sae_a)
    path_b = _resolve_align_sae_path(sae_b)
    if path_a is None or path_b is None:
        print("Usage: aquin sae align --sae-a <path> --sae-b <path> [--save <map.json>] [--max-features N]")
        print("       Dictionaries live under ~/.aquin/sae/user/<model>/<name>/sae_layerN.pt")
        sys.exit(1)
 
    _ensure_compute_env()
 
    from aquin.cli import _build_tool_ctx
    from aquin.compute.sae import SparseAutoencoder
    from aquin.compute.sae_diff import align_sae_decoders
    import torch
 
    ctx = _build_tool_ctx()
    from aquin.engine.sync_dispatch import require_active_session
 
    require_active_session(ctx, label="aquin sae align")
    sync_args = {"sae_a": str(path_a), "sae_b": str(path_b), "save": save_path or ""}
 
    from aquin.compute.device import resolve_compute_device
    device = resolve_compute_device()
    a = SparseAutoencoder.load(path_a, device=device)
    b = SparseAutoencoder.load(path_b, device=device)
    if not torch.isfinite(a.W_dec.data).all() or not torch.isfinite(b.W_dec.data).all():
        print("[sae align] decoder has NaNs; matching encoder directions instead", flush=True)
    if a.n_features != b.n_features or a.d_model != b.d_model:
        print(f"Warning: shape mismatch A=({a.n_features},{a.d_model}) B=({b.n_features},{b.d_model})")
 
    mapping = align_sae_decoders(
        a, b, max_features=int(max_f) if max_f else 256,
    )
    cosines = [m["cosine"] for m in mapping]
    mean_cos = sum(cosines) / len(cosines) if cosines else 0.0
    print(f"aligned {len(mapping)} feature pairs  mean cosine={mean_cos:.4f}")
 
    low = sorted(mapping, key=lambda x: x["cosine"])[:5]
    print("lowest cosine matches:")
    for row in low:
        print(f"  A[{row['feature_a']}] <-> B[{row['feature_b']}]  cos={row['cosine']:.4f}")
 
    result = {"pairs": mapping, "mean_cosine": round(mean_cos, 6), "n_pairs": len(mapping)}
    if save_path:
        Path(save_path).write_text(json.dumps(result, indent=2), encoding="utf-8")
        print(f"[sae align] wrote {save_path}")
    _sync_sae_result(ctx, "run_sae_align", sync_args, result)
 
 
def cmd_sae(args: list[str]) -> None:
    _OLD_EMBED_SUBS = {
        "features", "contrastive", "interp", "browser", "graph",
        "circuit", "steer", "absorption", "polysemy", "faithfulness",
    }
 
    if not args or args[0] in ("-h", "--help", "help"):
        print("Checkpoint SAE — train / align / publish on fine-tuned weights.")
        print("")
        print("Prerequisite: aquin load --model <id> && aquin load sae <model-l{n}>")
        print("  Example SAE id: llama-3.2-1b-l8")
        print("")
        print("Usage: aquin sae <train|align|catalog-metrics|publish> ...")
        print("")
        print("  train  --model <id> --layer <n> [--checkpoint <path>] [--quick] [--corpus <file>]")
        print("         [--activations <dir>] [--name <tag>] [--save <path>]")
        print("")
        print("  align  --sae-a <path> --sae-b <path> [--save <json>] [--max-features <n>]")
        print("")
        print("  catalog-metrics / publish  (public catalog pipeline)")
        print("  For base vs checkpoint feature deltas: aquin diff sae --checkpoint <path>")
        print("")
        print("Embedding SAE inspection:")
        print("  aquin check features | contrastive | browser | graph | circuit")
        print("  aquin check absorption | polysemy | faithfulness")
        print("  aquin interp | steer | decomp   (bare verbs)")
        print("")
        print("Docs: https://aquin.app/docs/checkpoint-sae")
        return
 
    sub = args[0]
    rest = args[1:]
 
    if sub in _OLD_EMBED_SUBS:
        if sub in ("steer", "interp"):
            print(f"Use: aquin {sub}", file=sys.stderr)
        else:
            print(f"Use: aquin check {sub}", file=sys.stderr)
        sys.exit(1)
 
    if sub == "train":
        cmd_sae_train(rest)
    elif sub == "align":
        cmd_sae_align(rest)
    elif sub in ("catalog-metrics", "catalog_metrics"):
        cmd_sae_catalog_metrics(rest)
    elif sub == "publish":
        from aquin.publish_public_sae import publish_public_sae
 
        publish_public_sae(rest)
    else:
        print(f"Unknown sae subcommand: {sub}")
        print("Run: aquin sae help")
        sys.exit(1)