activation_replay.py

Go to the documentation of this file.

#!/usr/bin/env python3
# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Replay/query saved activation captures by metadata."""
 
from __future__ import annotations
 
import json
from collections import defaultdict
from pathlib import Path
from typing import Any
 
import torch
 
 
def _read_json(path: Path) -> dict[str, Any]:
    data = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(data, dict):
        raise ValueError(f"Expected JSON object in {path}")
    return data
 
 
def load_capture_root(path: str | Path) -> dict[str, Any]:
    root = Path(path).expanduser().resolve()
    if root.is_file():
        root = root.parent
    manifest_path = root / "manifest.json"
    metadata_path = root / "metadata.json"
    summary_path = root / "summary.jsonl"
    if metadata_path.is_file():
        meta = _read_json(metadata_path)
        if manifest_path.is_file():
            man = _read_json(manifest_path)
            files = man.get("files") if isinstance(man.get("files"), dict) else {}
            meta_files = meta.get("files") if isinstance(meta.get("files"), dict) else {}
            if files or meta_files:
                merged = {**files, **meta_files}
                if isinstance(files.get("activations"), dict) and isinstance(meta_files.get("activations"), dict):
                    merged["activations"] = {**files["activations"], **meta_files["activations"]}
                meta["files"] = merged
            if not meta.get("probes") and man.get("probes"):
                meta["probes"] = man["probes"]
            if not meta.get("sae_features") and man.get("sae_features"):
                meta["sae_features"] = man["sae_features"]
    elif manifest_path.is_file():
        meta = _read_json(manifest_path)
    else:
        raise FileNotFoundError(f"No metadata.json or manifest.json in {root}")
 
    rows: list[dict[str, Any]] = []
    if summary_path.is_file():
        for line in summary_path.read_text(encoding="utf-8").splitlines():
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            if isinstance(row, dict):
                rows.append(row)
    return {
        "root": root,
        "metadata_path": metadata_path if metadata_path.is_file() else manifest_path,
        "manifest": meta,
        "rows": rows,
    }
 
 
def _flatten_probe_row(row: dict[str, Any], probe_lookup: dict[str, dict[str, Any]]) -> dict[str, Any]:
    probe_id = str(row.get("probe_id") or "")
    probe = probe_lookup.get(probe_id, {})
    probe_meta = probe.get("metadata") if isinstance(probe.get("metadata"), dict) else {}
    row_meta = row.get("metadata") if isinstance(row.get("metadata"), dict) else {}
    merged_meta = {**probe_meta, **row_meta}
    flat: dict[str, Any] = {
        "probe_id": probe_id,
        "layer": row.get("layer"),
        "d_model": row.get("d_model"),
        "text": probe.get("text"),
    }
    for key, val in merged_meta.items():
        flat[key] = val
    return flat
 
 
def _parse_filter(expr: str | None) -> list[tuple[str, str]]:
    if not expr:
        return []
    parts: list[tuple[str, str]] = []
    for chunk in expr.split(","):
        chunk = chunk.strip()
        if not chunk:
            continue
        if "=" not in chunk:
            raise ValueError(f"Bad filter {chunk!r}; use key=value,key2=value2")
        key, value = chunk.split("=", 1)
        parts.append((key.strip(), value.strip()))
    return parts
 
 
def _matches(row: dict[str, Any], filters: list[tuple[str, str]]) -> bool:
    for key, expected in filters:
        got = row.get(key)
        if got is None:
            return False
        if str(got) != expected:
            return False
    return True
 
 
def replay_capture(
    path: str | Path,
    *,
    filter_expr: str | None = None,
    group_by: str | None = None,
    limit: int | None = None,
) -> dict[str, Any]:
    loaded = load_capture_root(path)
    manifest = loaded["manifest"]
    probes = manifest.get("probes") if isinstance(manifest.get("probes"), list) else []
    probe_lookup = {
        str(p.get("id")): p for p in probes
        if isinstance(p, dict) and p.get("id") is not None
    }
    rows = [_flatten_probe_row(r, probe_lookup) for r in loaded["rows"]]
    filters = _parse_filter(filter_expr)
    filtered = [r for r in rows if _matches(r, filters)]
    if limit is not None and limit >= 0:
        filtered = filtered[:limit]
 
    result: dict[str, Any] = {
        "root": str(loaded["root"]),
        "metadata_path": str(loaded["metadata_path"]),
        "capture_id": manifest.get("capture_id"),
        "model_id": manifest.get("model_id"),
        "model_mode": manifest.get("model_mode"),
        "n_rows": len(rows),
        "n_filtered": len(filtered),
        "filter": filter_expr,
        "group_by": group_by,
        "rows": filtered,
    }
    if group_by:
        groups: dict[str, list[dict[str, Any]]] = defaultdict(list)
        for row in filtered:
            groups[str(row.get(group_by, "unknown"))].append(row)
        summary = []
        for key, items in sorted(groups.items(), key=lambda kv: (-len(kv[1]), kv[0])):
            layer_counts: dict[str, int] = defaultdict(int)
            for item in items:
                layer_counts[str(item.get("layer", "?"))] += 1
            summary.append({
                "group": key,
                "count": len(items),
                "layers": dict(sorted(layer_counts.items(), key=lambda kv: int(kv[0]) if kv[0].isdigit() else 999999)),
            })
        result["groups"] = summary
    return result
 
 
