cli.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited
# SPDX-License-Identifier: Apache-2.0
 
"""AQIT CLI: Recipe, train, EvalGate, inspect, SAE, patch."""
 
from __future__ import annotations
 
import sys
 
# Everything shipped in this tree is AQIT (Aquin framework is from-scratch later).
AQIT_VERBS = frozenset({
    "train",
    "load",
    "unload",
    "prompt",
    "prompting",
    "trace",
    "inspect",
    "feature",
    "feature-logits",
    "feature-neighbors",
    "extract-steer-vector",
    "diff",
    "sae",
    "capture-activations",
    "sweep",
    "activations",
    "find-feature",
    "sae-stats",
    "confidence-analysis",
    "check",
    "eval",
    "check-weights",
    "umap",
    "audit",
    "consistency-eval",
    "suppression-eval",
    "suppress-eval",
    "boundary-eval",
    "attention",
    "features",
    "attribution",
    "perturbation",
    "contrastive",
    "browser",
    "graph",
    "circuit",
    "absorption",
    "polysemy",
    "faithfulness",
    "replay",
    "list",
    "info",
    "red-team",
    "steer",
    "multi-steer",
    "simulate",
    "benchmark",
    "benchmarks",
})
 
 
def _print_help() -> None:
    print("AQIT  Aquin Interp Tooling")
    print("Recipe, train, EvalGate, inspect, SAE, patch. (Aquin framework is separate.)")
    print("")
    print("Usage: aqit <command> ...")
    print("")
    print("  train <recipe.yaml>     DataRevision → train → Checkpoint → EvalGate")
    print("  load model <id>         /  unload  /  prompt")
    print("  eval custom|consistency|suppress|boundary")
    print("  check attention|layer|perturbation|weight|confidence|trajectory")
    print("  diff weight|sae|residue")
    print("  trace / benchmark / --umap")
    print("  sae train|align    info sae / load sae")
    print("  feature locate|logit|neighbor")
    print("  steer / sweep / multi-steer")
    print("  activations capture|replay    features compare")
    print("  simulate / list simulation / replay simulation / compare simulation")
    print("")
    print("SDK:  import aqit")
    print("      aqit.train.run(\"recipe.yaml\")")
    print("      aqit.inspect.trace(...)")
    print("      aqit.patch.steer(...)")
 
 
def dispatch(argv: list[str]) -> None:
    """Run one AQIT command (argv includes the verb)."""
    if not argv or argv[0] in ("-h", "--help", "help"):
        _print_help()
        return
 
    cmd = argv[0]
    rest = argv[1:]
 
    from aquin import cli as aquin_cli
 
    if cmd == "train":
        from aquin.train_cli import cmd_train
 
        cmd_train(rest)
    elif cmd == "load":
        aquin_cli.cmd_load(rest)
    elif cmd == "unload":
        aquin_cli.cmd_unload(rest)
    elif cmd in ("prompt", "prompting"):
        aquin_cli.cmd_prompt(rest)
    elif cmd == "trace":
        aquin_cli.cmd_trace(rest)
    elif cmd == "inspect":
        print("Use: aqit trace", file=sys.stderr)
        sys.exit(1)
    elif cmd == "feature":
        from aquin.feature_cli import cmd_feature
 
        cmd_feature(rest)
    elif cmd == "feature-logits":
        print("Use: aqit feature logit", file=sys.stderr)
        sys.exit(1)
    elif cmd == "feature-neighbors":
        print("Use: aqit feature neighbor", file=sys.stderr)
        sys.exit(1)
    elif cmd == "extract-steer-vector":
        print("Use: aqit steer --feature_idx <n> --save <path>", file=sys.stderr)
        sys.exit(1)
    elif cmd == "replay":
        from aquin.replay_cli import cmd_replay
 
        cmd_replay(rest)
    elif cmd == "list":
        aquin_cli.cmd_list(rest)
    elif cmd == "info":
        aquin_cli.cmd_info(rest)
    elif cmd == "diff":
        from aquin.diff_cli import cmd_diff
 
        cmd_diff(rest)
    elif cmd == "sae":
        from aquin.sae_cli import cmd_sae
 
        cmd_sae(rest)
    elif cmd == "capture-activations":
        from aquin.capture_cli import cmd_capture_activations
 
        cmd_capture_activations(rest)
    elif cmd == "sweep":
        from aquin.sweep_cli import cmd_sweep
 
        cmd_sweep(rest)
    elif cmd == "activations":
        from aquin.activations_cli import cmd_activations
 
        cmd_activations(rest)
    elif cmd == "find-feature":
        print("Use: aqit feature locate", file=sys.stderr)
        sys.exit(1)
    elif cmd == "sae-stats":
        from aquin.sae_stats_cli import cmd_sae_stats
 
        cmd_sae_stats(rest)
    elif cmd == "confidence-analysis":
        print("Use: aqit check confidence", file=sys.stderr)
        sys.exit(1)
    elif cmd == "check":
        from aquin.check_cli import cmd_check
 
        cmd_check(rest)
    elif cmd == "eval":
        from aquin.eval_cli import cmd_eval
 
        cmd_eval(rest)
    elif cmd == "check-weights":
        print("Use: aqit check weight", file=sys.stderr)
        sys.exit(1)
    elif cmd == "umap":
        print("Use: aqit <cmd> --umap  (e.g. aqit feature locate --umap)", file=sys.stderr)
        sys.exit(1)
    elif cmd == "audit":
        print("Use: aqit eval consistency | suppress | boundary", file=sys.stderr)
        sys.exit(1)
    elif cmd == "consistency-eval":
        print("Use: aqit eval consistency", file=sys.stderr)
        sys.exit(1)
    elif cmd in ("suppression-eval", "suppress-eval"):
        print("Use: aqit eval suppress", file=sys.stderr)
        sys.exit(1)
    elif cmd == "boundary-eval":
        print("Use: aqit eval boundary", file=sys.stderr)
        sys.exit(1)
    elif cmd == "attention":
        print("Use: aqit check attention", file=sys.stderr)
        sys.exit(1)
    elif cmd == "features":
        from aquin.features_cli import cmd_features
 
        cmd_features(rest)
    elif cmd in (
        "attribution",
        "perturbation",
        "contrastive",
        "browser",
        "graph",
        "circuit",
        "absorption",
        "polysemy",
        "faithfulness",
    ):
        print(f"Use: aqit check {cmd}", file=sys.stderr)
        sys.exit(1)
    elif cmd in aquin_cli._active_cli_tool_map():
        aquin_cli.cmd_tool(cmd, rest)
    else:
        print(f"Unknown AQIT command: {cmd}")
        print("Run  aqit help")
        sys.exit(1)
 
 
def is_aqit_verb(cmd: str) -> bool:
    if cmd in AQIT_VERBS:
        return True
    from aquin.cli import _active_cli_tool_map
 
    return cmd in _active_cli_tool_map()
 
 
def main() -> None:
    dispatch(sys.argv[1:])