steer_vector.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""LAT steering-vector extract, persist, and load for SAE decoder directions."""
from __future__ import annotations
 
import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
 
import torch
 
SCHEMA_VERSION = 1
KIND = "sae_decoder_steer"
 
 
def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()
 
 
def _experiment_path(model_id: str) -> Path:
    slug = model_id.replace("/", "--")
    return Path.home() / ".aquin" / "experiments" / f"{slug}.json"
 
 
def lookup_probe_record(model_id: str, probe_id: str | None = None) -> dict[str, Any] | None:
    """Load persisted find-feature record from ~/.aquin/experiments/<model>.json."""
    path = _experiment_path(model_id)
    if not path.is_file():
        return None
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return None
    if not isinstance(data, dict):
        return None
    key = probe_id or "deception_feature"
    rec = data.get(key)
    return rec if isinstance(rec, dict) else None
 
 
def resolve_feature_direction(
    model_id: str,
    feature_idx: int,
    *,
    layer: int | None = None,
    apply_norm_std: bool = True,
) -> tuple[torch.Tensor, int, dict[str, Any]]:
    """Return unit-ready decoder direction (d_model,) and resolved layer."""
    from aquin.compute.feature_analysis import load_norm, load_sae
    from aquin.compute.model_loader import get_config
 
    cfg = get_config(model_id)
    resolved_layer = int(layer if layer is not None else cfg["sae_layer"])
    try:
        sae = load_sae(model_id, resolved_layer)
    except FileNotFoundError:
        from aquin.compute.model_loader import get_available_sae_layers, get_sae_layer
 
        available = get_available_sae_layers(model_id)
        fallback = int(get_sae_layer(model_id))
        if available:
            resolved_layer = fallback if fallback in available else int(available[0])
            sae = load_sae(model_id, resolved_layer)
        else:
            raise
    if feature_idx < 0 or feature_idx >= sae.W_dec.shape[0]:
        raise ValueError(f"feature_idx {feature_idx} out of range (0..{sae.W_dec.shape[0] - 1})")
 
    feat_dir = sae.W_dec[feature_idx].detach().float().clone()
    if not torch.isfinite(feat_dir).all():
        enc = sae.W_enc.data.detach().float()
        if enc.ndim == 2 and enc.shape[1] > feature_idx:
            feat_dir = enc[:, feature_idx].clone()
        feat_dir = torch.nan_to_num(feat_dir, nan=0.0, posinf=0.0, neginf=0.0)
    norm_stats: dict[str, Any] | None = None
    norm = load_norm(model_id, resolved_layer) if apply_norm_std else None
    if norm is not None and norm.get("std") is not None:
        # Keep std on feat_dir's device — SAE may be CUDA while a naive .cpu()
        # std would raise "cuda:0 and cpu" on the multiply below.
        std = norm["std"].detach().float().to(device=feat_dir.device)
        feat_dir = feat_dir * std
        mean = norm.get("mean")
        norm_stats = {
            "mean": mean.detach().float().cpu().tolist() if mean is not None else None,
            "std": std.detach().cpu().tolist(),
            "applied_to_vector": True,
        }
 
    meta = {
        "d_model": int(feat_dir.shape[-1]),
        "vector_l2_norm": round(float(feat_dir.norm().item()), 6),
        "norm": norm_stats,
    }
    return feat_dir, resolved_layer, meta
 
 
def build_steer_vector_payload(
    model_id: str,
    feature_idx: int,
    *,
    layer: int | None = None,
    feature_label: str | None = None,
    probe_id: str | None = None,
) -> dict[str, Any]:
    direction, resolved_layer, vec_meta = resolve_feature_direction(
        model_id, feature_idx, layer=layer,
    )
    probe_record = lookup_probe_record(model_id, probe_id)
    label = feature_label or f"F{feature_idx}"
 
    return {
        "schema_version": SCHEMA_VERSION,
        "kind": KIND,
        "model_id": model_id,
        "layer": resolved_layer,
        "feature_idx": int(feature_idx),
        "feature_label": label or f"F{feature_idx}",
        "probe_id": probe_id or ("deception_feature" if probe_record else None),
        "probe_record": probe_record,
        "d_model": vec_meta["d_model"],
        "vector_l2_norm": vec_meta["vector_l2_norm"],
        "norm": vec_meta.get("norm"),
        "vector": direction.cpu().tolist(),
        "created_at": _now_iso(),
        "source": "steer --save",
    }
 
 
def write_steer_vector_file(payload: dict[str, Any], output: str | Path) -> Path:
    out = Path(output).expanduser().resolve()
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    return out
 
 
def load_steer_vector_file(path: str | Path) -> dict[str, Any]:
    p = Path(path).expanduser().resolve()
    if not p.is_file():
        raise FileNotFoundError(f"Steering vector file not found: {p}")
    try:
        data = json.loads(p.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:
        raise ValueError(f"Invalid steering vector JSON: {p}") from exc
    if not isinstance(data, dict):
        raise ValueError(f"Steering vector file must be a JSON object: {p}")
    if data.get("kind") != KIND:
        raise ValueError(f"Unsupported steering vector kind: {data.get('kind')!r}")
    if "vector" not in data or not isinstance(data["vector"], list) or not data["vector"]:
        raise ValueError("Steering vector file missing non-empty 'vector' array")
    return data
 
 
def steer_vector_tensor(
    data: dict[str, Any],
    *,
    device: torch.device | str,
    dtype: torch.dtype,
) -> torch.Tensor:
    vec = torch.tensor(data["vector"], device=device, dtype=dtype)
    if vec.ndim != 1:
        raise ValueError("Steering vector must be a 1-D array")
    return vec
 
 
def extract_steer_vector(
    model_id: str,
    feature_idx: int,
    output: str | Path,
    *,
    layer: int | None = None,
    feature_label: str | None = None,
    probe_id: str | None = None,
) -> dict[str, Any]:
    payload = build_steer_vector_payload(
        model_id,
        feature_idx,
        layer=layer,
        feature_label=feature_label,
        probe_id=probe_id,
    )
    out_path = write_steer_vector_file(payload, output)
    return {
        "status": "done",
        "output_path": str(out_path),
        "model_id": payload["model_id"],
        "layer": payload["layer"],
        "feature_idx": payload["feature_idx"],
        "feature_label": payload["feature_label"],
        "probe_id": payload.get("probe_id"),
        "d_model": payload["d_model"],
        "vector_l2_norm": payload["vector_l2_norm"],
        "norm": payload.get("norm"),
        "source_model_id": payload["model_id"],
        "source_layer": payload["layer"],
        "source_feature_idx": payload["feature_idx"],
    }
 
 
def _generate_steered_response(
    model: Any,
    prompt: str,
    *,
    steer_vec: torch.Tensor,
    steer_strength: float,
    layer: int,
    max_new_tokens: int,
) -> str:
    from aquin.compute.causal_trace import _format_prompt
 
    def _steer_hook(value, hook):
        vec = steer_vec.to(device=value.device, dtype=value.dtype)
        return value + steer_strength * vec.unsqueeze(0).unsqueeze(0)
 
    hook_name = f"blocks.{layer}.hook_resid_post"
    fmt_prompt = _format_prompt(model, prompt)
    tokens = model.to_tokens(fmt_prompt)
    steered_tokens: list[int] = []
    cur = tokens
    for _ in range(max_new_tokens):
        with torch.no_grad():
            out = model.run_with_hooks(cur, fwd_hooks=[(hook_name, _steer_hook)])
        next_id = int(out[0, -1].argmax().item())
        if next_id == model.tokenizer.eos_token_id:
            break
        steered_tokens.append(next_id)
        cur = torch.cat([cur, torch.tensor([[next_id]], device=cur.device)], dim=1)
    return model.tokenizer.decode(steered_tokens, skip_special_tokens=True)
 
 
def run_steer_with_vector(
    *,
    model_id: str,
    prompt: str | None = None,
    steer_strength: float,
    layer: int | None = None,
    feature_idx: int | None = None,
    vector_path: str | Path | None = None,
    vector_data: dict[str, Any] | None = None,
    feature_label: str | None = None,
    max_new_tokens: int = 80,
    ctx: dict | None = None,
    args: dict | None = None,
) -> dict[str, Any]:
    """Steer generation using a live feature index or a saved vector file."""
    from aquin.compute.causal_trace import run_chat
    from aquin.compute.feature_analysis import format_feature_ref, resolve_feature_label
    from aquin.compute.model_loader import get_config, get_loaded_model, load_model, resolve_model_id
 
    args = args or {}
    do_eval = bool(args.get("eval"))
    explicit_prompt = args.get("prompt")
    if explicit_prompt is not None:
        explicit_prompt = str(explicit_prompt).strip() or None
 
    prompt_text = (prompt or "").strip() if prompt else ""
    if not prompt_text and explicit_prompt:
        prompt_text = explicit_prompt
 
    # --eval without an explicit --prompt: skip single-demo generation; suite is the output.
    run_single = bool(prompt_text) and (not do_eval or explicit_prompt is not None)
    if do_eval and not explicit_prompt:
        run_single = False
        prompt_text = ""
 
    model_id = resolve_model_id(model_id)
    model = get_loaded_model()
    if model is None:
        model = load_model(model_id)
 
    cfg = get_config(model_id)
    resolved_layer = layer
    steer_vec: torch.Tensor | None = None
    vec_meta: dict[str, Any] | None = None
 
    if vector_path or vector_data:
        vec_meta = vector_data or load_steer_vector_file(str(vector_path))
        steer_vec = steer_vector_tensor(
            vec_meta,
            device=model.W_E.device,
            dtype=model.W_E.dtype,
        )
        if resolved_layer is None:
            resolved_layer = int(vec_meta.get("layer", cfg["sae_layer"]))
        if feature_idx is None:
            feature_idx = int(vec_meta.get("feature_idx", 0))
        if not feature_label:
            feature_label = str(vec_meta.get("feature_label") or f"F{feature_idx}")
        if steer_vec.shape[-1] != model.cfg.d_model:
            return {
                "error": (
                    f"Vector d_model={steer_vec.shape[-1]} does not match model d_model={model.cfg.d_model}. "
                    "Use a vector extracted from a compatible checkpoint."
                ),
            }
    elif feature_idx is not None:
        direction, resolved_layer, _ = resolve_feature_direction(
            model_id, int(feature_idx), layer=layer,
        )
        steer_vec = direction.to(device=model.W_E.device, dtype=model.W_E.dtype)
        if not feature_label:
            # CLI paths (sweep/steer) pass no session ctx — keep a cheap display name.
            # Causal re-labeling is for UI sessions that provide ctx.
            if ctx:
                feature_label = resolve_feature_label(
                    int(feature_idx),
                    ctx=ctx,
                    args={**args, "model_id": model_id},
                    layer=resolved_layer,
                )
            else:
                feature_label = f"F{feature_idx}"
    else:
        return {"error": "Provide feature_idx or vector (path to saved LAT file)."}
 
    resolved_layer = int(resolved_layer if resolved_layer is not None else cfg["sae_layer"])
    assert steer_vec is not None
 
    baseline = ""
    steered_response = ""
    if run_single:
        try:
            baseline = run_chat(
                prompt_text, model_id=model_id, max_new_tokens=max_new_tokens, temperature=0.0,
            )
        except Exception as e:
            return {"error": f"Baseline generation failed: {e}"}
        try:
            steered_response = _generate_steered_response(
                model,
                prompt_text,
                steer_vec=steer_vec,
                steer_strength=steer_strength,
                layer=resolved_layer,
                max_new_tokens=max_new_tokens,
            )
        except Exception as e:
            steered_response = f"[steer error: {e}]"
 
    out: dict[str, Any] = {
        "feature_idx": int(feature_idx or 0),
        "feature_label": feature_label or f"F{feature_idx}",
        "feature_ref": format_feature_ref(int(feature_idx or 0), feature_label or f"F{feature_idx}"),
        "steer_strength": steer_strength,
        "layer": resolved_layer,
        "prompt": prompt_text,
        "baseline_response": baseline,
        "steered_response": steered_response,
        "demo": run_single,
    }
    if vec_meta:
        out["vector_path"] = str(vector_path) if vector_path else None
        out["vector_source_model_id"] = vec_meta.get("model_id")
        out["vector_source_layer"] = vec_meta.get("layer")
        out["vector_source_feature_idx"] = vec_meta.get("feature_idx")
 
    if do_eval:
        from aquin.compute.steer_eval import run_steer_probe_eval
 
        eval_tokens = int(args.get("eval_max_new_tokens") or min(max_new_tokens, 64))
        threshold = float(args.get("threshold") if args.get("threshold") is not None else 0.5)
 
        def _baseline(p: str) -> str:
            return run_chat(p, model_id=model_id, max_new_tokens=eval_tokens, temperature=0.0)
 
        def _steered(p: str) -> str:
            return _generate_steered_response(
                model,
                p,
                steer_vec=steer_vec,
                steer_strength=steer_strength,
                layer=resolved_layer,
                max_new_tokens=eval_tokens,
            )
 
        eval_result = run_steer_probe_eval(
            generate_baseline=_baseline,
            generate_steered=_steered,
            prompts=args.get("prompts") or args.get("prompts_path"),
            reference_answers=args.get("reference_answers"),
            threshold=threshold,
            max_probes=int(args.get("max_probes") or 50),
        )
        if eval_result.get("error"):
            return {"error": eval_result["error"], **{k: v for k, v in out.items() if k != "error"}}
        out["eval"] = eval_result
        # Seed demo fields from first probe when --eval ran without --prompt
        if not run_single and eval_result.get("probes"):
            first = eval_result["probes"][0]
            out["prompt"] = first.get("prompt") or ""
            out["baseline_response"] = (first.get("baseline") or {}).get("response") or ""
            out["steered_response"] = (first.get("steered") or {}).get("response") or ""
 
    return out