activations_cli.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""aquin activations capture | replay"""
 
from __future__ import annotations
 
import json
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("Recorded activation workflows.")
    print("")
    print("Usage: aquin activations <capture|replay> ...")
    print("")
    print("  capture   Record activations + metadata into a capture directory")
    print("            aquin activations capture --dir ./captures/run1 [--prompts probes.jsonl]")
    print("")
    print("  replay    Slice recorded activation metadata by filter/group")
    print("            aquin activations replay --from ./captures/run1 [--filter lang=hi] [--group-by layer]")
    print("")
    print("Notes:")
    print("  - capture always writes root metadata.json + manifest.json + summary.jsonl")
    print("  - replay reads saved metadata; it does not re-run the model")
    print("")
    print("Docs: https://aquin.app/docs/sae-training")
 
 
def _cmd_replay_activations(args: list[str]) -> None:
    if not args or args[0] in ("-h", "--help", "help"):
        print("Usage: aquin activations replay --from <capture_dir> [--filter key=value,...] [--group-by key] [--limit N]")
        sys.exit(0 if args and args[0] in ("-h", "--help", "help") else 1)
 
    from aquin.compute.activation_replay import replay_capture
 
    source = _parse_flag(args, "--from") or _parse_flag(args, "--dir")
    if not source and args and not args[0].startswith("-"):
        source = args[0]
    if not source:
        print("Error: --from <capture_dir> is required")
        sys.exit(1)
 
    filter_expr = _parse_flag(args, "--filter")
    group_by = _parse_flag(args, "--group-by")
    limit_raw = _parse_flag(args, "--limit")
    try:
        limit = int(limit_raw) if limit_raw else 20
    except ValueError:
        print("Error: --limit must be an integer")
        sys.exit(1)
 
    try:
        result = replay_capture(source, filter_expr=filter_expr, group_by=group_by, limit=limit)
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)
 
    print(f"[activations replay] root={result['root']}")
    print(
        f"[activations replay] model={result.get('model_id', '—')}  "
        f"mode={result.get('model_mode', '—')}  matched={result.get('n_filtered', 0)}/{result.get('n_rows', 0)}"
    )
    if result.get("filter"):
        print(f"[activations replay] filter={result['filter']}")
    if result.get("group_by"):
        print(f"[activations replay] group_by={result['group_by']}")
 
    groups = result.get("groups") or []
    if groups:
        print("")
        print("Groups")
        for row in groups:
            layers = ", ".join(f"L{k}:{v}" for k, v in (row.get("layers") or {}).items())
            print(f"  {row.get('group', 'unknown')}: {row.get('count', 0)}  {layers}".rstrip())
        return
 
    rows = result.get("rows") or []
    if not rows:
        print("No matching activation rows.")
        return
    print("")
    print("Rows")
    for row in rows:
        text = " ".join(str(row.get("text") or "").split())
        if len(text) > 72:
            text = text[:71] + "…"
        meta = []
        for key in ("label", "lang", "stressor", "group", "quant_run_id"):
            if row.get(key) is not None:
                meta.append(f"{key}={row[key]}")
        meta_s = "  ".join(meta)
        print(f"  {row.get('probe_id', '?')}  L{row.get('layer', '?')}  {meta_s}".rstrip())
        if text:
            print(f"    {text}")
 
 
def cmd_activations(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 == "capture":
        from aquin.capture_cli import cmd_capture_activations
 
        cmd_capture_activations(rest)
        return
 
    if sub == "replay":
        _cmd_replay_activations(rest)
        return
 
    print(f"Unknown activations subcommand: {sub}")
    print("Run: aquin activations --help")
    sys.exit(1)