sae.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.
 
"""SAE — load / train / align / stats / features."""
 
from __future__ import annotations
 
from typing import Any
 
from aquin.sdk._runtime import invoke
 
 
def info(sae_id: str, **kwargs: Any) -> dict[str, Any]:
    from aquin.cli import cmd_info
 
    args = [sae_id]
    if kwargs.get("json"):
        args.append("--json")
    cmd_info(["sae", *args])
    return {"ok": True, "sae_id": sae_id}
 
 
def load(sae_id: str | None = None, **kwargs: Any) -> dict[str, Any]:
    from aquin.cli import cmd_load
 
    argv: list[str] = ["sae"]
    if sae_id:
        argv.append(sae_id)
    if kwargs.get("user"):
        argv.extend(["--user", str(kwargs["user"])])
    if kwargs.get("path"):
        argv.extend(["--path", str(kwargs["path"])])
    if kwargs.get("layer") is not None:
        argv.extend(["--layer", str(kwargs["layer"])])
    cmd_load(argv)
    return {"ok": True}
 
 
def train(**kwargs: Any) -> dict[str, Any]:
    from aquin.sae_cli import cmd_sae_train
 
    argv: list[str] = []
    for key, flag in (
        ("model", "--model"),
        ("layer", "--layer"),
        ("layers", "--layers"),
        ("name", "--name"),
        ("activations", "--activations"),
        ("checkpoint", "--checkpoint"),
    ):
        if kwargs.get(key) is not None:
            argv.extend([flag, str(kwargs[key])])
    if kwargs.get("quick"):
        argv.append("--quick")
    if kwargs.get("balance"):
        argv.append("--balance")
    cmd_sae_train(argv)
    return {"ok": True}
 
 
def align(**kwargs: Any) -> dict[str, Any]:
    from aquin.sae_cli import cmd_sae_align
 
    argv: list[str] = []
    for key, flag in (("sae_a", "--sae-a"), ("sae_b", "--sae-b"), ("sae-a", "--sae-a"), ("sae-b", "--sae-b")):
        if kwargs.get(key) is not None:
            argv.extend([flag, str(kwargs[key])])
    cmd_sae_align(argv)
    return {"ok": True}
 
 
def stats(**kwargs: Any) -> dict[str, Any]:
    return invoke("run_sae_stats", kwargs, command="sae-stats")
 
 
def feature_locate(**kwargs: Any) -> dict[str, Any]:
    return invoke("run_find_feature", kwargs, command="feature locate")
 
 
def feature_logit(**kwargs: Any) -> dict[str, Any]:
    return invoke("get_feature_logits", kwargs, command="feature logit")
 
 
def feature_neighbor(**kwargs: Any) -> dict[str, Any]:
    return invoke("get_feature_neighbors", kwargs, command="feature neighbor")