find_feature_cli.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""aquin feature locate — rank SAE features for honest vs deceptive probes."""
 
import json
import sys
from pathlib import Path
 
from aquin.cli_flags import reject_legacy_output_flags
 
 
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_int_flag(args: list[str], name: str, default: int) -> int:
    raw = _parse_flag(args, name)
    if raw is None:
        return default
    try:
        return int(raw)
    except ValueError:
        print(f"Error: {name} must be an integer")
        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 _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 _build_find_feature_card(payload: dict, sync_args: dict) -> dict:
    return {
        "type": "findFeature",
        "data": {
            "modelId": payload.get("model_id") or sync_args.get("model_id"),
            "layer": payload.get("layer"),
            "scorer": payload.get("scorer"),
            "nHonest": payload.get("n_honest"),
            "nDeceptive": payload.get("n_deceptive"),
            "promptsPath": payload.get("prompts_path") or sync_args.get("prompts"),
            "checkpoint": payload.get("checkpoint") or sync_args.get("checkpoint"),
            "chosenFeatureIdx": payload.get("chosen_feature_idx"),
            "chosenDelta": payload.get("chosen_delta"),
            "direction": payload.get("direction"),
            "conditioning": payload.get("conditioning"),
            "behavior": payload.get("behavior"),
            "warning": payload.get("warning"),
            "persistedKey": payload.get("persisted_key"),
            "experimentPath": payload.get("experiment_path"),
            "rankings": payload.get("rankings") or [],
            "status": payload.get("status", "done"),
        },
    }
 
 
def _print_find_feature(payload: dict) -> None:
    print(f"model     : {payload.get('model_id')}")
    print(f"layer     : {payload.get('layer')}")
    print(f"scorer    : {payload.get('scorer')}")
    if payload.get("direction"):
        print(f"direction : {payload.get('direction')}")
    if payload.get("conditioning"):
        print(f"condition : {payload.get('conditioning')}")
    behavior = payload.get("behavior")
    if isinstance(behavior, dict):
        print(
            f"behavior  : {behavior.get('n_truthful', '?')} truthful · "
            f"{behavior.get('n_deceptive', '?')} deceptive · "
            f"{behavior.get('n_ambiguous', '?')} ambiguous "
            f"({behavior.get('n_generated', '?')} generated)"
        )
    print(f"probes    : {payload.get('n_honest')} honest · {payload.get('n_deceptive')} deceptive")
    if payload.get("prompts_path"):
        print(f"prompts   : {payload.get('prompts_path')}")
    chosen = payload.get("chosen_feature_idx")
    if chosen is not None:
        print(f"\nchosen    : feature {chosen}  Δ={payload.get('chosen_delta')}")
    if payload.get("persisted_key"):
        print(f"persisted : {payload.get('persisted_key')} → {payload.get('experiment_path')}")
    if payload.get("warning"):
        print(f"\nwarning   : {payload.get('warning')}")
    rankings = payload.get("rankings") or []
    if rankings:
        direction = payload.get("direction") or "both"
        rank_label = {
            "deceptive": "top features by Δ (deceptive > honest)",
            "honest": "top features by |Δ| (honest > deceptive)",
            "both": "top features by |Δ| (deceptive − honest)",
        }.get(direction, "top features")
        print(f"\n{rank_label} ({min(len(rankings), 10)} shown):")
        for i, row in enumerate(rankings[:10], 1):
            interp = row.get("interp_score")
            extra = f"  interp={interp:.2f}" if interp is not None else ""
            print(
                f"  {i:2}. f{row['feature_idx']:<5}  "
                f"honest={row['honest_mean']:.4f}  deceptive={row['deceptive_mean']:.4f}  "
                f"Δ={row['delta']:+.4f}{extra}"
            )
 
 
def cmd_find_feature(args: list[str]) -> None:
    if _has_flag(args, "--help") or _has_flag(args, "-h"):
        _print_help()
        return
 
    if _parse_flag(args, "--model") is not None:
        print("Error: feature locate uses the loaded session model only.")
        print("       Run: aquin load --model <id>")
        sys.exit(1)
 
    reject_legacy_output_flags(args)
    scorer = _parse_flag(args, "--scorer") or "deception"
    direction = _parse_flag(args, "--direction") or "both"
    conditioning = _parse_flag(args, "--conditioning") or "behavior"
    prompts = _parse_flag(args, "--prompts")
    layer_s = _parse_flag(args, "--layer")
    checkpoint = _parse_flag(args, "--checkpoint")
    persist = _parse_flag(args, "--persist")
    save_path = _parse_flag(args, "--save")
    top_k = _parse_int_flag(args, "--top", 20)
    benchmark_top = _parse_int_flag(args, "--benchmark-top", 0)
    want_umap = _has_flag(args, "--umap")
 
    if checkpoint and not Path(checkpoint).exists():
        print(f"Checkpoint not found: {checkpoint}")
        sys.exit(1)
 
    _ensure_compute_env()
 
    from aquin.cli import _build_tool_ctx
    from aquin.engine.sync_dispatch import require_active_session, sync_cli_result
 
    mid = _require_loaded_model_id()
    layer = int(layer_s) if layer_s else None
    ctx = _build_tool_ctx(model_id=mid)
    require_active_session(ctx, label="aquin feature locate")
 
    sync_args = {
        "model_id": mid,
        "scorer": scorer,
        "direction": direction,
        "conditioning": conditioning,
        "prompts": prompts or "",
        "layer": layer,
        "checkpoint": checkpoint or "",
        "top_k": top_k,
        "persist": persist or "",
    }
 
    openai_client = None
    if benchmark_top > 0:
        try:
            from aquin.compute.openai_client import get_openai_client
 
            openai_client = get_openai_client(ctx)
        except Exception:
            print("[feature locate] --benchmark-top ignored (OpenAI client unavailable)", flush=True)
            benchmark_top = 0
 
    session_id = ctx.get("session_id") or ctx.get("state", {}).get("session_id")
 
    try:
        from aquin.compute.find_feature import run_find_feature
 
        print(f"[feature locate] model={mid}  scorer={scorer}  direction={direction}  conditioning={conditioning}  top={top_k}")
        payload = run_find_feature(
            mid,
            scorer=scorer,
            prompts_path=prompts,
            layer=layer,
            checkpoint_path=checkpoint,
            top_k=top_k,
            direction=direction,
            conditioning=conditioning,
            benchmark_top=benchmark_top,
            persist_key=persist,
            session_id=str(session_id) if session_id else None,
            openai_client=openai_client,
        )
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)
 
    _print_find_feature(payload)
    if save_path:
        Path(save_path).write_text(json.dumps(payload, indent=2), encoding="utf-8")
        print(f"\n[feature locate] wrote {save_path}")
 
    card = _build_find_feature_card(payload, sync_args)
    sync_cli_result(ctx, "run_find_feature", sync_args, payload, card=card)
 
    if want_umap:
        from aquin.cli import _run_umap_followup
 
        _run_umap_followup(
            ctx,
            result=payload,
            tool_args=sync_args,
            ensure_model=mid,
            layer=layer,
        )
 
 
def _print_help() -> None:
    print("Rank SAE features that separate honest vs deceptive probes (LLM).")
    print("")
    print("Prerequisite: aquin load --model <id>")
    print("  aquin load sae <model-l{n}>")
    print("")
    print("Usage: aquin feature locate [--scorer deception] [--prompts <json|jsonl>]")
    print("       [--layer N] [--checkpoint <path>] [--top N] [--direction both|deceptive|honest]")
    print("       [--conditioning behavior|prompt] [--benchmark-top K] [--persist <key>] [--save <json>]")
    print("       [--umap]")
    print("")
    print("  --prompts is required (honest/deceptive JSON or JSONL)")
    print("  --conditioning behavior (default): generate completions, classify output, bucket by behavior")
    print("  --conditioning prompt: legacy static encoding on probe text only")
    print("  --direction both (default): rank by |Δ|; deceptive: only Δ>0; honest: only Δ<0")
    print("  --persist writes chosen feature to ~/.aquin/experiments/<model>.json + session memory")
    print("  --umap loads SAE UMAP projection after ranking (web explorer)")
    print("  Syncs findFeature card to the web orchestrator.")
    print("")
    print("Docs: https://aquin.app/docs/deception")