scripts/build_sae_data.py

"""Build an SAE-feature post-train dataset in the exact vec-bank format pretrain.py consumes.

Per feature: direction = unit(W_enc[:,f]) (encoder column, probe side — the same side eval_sae.py
rewards); targets = the feature's top --targets corpus windows (by per-window peak act) decoded to
text. Features are split (seeded, --eval-frac) into train/eval and ONLY train features enter
vecs.f32/records.jsonl — the eval half is reserved for cross-uplift (eval_sae.py --split
<out>/split.json). build_stats.json carries n_examples so scripts/rl.py --data-dir <out> also
works on this bank unchanged.

    python scripts/build_sae_data.py --out-dir data/sae --targets 3
"""
import argparse
import json
import os
import random

import numpy as np
import torch
from transformers import AutoTokenizer

from mxf.config import D_MODEL, MODEL
from mxf.sae import load_max_acts, load_sae


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--sae-path", default=None, help="ae.pt path; default: hf_hub_download")
    ap.add_argument("--maxacts-path", default=None, help="max-acts .pt path; default: hf_hub_download")
    ap.add_argument("--out-dir", default="data/sae")
    ap.add_argument("--n-features", type=int, default=0, help="0 = ALL alive features; else seeded sample")
    ap.add_argument("--targets", type=int, default=3, help="top corpus windows per feature")
    ap.add_argument("--min-act", type=float, default=0.0, help="drop features with corpus peak <= this")
    ap.add_argument("--eval-frac", type=float, default=0.5)
    ap.add_argument("--seed", type=int, default=0)
    a = ap.parse_args()
    os.makedirs(a.out_dir, exist_ok=True)
    rng = random.Random(a.seed)

    tok = AutoTokenizer.from_pretrained(MODEL)
    sae = load_sae(a.sae_path)                                # cpu fp32
    data = load_max_acts(a.maxacts_path)
    tokens, acts = data["max_tokens"], data["max_acts"]       # [F, N, L]
    assert acts.shape[0] == sae.d_sae, f"max-acts F={acts.shape[0]} != SAE F={sae.d_sae}"
    dataset_max = acts.amax(dim=(1, 2))                       # [F] corpus peak per feature

    alive = (dataset_max > a.min_act).nonzero(as_tuple=True)[0].tolist()
    if a.n_features and a.n_features < len(alive):
        alive = rng.sample(alive, a.n_features)
    rng.shuffle(alive)
    n_eval = int(len(alive) * a.eval_frac)
    ev, train = sorted(alive[:n_eval]), sorted(alive[n_eval:])
    json.dump({"seed": a.seed, "eval_frac": a.eval_frac, "min_act": a.min_act,
               "train": train, "eval": ev}, open(f"{a.out_dir}/split.json", "w"))
    print(f"{len(alive)} alive features (min_act {a.min_act}) -> {len(train)} train / {len(ev)} eval",
          flush=True)

    dirs = torch.nn.functional.normalize(sae.W_enc, dim=0)    # [d, F] unit encoder columns
    rows, skipped = [], 0
    for j, f in enumerate(train):
        peak = acts[f].amax(dim=-1)                           # [N] per-window peaks
        for w in peak.argsort(descending=True)[: a.targets].tolist():
            if peak[w] <= 0:
                break                                         # desc order: rest are dead too
            text = tok.decode(tokens[f, w].tolist(), skip_special_tokens=True).strip()
            if len(text) < 3:
                skipped += 1
                continue
            rows.append((f, text, peak[w].item()))
        if (j + 1) % 5000 == 0:
            print(f"{j + 1}/{len(train)} features, {len(rows)} records", flush=True)

    assert rows, "no records minted — check --min-act / max-acts file"
    vecs = np.memmap(f"{a.out_dir}/vecs.f32", dtype=np.float32, mode="w+",
                     shape=(len(rows), D_MODEL))
    with open(f"{a.out_dir}/records.jsonl", "w") as recs:
        for n, (f, text, act) in enumerate(rows):             # one row per (feature, target)
            vecs[n] = dirs[:, f].numpy()
            recs.write(json.dumps({"vec_idx": n, "target_text": text, "feature": f,
                                   "act": round(act, 3)}) + "\n")
    vecs.flush()
    stats = {"n_examples": len(rows), "n_train_features": len(train), "n_eval_features": len(ev),
             "targets_per_feature": a.targets, "skipped_short": skipped, "seed": a.seed}
    json.dump(stats, open(f"{a.out_dir}/build_stats.json", "w"), indent=1)
    print(f"BUILD_SAE_DATA_DONE {stats}", flush=True)


if __name__ == "__main__":
    main()