weight_diff_cli.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""aquin diff weight — base vs checkpoint weight delta analysis."""
 
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("Compare base vs fine-tuned checkpoint weights (LLM).")
    print("")
    print("Prerequisite: aquin load --model <id>")
    print("")
    print("Usage: aquin diff weight --checkpoint <path>")
    print("       [--name <label>] [--prompts <json|jsonl>] [--no-behavioral]")
    print("       [--save <path>]")
    print("")
    print("  --checkpoint      Fine-tuned .pt state_dict or HF save_pretrained directory.")
    print("  --name            Label for checkpoint in output and web card.")
    print("  --prompts         Probe file for behavioral model-diff (LLM only).")
    print("  --no-behavioral   Skip generation-based behavioral scores (faster).")
    print("  --save            Write schema_version=1 JSON export.")
    print("")
    print("Reports per-layer ||dW||, rank/collapse signals, merge verdict (pass/warn/fail),")
    print("and optional behavioral probes on LLM checkpoints.")
    print("")
    print("Example:")
    print("  aquin diff weight --checkpoint ~/runs/checkpoint.pt --save weight-diff.json")
    print("")
    print("Docs: https://aquin.app/docs/checkpoint-sae")
 
 
def cmd_weight_diff(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 weight 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}")
            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 weight")
 
    tool_args: dict[str, Any] = {
        "model_id": mid,
        "checkpoint": str(ckpt),
        "name": _parse_flag(args, "--name") or ckpt.stem,
        "save": _parse_flag(args, "--save"),
        "no_behavioral": _has_flag(args, "--no-behavioral"),
    }
    if prompts_arg:
        tool_args["prompts"] = prompts_arg
    try:
        print(f"[diff weight] model={mid}  checkpoint={ckpt.name}")
        result = dispatch_with_sync("run_merge_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("diff weight", result)
 
    if isinstance(result, dict):
        if result.get("error"):
            sys.exit(1)
        if result.get("mergeVerdict") == "fail":
            sys.exit(2)