user_sae.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Local user-trained SAE registry and active binding for inspect / steer."""
 
from __future__ import annotations
 
import json
from pathlib import Path
from typing import Any
 
USER_SAE_ROOT = Path.home() / ".aquin" / "sae" / "user"
ACTIVE_SAE_PATH = Path.home() / ".aquin" / "active_user_sae.json"
 
 
def _read_meta(sae_path: Path) -> dict[str, Any]:
    meta_path = sae_path.with_suffix(".meta.json")
    if not meta_path.exists():
        return {}
    try:
        return json.loads(meta_path.read_text(encoding="utf-8"))
    except Exception:
        return {}
 
 
def _layer_from_filename(path: Path) -> int | None:
    stem = path.stem
    if stem.startswith("sae_layer"):
        try:
            return int(stem.replace("sae_layer", ""))
        except ValueError:
            return None
    return None
 
 
def user_sae_dirs_for_model(model_id: str, *, embedding: bool = False) -> list[Path]:
    from aquin.compute.model_loader import resolve_model_id
 
    del embedding  # LLM-only registry
    slug = resolve_model_id(model_id)
    roots = [USER_SAE_ROOT / slug]
    return [r for r in roots if r.is_dir()]
 
 
def resolve_user_sae_path(
    model_id: str,
    name: str,
    layer: int | None = None,
    *,
    embedding: bool | None = None,
) -> Path:
    """Resolve ~/.aquin/sae/user/<model>/<name>/sae_layer{L}.pt"""
    from aquin.compute.model_loader import resolve_model_id
 
    del embedding
    safe = name.replace("/", "--").replace(" ", "_")
    roots = user_sae_dirs_for_model(model_id)
 
    candidates: list[Path] = []
    for root in roots:
        run_dir = root / safe
        if not run_dir.is_dir():
            continue
        if layer is not None:
            candidates.append(run_dir / f"sae_layer{layer}.pt")
        else:
            found = sorted(run_dir.glob("sae_layer*.pt"))
            if len(found) > 1:
                layers = ", ".join(str(_layer_from_filename(p)) for p in found)
                raise ValueError(
                    f"User SAE '{name}' has multiple layers ({layers}). Pass --layer <n>."
                )
            candidates.extend(found)
 
    for path in candidates:
        if path.is_file():
            return path
 
    slug = resolve_model_id(model_id)
    expected = USER_SAE_ROOT / slug / safe
    if layer is not None:
        expected_file = expected / f"sae_layer{layer}.pt"
    else:
        expected_file = expected / "sae_layer<N>.pt"
 
    hints: list[str] = [
        f"Session model resolves to: {slug}",
        f"Expected path: {expected_file}",
    ]
    all_runs = list_user_saes(slug)
    if all_runs:
        hints.append("On disk:")
        for row in all_runs[:8]:
            hints.append(f"  --user {row['name']} --layer {row.get('layer', '?')}  ({row['model_id']})")
    else:
        hints.append(f"No user SAEs under {USER_SAE_ROOT}/")
        hints.append("Train: aquin sae train --layer <n> --name my-run --quick")
 
    layer_hint = f" --layer {layer}" if layer is not None else ""
    raise FileNotFoundError(
        f"No user SAE '{name}' for this model{layer_hint}.\n" + "\n".join(hints)
    )
 
 
def resolve_path_sae(path: str | Path, *, model_id: str | None = None, layer: int | None = None) -> tuple[Path, str, int]:
    """Resolve explicit .pt path; infer model + layer from meta or args."""
    sae_path = Path(path).expanduser().resolve()
    if not sae_path.is_file():
        raise FileNotFoundError(f"SAE file not found: {sae_path}")
 
    from aquin.compute.torch_io import is_valid_sae_checkpoint_path, load_checkpoint
 
    if not is_valid_sae_checkpoint_path(sae_path):
        detail: str | None = None
        try:
            load_checkpoint(sae_path, map_location="cpu")
            detail = "the file loaded, but it does not contain SAE weight tensors"
        except Exception as exc:
            detail = str(exc).strip() or exc.__class__.__name__
        raise ValueError(
            "Invalid SAE checkpoint for --path.\n"
            f"  file: {sae_path}\n"
            "  expected: a non-empty PyTorch/safetensors SAE checkpoint containing encoder/decoder weights\n"
            f"  detail: {detail}\n"
            "Use a real SAE weights file, or re-download/re-export the checkpoint and try again."
        )
 
    meta = _read_meta(sae_path)
    resolved_layer = layer
    if resolved_layer is None:
        resolved_layer = meta.get("layer")
    if resolved_layer is None:
        resolved_layer = _layer_from_filename(sae_path)
    if resolved_layer is None:
        raise ValueError(f"Could not infer layer for {sae_path}. Pass --layer <n>.")
 
    resolved_model = model_id or meta.get("model_id")
    if not resolved_model:
        raise ValueError(f"Could not infer model for {sae_path}. Pass --model <id> or use a .meta.json.")
 
    from aquin.compute.model_loader import resolve_model_id
 
    resolved_model = resolve_model_id(resolved_model)
    return sae_path, resolved_model, int(resolved_layer)
 
 
def list_user_saes(model_id: str | None = None) -> list[dict[str, Any]]:
    """Scan ~/.aquin/sae/user for trained SAE checkpoints."""
    rows: list[dict[str, Any]] = []
    if not USER_SAE_ROOT.is_dir():
        return rows
 
    for model_dir in sorted(USER_SAE_ROOT.iterdir()):
        if not model_dir.is_dir():
            continue
        dir_name = model_dir.name
        if dir_name.startswith("embed-"):
            continue
        slug = dir_name
        if model_id:
            try:
                from aquin.compute.model_loader import resolve_model_id
 
                want = resolve_model_id(model_id)
                if slug != want and dir_name != model_id:
                    continue
            except ValueError:
                if slug != model_id and dir_name != model_id:
                    continue
 
        for run_dir in sorted(model_dir.iterdir()):
            if not run_dir.is_dir():
                continue
            for sae_file in sorted(run_dir.glob("sae_layer*.pt")):
                layer = _layer_from_filename(sae_file)
                meta = _read_meta(sae_file)
                rows.append(
                    {
                        "name": run_dir.name,
                        "model_id": meta.get("model_id") or slug,
                        "layer": layer if layer is not None else meta.get("layer"),
                        "path": str(sae_file),
                        "embedding": False,
                        "steps": meta.get("steps"),
                    }
                )
    return rows
 
 
def load_active_binding() -> dict[str, Any] | None:
    if not ACTIVE_SAE_PATH.exists():
        return None
    try:
        data = json.loads(ACTIVE_SAE_PATH.read_text(encoding="utf-8"))
        return data if isinstance(data, dict) else None
    except Exception:
        return None
 
 
def save_active_binding(
    *,
    model_id: str,
    layer: int,
    path: str | Path,
    name: str | None = None,
    embedding: bool = False,
) -> dict[str, Any]:
    binding = {
        "model_id": model_id,
        "layer": int(layer),
        "path": str(Path(path).resolve()),
        "name": name,
        "embedding": False,
    }
    del embedding
    ACTIVE_SAE_PATH.parent.mkdir(parents=True, exist_ok=True)
    ACTIVE_SAE_PATH.write_text(json.dumps(binding, indent=2), encoding="utf-8")
    return binding
 
 
def clear_active_binding() -> None:
    if ACTIVE_SAE_PATH.exists():
        ACTIVE_SAE_PATH.unlink(missing_ok=True)
 
 
def get_active_user_sae_path(model_id: str, layer: int) -> Path | None:
    """Return active user SAE path when it matches model + layer."""
    binding = load_active_binding()
    if not binding:
        return None
    try:
        bound_layer = int(binding["layer"])
    except (KeyError, TypeError, ValueError):
        return None
    if bound_layer != int(layer):
        return None
 
    bound_model = binding.get("model_id")
    if not bound_model:
        return None
 
    from aquin.compute.model_loader import resolve_model_id
 
    try:
        if resolve_model_id(model_id) != resolve_model_id(str(bound_model)):
            return None
    except ValueError:
        if str(bound_model).lower() != str(model_id).lower():
            return None
 
    path = Path(str(binding.get("path", "")))
    return path if path.is_file() else None
 
 
def activate_user_sae(
    *,
    model_id: str,
    layer: int,
    path: Path,
    name: str | None = None,
) -> dict[str, Any]:
    """Register active user SAE and warm the in-process cache."""
    from aquin.compute.model_loader import resolve_model_id
    from aquin.compute import feature_analysis
 
    short = resolve_model_id(model_id)
 
    binding = save_active_binding(
        model_id=short,
        layer=layer,
        path=path,
        name=name,
        embedding=False,
    )
 
    feature_analysis._sae_cache.pop((short, layer), None)
    feature_analysis.load_sae(short, layer)
 
    return binding