find_feature.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Rank SAE features that separate honest vs deceptive probe sets (deception scorer)."""
 
from __future__ import annotations
 
import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
 
import torch
 
from aquin.compute.activation_capture import resolve_prompts_path
 
HONEST_LABELS = frozenset({"honest", "truthful", "true", "control", "baseline"})
DECEPTIVE_LABELS = frozenset({"deceptive", "deception", "lie", "lying", "dishonest", "false"})
TEXT_KEYS = ("instruction", "prompt", "text", "content", "honest", "deceptive")
VALID_DIRECTIONS = frozenset({"both", "deceptive", "honest"})
 
 
def normalize_direction(direction: str | None) -> str:
    d = (direction or "both").strip().lower()
    if d not in VALID_DIRECTIONS:
        raise ValueError(f"Unknown direction {direction!r}; supported: both, deceptive, honest")
    return d
 
 
def _resolve_deception_fixture() -> Path:
    """Bundled probes were removed. Callers must pass prompts=."""
    raise FileNotFoundError(
        "No default deception probes. Pass prompts= with an honest/deceptive JSON/JSONL path."
    )
 
 
def _row_text(row: dict[str, Any], *keys: str) -> str:
    for key in keys:
        val = row.get(key)
        if isinstance(val, str) and val.strip():
            return val.strip()
    return ""
 
 
def load_deception_probes(path: str | Path | None) -> tuple[list[str], list[str], dict[str, Any]]:
    """
    Load honest and deceptive text lists from JSON/JSONL.
 
    Supports:
      - paired rows: honest + deceptive fields on one object
      - labeled rows: instruction/prompt/text + label honest|deceptive
    """
    p = Path(path) if path else _resolve_deception_fixture()
    if not p.is_file():
        resolved = resolve_prompts_path(p)
        if resolved is None:
            raise FileNotFoundError(f"Prompts file not found: {p}")
        p = resolved
 
    text = p.read_text(encoding="utf-8").strip()
    if p.suffix == ".jsonl":
        rows = [json.loads(line) for line in text.splitlines() if line.strip()]
    else:
        data = json.loads(text)
        rows = data if isinstance(data, list) else [data]
 
    honest: list[str] = []
    deceptive: list[str] = []
    n_pairs = 0
    n_labeled = 0
 
    for row in rows:
        if not isinstance(row, dict):
            continue
        h = _row_text(row, "honest", "instruction_honest", "true", "control")
        d = _row_text(row, "deceptive", "instruction_deceptive", "false", "lie")
        if h and d:
            honest.append(h)
            deceptive.append(d)
            n_pairs += 1
            continue
 
        label = str(row.get("label") or row.get("class") or row.get("condition") or "").strip().lower()
        body = _row_text(row, *TEXT_KEYS)
        if not body or not label:
            continue
        if label in HONEST_LABELS:
            honest.append(body)
            n_labeled += 1
        elif label in DECEPTIVE_LABELS:
            deceptive.append(body)
            n_labeled += 1
 
    if not honest or not deceptive:
        raise ValueError(
            f"No honest/deceptive probe pairs in {p}. "
            "Use paired rows (honest+deceptive) or labeled rows (label: honest|deceptive)."
        )
 
    meta = {
        "prompts_path": str(p),
        "n_honest": len(honest),
        "n_deceptive": len(deceptive),
        "n_pairs": n_pairs,
        "n_labeled": n_labeled,
    }
    return honest, deceptive, meta
 
 
def _rank_features(
    honest_acts: torch.Tensor,
    deceptive_acts: torch.Tensor,
    *,
    top_k: int,
    direction: str = "both",
) -> list[dict[str, Any]]:
    """honest_acts, deceptive_acts: (n_features,) mean activations."""
    direction = normalize_direction(direction)
    delta = deceptive_acts - honest_acts
    abs_delta = delta.abs()
 
    if direction == "deceptive":
        mask = delta > 0
        rank_scores = delta
    elif direction == "honest":
        mask = delta < 0
        rank_scores = -delta
    else:
        mask = torch.ones_like(delta, dtype=torch.bool)
        rank_scores = abs_delta
 
    if not bool(mask.any().item()):
        return []
 
    valid_idx = mask.nonzero(as_tuple=False).squeeze(-1)
    valid_scores = rank_scores[mask]
    k = min(top_k, int(valid_scores.shape[0]))
    local_top = valid_scores.topk(k).indices
    top_idx = valid_idx[local_top].tolist()
 
    rows = [
        {
            "feature_idx": int(i),
            "honest_mean": round(float(honest_acts[i].item()), 6),
            "deceptive_mean": round(float(deceptive_acts[i].item()), 6),
            "delta": round(float(delta[i].item()), 6),
            "abs_delta": round(float(abs_delta[i].item()), 6),
        }
        for i in top_idx
    ]
    if direction == "deceptive":
        rows.sort(key=lambda r: r["delta"], reverse=True)
    elif direction == "honest":
        rows.sort(key=lambda r: r["delta"])
    else:
        rows.sort(key=lambda r: r["abs_delta"], reverse=True)
    return rows
 
 
def _find_feature_warning(
    direction: str,
    chosen: dict[str, Any] | None,
    *,
    persisted: bool,
    conditioning: str = "behavior",
    behavior_meta: dict[str, Any] | None = None,
) -> str | None:
    direction = normalize_direction(direction)
    warnings: list[str] = []
 
    if conditioning == "behavior" and behavior_meta:
        n_truthful = int(behavior_meta.get("n_truthful") or 0)
        n_deceptive = int(behavior_meta.get("n_deceptive") or 0)
        n_ambiguous = int(behavior_meta.get("n_ambiguous") or 0)
        n_generated = int(behavior_meta.get("n_generated") or 0)
        if n_truthful < 2 or n_deceptive < 2:
            warnings.append(
                f"Small behavior buckets (truthful={n_truthful}, deceptive={n_deceptive}). "
                "Add probes or tune references."
            )
        if n_generated and n_ambiguous > n_generated // 2:
            warnings.append(
                f"{n_ambiguous}/{n_generated} completions were ambiguous; "
                "classification may be weak."
            )
    elif conditioning == "prompt":
        warnings.append(
            "Using prompt conditioning (static probe text). "
            "Default behavior mode generates completions and buckets by observed output."
        )
 
    if chosen is None:
        if direction == "deceptive":
            warnings.append(
                "No features with positive delta (deceptive > honest). "
                "Try --direction both, another layer, or custom probes."
            )
        elif direction == "honest":
            warnings.append(
                "No features with negative delta (honest > deceptive). "
                "Try --direction both, another layer, or custom probes."
            )
        else:
            warnings.append("No separating features found.")
        return " ".join(warnings) if warnings else None
 
    delta = chosen.get("delta")
    if delta is not None and direction == "both" and float(delta) < 0:
        msg = (
            "Top feature has negative delta (honest activates more than deceptive). "
            "Use --direction deceptive for deception-biased ranking."
        )
        if persisted:
            msg = (
                "Persisted feature has negative delta (honest > deceptive). "
                "Use --direction deceptive to pick deception-biased features."
            )
        warnings.append(msg)
 
    return " ".join(warnings) if warnings else None
 
 
def persist_experiment_record(
    model_id: str,
    key: str,
    record: dict[str, Any],
    *,
    session_id: str | None = None,
) -> str:
    """Write canonical feature record to ~/.aquin/experiments/<model>.json and optional session memory."""
    exp_dir = Path.home() / ".aquin" / "experiments"
    exp_dir.mkdir(parents=True, exist_ok=True)
    slug = model_id.replace("/", "--")
    path = exp_dir / f"{slug}.json"
    existing: dict[str, Any] = {}
    if path.is_file():
        try:
            existing = json.loads(path.read_text(encoding="utf-8"))
        except Exception:
            existing = {}
    if not isinstance(existing, dict):
        existing = {}
    existing[key] = record
    existing["updated_at"] = datetime.now(timezone.utc).isoformat()
    path.write_text(json.dumps(existing, indent=2), encoding="utf-8")
 
    if session_id:
        from aquin.engine.session_memory_store import patch_local_memory
 
        patch_local_memory(session_id, key, record.get("feature_idx"))
        patch_local_memory(session_id, f"{key}_record", record)
 
    return str(path)
 
 
def run_find_feature(
    model_id: str,
    *,
    scorer: str = "deception",
    prompts_path: str | Path | None = None,
    layer: int | None = None,
    checkpoint_path: str | Path | None = None,
    top_k: int = 20,
    direction: str = "both",
    conditioning: str = "behavior",
    max_new_tokens: int = 64,
    temperature: float = 0.0,
    benchmark_top: int = 0,
    persist_key: str | None = None,
    session_id: str | None = None,
    openai_client: Any | None = None,
) -> dict[str, Any]:
    direction = normalize_direction(direction)
    from aquin.compute.deception_behavior import normalize_conditioning
 
    conditioning = normalize_conditioning(conditioning)
    if scorer != "deception":
        raise ValueError(f"Unknown scorer {scorer!r}; supported: deception")
 
    from aquin.compute.deception_behavior import collect_behavior_encode_texts
    from aquin.compute.feature_analysis import load_sae
    from aquin.compute.model_loader import get_sae_layer, resolve_model_id
    from aquin.compute.sae_diff import load_tl_from_checkpoint, mean_sae_activations
 
    mid = resolve_model_id(model_id)
    sae_layer = int(layer if layer is not None else get_sae_layer(mid))
    behavior_meta: dict[str, Any] | None = None
 
    model = load_tl_from_checkpoint(mid, checkpoint_path)
    sae = load_sae(mid, sae_layer)
 
    if conditioning == "behavior":
        honest, deceptive, behavior_meta = collect_behavior_encode_texts(
            prompts_path,
            mid,
            model,
            checkpoint_path=checkpoint_path,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
        )
        if not honest or not deceptive:
            del model
            n_gen = int((behavior_meta or {}).get("n_generated") or 0)
            n_amb = int((behavior_meta or {}).get("n_ambiguous") or 0)
            preview = ""
            samples = (behavior_meta or {}).get("samples") or []
            if samples:
                first = samples[0]
                preview = (
                    f" Example: prompt={first.get('probe_id')!r} "
                    f"behavior={first.get('behavior')!r} "
                    f"response={first.get('response_preview')!r}."
                )
            raise ValueError(
                "Behavior conditioning needs at least one truthful and one deceptive completion. "
                f"Got truthful={len(honest)}, deceptive={len(deceptive)} "
                f"(generated={n_gen}, ambiguous={n_amb}).{preview} "
                "Use probes with honest_reference/deceptive_reference, "
                "or pass --conditioning prompt to use static honest/deceptive text."
            )
        probe_meta = {
            "prompts_path": behavior_meta.get("prompts_path"),
            "n_honest": len(honest),
            "n_deceptive": len(deceptive),
            "n_pairs": 0,
            "n_labeled": 0,
        }
    else:
        honest, deceptive, probe_meta = load_deception_probes(prompts_path)
        behavior_meta = None
 
    honest_mean = mean_sae_activations(model, sae, sae_layer, honest, mid)
    deceptive_mean = mean_sae_activations(model, sae, sae_layer, deceptive, mid)
    rankings = _rank_features(honest_mean, deceptive_mean, top_k=top_k, direction=direction)
 
    chosen = rankings[0] if rankings else None
    chosen_idx = int(chosen["feature_idx"]) if chosen else None
 
    if benchmark_top > 0 and rankings and openai_client is not None:
        from aquin.compute.interp_score import run_interp_score
 
        probe = deceptive[0] if deceptive else honest[0]
        for row in rankings[: max(1, min(benchmark_top, len(rankings)))]:
            fidx = int(row["feature_idx"])
            try:
                interp = run_interp_score(
                    fidx, probe, model, sae, openai_client,
                    model_id=mid, layer=sae_layer,
                )
                row["interp_score"] = interp.get("score")
                row["purity_score"] = interp.get("purity_score")
            except Exception as exc:
                row["interp_error"] = str(exc)
 
        def _combined(row: dict[str, Any]) -> float:
            if direction == "deceptive":
                base = float(row.get("delta") or 0)
            elif direction == "honest":
                base = float(-(row.get("delta") or 0))
            else:
                base = float(row.get("abs_delta") or 0)
            interp = row.get("interp_score")
            if interp is None:
                return base
            return base + 0.05 * float(interp)
 
        rankings.sort(key=_combined, reverse=True)
        chosen = rankings[0]
        chosen_idx = int(chosen["feature_idx"])
 
    del model
 
    warning = _find_feature_warning(
        direction,
        chosen,
        persisted=bool(persist_key),
        conditioning=conditioning,
        behavior_meta=behavior_meta,
    )
 
    payload: dict[str, Any] = {
        "type": "findFeature",
        "status": "done",
        "scorer": scorer,
        "direction": direction,
        "conditioning": conditioning,
        "model_id": mid,
        "layer": sae_layer,
        "checkpoint": str(checkpoint_path) if checkpoint_path else None,
        "n_honest": probe_meta["n_honest"],
        "n_deceptive": probe_meta["n_deceptive"],
        "prompts_path": probe_meta.get("prompts_path"),
        "chosen_feature_idx": chosen_idx,
        "chosen_delta": chosen.get("delta") if chosen else None,
        "warning": warning,
        "rankings": rankings,
    }
    if behavior_meta:
        payload["behavior"] = {
            "n_generated": behavior_meta.get("n_generated"),
            "n_truthful": behavior_meta.get("n_truthful"),
            "n_deceptive": behavior_meta.get("n_deceptive"),
            "n_ambiguous": behavior_meta.get("n_ambiguous"),
            "intent_behavior_match": behavior_meta.get("intent_behavior_match"),
            "samples": behavior_meta.get("samples"),
        }
 
    if persist_key and chosen_idx is not None:
        record = {
            "feature_idx": chosen_idx,
            "layer": sae_layer,
            "model_id": mid,
            "scorer": scorer,
            "direction": direction,
            "conditioning": conditioning,
            "delta": chosen.get("delta"),
            "abs_delta": chosen.get("abs_delta"),
            "prompts_path": probe_meta.get("prompts_path"),
            "chosen_at": datetime.now(timezone.utc).isoformat(),
        }
        exp_path = persist_experiment_record(mid, persist_key, record, session_id=session_id)
        payload["persisted_key"] = persist_key
        payload["experiment_path"] = exp_path
 
    return payload