capture_cli.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""aquin capture-activations"""
 
from __future__ import annotations
 
import sys
from pathlib import Path
 
DEFAULT_PROBE_COUNT = 6
 
 
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 _has_flag(args: list[str], name: str) -> bool:
    return name in args
 
 
def _parse_count(args: list[str]) -> int:
    raw = _parse_flag(args, "--count")
    if raw is None:
        return DEFAULT_PROBE_COUNT
    try:
        return max(1, min(int(raw), 64))
    except ValueError:
        print("Error: --count must be an integer between 1 and 64")
        sys.exit(1)
 
 
def _ensure_compute_env() -> None:
    from aquin.compute.loader_shim import apply as _shim_apply
    from aquin.engine.local_server import start as _start_local_server
 
    _shim_apply()
    _start_local_server()
 
 
def _build_capture_card(payload: dict, sync_args: dict) -> dict:
    probes_source = payload.get("probes_source") or sync_args.get("probes_source")
    return {
        "type": "activationCapture",
        "data": {
            "modelId": payload.get("model_id") or sync_args.get("model_id"),
            "modelMode": payload.get("model_mode"),
            "nProbes": payload.get("n_probes"),
            "layers": payload.get("layers") or [],
            "position": payload.get("position") or sync_args.get("position"),
            "encodeSae": bool(payload.get("encode_sae") or sync_args.get("encode_sae")),
            "dModel": payload.get("d_model"),
            "outputDir": payload.get("output_dir") or sync_args.get("dir"),
            "manifestPath": payload.get("manifest_path"),
            "checkpoint": payload.get("checkpoint") or sync_args.get("checkpoint") or None,
            "captureName": payload.get("capture_id") or sync_args.get("name"),
            "probeSamples": payload.get("probe_samples") or [],
            "saeFeatures": payload.get("sae_features"),
            "probesSource": probes_source,
            "topic": payload.get("topic") or sync_args.get("topic"),
            "status": payload.get("status", "done"),
        },
    }
 
 
def _require_loaded_model_id() -> str:
    from aquin.compute.model_loader import get_active_model_id, resolve_model_id
 
    active = (get_active_model_id() or "").strip()
    if not active:
        print("Error: no model loaded. Run: aquin load --model <id>")
        sys.exit(1)
    try:
        return resolve_model_id(active)
    except ValueError as e:
        print(f"Error: {e}")
        sys.exit(1)
 
def cmd_capture_activations(args: list[str]) -> None:
    if _has_flag(args, "--help") or _has_flag(args, "-h"):
        _print_help()
        return
 
    requested_model = _parse_flag(args, "--model")
    if requested_model is not None and _parse_flag(args, "--checkpoint") is None:
        print("Error: capture-activations uses the loaded session model only.")
        print("       Pass --model <id> only when using --checkpoint <path>.")
        sys.exit(1)
 
    prompts_path = _parse_flag(args, "--prompts")
    out_dir_arg = _parse_flag(args, "--dir") or _parse_flag(args, "--output")
    layers_spec = _parse_flag(args, "--layers")
    checkpoint = _parse_flag(args, "--checkpoint")
    name = _parse_flag(args, "--name")
    topic = _parse_flag(args, "--topic")
    position = _parse_flag(args, "--position") or "last"
    granularity = _parse_flag(args, "--granularity") or "prompt"
    balance_group = _parse_flag(args, "--group")
    sae_layer_s = _parse_flag(args, "--sae-layer")
    encode_sae = _has_flag(args, "--encode-sae")
    balance = _has_flag(args, "--balance")
    count = _parse_count(args)
 
    if not out_dir_arg:
        print("Error: --dir <path> is required (alias: --output <path>).", file=sys.stderr)
        print("       Example: aquin capture-activations --dir ./captures/my-run", file=sys.stderr)
        sys.exit(1)
 
    if position not in ("last", "mean"):
        print("Error: --position must be last or mean")
        sys.exit(1)
    if granularity not in ("prompt", "token"):
        print("Error: --granularity must be prompt or token")
        sys.exit(1)
 
    _ensure_compute_env()
 
    from aquin.cli import _build_tool_ctx
    from aquin.compute.activation_capture import (
        llm_layer_count,
        parse_layers,
        resolve_capture_model_id,
        resolve_probes_for_capture,
        run_capture_activations,
    )
    from aquin.compute.model_loader import resolve_model_id
    from aquin.engine.sync_dispatch import require_active_session, sync_cli_result
 
    try:
        if requested_model:
            mid = resolve_model_id(requested_model)
            model_mode = "llm"
        else:
            mid, model_mode = resolve_capture_model_id(_require_loaded_model_id())
    except ValueError as e:
        print(f"Error: {e}")
        sys.exit(1)
 
    if checkpoint and not Path(checkpoint).exists():
        print(f"Checkpoint not found: {checkpoint}")
        sys.exit(1)
 
    out_dir = Path(out_dir_arg)
    out_dir.mkdir(parents=True, exist_ok=True)
 
    try:
        probes, probe_meta = resolve_probes_for_capture(
            model_id=mid,
            model_mode=model_mode,
            prompts_path=prompts_path,
            count=count,
            topic=topic,
            balance=balance,
            balance_group=balance_group,
            output_dir=out_dir if not prompts_path else None,
        )
    except (FileNotFoundError, ValueError) as e:
        print(f"Error: {e}")
        sys.exit(1)
 
    try:
        n_layers = llm_layer_count(mid, checkpoint)
        layer_list = parse_layers(layers_spec, n_layers)
    except ValueError as e:
        print(f"Error: {e}")
        sys.exit(1)
 
    sae_layer = int(sae_layer_s) if sae_layer_s else None
    ckpt_name = name or (Path(checkpoint).stem if checkpoint else "base")
 
    ctx = _build_tool_ctx(model_id=mid)
    require_active_session(ctx, label="aquin capture-activations")
    sync_args = {
        "model_id": mid,
        "model_mode": model_mode,
        "prompts": probe_meta.get("prompts_path") or probe_meta.get("generated_probes_path") or "",
        "probes_source": probe_meta.get("probes_source"),
        "topic": probe_meta.get("topic") or topic,
        "count": count,
        "dir": str(out_dir),
        "layers": layers_spec or "all",
        "checkpoint": checkpoint or "",
        "name": ckpt_name,
        "position": position,
        "granularity": granularity,
        "encode_sae": encode_sae,
        "balance": balance,
        "group": balance_group or "",
    }
 
    print(
        f"[capture] mode={model_mode}  model={mid}  probes={len(probes)}  "
        f"source={probe_meta.get('probes_source')}  layers={layer_list}  "
        f"position={position}  granularity={granularity}  balance={balance}"
        f"{f' group={balance_group}' if balance_group else ''}  → {out_dir}"
    )
    try:
        payload = run_capture_activations(
            mid,
            probes,
            out_dir,
            model_mode=model_mode,
            layers=layer_list,
            checkpoint_path=checkpoint,
            checkpoint_name=ckpt_name if checkpoint else None,
            position=position,  # type: ignore[arg-type]
            granularity=granularity,  # type: ignore[arg-type]
            encode_sae=encode_sae,
            sae_layer=sae_layer,
            capture_name=ckpt_name,
            manifest_extras=probe_meta,
        )
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)
 
    print(f"[capture] wrote {payload['manifest_path']}")
    if payload.get("metadata_path"):
        print(f"[capture] wrote {payload['metadata_path']}")
    for layer, rel in payload.get("activation_files", {}).items():
        print(
            f"[capture]   layer {layer}: {rel}  "
            f"shape=({payload['n_probes']}, {payload.get('d_model', 'd_model')})"
        )
    if payload.get("sae_features"):
        sf = payload["sae_features"]
        print(f"[capture]   sae layer {sf['layer']}: n_features={sf['n_features']}")
 
    card = _build_capture_card(payload, sync_args)
    sync_cli_result(ctx, "run_capture_activations", sync_args, payload, card=card)
 
 
def _print_help() -> None:
    print("Batch activation capture for labeled probe sets (LLM).")
    print("")
    print("Prerequisite: aquin load --model <id>")
    print("")
    print("Usage: aquin capture-activations --dir <path>")
    print("       (alias: --output <path>)")
    print("       [--prompts <json|jsonl>]  (omit to auto-generate probes)")
    print("       [--count N] [--topic <text>]  (default count=6; LLM generates, embed uses templates)")
    print("       [--layers all|8,15] [--checkpoint <path>]  (LLM only)")
    print("       [--position last|mean] [--granularity prompt|token] [--encode-sae] [--sae-layer N] [--name <label>] [--balance] [--group <field>]")
    print("")
    print("  Without --prompts: loaded LLM writes probe lines.")
    print("  Generated probes saved to <dir>/probes.jsonl.")
    print("  Always writes root manifest.json + metadata.json + summary.jsonl.")
    print("  --granularity token saves per-token activations + token_spans.jsonl.")
    print("  --balance samples evenly across metadata groups (label, stressor, lang, group) when present.")
    print("  --group chooses the metadata field used with --balance.")
    print("  Syncs activationCapture card.")
    print("")
    print("Docs: https://aquin.app/docs/sae-training")