confidence_analysis.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Standalone confidence pattern analysis over probe datasets (LLM)."""
 
from __future__ import annotations
 
from typing import Any
 
import torch
import torch.nn.functional as F
 
from aquin.compute.calibration import ece_proxy
from aquin.compute.sae_stats import _llm_probe_acts, load_probes
 
_META_KEYS = ("id", "stressor", "lang", "quant_run_id", "condition", "label", "group")
DEFAULT_THRESHOLD = 0.40
_BASELINE_TAG = "baseline"
 
 
def _prompt_confidence_metrics(model: Any, text: str) -> dict[str, Any]:
    tokens = model.to_tokens(text[:512])
    with torch.no_grad():
        logits = model(tokens)
    probs = F.softmax(logits[0].float(), dim=-1)
    log_probs = F.log_softmax(logits[0].float(), dim=-1)
    greedy = logits[0].argmax(dim=-1)
    idx = torch.arange(len(greedy), device=greedy.device)
    token_confs = probs[idx, greedy]
    entropies = -(probs * log_probs).sum(dim=-1)
    mean_conf = float(token_confs.mean().item())
    last_dist = probs[-1]
    return {
        "mean_confidence": round(mean_conf, 4),
        "max_prob": round(float(max(token_confs.max().item(), last_dist.max().item())), 4),
        "entropy": round(float(entropies.mean().item()), 4),
        "ece_proxy": round(abs(mean_conf - 0.5), 4),
        "last_token_entropy": round(float(-(last_dist * last_dist.clamp(min=1e-10).log()).sum().item()), 4),
    }
 
 
def _llm_sae_join_row(model: Any, model_id: str, layer: int, text: str) -> dict[str, Any]:
    from aquin.compute.feature_analysis import load_sae
 
    try:
        sae = load_sae(model_id, layer)
    except Exception as exc:
        return {"sae_available": False, "sae_error": str(exc)}
    acts = _llm_probe_acts(model, sae, model_id, layer, text)
    l0 = float((acts > 1e-6).sum().item())
    top_idx = int(acts.argmax().item())
    return {
        "sae_available": True,
        "sae_layer": layer,
        "mean_l0": round(l0, 4),
        "top_feature_idx": top_idx,
        "top_feature_act": round(float(acts[top_idx].item()), 6),
    }
 
 
def _stressor_summary(rows: list[dict[str, Any]], *, baseline: str = _BASELINE_TAG) -> list[dict[str, Any]]:
    groups: dict[str, list[dict[str, Any]]] = {}
    for row in rows:
        key = str(row.get("stressor") or row.get("label") or "unknown")
        groups.setdefault(key, []).append(row)
 
    base_rows = groups.get(baseline, [])
    base_conf = (
        sum(r["mean_confidence"] for r in base_rows) / len(base_rows)
        if base_rows else None
    )
    base_ent = (
        sum(r["entropy"] for r in base_rows) / len(base_rows)
        if base_rows else None
    )
 
    out: list[dict[str, Any]] = []
    for stressor, items in sorted(groups.items()):
        confs = [r["mean_confidence"] for r in items]
        ents = [r["entropy"] for r in items]
        entry = {
            "stressor": stressor,
            "n_probes": len(items),
            "mean_confidence": round(sum(confs) / len(confs), 4),
            "mean_entropy": round(sum(ents) / len(ents), 4),
            "ece_proxy": ece_proxy(confs),
            "low_confidence_count": sum(1 for r in items if r.get("low_confidence")),
        }
        if base_conf is not None and stressor != baseline:
            entry["confidence_delta"] = round(entry["mean_confidence"] - base_conf, 4)
        if base_ent is not None and stressor != baseline:
            entry["entropy_delta"] = round(entry["mean_entropy"] - base_ent, 4)
        out.append(entry)
    return out
 
 
def _finalize_payload(
    *,
    mode: str,
    model_id: str,
    probes: list[dict[str, Any]],
    per_probe: list[dict[str, Any]],
    threshold: float,
    join_sae: bool,
    sae_layer: int | None,
) -> dict[str, Any]:
    confs = [r["mean_confidence"] for r in per_probe]
    stressor_summary = _stressor_summary(per_probe)
    heatmap_rows = [s["stressor"] for s in stressor_summary]
    heatmap_values = [
        [s["mean_confidence"], s["mean_entropy"], s["ece_proxy"]]
        for s in stressor_summary
    ]
 
    return {
        "schema_version": 1,
        "model_id": model_id,
        "mode": mode,
        "n_probes": len(per_probe),
        "threshold": threshold,
        "aggregate_ece_proxy": ece_proxy(confs),
        "mean_confidence": round(sum(confs) / len(confs), 4) if confs else 0.0,
        "low_confidence_count": sum(1 for r in per_probe if r["low_confidence"]),
        "join_sae": join_sae,
        "sae_layer": sae_layer if join_sae else None,
        "probes": per_probe,
        "stressor_summary": stressor_summary,
        "heatmap": {
            "metric": "confidence_patterns",
            "rows": heatmap_rows,
            "cols": ["mean_confidence", "mean_entropy", "ece_proxy"],
            "values": heatmap_values,
        },
    }
 
 
def run_confidence_analysis(
    model: Any,
    model_id: str,
    probes: list[dict[str, Any]],
    *,
    threshold: float = DEFAULT_THRESHOLD,
    join_sae: bool = False,
    sae_layer: int | None = None,
) -> dict[str, Any]:
    from aquin.compute.model_loader import get_sae_layer
 
    resolved_layer = int(sae_layer if sae_layer is not None else get_sae_layer(model_id))
 
    per_probe: list[dict[str, Any]] = []
    for probe in probes:
        text = probe["text"]
        metrics = _prompt_confidence_metrics(model, text)
        row: dict[str, Any] = {
            "id": probe.get("id"),
            "text": text,
            **{k: probe[k] for k in _META_KEYS if k in probe and k != "id"},
            **metrics,
            "low_confidence": metrics["mean_confidence"] < threshold,
        }
        if join_sae:
            row.update(_llm_sae_join_row(model, model_id, resolved_layer, text))
        per_probe.append(row)
 
    return _finalize_payload(
        mode="llm",
        model_id=model_id,
        probes=probes,
        per_probe=per_probe,
        threshold=threshold,
        join_sae=join_sae,
        sae_layer=resolved_layer,
    )
 
 
def run_confidence_analysis_from_args(args: dict[str, Any]) -> dict[str, Any]:
    from aquin.compute.model_loader import (
        get_loaded_model,
        load_model,
        resolve_model_id,
    )
 
    model_id = args.get("model_id") or "llama-3.2-1b"
    prompts_raw = args.get("prompts")
    if prompts_raw is None or prompts_raw == "":
        return {
            "error": (
                "Missing prompts. Pass a file path to JSON/JSONL, a JSON array of "
                'strings, e.g. \'["The capital of France is", "2+2 equals"]\', or an '
                "array of {text: ...} objects."
            )
        }
 
    try:
        probes = _coerce_probes(prompts_raw)
    except (OSError, ValueError, TypeError) as exc:
        return {"error": str(exc)}
 
    join_sae = bool(args.get("join_sae") or args.get("join-sae"))
    layer_raw = args.get("layer") or args.get("sae_layer")
    sae_layer = int(layer_raw) if layer_raw is not None else None
    threshold = float(args.get("threshold") or DEFAULT_THRESHOLD)
 
    try:
        model_id = resolve_model_id(str(model_id))
    except ValueError as exc:
        return {"error": str(exc)}
 
    model = get_loaded_model()
    if model is None:
        try:
            model = load_model(model_id)
        except Exception as exc:
            return {"error": str(exc)}
 
    return run_confidence_analysis(
        model,
        model_id,
        probes,
        threshold=threshold,
        join_sae=join_sae,
        sae_layer=sae_layer,
    )
 
 
def _coerce_probes(prompts_raw: Any) -> list[dict[str, Any]]:
    """Accept file path, JSON array string, list of strings, or list of probe objects."""
    if isinstance(prompts_raw, list):
        return _rows_to_probes(prompts_raw)
 
    if isinstance(prompts_raw, dict):
        if isinstance(prompts_raw.get("probes"), list):
            return _rows_to_probes(prompts_raw["probes"])
        text = prompts_raw.get("text") or prompts_raw.get("prompt")
        if text:
            return [{"id": "probe_0", "text": str(text)}]
        raise ValueError("prompts object needs probes[] or text/prompt")
 
    s = str(prompts_raw).strip()
    if not s:
        raise ValueError("prompts is empty")
 
    # Inline JSON list / object
    if s[0] in "[{":
        import json
 
        try:
            parsed = json.loads(s)
        except json.JSONDecodeError as exc:
            raise ValueError(f"prompts looks like JSON but failed to parse: {exc}") from exc
        return _coerce_probes(parsed)
 
    # Single bare string that is not an existing path → treat as one probe
    # Prefer file when path exists (legacy CLI).
    from pathlib import Path
 
    try:
        return load_probes(s)
    except (OSError, ValueError, FileNotFoundError):
        p = Path(s).expanduser()
        if p.is_file():
            raise
        # newline-separated mini list
        lines = [ln.strip() for ln in s.splitlines() if ln.strip()]
        if len(lines) > 1:
            return _rows_to_probes(lines)
        return [{"id": "probe_0", "text": s}]
 
 
def _rows_to_probes(rows: list[Any]) -> list[dict[str, Any]]:
    probes: list[dict[str, Any]] = []
    for i, row in enumerate(rows):
        if isinstance(row, str):
            text = row.strip()
            if text:
                probes.append({"id": f"probe_{i}", "text": text})
            continue
        if not isinstance(row, dict):
            raise ValueError(f"Probe row {i} must be string or object")
        text = (
            row.get("text")
            or row.get("prompt")
            or row.get("input")
            or row.get("sentence")
        )
        if not text:
            raise ValueError(f"Probe row {i} missing text/prompt")
        out: dict[str, Any] = {"id": row.get("id") or f"probe_{i}", "text": str(text)}
        for key in _META_KEYS:
            if row.get(key) is not None:
                out[key] = row[key]
        probes.append(out)
    if not probes:
        raise ValueError("No probes after parsing prompts")
    return probes