interp_score.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.
 
"""
Ingested from inspection-backend/interp_score.py.
Import adaptation only: sae/feature_analysis/model_config -> aquin.compute.*.
"""
from __future__ import annotations
 
import torch
import torch.nn.functional as F
 
from transformer_lens import HookedTransformer
 
from aquin.compute.sae import SparseAutoencoder
from aquin.compute.feature_analysis import (
    get_causal_label,
    load_sae,
    load_norm,
    normalize,
    _sae_feature_ablate_hook,
)
from aquin.compute.model_loader import get_config
 
from aquin.compute.device import resolve_compute_device, synchronize_device
 
DEVICE = resolve_compute_device()
N_SAMPLES = 10
 
 
def _generate_sentences(label: str, client, n: int = N_SAMPLES) -> dict:
    from aquin.compute.llm_json import chat_json_completion, sentence_lists_from_llm
 
    prompt = f"""You are helping evaluate a sparse autoencoder feature with the label: "{label}"
 
Generate {n} short sentences (5-15 words each) where this feature SHOULD fire strongly,
and {n} short sentences where this feature should NOT fire.
 
Reply ONLY with a JSON object in this exact format, no markdown:
{{
  "positive": ["sentence1", "sentence2", ...],
  "negative": ["sentence1", "sentence2", ...]
}}"""
 
    resp = chat_json_completion(
        client,
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        max_tokens=800,
        temperature=0.7,
    )
    raw = resp.choices[0].message.content.strip()
    positive, negative = sentence_lists_from_llm(raw)
    return {"positive": positive, "negative": negative}
 
 
def _get_feature_activation(sentence: str, feature_idx: int, model: HookedTransformer, sae: SparseAutoencoder, model_id: str = "llama-3.2-1b", layer: int | None = None) -> float:
    cfg = get_config(model_id)
    resolved_layer = layer if layer is not None else cfg["sae_layer"]
    tokens = model.to_tokens(sentence)
    with torch.no_grad():
        _, cache = model.run_with_cache(
            tokens,
            names_filter=f"blocks.{resolved_layer}.hook_resid_post",
            return_type=None,
        )
    resid = cache[f"blocks.{resolved_layer}.hook_resid_post"][0]
    acts = sae.encode(normalize(resid, model_id, resolved_layer))
    return float(acts[:, feature_idx].max().item())
 
 
def _cohen_d_score(pos_acts: list[float], neg_acts: list[float]) -> float:
    if not pos_acts or not neg_acts:
        return 0.0
    pos_t = torch.tensor(pos_acts, dtype=torch.float32)
    neg_t = torch.tensor(neg_acts, dtype=torch.float32)
    mu_pos = pos_t.mean().item()
    mu_neg = neg_t.mean().item()
    pooled_std = torch.cat([pos_t, neg_t]).std().item() + 1e-6
    raw = (mu_pos - mu_neg) / pooled_std
    return round(float(max(0.0, min(1.0, raw))), 4)
 
 
def _feature_purity_score(sentences: list[str], client) -> float | None:
    if not sentences or client is None:
        return None
 
    try:
        resp = client.embeddings.create(
            model="text-embedding-3-large",
            input=sentences[:10],
        )
        vecs = torch.tensor(
            [e.embedding for e in resp.data], dtype=torch.float32
        )
        vecs = F.normalize(vecs, dim=-1)
        sim_matrix = vecs @ vecs.T
        n = vecs.shape[0]
        if n < 2:
            return None
        upper_mask = torch.ones(n, n, dtype=torch.bool).triu(diagonal=1)
        mean_sim = sim_matrix[upper_mask].mean().item()
        purity = (mean_sim + 1.0) / 2.0
        return round(float(purity), 4)
    except Exception as e:
        print(f"[purity] embedding error: {e}", flush=True)
        return None
 
 
def _finite(x: float | None, default: float | None = None) -> float | None:
    import math
 
    if x is None:
        return default
    try:
        f = float(x)
    except (TypeError, ValueError):
        return default
    if math.isnan(f) or math.isinf(f):
        return default
    return f
 
 
def _kl_div(p: torch.Tensor, q: torch.Tensor) -> float:
    p = p.clamp(min=1e-10)
    q = q.clamp(min=1e-10)
    return float((p * (p / q).log()).sum().item())
 
 
def _entropy(p: torch.Tensor) -> float:
    p = p.clamp(min=1e-10)
    return float(-(p * p.log()).sum().item())
 
 
