activation_capture.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Batch activation capture for labeled probe sets (SAE training, diff, analysis)."""
 
from __future__ import annotations
 
import json
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Literal
 
import torch
 
PROBE_TEXT_KEYS = ("instruction", "prompt", "text", "content", "response")
RESERVED_PROBE_KEYS = frozenset({"id", *PROBE_TEXT_KEYS})
 
Position = Literal["last", "mean"]
Granularity = Literal["prompt", "token"]
ModelMode = Literal["llm"]
MAX_PROBE_COUNT = 64
BALANCE_PRIORITY = ("label", "stressor", "lang", "group")
 
 
def resolve_prompts_path(path: str | Path) -> Path | None:
    """Resolve prompts file from cwd or repo parent (e.g. when run from cli/)."""
    raw = Path(path)
    if raw.is_file():
        return raw.resolve()
    candidates = [
        Path.cwd() / raw,
        Path.cwd().parent / raw,
        Path.cwd().parent.parent / raw,
    ]
    for candidate in candidates:
        if candidate.is_file():
            return candidate.resolve()
    return None
 
 
def write_probes_jsonl(path: Path, probes: list[dict[str, Any]]) -> Path:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as fh:
        for probe in probes:
            row = {"id": probe["id"], "instruction": probe["text"], **(probe.get("metadata") or {})}
            fh.write(json.dumps(row, ensure_ascii=False) + "\n")
    return path
 
 
def _balance_key(probe: dict[str, Any], group: str | None = None) -> str | None:
    meta = probe.get("metadata") if isinstance(probe.get("metadata"), dict) else {}
    if group:
        val = meta.get(group)
        if val is not None and str(val).strip():
            return f"{group}:{val}"
        return None
    for key in BALANCE_PRIORITY:
        val = meta.get(key)
        if val is not None and str(val).strip():
            return f"{key}:{val}"
    return None
 
 
def balance_probes(
    probes: list[dict[str, Any]],
    *,
    group: str | None = None,
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    groups: dict[str, list[dict[str, Any]]] = {}
    for probe in probes:
        key = _balance_key(probe, group=group)
        if key is None:
            return probes, {
                "balanced": False,
                "balance_reason": "no_metadata_groups",
                **({"balance_group": group} if group else {}),
            }
        groups.setdefault(key, []).append(probe)
    if len(groups) <= 1:
        return probes, {
            "balanced": False,
            "balance_reason": "single_group",
            "balance_groups": len(groups),
            **({"balance_group": group} if group else {}),
        }
 
    target = min(len(items) for items in groups.values())
    balanced: list[dict[str, Any]] = []
    group_counts: dict[str, int] = {}
    for key in sorted(groups):
        kept = groups[key][:target]
        balanced.extend(kept)
        group_counts[key] = len(kept)
    return balanced, {
        "balanced": True,
        "balance_groups": len(groups),
        "balance_per_group": target,
        "balance_counts": group_counts,
        **({"balance_group": group} if group else {}),
    }
 
 
def generate_llm_probes(
    model_id: str,
    count: int,
    *,
    topic: str = "general knowledge, reasoning, and instructions",
) -> list[dict[str, Any]]:
    """Use the loaded LLM to generate diverse probe strings."""
    from aquin.compute.causal_trace import run_chat
    from aquin.compute.model_loader import get_loaded_model, load_model
    from aquin.compute.sae_diff import DEFAULT_PROMPTS
 
    model = get_loaded_model() or load_model(model_id)
    n = max(1, min(int(count), MAX_PROBE_COUNT))
    topic = " ".join(str(topic).split()) or "general knowledge"
 
    prompt = (
        f"Generate exactly {n} diverse short English probe prompts for language model "
        f"interpretability testing. Mix factual completions, reasoning, and instruction-style "
        f"prompts. Theme: {topic}. Output one probe per line. No numbering, bullets, or explanations."
    )
    text = run_chat(
        prompt,
        model_id=model_id,
        max_new_tokens=min(n * 28, 900),
        temperature=0.85,
        model=model,
    )
 
    lines: list[str] = []
    for raw in text.splitlines():
        line = raw.strip().lstrip("-•*0123456789.) ").strip('"').strip("'").strip()
        if len(line) < 8 or line.lower().startswith(("here are", "sure", "generate", "probe")):
            continue
        lines.append(line[:512])
 
    probes: list[dict[str, Any]] = [
        {
            "id": f"g{i}",
            "text": line,
            "metadata": {"source": "generated", "topic": topic, "group": "factual" if i % 2 == 0 else "instruction"},
        }
        for i, line in enumerate(lines[:n])
    ]
 
    if len(probes) < n:
        for j, fallback in enumerate(DEFAULT_PROMPTS):
            if len(probes) >= n:
                break
            probes.append({
                "id": f"d{j}",
                "text": fallback,
                "metadata": {"source": "default", "topic": topic, "group": "factual" if j % 2 == 0 else "instruction"},
            })
 
    print(f"[capture] generated {len(probes[:n])} LLM probes (topic={topic!r})", flush=True)
    return probes[:n]
 
 
 
def resolve_probes_for_capture(
    *,
    model_id: str,
    model_mode: ModelMode,
    prompts_path: str | Path | None,
    count: int,
    topic: str | None,
    balance: bool = False,
    balance_group: str | None = None,
    output_dir: Path | None = None,
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    """Load probes from file or generate when --prompts omitted."""
    meta: dict[str, Any] = {
        "probes_source": "file",
        "requested_count": count,
        "topic": topic,
    }
 
    if prompts_path:
        resolved = resolve_prompts_path(prompts_path)
        if resolved is None:
            raise FileNotFoundError(
                f"Prompts file not found: {prompts_path} "
                f"(cwd={Path.cwd()}). Omit --prompts to auto-generate with --count and --topic."
            )
        meta["prompts_path"] = str(resolved)
        probes = load_probes(resolved)
        if count and len(probes) > count:
            probes = probes[:count]
        if balance:
            probes, bal_meta = balance_probes(probes, group=balance_group)
            meta.update(bal_meta)
        return probes, meta
 
    theme = topic or "general knowledge, reasoning, and instructions"
    meta["probes_source"] = "generated"
    probes = generate_llm_probes(model_id, count, topic=theme)
 
    if output_dir is not None:
        out_path = write_probes_jsonl(output_dir / "probes.jsonl", probes)
        meta["generated_probes_path"] = str(out_path)
        print(f"[capture] wrote {out_path}", flush=True)
 
    if balance:
        probes, bal_meta = balance_probes(probes, group=balance_group)
        meta.update(bal_meta)
 
    return probes, meta
 
 
def load_probes(path: str | Path | None, *, fallback: list[str] | None = None) -> list[dict[str, Any]]:
    """Load probes with optional metadata from JSON/JSONL."""
    if not path:
        texts = fallback or []
        return [{"id": f"p{i}", "text": t, "metadata": {}} for i, t in enumerate(texts)]
 
    p = Path(path)
    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]
 
    out: list[dict[str, Any]] = []
    for i, row in enumerate(rows):
        if isinstance(row, str):
            out.append({"id": f"p{i}", "text": row.strip(), "metadata": {}})
            continue
        if not isinstance(row, dict):
            continue
        text_val = ""
        for key in PROBE_TEXT_KEYS:
            val = row.get(key)
            if isinstance(val, str) and val.strip():
                text_val = val.strip()
                break
        if not text_val:
            continue
        probe_id = str(row.get("id") or f"p{i}")
        metadata = {k: v for k, v in row.items() if k not in RESERVED_PROBE_KEYS}
        out.append({"id": probe_id, "text": text_val, "metadata": metadata})
 
    if not out:
        raise ValueError(f"No probes found in {p}")
    return out
 
 
def parse_layers(spec: str | None, n_layers: int) -> list[int]:
    if not spec or spec.strip().lower() == "all":
        return list(range(n_layers))
    layers: list[int] = []
    for part in spec.split(","):
        part = part.strip()
        if not part:
            continue
        layer = int(part)
        if layer < 0 or layer >= n_layers:
            raise ValueError(f"Layer {layer} out of range [0, {n_layers - 1}]")
        layers.append(layer)
    if not layers:
        raise ValueError("No layers specified")
    return sorted(set(layers))
 
 
def resolve_capture_model_id(model_id: str) -> tuple[str, ModelMode]:
    from aquin.compute.model_loader import resolve_model_id
 
    return resolve_model_id(model_id), "llm"
 
 
def llm_layer_count(model_id: str, checkpoint_path: str | Path | None = None) -> int:
    """Catalog layer count — never load weights just to read n_layers (OOM on T4)."""
    from aquin.compute.model_loader import get_config, resolve_model_id
 
    _ = checkpoint_path  # checkpoint cannot change architecture layer count
    return int(get_config(resolve_model_id(model_id))["n_layers"])
 
 
 
def _pool_activation(tensor: torch.Tensor, position: Position) -> torch.Tensor:
    """tensor: (seq, d_model) -> (d_model,)"""
    if position == "last":
        return tensor[-1].float().cpu()
    return tensor.mean(dim=0).float().cpu()
 
 
def _forward_llm_layers(
    model: Any,
    text: str,
    layers: list[int],
    *,
    position: Position,
    max_chars: int = 512,
) -> dict[int, torch.Tensor]:
    tokens = model.to_tokens(text[:max_chars])
    hooks = {f"blocks.{layer}.hook_resid_post" for layer in layers}
    with torch.no_grad():
        _, cache = model.run_with_cache(
            tokens,
            names_filter=lambda name: name in hooks,
            return_type=None,
        )
    out: dict[int, torch.Tensor] = {}
    for layer in layers:
        key = f"blocks.{layer}.hook_resid_post"
        if key in cache:
            out[layer] = _pool_activation(cache[key][0], position)
    return out
 
 
def _forward_llm_layers_token(
    model: Any,
    text: str,
    layers: list[int],
    *,
    max_chars: int = 512,
) -> tuple[dict[int, torch.Tensor], list[str]]:
    tokens = model.to_tokens(text[:max_chars])
    hooks = {f"blocks.{layer}.hook_resid_post" for layer in layers}
    with torch.no_grad():
        _, cache = model.run_with_cache(
            tokens,
            names_filter=lambda name: name in hooks,
            return_type=None,
        )
    token_ids = tokens[0].tolist()
    token_texts = [model.tokenizer.decode([tid], skip_special_tokens=False) for tid in token_ids]
    out: dict[int, torch.Tensor] = {}
    for layer in layers:
        key = f"blocks.{layer}.hook_resid_post"
        if key in cache:
            out[layer] = cache[key][0].float().cpu()
    return out, token_texts
 
 
 
def _write_capture_artifacts(
    *,
    out_root: Path,
    probes: list[dict[str, Any]],
    layer_list: list[int],
    per_layer: dict[int, list[torch.Tensor]],
    summary_rows: list[dict[str, Any]],
    model_id: str,
    model_mode: ModelMode,
    d_model: int,
    checkpoint_path: str | Path | None,
    checkpoint_name: str | None,
    position: Position,
    encode_sae: bool,
    capture_name: str | None,
    sae_file: str | None,
    sae_features: dict[str, Any] | None,
    manifest_extras: dict[str, Any] | None = None,
    granularity: Granularity = "prompt",
    token_spans: list[dict[str, Any]] | None = None,
) -> dict[str, Any]:
    (out_root / "layers").mkdir(parents=True, exist_ok=True)
    activation_files: dict[str, str] = {}
    for layer, vecs in per_layer.items():
        if not vecs:
            continue
        rel = f"layers/layer_{layer}.pt"
        stacked = torch.cat(vecs, dim=0) if granularity == "token" else torch.stack(vecs)
        torch.save(stacked, out_root / rel)
        activation_files[str(layer)] = rel
 
    capture_id = capture_name or uuid.uuid4().hex[:12]
    manifest: dict[str, Any] = {
        "schema_version": 1,
        "capture_id": capture_id,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "model_id": model_id,
        "model_mode": model_mode,
        "checkpoint": str(checkpoint_path) if checkpoint_path else None,
        "checkpoint_name": checkpoint_name,
        "position": position,
        "granularity": granularity,
        "encode_sae": encode_sae,
        "n_probes": len(probes),
        "layers": layer_list,
        "d_model": d_model,
        "output_dir": str(out_root.resolve()),
        "files": {
            "summary": "summary.jsonl",
            "activations": activation_files,
        },
        "probes": probes,
    }
    if manifest_extras:
        manifest.update(manifest_extras)
    if sae_file:
        manifest["files"]["sae_features"] = sae_file
        manifest["sae_features"] = sae_features
    if granularity == "token":
        manifest["files"]["token_spans"] = "token_spans.jsonl"
 
    metadata_payload: dict[str, Any] = {
        "schema_version": 1,
        "kind": "activation_capture_metadata",
        "capture_id": capture_id,
        "created_at": manifest["created_at"],
        "model_id": model_id,
        "model_mode": model_mode,
        "checkpoint": str(checkpoint_path) if checkpoint_path else None,
        "checkpoint_name": checkpoint_name,
        "position": position,
        "granularity": granularity,
        "encode_sae": encode_sae,
        "n_probes": len(probes),
        "layers": layer_list,
        "d_model": d_model,
        "output_dir": str(out_root.resolve()),
        "files": manifest["files"],
        "probes": probes,
    }
    if manifest_extras:
        metadata_payload.update(manifest_extras)
    if sae_features:
        metadata_payload["sae_features"] = sae_features
 
    (out_root / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    (out_root / "metadata.json").write_text(json.dumps(metadata_payload, indent=2), encoding="utf-8")
    with (out_root / "summary.jsonl").open("w", encoding="utf-8") as fh:
        for row in summary_rows:
            fh.write(json.dumps(row) + "\n")
    if granularity == "token":
        with (out_root / "token_spans.jsonl").open("w", encoding="utf-8") as fh:
            for row in (token_spans or []):
                fh.write(json.dumps(row, ensure_ascii=False) + "\n")
 
    probe_samples = [
        {
            "id": p["id"],
            **({"label": p["metadata"]["label"]} if p.get("metadata", {}).get("label") else {}),
            **({"lang": p["metadata"]["lang"]} if p.get("metadata", {}).get("lang") else {}),
        }
        for p in probes[:6]
    ]
 
    return {
        "capture_id": capture_id,
        "model_id": model_id,
        "model_mode": model_mode,
        "n_probes": len(probes),
        "layers": layer_list,
        "position": position,
        "granularity": granularity,
        "encode_sae": encode_sae,
        "d_model": d_model,
        "output_dir": str(out_root.resolve()),
        "manifest_path": str((out_root / "manifest.json").resolve()),
        "metadata_path": str((out_root / "metadata.json").resolve()),
        "activation_files": activation_files,
        "sae_features": sae_features,
        "probe_samples": probe_samples,
        "checkpoint": str(checkpoint_path) if checkpoint_path else None,
        "checkpoint_name": checkpoint_name,
        "status": "done",
        **(manifest_extras or {}),
    }
 
 
def _run_capture_llm(
    model_id: str,
    probes: list[dict[str, Any]],
    output_dir: str | Path,
    *,
    layers: list[int],
    checkpoint_path: str | Path | None,
    checkpoint_name: str | None,
    position: Position,
    encode_sae: bool,
    sae_layer: int | None,
    capture_name: str | None,
    manifest_extras: dict[str, Any] | None = None,
    granularity: Granularity = "prompt",
) -> dict[str, Any]:
    from aquin.compute.device import empty_device_cache
    from aquin.compute.feature_analysis import load_sae, normalize
    from aquin.compute.model_loader import get_sae_layer
    from aquin.compute.sae_diff import load_tl_from_checkpoint
 
    # Checkpoint path builds a second full model — free any resident weights first (T4).
    if checkpoint_path is not None:
        from aquin.compute.model_runtime import unload_weights
 
        unload_weights(clear_active=False)
        empty_device_cache()
    model = load_tl_from_checkpoint(model_id, checkpoint_path)
    d_model = int(model.cfg.d_model)
    out_root = Path(output_dir)
    (out_root / "layers").mkdir(parents=True, exist_ok=True)
 
    per_layer: dict[int, list[torch.Tensor]] = {layer: [] for layer in layers}
    summary_rows: list[dict[str, Any]] = []
    token_spans: list[dict[str, Any]] = []
 
    if granularity == "token" and encode_sae:
        raise ValueError("--encode-sae is not yet supported with --granularity token")
 
    for probe in probes:
        if granularity == "token":
            acts, token_texts = _forward_llm_layers_token(model, probe["text"], layers)
        else:
            acts = _forward_llm_layers(model, probe["text"], layers, position=position)
            token_texts = []
        row_base = {"probe_id": probe["id"], "metadata": probe.get("metadata") or {}}
        for layer, vec in acts.items():
            per_layer[layer].append(vec)
            if granularity == "token":
                start = sum(t.shape[0] for t in per_layer[layer][:-1])
                end = start + int(vec.shape[0])
                summary_rows.append({**row_base, "layer": layer, "n_tokens": int(vec.shape[0]), "d_model": int(vec.shape[1])})
                token_spans.append({
                    **row_base,
                    "layer": layer,
                    "token_start": start,
                    "token_end": end,
                    "n_tokens": int(vec.shape[0]),
                    "tokens": token_texts,
                })
            else:
                summary_rows.append({**row_base, "layer": layer, "d_model": int(vec.shape[0])})
 
    sae_file = None
    sae_features = None
    if encode_sae:
        enc_layer = sae_layer if sae_layer is not None else get_sae_layer(model_id)
        if enc_layer not in layers:
            raise ValueError(
                f"--encode-sae requires layer {enc_layer} in captured layers ({layers})"
            )
        sae = load_sae(model_id, layer=enc_layer)
        encoded_rows: list[torch.Tensor] = []
        for probe in probes:
            acts = _forward_llm_layers(model, probe["text"], [enc_layer], position=position)
            resid = acts[enc_layer].to(next(sae.parameters()).device)
            encoded = sae.encode(normalize(resid.unsqueeze(0), model_id, enc_layer))[0]
            encoded_rows.append(encoded.float().cpu())
        sae_stack = torch.stack(encoded_rows)
        sae_rel = f"sae/sae_layer_{enc_layer}.pt"
        (out_root / "sae").mkdir(parents=True, exist_ok=True)
        torch.save(sae_stack, out_root / sae_rel)
        sae_file = sae_rel
        sae_features = {
            "layer": enc_layer,
            "n_features": int(sae_stack.shape[1]),
            "n_probes": int(sae_stack.shape[0]),
        }
 
    del model
    empty_device_cache()
    return _write_capture_artifacts(
        out_root=out_root,
        probes=probes,
        layer_list=layers,
        per_layer=per_layer,
        summary_rows=summary_rows,
        model_id=model_id,
        model_mode="llm",
        d_model=d_model,
        checkpoint_path=checkpoint_path,
        checkpoint_name=checkpoint_name,
        position=position,
        encode_sae=encode_sae,
        capture_name=capture_name,
        sae_file=sae_file,
        sae_features=sae_features,
        manifest_extras=manifest_extras,
        granularity=granularity,
        token_spans=token_spans,
    )
 
 
 
def run_capture_activations(
    model_id: str,
    probes: list[dict[str, Any]],
    output_dir: str | Path,
    *,
    model_mode: ModelMode | None = None,
    layers: list[int] | None = None,
    checkpoint_path: str | Path | None = None,
    checkpoint_name: str | None = None,
    position: Position = "last",
    encode_sae: bool = False,
    sae_layer: int | None = None,
    capture_name: str | None = None,
    manifest_extras: dict[str, Any] | None = None,
    granularity: Granularity = "prompt",
) -> dict[str, Any]:
    mid, _mode = resolve_capture_model_id(model_id)
    del model_mode  # LLM-only capture
 
    from aquin.compute.model_loader import get_config
 
    n_layers = int(get_config(mid)["n_layers"])
    layer_list = layers if layers is not None else list(range(n_layers))
 
    return _run_capture_llm(
        mid,
        probes,
        output_dir,
        layers=layer_list,
        checkpoint_path=checkpoint_path,
        checkpoint_name=checkpoint_name,
        position=position,
        encode_sae=encode_sae,
        sae_layer=sae_layer,
        capture_name=capture_name,
        manifest_extras=manifest_extras,
        granularity=granularity,
    )