activation_store.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Resolve saved activation directories for SAE training (chunks, captures, prior runs)."""
 
from __future__ import annotations
 
import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
 
import torch
 
CHUNK_SIZE = 50_000
MANIFEST_NAME = "manifest.json"
 
 
def find_chunk_paths(acts_dir: Path) -> list[Path]:
    return sorted(acts_dir.glob("chunk_*.pt"))
 
 
def read_manifest(acts_dir: Path) -> dict[str, Any] | None:
    path = acts_dir / MANIFEST_NAME
    if not path.is_file():
        return None
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
        return data if isinstance(data, dict) else None
    except Exception:
        return None
 
 
def write_collect_manifest(
    acts_dir: Path,
    *,
    model_id: str,
    layer: int,
    model_mode: str,
    d_model: int,
    source: str,
    n_vectors: int,
    checkpoint_path: str | None = None,
    corpus_path: str | None = None,
) -> Path:
    acts_dir.mkdir(parents=True, exist_ok=True)
    chunks = [p.name for p in find_chunk_paths(acts_dir)]
    manifest = {
        "schema_version": 1,
        "source": source,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "model_id": model_id,
        "layer": layer,
        "model_mode": model_mode,
        "d_model": d_model,
        "n_vectors": n_vectors,
        "checkpoint": checkpoint_path,
        "corpus": corpus_path,
        "files": {
            "chunks": chunks,
            "norm": "norm.pt" if (acts_dir / "norm.pt").is_file() else None,
        },
    }
    out = acts_dir / MANIFEST_NAME
    out.write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    return out
 
 
def _split_save_chunks(stacked: torch.Tensor, out_dir: Path) -> int:
    out_dir.mkdir(parents=True, exist_ok=True)
    n = int(stacked.shape[0])
    chunk_i = 0
    for idx in range(0, n, CHUNK_SIZE):
        torch.save(stacked[idx : idx + CHUNK_SIZE], out_dir / f"chunk_{chunk_i}.pt")
        chunk_i += 1
    return n
 
 
def import_capture_to_chunks(capture_dir: Path, layer: int, out_dir: Path) -> int:
    manifest = read_manifest(capture_dir)
    if not manifest:
        raise ValueError(f"No manifest.json in {capture_dir}")
 
    rel = (manifest.get("files") or {}).get("activations", {}).get(str(layer))
    if not rel:
        layers = list((manifest.get("files") or {}).get("activations", {}).keys())
        hint = f" Available layers: {', '.join(layers)}" if layers else ""
        raise ValueError(f"Capture has no activations for layer {layer}.{hint}")
 
    tensor_path = capture_dir / rel
    if not tensor_path.is_file():
        raise FileNotFoundError(f"Capture activation file missing: {tensor_path}")
 
    stacked = torch.load(tensor_path, map_location="cpu", weights_only=False)
    if stacked.dim() != 2:
        raise ValueError(f"Expected 2D activations [N, d_model], got shape {tuple(stacked.shape)}")
 
    n = _split_save_chunks(stacked, out_dir)
    write_collect_manifest(
        out_dir,
        model_id=str(manifest.get("model_id") or ""),
        layer=layer,
        model_mode=str(manifest.get("model_mode") or "llm"),
        d_model=int(stacked.shape[1]),
        source="capture",
        n_vectors=n,
        checkpoint_path=manifest.get("checkpoint"),
    )
    return n
 
 
def import_capture_to_chunks_balanced(
    capture_dir: Path,
    layer: int,
    out_dir: Path,
    *,
    group: str | None = None,
) -> int:
    manifest = read_manifest(capture_dir)
    if not manifest:
        raise ValueError(f"No manifest.json in {capture_dir}")
 
    rel = (manifest.get("files") or {}).get("activations", {}).get(str(layer))
    if not rel:
        layers = list((manifest.get("files") or {}).get("activations", {}).keys())
        hint = f" Available layers: {', '.join(layers)}" if layers else ""
        raise ValueError(f"Capture has no activations for layer {layer}.{hint}")
 
    tensor_path = capture_dir / rel
    if not tensor_path.is_file():
        raise FileNotFoundError(f"Capture activation file missing: {tensor_path}")
 
    stacked = torch.load(tensor_path, map_location="cpu", weights_only=False)
    if stacked.dim() != 2:
        raise ValueError(f"Expected 2D activations [N, d_model], got shape {tuple(stacked.shape)}")
 
    probes = manifest.get("probes") if isinstance(manifest.get("probes"), list) else []
    if not probes:
        return import_capture_to_chunks(capture_dir, layer, out_dir)
 
    from aquin.compute.activation_capture import balance_probes
 
    balanced, meta = balance_probes([p for p in probes if isinstance(p, dict)], group=group)
    if not meta.get("balanced"):
        return import_capture_to_chunks(capture_dir, layer, out_dir)
 
    keep_ids = {str(p.get("id")) for p in balanced if p.get("id") is not None}
    keep_indices = [
        i for i, p in enumerate(probes)
        if isinstance(p, dict) and str(p.get("id")) in keep_ids
    ]
    if not keep_indices:
        return import_capture_to_chunks(capture_dir, layer, out_dir)
 
    subset = stacked[torch.tensor(keep_indices, dtype=torch.long)]
    n = _split_save_chunks(subset, out_dir)
    write_collect_manifest(
        out_dir,
        model_id=str(manifest.get("model_id") or ""),
        layer=layer,
        model_mode=str(manifest.get("model_mode") or "llm"),
        d_model=int(subset.shape[1]),
        source="capture-balanced",
        n_vectors=n,
        checkpoint_path=manifest.get("checkpoint"),
    )
    return n
 
 
def resolve_training_acts_dir(
    path: str | Path,
    layer: int,
    *,
    balance: bool = False,
    balance_group: str | None = None,
) -> Path:
    """
    Return a directory containing chunk_*.pt for training.
 
    Accepts:
    - a collect dir (chunk_*.pt + optional manifest.json)
    - a prior train run dir (…/my-run/_acts_layer{N})
    - a capture-activations output dir (manifest.json + layers/layer_{N}.pt)
    """
    root = Path(path).expanduser().resolve()
    if not root.exists():
        raise FileNotFoundError(f"Activations path not found: {root}")
 
    if find_chunk_paths(root):
        return root
 
    sub = root / f"_acts_layer{layer}"
    if find_chunk_paths(sub):
        return sub
 
    if (root / MANIFEST_NAME).is_file():
        manifest = read_manifest(root)
        if manifest and manifest.get("files", {}).get("activations"):
            suffix = f"_balanced_{balance_group}" if balance and balance_group else ("_balanced" if balance else "")
            cache = root / f"_train_cache_layer{layer}{suffix}"
            if find_chunk_paths(cache):
                return cache
            if balance:
                n = import_capture_to_chunks_balanced(root, layer, cache, group=balance_group)
                print(f"[sae-train] imported {n:,} balanced vectors from capture → {cache}", flush=True)
            else:
                n = import_capture_to_chunks(root, layer, cache)
                print(f"[sae-train] imported {n:,} vectors from capture → {cache}", flush=True)
            return cache
 
    raise FileNotFoundError(
        f"No training activations for layer {layer} under {root}.\n"
        f"Expected chunk_*.pt, _acts_layer{layer}/, or capture manifest with layers/layer_{layer}.pt"
    )
 
 
def count_vectors(acts_dir: Path) -> int:
    total = 0
    for p in find_chunk_paths(acts_dir):
        t = torch.load(p, map_location="cpu", weights_only=False)
        total += int(t.shape[0])
    return total
 
 
def compute_norm(acts_dir: Path, d_model: int) -> tuple[torch.Tensor, torch.Tensor]:
    norm_path = acts_dir / "norm.pt"
    if norm_path.exists():
        d = torch.load(norm_path, map_location="cpu", weights_only=False)
        return d["mean"], d["std"]
 
    chunks = find_chunk_paths(acts_dir)
    if not chunks:
        raise FileNotFoundError(f"No chunk_*.pt in {acts_dir}")
 
    running_mean = torch.zeros(d_model)
    n_total = 0
    for p in chunks:
        c = torch.load(p, map_location="cpu", weights_only=False)
        running_mean += c.sum(0)
        n_total += c.shape[0]
    mean = running_mean / max(n_total, 1)
 
    running_var = torch.zeros(d_model)
    for p in chunks:
        c = torch.load(p, map_location="cpu", weights_only=False)
        running_var += ((c - mean) ** 2).sum(0)
    std = (running_var / max(n_total, 1)).sqrt().clamp(min=1e-6)
    torch.save({"mean": mean, "std": std}, norm_path)
    return mean, std
 
 
def validate_acts_manifest(acts_dir: Path, *, model_id: str, layer: int) -> None:
    manifest = read_manifest(acts_dir)
    if not manifest:
        return
    man_layer = manifest.get("layer")
    if man_layer is not None and int(man_layer) != int(layer):
        raise ValueError(
            f"Activations manifest is for layer {man_layer}, but --layer {layer} was passed."
        )
    man_model = manifest.get("model_id")
    if not man_model:
        return
    try:
        from aquin.compute.model_loader import resolve_model_id
 
        want = resolve_model_id(model_id)
        got = resolve_model_id(str(man_model))
        if want != got:
            print(
                f"[sae-train] warning: activations model_id={got!r} differs from session model {want!r}",
                flush=True,
            )
    except ValueError:
        pass