confidence_analysis_cli.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""aquin check confidence — per-probe confidence + optional SAE join."""
 
import sys
 
from aquin.cli_flags import reject_legacy_output_flags
from pathlib import Path
from typing import Any
 
 
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 _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 _print_help() -> None:
    print("Measure confidence patterns over a probe dataset (LLM).")
    print("")
    print("Prerequisite: aquin load --model <id>")
    print("  Token logits (mean greedy prob, entropy)")
    print("")
    print("Usage: aquin check confidence --prompts <json|jsonl>")
    print("       [--threshold <0-1>]")
    print("       [--join-sae] [--layer <n>]")
    print("       [--save <path>]")
    print("       [--check]")
    print("")
    print("  --prompts     Probe file (text + optional id, stressor, lang, quant_run_id).")
    print("  --threshold   Low-confidence cutoff (default: 0.40).")
    print("  --join-sae    Attach SAE mean L0 + top feature per probe.")
    print("  --layer       SAE layer for join (default: model sae_layer).")
    print("  --save        Write schema_version=1 JSON export.")
    print("  --check       Save confidence-analysis-check.json and confidence-analysis-check.png in cwd.")
    print("")
    print("Example (LLM):")
    print("  aquin check confidence --prompts probes.jsonl --join-sae --layer 8 --check")
    print("")
    print("Docs: https://aquin.app/docs/evals/llm")
 
 
def cmd_confidence_analysis(args: list[str]) -> None:
    if _has_flag(args, "--help") or _has_flag(args, "-h"):
        _print_help()
        return
 
    reject_legacy_output_flags(args)
 
    prompts = _parse_flag(args, "--prompts")
    if not prompts:
        print("Error: --prompts is required.")
        print("")
        _print_help()
        sys.exit(1)
 
    from aquin.compute.activation_capture import resolve_prompts_path
 
    resolved = resolve_prompts_path(Path(prompts).expanduser())
    if resolved is None:
        print(f"Error: probe file not found: {prompts}")
        print("       Tried cwd and repo parent.")
        sys.exit(1)
    prompts = str(resolved)
 
    if _parse_flag(args, "--model") is not None:
        print("Error: check confidence uses the loaded session model only.")
        sys.exit(1)
 
    _ensure_compute_env()
 
    from aquin.cli import _build_tool_ctx
    from aquin.engine.sync_dispatch import dispatch_with_sync, require_active_session
 
    mid = _require_loaded_model_id()
    ctx = _build_tool_ctx(model_id=mid)
    require_active_session(ctx, label="aquin check confidence")
 
    tool_args: dict[str, Any] = {
        "model_id": mid,
        "prompts": prompts,
        "threshold": _parse_flag(args, "--threshold"),
        "join_sae": _has_flag(args, "--join-sae"),
        "layer": _parse_flag(args, "--layer"),
        "save": _parse_flag(args, "--save"),
    }
    do_check = _has_flag(args, "--check")
 
    try:
        print(
            f"[check confidence] model={mid}  probes={prompts}"
            f"  join_sae={tool_args['join_sae']}"
        )
        result = dispatch_with_sync("run_confidence_analysis", tool_args, ctx)
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)
 
    from aquin.cli_output import print_tool_result
 
    print_tool_result("check confidence", result)
 
    if do_check:
        if isinstance(result, dict) and result.get("error"):
            print(f"[check confidence --check] skipped save: {result['error']}", file=sys.stderr)
        else:
            import os
            from aquin.confidence_analysis_check import write_confidence_analysis_check
 
            try:
                json_path, png_path = write_confidence_analysis_check(
                    result, tool_name="run_confidence_analysis", cwd=os.getcwd(),
                )
                print(f"Saved {json_path}")
                print(f"Saved {png_path}")
            except Exception as exc:
                print(f"[check confidence --check] failed to save files: {exc}", file=sys.stderr)
                import traceback
                traceback.print_exc()
                sys.exit(1)
 
    if isinstance(result, dict) and result.get("error"):
        sys.exit(1)