"""Stage 3: mint (direction, target_texts) pairs — with REAL, VALIDATED probes. For each sampled cluster pair (A,B) we TRAIN a logistic-regression probe on a train split of the READ_LAYER residuals and only keep the direction if it actually works: - held-out separation AUC(A vs B) >= --min-auc (else the "concept direction" is meaningless) - cluster-A val residuals project HIGH on the unit probe and clear cluster-B by a margin (else the direction isn't something the corpus strongly exhibits → nothing to maximize). The kept unit direction is the conditioning vector; targets = centroid-closest cluster-A texts. We log the AUC / projection-margin distributions and the fraction of pairs dropped. OUT-OF-CORE: residuals come straight from the clustering stage's emb.f32 ([n_docs, d] memmap, row-aligned with assign.npy / texts.jsonl — no separate resid cache). Cluster member rows are gathered on demand through a bounded LRU (--cache-gb), texts through a byte-offset index, so RAM never scales with corpus size. Probes are GPU-batched logistic regressions: --probe-batch pairs fit in parallel per exact Newton/IRLS run (sklearn's objective: sum-BCE + 0.5/C·||w||²) — replaces sklearn-per-pair, which is far too slow at millions of pairs. python scripts/build_data.py --acts-dir data/actclusters --n-examples 8000000 \ --min-auc 0.9 --min-margin 1.0 --hard-neg-k 16 """ import argparse import json import os import numpy as np import torch from mxf.config import D_MODEL, BuildDataConfig class TextBank: """Random access into texts.jsonl by row via a byte-offset index (corpus never in RAM).""" def __init__(self, path): self.f = open(path, "rb") off = [0] for line in self.f: off.append(off[-1] + len(line)) self.off = np.asarray(off[:-1], dtype=np.int64) def __getitem__(self, i): self.f.seek(self.off[i]) return json.loads(self.f.readline())["t"] class ResidLRU: """Bounded LRU over per-cluster member-row gathers from the emb memmap.""" def __init__(self, emb, budget_gb): self.emb, self.d, self.used, self.max = emb, {}, 0, int(budget_gb * 1e9) def get(self, c, rows): v = self.d.pop(c, None) if v is None: v = np.asarray(self.emb[rows]) self.used += v.nbytes while self.used > self.max and self.d: self.used -= self.d.pop(next(iter(self.d))).nbytes self.d[c] = v # (re)insert = most recent return v @torch.no_grad() def fit_probes_gpu(feats, C, iters, device): """One L2 logistic regression per (Atr, Btr, Ava, Bva) tuple, ALL fitted in parallel on GPU. Exact Newton/IRLS in the n-dim dual: with n<=~128 samples << d=4096, w* = Xᵀα (representer), and Woodbury keeps every solve at [P, n, n] fp64 — matches sklearn's optimum of sum-BCE + 0.5/C·||w||² (+ ~free intercept, δ=1e-4) in ~10 iterations. Padded batch + masks; val AUC / medians batched too. Returns per pair (unit_w, val_auc, projA_med, projB_med) or None (degenerate).""" P, d = len(feats), feats[0][0].shape[1] nt = max(len(A) + len(B) for A, B, _, _ in feats) va = max(1, max(len(v) for _, _, v, _ in feats)); vb = max(1, max(len(v) for _, _, _, v in feats)) X = np.zeros((P, nt, d), np.float32); y = np.zeros((P, nt), np.float64); m = np.zeros((P, nt), np.float64) Xa = np.zeros((P, va, d), np.float32); ma = np.zeros((P, va), bool) Xb = np.zeros((P, vb, d), np.float32); mb = np.zeros((P, vb), bool) for i, (A, B, Av, Bv) in enumerate(feats): # pad on CPU, one H2D per tensor (not per pair) X[i, : len(A)] = A; X[i, len(A) : len(A) + len(B)] = B y[i, : len(A)] = 1.0; m[i, : len(A) + len(B)] = 1.0 Xa[i, : len(Av)] = Av; ma[i, : len(Av)] = True Xb[i, : len(Bv)] = Bv; mb[i, : len(Bv)] = True X, y, m, Xa, ma, Xb, mb = (torch.from_numpy(t).to(device) for t in (X, y, m, Xa, ma, Xb, mb)) G = (C * (X @ X.transpose(1, 2))).double() + 1e4 # X̃D⁻¹X̃ᵀ with X̃=[X,1], D=diag(1/C..,δ) eye = torch.eye(nt, device=device, dtype=torch.float64) z = torch.zeros(P, nt, device=device, dtype=torch.float64) c = torch.zeros_like(z) for _ in range(iters): # damped-free Newton: θ' = -D⁻¹X̃ᵀc p = torch.sigmoid(z) r = (p - y) * m sh = ((p * (1 - p)).clamp(min=1e-12) * m).sqrt() q = (G @ r.unsqueeze(-1)).squeeze(-1) + z # X̃D⁻¹∇: pads contribute 0 (r,c=0 there) A_ = eye + sh.unsqueeze(-1) * G * sh.unsqueeze(-2) v = torch.linalg.solve(A_, (sh * q).unsqueeze(-1)).squeeze(-1) c = r - sh * v z = -(G @ c.unsqueeze(-1)).squeeze(-1) w = -C * torch.einsum("pn,pnd->pd", c.float(), X) nrm = w.norm(dim=1) wu = w / nrm[:, None].clamp(min=1e-12) pa = torch.einsum("pvd,pd->pv", Xa, wu).masked_fill(~ma, torch.nan) pb = torch.einsum("pvd,pd->pv", Xb, wu).masked_fill(~mb, torch.nan) both = ma[:, :, None] & mb[:, None, :] gt = (pa[:, :, None] > pb[:, None, :]).double() + 0.5 * (pa[:, :, None] == pb[:, None, :]).double() auc = (gt.where(both, 0.0)).sum((1, 2)) / both.sum((1, 2)).clamp(min=1) medA, medB = pa.nanquantile(0.5, dim=1), pb.nanquantile(0.5, dim=1) nrm, WU, auc, medA, medB = (t.cpu().numpy() for t in (nrm, wu, auc, medA, medB)) return [None if (nrm[i] < 1e-8 or not len(f[2]) or not len(f[3])) else (WU[i], float(auc[i]), float(medA[i]), float(medB[i])) for i, f in enumerate(feats)] def main(): cfg = BuildDataConfig() ap = argparse.ArgumentParser() ap.add_argument("--acts-dir", default="data/actclusters", help="embed_cluster_acts out-dir (emb.f32 / assign.npy / centroids.npy / texts.jsonl)") ap.add_argument("--out-dir", default=cfg.out_dir) ap.add_argument("--n-examples", type=int, default=cfg.n_examples) ap.add_argument("--targets", type=int, default=cfg.targets_per_example) ap.add_argument("--min-auc", type=float, default=0.9, help="drop pairs the probe can't separate") ap.add_argument("--min-margin", type=float, default=1.0, help="min (projA_med - projB_med) in resid units: cluster must activate the probe") ap.add_argument("--val-frac", type=float, default=0.3) ap.add_argument("--probe-c", type=float, default=cfg.probe_c) ap.add_argument("--hard-neg-k", type=int, default=0, help="0 = random B (trivially-separable, cartoon directions). >0 = draw B from " "A's k nearest clusters by centroid → subtle, information-rich probe directions.") ap.add_argument("--members-cap", type=int, default=64, help="residual rows per cluster used to fit/val the probe") ap.add_argument("--probe-batch", type=int, default=512, help="cluster pairs fitted per GPU LR batch") ap.add_argument("--probe-iters", type=int, default=10, help="Newton/IRLS iterations per LR fit") ap.add_argument("--cache-gb", type=float, default=32.0, help="LRU budget for member-residual gathers") ap.add_argument("--device", default="cuda") a = ap.parse_args() os.makedirs(a.out_dir, exist_ok=True) rng = np.random.default_rng(cfg.seed) meta = json.load(open(f"{a.acts_dir}/meta.json")) K = meta["clusters"] emb = np.memmap(f"{a.acts_dir}/emb.f32", dtype=np.float32, mode="r", shape=(meta["n_docs"], meta["d"])) # -1 is illegal in memmap shapes assign = np.load(f"{a.acts_dir}/assign.npy") cent = np.load(f"{a.acts_dir}/centroids.npy") texts = TextBank(f"{a.acts_dir}/texts.jsonl") # group doc rows by cluster CSR-style: rows(c) = order[startx[c]:startx[c+1]] (no 1M py lists) order = np.argsort(assign, kind="stable") startx = np.zeros(K + 1, dtype=np.int64) np.cumsum(np.bincount(assign, minlength=K), out=startx[1:]) # per usable cluster (>=8 members): shuffled member rows (probe fit/val, <= members-cap), # train/val split, target rows = centroid-closest among docs NOT in the probe-fit split — # targets come only from docs the probe never saw, so direction↔text can't be memorized mrows, split, tgt_rows, mus = {}, {}, {}, {} for c in range(K): rows = order[startx[c] : startx[c + 1]] if len(rows) < 8: continue perm = rows[rng.permutation(len(rows))] mem = perm[: a.members_cap] s = max(2, int(len(mem) * (1 - a.val_frac))) cand = np.sort(perm[s:]) # all cluster docs minus probe-fit rows if not len(cand): continue dctr = np.linalg.norm(np.asarray(emb[cand]) - cent[c], axis=1) mrows[c], split[c], tgt_rows[c] = mem, s, cand[np.argsort(dctr)[:32]].tolist() if a.hard_neg_k > 0: mus[c] = np.asarray(emb[np.sort(mem[:s])]).mean(0) if len(mrows) % 50_000 == 0: print(f" prepped {len(mrows)} clusters", flush=True) pool = np.array(sorted(mrows)) nid = {int(c): i for i, c in enumerate(pool)} print(f"{len(pool)} clusters with >=8 members", flush=True) # hard-negative pairing: nearest clusters in QWEN3 layer-27 residual space (train-split mean # resid) — the probe/injection/reward live in Qwen3 residuals, so the genuinely hard negatives # are pairs the model itself struggles to linearly separate. GPU brute-force top-k. nearest = None if a.hard_neg_k > 0: mu = torch.from_numpy(np.stack([mus[c] for c in pool])).to(a.device) mn, mub = (mu * mu).sum(1), mu.to(torch.bfloat16) nearest = np.empty((len(pool), min(a.hard_neg_k, len(pool) - 1)), dtype=np.int64) cb = max(64, int(4e9 / (len(pool) * 4))) for s0 in range(0, len(pool), cb): d2 = mn[None] - 2 * (mub[s0 : s0 + cb] @ mub.T).float() d2[torch.arange(len(d2)), torch.arange(s0, s0 + len(d2))] = torch.inf # exclude self nearest[s0 : s0 + len(d2)] = d2.topk(nearest.shape[1], dim=1, largest=False).indices.cpu().numpy() nearest = pool[nearest] # pool index -> cluster id del mu, mub, mn, d2 torch.cuda.empty_cache() lru = ResidLRU(emb, a.cache_gb) vec_bank = np.memmap(f"{a.out_dir}/vecs.f32", dtype=np.float32, mode="w+", shape=(a.n_examples, D_MODEL)) recs = open(f"{a.out_dir}/records.jsonl", "w") aucs, margins = [], [] n, tried, dropped = 0, 0, 0 while n < a.n_examples: pairs, feats = [], [] while len(pairs) < a.probe_batch: # sample a GPU batch of cluster pairs A = int(rng.choice(pool)) B = int(rng.choice(nearest[nid[A]])) if a.hard_neg_k > 0 else int(rng.choice(pool)) if B == A: continue RA, RB = lru.get(A, mrows[A]), lru.get(B, mrows[B]) sA, sB = split[A], split[B] pairs.append((A, B)) feats.append((RA[:sA], RB[:sB], RA[sA:], RB[sB:])) for (A, B), pr in zip(pairs, fit_probes_gpu(feats, a.probe_c, a.probe_iters, a.device)): tried += 1 if tried % 25_000 == 0: med = (f"auc med {np.median(aucs):.3f} margin med {np.median(margins):.2f}" if aucs else "none kept yet") print(f"minted {n}/{a.n_examples} | kept {tried-dropped}/{tried} pairs " f"(drop {dropped/tried:.0%}) | {med}", flush=True) if pr is None: dropped += 1; continue wu, auc, pA, pB = pr if auc < a.min_auc or (pA - pB) < a.min_margin: # weak probe OR cluster doesn't activate it dropped += 1; continue aucs.append(auc); margins.append(pA - pB) for tr in rng.permutation(tgt_rows[A])[: a.targets]: if n >= a.n_examples: break vec_bank[n] = wu recs.write(json.dumps({"vec_idx": n, "target_text": texts[int(tr)][:1200], "cluster": A, "val_auc": round(auc, 3), "proj_margin": round(pA - pB, 2)}) + "\n") n += 1 if n >= a.n_examples: break recs.close(); vec_bank.flush() stats = {"n_examples": n, "pairs_tried": tried, "pairs_dropped": dropped, "drop_frac": dropped / max(tried, 1), "auc_median": float(np.median(aucs)), "auc_p10": float(np.percentile(aucs, 10)), "margin_median": float(np.median(margins)), "min_auc": a.min_auc, "min_margin": a.min_margin} json.dump(stats, open(f"{a.out_dir}/build_stats.json", "w"), indent=1) print(f"BUILD_DATA_DONE {stats}", flush=True) if __name__ == "__main__": main()