residual_drift_cli.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""aquin diff residue — base vs checkpoint per-layer activation drift."""
 
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 _resolve_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("Per-layer activation drift between catalog base and a fine-tuned checkpoint.")
    print("")
    print("Prerequisite: aquin load --model <id>")
    print("")
    print("Usage: aquin diff residue --checkpoint <path>")
    print("       [--prompts <json|jsonl>] [--name <label>] [--save <path>]")
    print("")
    print("  --checkpoint  Fine-tuned .pt state_dict or HF save_pretrained directory.")
    print("  --prompts     Probe strings (default: built-in short prompts).")
    print("  --name        Label for checkpoint in output and web card.")
    print("  --save        Write schema_version=1 JSON export.")
    print("")
    print("LLM: last-token hook_resid_post cosine distance per layer.")
    print("Embedding: mean-pooled hidden-state cosine distance per encoder layer.")
    print("")
    print("Example:")
    print("  aquin diff residue --checkpoint ~/runs/checkpoint.pt \\")
    print("    --prompts probes.json --save drift.json")
    print("")
    print("Docs: https://aquin.app/docs/checkpoint-sae")
 
 
def cmd_residual_drift(args: list[str]) -> None:
    if _has_flag(args, "--help") or _has_flag(args, "-h"):
        _print_help()
        return
 
    reject_legacy_output_flags(args)
 
    checkpoint = _parse_flag(args, "--checkpoint")
    if not checkpoint:
        print("Error: --checkpoint is required.")
        print("")
        _print_help()
        sys.exit(1)
 
    ckpt = Path(checkpoint).expanduser()
    if not ckpt.exists():
        print(f"Error: checkpoint not found: {checkpoint}")
        sys.exit(1)
 
    if _parse_flag(args, "--model") is not None:
        print("Error: diff residue uses the loaded session model only.")
        sys.exit(1)
 
    prompts_arg = _parse_flag(args, "--prompts")
    if prompts_arg:
        from aquin.compute.activation_capture import resolve_prompts_path
 
        resolved = resolve_prompts_path(Path(prompts_arg).expanduser())
        if resolved is None:
            print(f"Error: probe file not found: {prompts_arg}")
            print("       Pass a probe file that exists on disk.")
            sys.exit(1)
        prompts_arg = str(resolved)
 
    _ensure_compute_env()
 
    from aquin.cli import _build_tool_ctx
    from aquin.engine.sync_dispatch import dispatch_with_sync, require_active_session
 
    mid = _resolve_model_id()
    ctx = _build_tool_ctx(model_id=mid)
    require_active_session(ctx, label="aquin diff residue")
 
    tool_args: dict[str, Any] = {
        "model_id": mid,
        "checkpoint": str(ckpt),
        "name": _parse_flag(args, "--name") or ckpt.stem,
        "save": _parse_flag(args, "--save"),
    }
    if prompts_arg:
        tool_args["prompts"] = prompts_arg
    try:
        print(f"[diff residue] model={mid}  checkpoint={ckpt.name}")
        result = dispatch_with_sync("run_residual_drift", 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("diff residue", result)
 
    if isinstance(result, dict) and result.get("error"):
        sys.exit(1)