COHORT_FIELDS = ("group", "label", "lang", "stressor", "source")
 
 
def _probe_groups(probes: list[dict[str, Any]], field: str) -> dict[str, list[int]]:
    grouped: dict[str, list[int]] = defaultdict(list)
    for i, probe in enumerate(probes):
        if not isinstance(probe, dict):
            continue
        meta = probe.get("metadata") if isinstance(probe.get("metadata"), dict) else {}
        grouped[str(meta.get(field, "unknown"))].append(i)
    return grouped
 
 
def _usable_groups(grouped: dict[str, list[int]]) -> bool:
    keys = [k for k, idxs in grouped.items() if idxs]
    return len(keys) >= 2
 
 
def resolve_compare_group(
    probes: list[dict[str, Any]],
    requested: str,
) -> tuple[str, dict[str, list[int]], bool]:
    """Pick a metadata field with ≥2 cohorts. Last resort: split probes in half."""
    grouped = _probe_groups(probes, requested)
    if _usable_groups(grouped):
        return requested, grouped, False
    for field in COHORT_FIELDS:
        if field == requested:
            continue
        candidate = _probe_groups(probes, field)
        if _usable_groups(candidate):
            return field, candidate, False
    n = len(probes)
    if n < 2:
        raise ValueError(
            f"Need at least 2 groups for {requested!r}; found {len(grouped)}. "
            "Recapture with labeled probes (metadata group/label/lang)."
        )
    mid = max(1, n // 2)
    inferred = {"a": list(range(0, mid)), "b": list(range(mid, n))}
    return requested, inferred, True
 
 
def _write_json(path: Path, data: dict[str, Any]) -> None:
    path.write_text(json.dumps(data, indent=2), encoding="utf-8")
 
 
def encode_capture_sae_features(
    path: str | Path,
    *,
    sae_layer: int | None = None,
    sae: Any | None = None,
) -> dict[str, Any]:
    """Encode residual activations in a capture with the SAE for that layer."""
    loaded = load_capture_root(path)
    manifest = loaded["manifest"]
    root = loaded["root"]
    files = manifest.get("files") if isinstance(manifest.get("files"), dict) else {}
    sae_rel = files.get("sae_features")
    if sae_rel and (root / str(sae_rel)).is_file():
        return {
            "root": str(root),
            "encoded": False,
            "sae_features": manifest.get("sae_features"),
        }
 
    if str(manifest.get("granularity") or "prompt") == "token":
        raise ValueError("Cannot encode SAE features on a token-level capture.")
 
    model_id = str(manifest.get("model_id") or "").strip()
    if not model_id:
        raise ValueError("Capture has no model_id; cannot encode SAE features.")
 
    activations = files.get("activations") if isinstance(files.get("activations"), dict) else {}
    enc_layer = int(sae_layer) if sae_layer is not None else None
    if enc_layer is None:
        try:
            from aquin.compute.model_loader import get_sae_layer
 
            enc_layer = int(get_sae_layer(model_id))
        except Exception:
            enc_layer = None
    if enc_layer is None or str(enc_layer) not in activations:
        try:
            from aquin.compute.model_loader import get_available_sae_layers
 
            for layer in get_available_sae_layers(model_id):
                if str(layer) in activations:
                    enc_layer = int(layer)
                    break
        except Exception:
            pass
    if enc_layer is None or str(enc_layer) not in activations:
        raise ValueError(
            "Capture has no residual activations for a loadable SAE layer. "
            "Re-run activations.capture with encode_sae=true (pythia-70m: layer 3)."
        )
 
    resid_path = root / str(activations[str(enc_layer)])
    if not resid_path.is_file():
        raise FileNotFoundError(f"Missing residual activations: {resid_path}")
    resid = torch.load(resid_path, map_location="cpu", weights_only=False)
    if not isinstance(resid, torch.Tensor) or resid.dim() != 2:
        raise ValueError(
            f"Expected residual matrix [N, d], got {type(resid)} "
            f"{tuple(resid.shape) if isinstance(resid, torch.Tensor) else ''}"
        )
 
    if sae is None:
        from aquin.compute.feature_analysis import load_sae
 
        sae = load_sae(model_id, layer=int(enc_layer))
    from aquin.compute.feature_analysis import normalize
 
    device = next(sae.parameters()).device
    with torch.no_grad():
        encoded = sae.encode(normalize(resid.to(device), model_id, int(enc_layer))).float().cpu()
 
    sae_rel = f"sae/sae_layer_{int(enc_layer)}.pt"
    (root / "sae").mkdir(parents=True, exist_ok=True)
    torch.save(encoded, root / sae_rel)
    sae_features = {
        "layer": int(enc_layer),
        "n_features": int(encoded.shape[1]),
        "n_probes": int(encoded.shape[0]),
    }
    files = dict(files)
    files["sae_features"] = sae_rel
    manifest["files"] = files
    manifest["sae_features"] = sae_features
    manifest["encode_sae"] = True
 
    meta_path = loaded["metadata_path"]
    _write_json(meta_path, manifest)
    other = root / ("manifest.json" if meta_path.name == "metadata.json" else "metadata.json")
    if other.is_file():
        other_data = _read_json(other)
        other_files = other_data.get("files") if isinstance(other_data.get("files"), dict) else {}
        other_files["sae_features"] = sae_rel
        other_data["files"] = other_files
        other_data["sae_features"] = sae_features
        other_data["encode_sae"] = True
        _write_json(other, other_data)
 
    return {"root": str(root), "encoded": True, "sae_features": sae_features}
 
 
def compare_capture_features(
    path: str | Path,
    *,
    group: str,
    top_k: int = 10,
    encode_sae: bool = True,
    sae_layer: int | None = None,
    sae: Any | None = None,
) -> dict[str, Any]:
    loaded = load_capture_root(path)
    manifest = loaded["manifest"]
    probes = manifest.get("probes") if isinstance(manifest.get("probes"), list) else []
    files = manifest.get("files") if isinstance(manifest.get("files"), dict) else {}
    sae_rel = files.get("sae_features")
    sae_path = loaded["root"] / str(sae_rel) if sae_rel else None
    if not sae_rel or sae_path is None or not sae_path.is_file():
        if not encode_sae:
            raise ValueError("Capture has no SAE feature matrix. Re-run capture with --encode-sae.")
        encode_capture_sae_features(path, sae_layer=sae_layer, sae=sae)
        loaded = load_capture_root(path)
        manifest = loaded["manifest"]
        probes = manifest.get("probes") if isinstance(manifest.get("probes"), list) else []
        files = manifest.get("files") if isinstance(manifest.get("files"), dict) else {}
        sae_rel = files.get("sae_features")
        sae_path = loaded["root"] / str(sae_rel) if sae_rel else None
        if not sae_rel or sae_path is None or not sae_path.is_file():
            raise ValueError("Capture has no SAE feature matrix. Re-run capture with --encode-sae.")
 
    feats = torch.load(sae_path, map_location="cpu", weights_only=False)
    if feats.dim() != 2:
        raise ValueError(f"Expected SAE features [N, F], got shape {tuple(feats.shape)}")
    if probes and len(probes) != int(feats.shape[0]):
        raise ValueError(f"Probe count {len(probes)} does not match SAE rows {int(feats.shape[0])}")
    if not probes:
        raise ValueError("Capture has no probes; features.compare needs a labeled activation capture.")
 
    group, grouped, inferred = resolve_compare_group(probes, group)
 
    group_means: dict[str, torch.Tensor] = {}
    for key, idxs in grouped.items():
        if not idxs:
            continue
        group_means[key] = feats[torch.tensor(idxs, dtype=torch.long)].mean(dim=0)
 
    keys = sorted(group_means)
    stacked = torch.stack([group_means[k] for k in keys], dim=0)
    spread = stacked.max(dim=0).values - stacked.min(dim=0).values
    k = min(int(top_k), int(spread.shape[0]))
    top_idx = spread.topk(k).indices.tolist()
 
    rows = []
    for idx in top_idx:
        per_group = {key: round(float(group_means[key][idx].item()), 6) for key in keys}
        rows.append({
            "feature_idx": int(idx),
            "spread": round(float(spread[idx].item()), 6),
            "groups": per_group,
        })
 
    return {
        "root": str(loaded["root"]),
        "capture_id": manifest.get("capture_id"),
        "model_id": manifest.get("model_id"),
        "group": group,
        "n_groups": len(group_means),
        "group_sizes": {key: len(grouped[key]) for key in keys},
        "sae_layer": (manifest.get("sae_features") or {}).get("layer"),
        "n_features": int(feats.shape[1]),
        "top_k": k,
        "group_inferred": inferred,
        "features": rows,
    }