1#!/usr/bin/env python2"""Stage C: counterfactual (paired intervention) label collection with a FIXED editor checkpoint.34For each batch: encode; reach a state s by applying r_pre ~ U{0..max_prefix} random forced actions; evaluate5the 4 logged candidate actions (policy preference if a controller is given, competitive alternative, 2 random;6singles/pairs balanced) from the SAME state and store g(s,a;y) with log-propensities. On every `audit_every`-th7batch all 24 actions are evaluated (oracle audit). The dataset stores what is needed to REGENERATE each state8deterministically (tokens, prefix actions) rather than the slot tensors. Partitions never share a state.910Usage: python -m training.collect_interventions --config configs/experiments/labels_pilot.yaml [--device mps]11"""12from __future__ import annotations1314import argparse, json, sys, time15from pathlib import Path1617import numpy as np18import torch1920sys.path.insert(0, str(Path(__file__).resolve().parents[1]))21from prisme.counterfactual import all_action_gains, paired_gains, sample_logged_actions, state_loss # noqa: E40222from prisme.flops import action_cost_proxy # noqa: E40223from training.common import Run, TokenBin, load_checkpoint, load_yaml, pick_device, seed_all, sync # noqa: E40224from training.train_editor import build_model, sample_forced_actions # noqa: E40225from data.synthetic_dataset import AnyTokenizer, SyntheticBatches, load_items # noqa: E402262728@torch.no_grad()29def collect(model, cfg, batches, n_batches, rng, device, max_prefix, audit_every, use_policy, cost):30 recs = []31 for bi in range(n_batches):32 x, y = next(batches)33 B, T = x.shape34 ctx = model.encode(x)35 r_pre = int(rng.integers(0, max_prefix + 1))36 R_total = int(rng.choice([1, 2, 4, 8]))37 R_total = max(R_total, r_pre + 1)38 prefix, _ = sample_forced_actions(cfg, B, T, [r_pre], [1.0], rng, device) if r_pre > 0 else (torch.full((B, T, cfg.r_max), cfg.stop_action, device=device, dtype=torch.long), None)39 S = ctx.S040 edited = torch.zeros(B, T, dtype=torch.bool, device=device)41 for r in range(r_pre):42 a = prefix[:, :, r]43 S, _ = model.apply_action(S, ctx, a, r, R_total - r - 1)44 edited |= a != cfg.stop_action45 step = torch.full((B, T), r_pre, device=device, dtype=torch.long)46 remaining = torch.full((B, T), R_total - r_pre - 1, device=device, dtype=torch.long)47 pref = None48 if use_policy:49 preds = model.controller(ctx.ctrl_in, S, step, remaining, ctx.entropy)50 q = model.controller.q_values(preds, cost)51 pref = q.argmax(-1)52 gen = torch.Generator(device="cpu").manual_seed(int(rng.integers(0, 2**31)))53 acts, logp = sample_logged_actions(cfg, B, T, pref.cpu() if pref is not None else None, generator=gen, device="cpu")54 acts = acts.to(device); logp = logp.to(device)55 gains = paired_gains(model, S, ctx, y, edited, acts, step, remaining)56 base = state_loss(model, S, ctx, y, edited)57 rec = {"tokens": x.cpu().to(torch.int32), "targets": y.cpu().to(torch.int32), "prefix": prefix[:, :, :max(r_pre, 1)].cpu().to(torch.int16),58 "r_pre": r_pre, "R_total": R_total, "cand": acts.cpu().to(torch.int16), "gain": gains.cpu(), "logp": logp.cpu(),59 "base_loss": base.cpu(), "valid": (y != -100).cpu()}60 if audit_every and bi % audit_every == 0:61 rec["all_gains"] = all_action_gains(model, S, ctx, y, edited, step, remaining).cpu()62 recs.append(rec)63 if bi % 10 == 0:64 v = rec["valid"]65 print(f"batch {bi}/{n_batches} r_pre={r_pre} mean gain {float(gains[v].mean()):.5f} frac>0 {float((gains[v] > 0).float().mean()):.3f}", flush=True)66 return recs676869def main(argv=None) -> int:70 ap = argparse.ArgumentParser(description=__doc__)71 ap.add_argument("--config", required=True)72 ap.add_argument("--device", default="auto")73 ap.add_argument("--seed", type=int, default=None)74 ap.add_argument("--run-root", default="experiments/runs")75 args = ap.parse_args(argv)76 exp = load_yaml(args.config)77 seed = args.seed if args.seed is not None else exp.get("seed", 1)78 seed_all(seed); device = pick_device(args.device)79 cfg, model = build_model(exp)80 model = model.to(device).eval()81 meta = load_checkpoint(exp["checkpoint"], model, strict=False)82 ck_sha = meta.get("weights_sha256")83 tr = exp["collect"]84 data = exp["data"]85 run = Run(exp.get("name", "labels"), {"experiment": exp, "model": cfg.to_dict()}, root=args.run_root, seed=seed)86 run.write_manifest(device=str(device), checkpoint_parent=exp["checkpoint"], checkpoint_sha256=ck_sha)87 cost = torch.tensor(action_cost_proxy(cfg, tr["seq_len"]), device=device)88 parts = {}89 for part, spec in exp["partitions"].items(): # e.g. train: {split: train, n_batches: 200}, calib: {split: calib, ...}90 rng = np.random.default_rng(seed + hash(part) % 1000)91 if data["kind"] == "tokens":92 tb = TokenBin(spec["bin"], tr["seq_len"]); it = tb.batches(tr["batch_size"], seed + spec.get("stream_seed", 0), device)93 else:94 tok = AnyTokenizer(data.get("tokenizer", "char"))95 sb = SyntheticBatches(load_items(Path(data["splits_dir"]) / f"{spec['split']}.jsonl"), tok, tr["seq_len"])96 it = sb.random_batches(tr["batch_size"], seed + spec.get("stream_seed", 0), device, data.get("loss_span", "answer"))97 t0 = time.time()98 recs = collect(model, cfg, it, spec["n_batches"], rng, device, tr.get("max_prefix", 3), tr.get("audit_every", 5),99 exp.get("use_policy", False), cost)100 sync(device)101 path = run.dir / f"labels_{part}.pt"102 torch.save({"records": recs, "checkpoint_sha256": ck_sha, "config": cfg.to_dict(), "part": part, "seed": seed}, path)103 g = torch.cat([r["gain"][r["valid"]] for r in recs]); n_states = sum(int(r["valid"].sum()) for r in recs)104 parts[part] = {"file": str(path), "n_states": n_states, "n_labels": int(g.numel()), "mean_gain": float(g.mean()),105 "frac_positive": float((g > 0).float().mean()), "frac_harm": float((g < 0).float().mean()),106 "seconds": round(time.time() - t0, 1), "n_audit_batches": sum("all_gains" in r for r in recs)}107 print(part, parts[part], flush=True)108 run.finish(partitions=parts)109 print("run", run.run_id)110 return 0111112113if __name__ == "__main__":114 raise SystemExit(main())115