Download reader/read_shuffle_cost.py from RonBuildsAI/sleep-form-rescue-count: direct link, hf CLI and curl.
- Browser
- Download file 11.8 kB
-
https://huggingface.co/RonBuildsAI/sleep-form-rescue-count/resolve/main/reader/read_shuffle_cost.py
- Command line
-
hf download hf://RonBuildsAI/sleep-form-rescue-count/reader/read_shuffle_cost.py
-
curl -L -o read_shuffle_cost.py https://huggingface.co/RonBuildsAI/sleep-form-rescue-count/resolve/main/reader/read_shuffle_cost.py
11.8 kB
| # SPDX-License-Identifier: BSD-3-Clause-Clear | |
| # Copyright (c) 2026 Ronald Lee Hampton. All rights reserved. | |
| # No express or implied licenses to any party's patent rights are granted by this license. | |
| # | |
| # read_shuffle_cost.py: re-measures the shuffle cost of every shipped brain and re-counts the switchboards. | |
| # | |
| # Usage (from the package folder): | |
| # python -I reader/read_shuffle_cost.py --brains brains --data data --expect results/per_brain_record.csv | |
| # Needs: Python 3, numpy, torch, safetensors. CPU only, one thread, about a minute for the 80 brains. No network. | |
| # Side effects: none; it prints only (redirect the output to keep it). | |
| # | |
| # THE MEASURE. Each brain is run on the 256 held-out windows (66 tokens each: two context words, then 64 targets), 16,384 positions in | |
| # one batch. The reasoning core's quantized output (qk, one 32-number vector per position) is then shuffled across positions with | |
| # numpy default_rng(0).permutation(16384), everything else held fixed, and the network is run again. The SHUFFLE COST is the mean, over | |
| # the positions whose target is not one of the 40 excluded ids, of (cross-entropy shuffled minus cross-entropy as trained), in nats. | |
| # THE LINE (fixed before the run): a core is a SWITCHBOARD if its shuffle cost is below 0.025 nats, else a CONTEXT WORKER. | |
| # A RESCUE is a seed whose plain (ST-0) core is a switchboard and whose sleep-form (S-SLEEP) core is a context worker; a HARM is the | |
| # reverse. THE BAR (fixed before the run): S-SLEEP ends with fewer switchboards than ST-0, AND at least 5 seeds are discordant, AND the | |
| # exact one-sided sign test on the discordant seeds gives p at or below 0.05. | |
| # | |
| # This file holds only the forward pass the measure needs (no training code). Tensor names keep the run's own prefixes: | |
| # tiny.* the routed base network: two word embeddings summed, a router to one of 48 experts (top 1, its score nudged by the pull | |
| # 0.4 x (mean size minus the expert's size), sizes in "y"), each expert a two-layer block, a layer norm and a shared head; | |
| # rc.* the reasoning core: written from the integer part's hidden state, one gated pass, quantized to whole numbers in -127..127; | |
| # cc.* the small integer part with ternary weights: reads the two context words' ids as 12 bits each, and writes a gain on the | |
| # base network's expert output from its own hidden state and the core's quantized output (zero gain at the start). | |
| # cc.Wv and cc.bv, a two-output head the measure does not use, were removed from the shipped files (results/head_removed_record.csv). | |
| # Every operation below is written in the same order as the run's own code so that the numbers can match it bit for bit on the same | |
| # software and hardware; on other hardware or torch versions the last digits may differ (the --expect check reports by how much). | |
| import sys, os, csv, glob, math, hashlib, argparse | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| V, D, T, CTX, NBITS = 4096, 32, 64, 2, 12 | |
| S_PULL, RHO = 0.4, 0.5 | |
| LINE = 0.025 | |
| FORMS = ("ST-0", "S-SLEEP") | |
| P_BAR, MIN_DISC = 0.05, 5 | |
| BITW = 2 ** torch.arange(NBITS) | |
| def sha256_file(p): | |
| h = hashlib.sha256() | |
| with open(p, "rb") as f: | |
| for blk in iter(lambda: f.read(1 << 20), b""): | |
| h.update(blk) | |
| return h.hexdigest() | |
| def sign_p(r, n): | |
| """exact one-sided p of at least r rescues among n discordant seeds, each a fair coin.""" | |
| return float(sum(math.comb(n, k) for k in range(r, n + 1)) / 2 ** n) if n else 1.0 | |
| def tern(w): | |
| return (w > 0.5).float() - (w < -0.5).float() | |
| def ste_tern(w): # the training-time expression, kept so the floating-point result is the same | |
| return w + (tern(w) - w) | |
| def ste_round(x): | |
| return x + (torch.round(x) - x) | |
| def ste_floor(x): | |
| return x + (torch.floor(x) - x) | |
| def rms(x, w): | |
| return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * w | |
| def msg(xs): | |
| """the two context words' token ids as 12 bits each, +1/-1, newest word first.""" | |
| parts = [((x[:, None] // BITW) % 2).float() * 2 - 1 for x in xs[:2]] | |
| return torch.cat(parts, dim=-1) | |
| def hidden(P, xs): | |
| """the integer part's hidden state: ternary weights, whole-number biases, clamps; every value a whole number.""" | |
| m = msg(xs) | |
| a1 = m @ ste_tern(P["cc.W1"]).t() + ste_round(P["cc.b1"] * 1.0) | |
| h1 = torch.clamp(a1, 0.0, 31.0) | |
| a2 = h1 @ ste_tern(P["cc.W2"]).t() + ste_round(P["cc.b2"] * 16.0) | |
| return torch.clamp(ste_floor(a2 / 4.0), 0.0, 127.0) | |
| def pull_bias(P): | |
| y = P["y"].numpy() | |
| return torch.tensor(S_PULL * (y.mean() - y), dtype=torch.float32) | |
| def forward(P, xs, bias, h2, perm=None): | |
| """logits at every position; with perm, the reasoning core's quantized output is shuffled across positions first.""" | |
| # the routed base network | |
| x = F.embedding(xs[0], P["tiny.emb.0.weight"]) | |
| x = x + F.embedding(xs[1], P["tiny.emb.1.weight"]) | |
| r = F.linear(x, P["tiny.router.weight"], P["tiny.router.bias"]) | |
| score = r + bias | |
| p = F.softmax(r, dim=-1) | |
| sel = torch.argmax(score, dim=-1) | |
| g = p.gather(1, sel[:, None]).squeeze(1) | |
| h = F.gelu(torch.bmm(x[:, None, :], P["tiny.w1"][sel])).squeeze(1) | |
| y = torch.bmm(h[:, None, :], P["tiny.w2"][sel]).squeeze(1) | |
| contrib = g[:, None] * y | |
| # the reasoning core, written from the integer part's hidden state | |
| K0 = (1 - RHO) * torch.tanh(rms(h2 / 127.0, P["rc.nwc.weight"]) @ P["rc.wkc"]) | |
| z = rms(K0, P["rc.ns.weight"]) | |
| u = F.silu(z @ P["rc.w1"]) * (z @ P["rc.w3"]) | |
| K2 = RHO * K0 + (1 - RHO) * torch.tanh(u @ P["rc.wh"]) | |
| qk = torch.clamp(ste_round(K2 * 127.0), -127.0, 127.0) | |
| if perm is not None: | |
| qk = qk[perm] # THE SHUFFLE: the core's output moved across positions, nothing else changed | |
| # the integer part's write: a gain on the expert output | |
| c = h2 @ ste_tern(P["cc.Wcg"]).t() | |
| c = c + qk @ ste_tern(P["cc.Wck"]).t() | |
| gain = 2.0 * torch.sigmoid(c / 256.0) - 1.0 | |
| o = x + contrib * (1.0 + gain) | |
| o = F.layer_norm(o, (D,), P["tiny.norm.weight"], P["tiny.norm.bias"], 1e-5) | |
| return F.linear(o, P["tiny.head.weight"], P["tiny.head.bias"]) | |
| def heldout(data_dir): | |
| win = np.load(os.path.join(data_dir, "heldout_windows_uint16.npy"), allow_pickle=False).astype(np.int64) | |
| if win.shape != (256, T + CTX): | |
| raise SystemExit("REFUSED: held-out windows have shape %s" % (win.shape,)) | |
| xs = [torch.from_numpy(np.ascontiguousarray(win[:, CTX - 1 - k:T + CTX - 1 - k])).reshape(-1) for k in range(CTX)] | |
| tgt = torch.from_numpy(np.ascontiguousarray(win[:, CTX:])).reshape(-1) | |
| ids = [int(l) for l in open(os.path.join(data_dir, "excluded_target_ids.txt"), encoding="utf-8") if l.strip() and not l.startswith("#")] | |
| excl = torch.zeros(V, dtype=torch.bool) | |
| excl[torch.tensor(ids)] = True | |
| allowed = (~excl[tgt]).numpy() | |
| return win, xs, tgt, allowed, ids | |
| def main(a): | |
| from safetensors.torch import load_file | |
| torch.set_num_threads(1) | |
| win, xs, tgt, allowed, ids = heldout(a.data) | |
| perm = torch.from_numpy(np.random.default_rng(0).permutation(len(tgt))) | |
| print("held-out windows %d, positions %d, scored positions %d, excluded ids %d" % (win.shape[0], len(tgt), int(allowed.sum()), len(ids))) | |
| print("torch %s, numpy %s, threads %d" % (torch.__version__, np.__version__, torch.get_num_threads())) | |
| paths = sorted(glob.glob(os.path.join(a.brains, "*.safetensors"))) | |
| per = {} | |
| for path in paths: | |
| name = os.path.basename(path)[:-len(".safetensors")] | |
| form, s = name.split("_")[0], int(name.split("_")[1][1:]) | |
| P = load_file(path) | |
| with torch.no_grad(): | |
| bias = pull_bias(P) | |
| h2 = hidden(P, xs) | |
| lg_on = forward(P, xs, bias, h2) | |
| lg_sh = forward(P, xs, bias, h2, perm) | |
| ce_on = F.cross_entropy(lg_on, tgt, reduction="none").double().numpy() | |
| ce_sh = F.cross_entropy(lg_sh, tgt, reduction="none").double().numpy() | |
| cost = float((ce_sh - ce_on)[allowed].mean()) | |
| per[(form, s)] = {"name": name, "cost": cost, "label": "SWITCHBOARD" if cost < LINE else "CONTEXT WORKER", | |
| "ce_on": float(ce_on[allowed].mean()), | |
| "logits_sha256": hashlib.sha256(lg_on.contiguous().numpy().tobytes()).hexdigest(), | |
| "file_sha256": sha256_file(path)} | |
| print("%-8s s%d shuffle cost %.6f %-14s (held-out loss as trained %.6f)" % (form, s, cost, per[(form, s)]["label"], per[(form, s)]["ce_on"])) | |
| seeds = sorted({s for (_f, s) in per}) | |
| if not seeds or not all((f, s) in per for f in FORMS for s in seeds): | |
| print("REFUSED: no brains, or a seed is missing one of its two brains") | |
| return 2 | |
| sb = {f: [s for s in seeds if per[(f, s)]["label"] == "SWITCHBOARD"] for f in FORMS} | |
| res = [s for s in seeds if per[("ST-0", s)]["label"] == "SWITCHBOARD" and per[("S-SLEEP", s)]["label"] == "CONTEXT WORKER"] | |
| harm = [s for s in seeds if per[("ST-0", s)]["label"] == "CONTEXT WORKER" and per[("S-SLEEP", s)]["label"] == "SWITCHBOARD"] | |
| n = len(res) + len(harm) | |
| p = sign_p(len(res), n) | |
| passed = len(sb["S-SLEEP"]) < len(sb["ST-0"]) and n >= MIN_DISC and p <= P_BAR | |
| print("") | |
| print("seeds %d (%d to %d); THE LINE: switchboard below %.3f nats of shuffle cost" % (len(seeds), seeds[0], seeds[-1], LINE)) | |
| for f in FORMS: | |
| print(" %-8s switchboards %d of %d: %s" % (f, len(sb[f]), len(seeds), sb[f])) | |
| print(" rescues %d %s; harms %d %s; discordant %d; exact one-sided sign test p %.5f" % (len(res), res, len(harm), harm, n, p)) | |
| print("THE BAR (S-SLEEP fewer, at least %d discordant, p <= %.2f): %s" % (MIN_DISC, P_BAR, "CONFIRMED" if passed else "NOT CONFIRMED")) | |
| near = sorted((abs(v["cost"] - LINE), "%s s%d" % k, round(v["cost"], 6)) for k, v in per.items())[:3] | |
| print(" nearest the line: %s" % [(k, c) for _d, k, c in near]) | |
| rc = 0 | |
| if a.expect: | |
| rows = {(r["form"], int(r["seed"])): r for r in csv.DictReader(open(a.expect, encoding="utf-8", newline=""))} | |
| keys = [k for k in per if k in rows] | |
| lab_eq = sum(1 for k in keys if rows[k]["label"] == per[k]["label"]) | |
| file_eq = sum(1 for k in keys if rows[k]["safetensors_sha256"] == per[k]["file_sha256"]) | |
| cost_eq = sum(1 for k in keys if float(rows[k]["shuffle_cost"]) == per[k]["cost"]) | |
| ce_eq = sum(1 for k in keys if float(rows[k]["ce_on_allowed_mean"]) == per[k]["ce_on"]) | |
| lg_eq = sum(1 for k in keys if rows[k]["logits_sha256_on"] == per[k]["logits_sha256"]) | |
| dmax = max(abs(float(rows[k]["shuffle_cost"]) - per[k]["cost"]) for k in keys) if keys else float("nan") | |
| cmax = max(abs(float(rows[k]["ce_on_allowed_mean"]) - per[k]["ce_on"]) for k in keys) if keys else float("nan") | |
| print("") | |
| print("AGAINST %s (%d rows):" % (os.path.basename(a.expect), len(rows))) | |
| print(" brain files equal to the record's sha256: %d of %d" % (file_eq, len(per))) | |
| print(" labels equal: %d of %d" % (lab_eq, len(per))) | |
| print(" shuffle costs equal bit for bit: %d of %d; largest absolute difference %.3g nats" % (cost_eq, len(per), dmax)) | |
| print(" held-out loss as trained equal bit for bit: %d of %d; largest absolute difference %.3g nats" % (ce_eq, len(per), cmax)) | |
| print(" logits as trained equal bit for bit (sha256): %d of %d" % (lg_eq, len(per))) | |
| if lab_eq != len(per) or file_eq != len(per) or len(rows) != len(per): | |
| rc = 1 | |
| return rc | |
| if __name__ == "__main__": | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--brains", default="brains") | |
| ap.add_argument("--data", default="data") | |
| ap.add_argument("--expect", default="") | |
| sys.exit(main(ap.parse_args())) | |