def run_mui_score(
    feature_idx: int,
    prompt: str,
    model: HookedTransformer,
    sae: SparseAutoencoder,
    n_positions: int = 8,
    model_id: str = "llama-3.2-1b",
    layer: int | None = None,
) -> dict:
    cfg = get_config(model_id)
    sae_layer = layer if layer is not None else cfg["sae_layer"]
    tokens = model.to_tokens(prompt)
 
    with torch.no_grad():
        baseline_logits, cache = model.run_with_cache(
            tokens,
            names_filter=f"blocks.{sae_layer}.hook_resid_post",
        )
 
    resid = cache[f"blocks.{sae_layer}.hook_resid_post"][0]
    acts = sae.encode(normalize(resid, model_id, sae_layer))
 
    feat_acts = acts[:, feature_idx]
    top_k = min(n_positions, feat_acts.shape[0])
    top_positions = feat_acts.topk(top_k).indices.tolist()
    top_positions = [p for p in top_positions if feat_acts[p].item() > 0.01]
 
    if not top_positions:
        return {
            "feature_idx": feature_idx,
            "score": 0.0,
            "per_position": [],
            "mean_kl": 0.0,
            "baseline_entropy": None,
        }
 
    baseline_probs = torch.softmax(baseline_logits[0, -1], dim=-1)
    baseline_H = _entropy(baseline_probs)
 
    per_position = []
    kl_vals = []
 
    for pos in top_positions:
        act_val = feat_acts[pos].item()
 
        def ablate(value, hook, pos=pos):
            return _sae_feature_ablate_hook(
                value, pos, feature_idx, sae, model_id, sae_layer,
            )
 
        with torch.no_grad():
            abl_logits = model.run_with_hooks(
                tokens,
                fwd_hooks=[(f"blocks.{sae_layer}.hook_resid_post", ablate)]
            )
 
        abl_probs = torch.softmax(abl_logits[0, -1], dim=-1)
        kl = _finite(_kl_div(baseline_probs, abl_probs), 0.0) or 0.0
        kl_vals.append(kl)
 
        context = model.tokenizer.decode(
            tokens[0, max(0, pos - 3):pos + 4].tolist()
        )
        per_position.append({
            "position": pos,
            "activation": round(act_val, 3),
            "kl_divergence": round(kl, 4),
            "context": context,
        })
 
    mean_kl = sum(kl_vals) / len(kl_vals) if kl_vals else 0.0
    baseline_H = _finite(baseline_H, 0.0) or 0.0
    mean_kl = _finite(mean_kl, 0.0) or 0.0
    mui = _finite(min(mean_kl / max(baseline_H, 1e-6), 1.0), 0.0) or 0.0
    mui = round(float(mui), 4)
 
    synchronize_device()
 
    return {
        "feature_idx": feature_idx,
        "score": mui,
        "per_position": sorted(per_position, key=lambda x: x["kl_divergence"], reverse=True),
        "mean_kl": round(mean_kl, 4),
        "baseline_entropy": round(baseline_H, 4),
    }
 
 
def run_interp_score(
    feature_idx: int,
    prompt: str,
    model: HookedTransformer,
    sae: SparseAutoencoder,
    client,
    n_samples: int = N_SAMPLES,
    model_id: str = "llama-3.2-1b",
    layer: int | None = None,
) -> dict:
    label = get_causal_label(feature_idx, prompt, model, sae, client, model_id=model_id, layer=layer)
    mui_result = run_mui_score(feature_idx, prompt, model, sae, model_id=model_id, layer=layer)
 
    if client is None:
        return {
            "feature_idx": feature_idx,
            "label": label,
            "score": None,
            "purity_score": None,
            "mui_score": mui_result["score"],
            "mui_per_position": mui_result["per_position"],
            "mui_mean_kl": mui_result["mean_kl"],
            "baseline_entropy": mui_result["baseline_entropy"],
            "error": "OpenAI not available. Set OPENAI_API_KEY on this machine.",
            "positive_examples": [],
            "negative_examples": [],
            "positive_mean": None,
            "negative_mean": None,
        }
 
    try:
        sentences = _generate_sentences(label, client, n=n_samples)
    except Exception as e:
        from aquin.compute.llm_json import llm_sentence_generation_error
        print(f"[interp] sentence generation: {e}", flush=True)
        return {
            "feature_idx": feature_idx,
            "label": label,
            "score": None,
            "purity_score": None,
            "mui_score": mui_result["score"],
            "mui_per_position": mui_result["per_position"],
            "mui_mean_kl": mui_result["mean_kl"],
            "baseline_entropy": mui_result["baseline_entropy"],
            "error": llm_sentence_generation_error(),
            "positive_examples": [],
            "negative_examples": [],
            "positive_mean": None,
            "negative_mean": None,
        }
 
    positive_results = []
    for sent in sentences.get("positive", []):
        try:
            act = _get_feature_activation(sent, feature_idx, model, sae, model_id=model_id, layer=layer)
            positive_results.append({"sentence": sent, "activation": round(act, 4)})
        except Exception as e:
            print(f"[interp_score] positive sentence failed: {e}", flush=True)
 
    negative_results = []
    for sent in sentences.get("negative", []):
        try:
            act = _get_feature_activation(sent, feature_idx, model, sae, model_id=model_id, layer=layer)
            negative_results.append({"sentence": sent, "activation": round(act, 4)})
        except Exception as e:
            print(f"[interp_score] negative sentence failed: {e}", flush=True)
 
    pos_acts = [r["activation"] for r in positive_results]
    neg_acts = [r["activation"] for r in negative_results]
 
    score      = _cohen_d_score(pos_acts, neg_acts)
    pos_mean   = round(sum(pos_acts) / len(pos_acts), 4) if pos_acts else None
    neg_mean   = round(sum(neg_acts) / len(neg_acts), 4) if neg_acts else None
 
    purity_score = _feature_purity_score(
        [r["sentence"] for r in positive_results], client
    )
 
    synchronize_device()
 
    return {
        "feature_idx": feature_idx,
        "label": label,
        "score": score,
        "purity_score": purity_score,
        "mui_score": mui_result["score"],
        "mui_per_position": mui_result["per_position"],
        "mui_mean_kl": mui_result["mean_kl"],
        "baseline_entropy": mui_result["baseline_entropy"],
        "positive_mean": pos_mean,
        "negative_mean": neg_mean,
        "positive_examples": sorted(positive_results, key=lambda x: x["activation"], reverse=True),
        "negative_examples": sorted(negative_results, key=lambda x: x["activation"], reverse=True),
    }