inspect.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
"""Inspect — diffs, traces, layer/attention/weight checks."""
 
from __future__ import annotations
 
from pathlib import Path
from typing import Any
 
from aquin.sdk._runtime import invoke, invoke_direct, prepare_compute
 
 
def diff_weight(**kwargs: Any) -> dict[str, Any]:
    return invoke("run_merge_analysis", kwargs, command="diff weight")
 
 
def diff_sae(
    *,
    checkpoint: str | Path,
    prompts: Any = None,
    model_id: str | None = None,
    layer: int | None = None,
    sae_path: str | Path | None = None,
    name: str | None = None,
    **kwargs: Any,
) -> dict[str, Any]:
    """Compare SAE feature means: base vs checkpoint."""
    from aquin.compute.model_loader import get_active_model_id, resolve_model_id
    from aquin.compute.sae_diff import load_prompts, run_sae_diff
 
    mid = resolve_model_id(model_id or get_active_model_id() or "")
    prepare_compute(model_id=mid)
 
    if prompts is None:
        raise ValueError("prompts= is required (list[str] or path).")
    if isinstance(prompts, (str, Path)):
        prompt_list = load_prompts(str(prompts))
    else:
        prompt_list = [str(p) for p in prompts]
 
    return invoke_direct(
        run_sae_diff,
        mid,
        prompt_list,
        command="diff sae",
        tool_name="run_sae_diff",
        target_checkpoint=str(checkpoint),
        checkpoint_name=name or Path(str(checkpoint)).stem,
        layer=layer,
        sae_path=sae_path,
        **kwargs,
    )
 
 
def diff_residue(**kwargs: Any) -> dict[str, Any]:
    return invoke("run_residual_drift", kwargs, command="diff residue")
 
 
def trajectory(**kwargs: Any) -> dict[str, Any]:
    return invoke("run_trajectory_analysis", kwargs, command="check trajectory")
 
 
def attention(**kwargs: Any) -> dict[str, Any]:
    return invoke("run_attention_routing", kwargs, command="check attention")
 
 
def layer(**kwargs: Any) -> dict[str, Any]:
    return invoke("run_layer_analysis", kwargs, command="check layer")
 
 
def perturbation(**kwargs: Any) -> dict[str, Any]:
    return invoke("run_perturbation_sensitivity", kwargs, command="check perturbation")
 
 
def weight(**kwargs: Any) -> dict[str, Any]:
    return invoke("check_weights", kwargs, command="check weight")
 
 
def trace(**kwargs: Any) -> dict[str, Any]:
    return invoke("run_full_inspection", kwargs, command="trace")
 
 
def benchmark(**kwargs: Any) -> dict[str, Any]:
    return invoke("run_benchmarks_on_top_feature", kwargs, command="benchmark")
 
 
def umap(**kwargs: Any) -> dict[str, Any]:
    return invoke("ensure_umap_loaded", kwargs, command="umap")