sweep_cli.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""aquin sweep — steering sweep across strengths."""
 
from __future__ import annotations
 
import sys
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 parse_strengths(raw: str | None) -> list[float]:
    text = (raw or "").strip()
    if not text:
        return [-10.0, -5.0, 0.0, 5.0, 10.0]
    out: list[float] = []
    for part in text.split(","):
        part = part.strip()
        if not part:
            continue
        out.append(float(part))
    if not out:
        raise ValueError("No strengths parsed")
    return out
 
 
def _print_help() -> None:
    print("Sweep steering strength over a fixed feature/vector.")
    print("")
    print("Usage: aquin sweep (--feature_idx <n> | --vector <path>) [--strengths <csv>]")
    print("       [--prompt <text>] [--eval] [--prompts <json|jsonl>] [--reference_answers <json>]")
    print("")
    print("Examples:")
    print('  aquin sweep --feature_idx 42 --prompt "Explain photosynthesis"')
    print("  aquin sweep --vector vec.json --eval --prompts probes.jsonl")
 
 
def cmd_sweep(args: list[str]) -> None:
    if not args or args[0] in ("-h", "--help", "help"):
        _print_help()
        return
 
    from aquin.compute.loader_shim import apply as _shim_apply
    from aquin.engine.local_server import start as _start_local_server
    from aquin.compute.model_loader import (
        get_active_model_id,
        resolve_model_id,
    )
    from aquin.compute.steer_vector import run_steer_with_vector
 
    feature_idx_raw = _parse_flag(args, "--feature_idx")
    vector = _parse_flag(args, "--vector")
    if feature_idx_raw is None and not vector:
        print("Error: pass --feature_idx <n> or --vector <path>")
        sys.exit(1)
    try:
        strengths = parse_strengths(_parse_flag(args, "--strengths"))
    except ValueError as e:
        print(f"Error: {e}")
        sys.exit(1)
 
    active = (get_active_model_id() or "").strip()
    if not active:
        print("Error: no model loaded. Run: aquin load model <id>")
        sys.exit(1)
    _shim_apply()
    _start_local_server()
 
    model_id = resolve_model_id(active)
    feature_idx = int(feature_idx_raw) if feature_idx_raw is not None else None
    prompt = _parse_flag(args, "--prompt")
    layer_raw = _parse_flag(args, "--layer")
    layer = int(layer_raw) if layer_raw else None
    max_new_tokens_raw = _parse_flag(args, "--max_new_tokens")
    max_new_tokens = int(max_new_tokens_raw) if max_new_tokens_raw else 80
    do_eval = _has_flag(args, "--eval")
    prompts = _parse_flag(args, "--prompts")
    refs = _parse_flag(args, "--reference_answers")
    threshold = _parse_flag(args, "--threshold")
    max_probes = _parse_flag(args, "--max_probes")
 
    rows: list[dict[str, Any]] = []
    feature_label = None
    resolved_layer = layer
    for strength in strengths:
        tool_args: dict[str, Any] = {
            "prompt": prompt,
            "eval": do_eval,
            "prompts": prompts,
            "reference_answers": refs,
        }
        if threshold is not None:
            tool_args["threshold"] = float(threshold)
        if max_probes is not None:
            tool_args["max_probes"] = int(max_probes)
        result = run_steer_with_vector(
            model_id=model_id,
            prompt=prompt,
            steer_strength=float(strength),
            layer=resolved_layer,
            feature_idx=feature_idx,
            vector_path=vector,
            feature_label=feature_label,
            max_new_tokens=max_new_tokens,
            args=tool_args,
        )
        if result.get("error"):
            print(f"Error at strength {strength}: {result['error']}")
            sys.exit(1)
        feature_label = str(result.get("feature_label") or feature_label or "")
        resolved_layer = int(result.get("layer") or resolved_layer or 0)
        row: dict[str, Any] = {
            "strength": float(strength),
            "prompt": result.get("prompt"),
            "steered_response": result.get("steered_response"),
        }
        eval_data = result.get("eval") if isinstance(result.get("eval"), dict) else None
        if eval_data:
            row["pass_rate"] = (eval_data.get("steered") or {}).get("pass_rate")
            row["delta_pass_rate"] = eval_data.get("delta_pass_rate")
            row["mode"] = eval_data.get("mode")
        rows.append(row)
 
    print(
        f"[sweep] model={model_id}  feature={feature_idx if feature_idx is not None else vector}  "
        f"layer={resolved_layer if resolved_layer is not None else '—'}  n={len(rows)}"
    )
    if do_eval:
        print("")
        print("strength   pass   delta   mode")
        for row in rows:
            pr = row.get("pass_rate")
            dpr = row.get("delta_pass_rate")
            pr_s = "—" if pr is None else f"{100.0 * float(pr):5.1f}%"
            dpr_s = "—" if dpr is None else f"{100.0 * float(dpr):+5.1f}%"
            print(f"{row['strength']:>7.2f}   {pr_s:>6}  {dpr_s:>7}   {row.get('mode', '—')}")
    else:
        print("")
        print("strength   response")
        for row in rows:
            resp = " ".join(str(row.get("steered_response") or "").split())
            if len(resp) > 88:
                resp = resp[:87] + "…"
            print(f"{row['strength']:>7.2f}   {resp}")