Buckets:
| """cwm3/probes/dump_f_hat.py — F_hat(t+k), k = 1..5, of ONE E0 run on the PROBE rows (results/probes/rows_{regime}_{split}.npz), built through the trainer's own | |
| data path (the run's pack version / side stream / token bank; train/batch.prep_E with need_targets=False) for the future_vs_predicted read-out protocol | |
| (innovation 3 b2; probes/future_vs_predicted.py --space v2). Read-only w.r.t. the run. GPU when >= 8 GB VRAM is free, else CPU. | |
| -> results/probes/fhat/<run>_{regime}_{split}.npz {F_hat f16 [n,5,256], rows int64 [n], run, best, pack_version} | |
| python probes/dump_f_hat.py --run E0_pooled_P_s0_v2 [--sets manual/test,engaged/test] | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| import sys | |
| import time | |
| from pathlib import Path | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| import numpy as np # noqa: E402 | |
| import torch # noqa: E402 | |
| sys.path.insert(0, str(Path(__file__).resolve().parent)) | |
| import common as C # noqa: E402 | |
| import cwm3_dataset as ds # noqa: E402 | |
| import batch as bt # noqa: E402 | |
| import train_cwm3 as TC # noqa: E402 (model_kwargs: the run's own module layout incl. opt-in heads) | |
| from models import build_branch # noqa: E402 | |
| OUT = C.OUT / "fhat" | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--run", required=True) | |
| ap.add_argument("--sets", default="manual/test,engaged/test") | |
| ap.add_argument("--force", action="store_true") | |
| a = ap.parse_args() | |
| torch.set_num_threads(4) | |
| OUT.mkdir(parents=True, exist_ok=True) | |
| rd = C.RUNS / a.run | |
| cfgj = json.loads((rd / "config.json").read_text()) | |
| cfg, args = cfgj["cfg"], cfgj["args"] | |
| assert cfg["branch"] == "E", cfg["branch"] | |
| ds.set_pack_version(cfg.get("pack_version", args.get("pack_version", "v1"))) | |
| best = (rd / "BEST").read_text().strip() | |
| dev = C.torch_device() | |
| z = torch.load(rd / best, map_location="cpu", weights_only=False) | |
| model = build_branch("E", int(args["seed"]), **TC.model_kwargs(cfg)) | |
| model.load_state_dict(z["model"]) | |
| model = model.to(dev).eval() | |
| st = {k: np.asarray(v) for k, v in np.load(rd / "stats.npz").items()} | |
| side = cfg.get("side", "vf256") | |
| bank = ds.load_bank(need_tok=cfg.get("tokens", "region") in ("region", "spat256"), need_pooled=cfg.get("tokens") == "pooled") | |
| for rs in a.sets.split(","): | |
| regime, split = rs.split("/") | |
| outp = OUT / f"{a.run}_{regime}_{split}.npz" | |
| if outp.exists() and not a.force: | |
| print(f"{outp.name} exists — skipped", flush=True) | |
| continue | |
| t0 = time.time() | |
| rows = C.probe_rows(regime, split, verbose=lambda m: None)["rows"] | |
| raw = ds.load_regimes([regime], split, frozenset({"tokens", side}), verbose=lambda m: None) | |
| raw = ds.subset(raw, rows) | |
| n = len(rows) | |
| store = bt.GpuStore(raw, bt.KEYS["E"], dev, bank) | |
| F_hat = np.zeros((n, 5, 256), np.float16) | |
| bs = 4096 | |
| with torch.no_grad(): | |
| for i in range(0, n, bs): | |
| idx = torch.arange(i, min(i + bs, n), device=dev) | |
| b, _ = bt.prep_E(store, idx, st, cfg, need_targets=False) | |
| F_hat[i:i + len(idx)] = model(b)["F"].float().cpu().numpy().astype(np.float16) | |
| np.savez(outp, F_hat=F_hat, rows=rows, run=np.array(a.run), best=np.array(best), pack_version=np.array(ds.pack_version()), device=np.array(str(dev)), | |
| vf_ok=raw[side + "_ok"], tok_ok16=(raw["tok_idx"] >= 0).all(1)) | |
| print(f"{outp.name}: n={n:,} pack {ds.pack_version()} side {side} device {dev} {time.time() - t0:.0f}s", flush=True) | |
| del store, raw | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 3.66 kB
- Xet hash:
- 25868c11f16f864ee6bb71a9473eb313463a963aa77fea757b08540799e1c072
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.