features_cli.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""aquin features compare"""
 
from __future__ import annotations
 
import sys
 
 
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 _print_help() -> None:
    print("Feature-level utilities over saved captures.")
    print("")
    print("Usage: aquin features compare --from <capture_dir> [--group <field>] [--top N] [--encode-sae] [--sae-layer N]")
    print("")
    print("Notes:")
    print("  - compare reads saved capture metadata + SAE feature matrix")
    print("  - if the capture has residuals but no SAE matrix, encodes with the loaded SAE")
    print("  - --group is a probe metadata field (group, label, lang, …); default group")
 
 
def _cmd_features_compare(args: list[str]) -> None:
    if not args or args[0] in ("-h", "--help", "help"):
        print("Usage: aquin features compare --from <capture_dir> [--group <field>] [--top N] [--encode-sae] [--sae-layer N]")
        sys.exit(0 if args and args[0] in ("-h", "--help", "help") else 1)
 
    from aquin.compute.activation_replay import compare_capture_features
 
    source = _parse_flag(args, "--from") or _parse_flag(args, "--dir")
    group = _parse_flag(args, "--group") or "group"
    top_raw = _parse_flag(args, "--top")
    sae_layer_raw = _parse_flag(args, "--sae-layer")
    encode_sae = "--no-encode-sae" not in args
    if not source:
        print("Error: --from <capture_dir> is required")
        sys.exit(1)
    try:
        top_k = int(top_raw) if top_raw else 10
    except ValueError:
        print("Error: --top must be an integer")
        sys.exit(1)
    try:
        sae_layer = int(sae_layer_raw) if sae_layer_raw else None
    except ValueError:
        print("Error: --sae-layer must be an integer")
        sys.exit(1)
 
    try:
        result = compare_capture_features(
            source,
            group=group,
            top_k=top_k,
            encode_sae=encode_sae,
            sae_layer=sae_layer,
        )
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)
 
    print(
        f"[features compare] model={result.get('model_id', '—')}  "
        f"group={result.get('group', '—')}  groups={result.get('n_groups', 0)}  "
        f"sae_layer={result.get('sae_layer', '—')}"
    )
    sizes = result.get("group_sizes") or {}
    if sizes:
        print("[features compare] sizes: " + "  ".join(f"{k}={v}" for k, v in sizes.items()))
    if result.get("group_inferred"):
        print("[features compare] probes had no cohort metadata; split into a/b")
    print("")
    for row in result.get("features") or []:
        groups = row.get("groups") or {}
        group_s = "  ".join(f"{k}={v:.4f}" for k, v in groups.items())
        print(f"f{row.get('feature_idx', '?')}  spread={row.get('spread', 0):.4f}  {group_s}")
 
 
def cmd_features(args: list[str]) -> None:
    if not args or args[0] in ("-h", "--help", "help"):
        _print_help()
        return
    sub = args[0]
    rest = args[1:]
    if sub == "compare":
        _cmd_features_compare(rest)
        return
    print(f"Unknown features subcommand: {sub}")
    print("Run: aquin features --help")
    sys.exit(1)