SPIRA v2: entropy-patch tokenization + hierarchical dopamine-gated spiking cortex + hippocampal one-shot memory with replay consolidation
032a961 verified | #!/usr/bin/env python3 | |
| """SPIRA evaluation harness v2 — proofs + hierarchy ablation + replay. | |
| Produces artifacts/metrics.json with: | |
| 0. vocab coverage : readout-vocabulary coverage curve (Zipf knee) | |
| 1. compression : EPT vs byte-level BPE (matched vocab) vs gzip-9 | |
| 2. prediction : unigram / bigram / SNN (fast-only) / hierarchical | |
| cortex / AM / fused, on held-out test + a slice of the | |
| TRAINED region (memorization reference) | |
| 3. one-shot : single-pass teaching of an unseen book -> exact recall | |
| 3b. replay : hippocampal replay -> cortical consolidation curve, | |
| AM integrity after replay, and interference on test | |
| 4. correction : recall vs number of corrupted context lanes | |
| 5. generation : samples, speed, synaptic-op accounting | |
| Sections are resumable: completed sections are skipped when metrics.json | |
| already contains them, so the harness can be re-invoked across process | |
| boundaries (each call loads the checkpoint fresh; the replay experiment is | |
| self-contained in one section). | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| import time | |
| import numpy as np | |
| from .baselines import (BigramBaseline, UnigramBaseline, bpe_encode, | |
| dict_bits, gzip_bits_per_byte, stream_entropy_bits) | |
| from .ept import EPT | |
| from .hashing import fnv1a64 | |
| from .model import SPIRA | |
| def distinct_n(text: str, n: int) -> float: | |
| toks = text.split() | |
| if len(toks) < n: | |
| return 0.0 | |
| grams = [" ".join(toks[i : i + n]) for i in range(len(toks) - n + 1)] | |
| return len(set(grams)) / len(grams) | |
| def dump(metrics: dict, path: str) -> None: | |
| with open(path, "w") as f: | |
| json.dump(metrics, f, indent=2, ensure_ascii=False) | |
| def main() -> None: | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--model", required=True) | |
| ap.add_argument("--train", required=True) | |
| ap.add_argument("--test", required=True) | |
| ap.add_argument("--oneshot", required=True) | |
| ap.add_argument("--out", required=True) | |
| ap.add_argument("--eval-positions", type=int, default=60_000) | |
| ap.add_argument("--replay-tokens", type=int, default=100_000) | |
| ap.add_argument("--replay-stages", type=int, default=4) | |
| args = ap.parse_args() | |
| train = open(args.train, "rb").read() | |
| test = open(args.test, "rb").read() | |
| oneshot = open(args.oneshot, "rb").read() | |
| model = SPIRA.load(args.model) | |
| metrics: dict = {} | |
| if os.path.exists(args.out): | |
| try: | |
| metrics = json.load(open(args.out)) | |
| print(f"(resuming: {len(metrics)} section(s) already present)") | |
| except Exception: | |
| metrics = {} | |
| # restore the train id stream for baselines / corruption / coverage | |
| t0 = time.perf_counter() | |
| train_patches = model.ept.segment(train) | |
| train_ids = np.fromiter((fnv1a64(p) for p in train_patches), dtype=np.uint64, | |
| count=len(train_patches)) | |
| model._ids_train = train_ids | |
| print(f"prep: segmented train ({len(train_ids)} patches) in {time.perf_counter()-t0:.1f}s") | |
| if model.core.hier: | |
| cal = model.calibrate_fusion(train[:60_000]) | |
| print(f"prep: fusion calibration sigma_fast={cal['sigma_fast']:.4f} " | |
| f"sigma_slow={cal['sigma_slow']:.4f} scale={cal['scale']:.6f} " | |
| f"beta={model.core.beta:.4f}") | |
| if "fusion_calibration" not in metrics: | |
| metrics["fusion_calibration"] = {**cal, "beta": float(model.core.beta)} | |
| dump(metrics, args.out) | |
| if "model_meta" not in metrics: | |
| metrics.update({"model_meta": model.meta, "config": vars(model.cfg)}) | |
| dump(metrics, args.out) | |
| # ============================================ 0. vocab coverage curve ===== | |
| if "vocab_coverage_curve" not in metrics: | |
| print("== [0/7] readout-vocabulary coverage curve ==") | |
| uniq, cnt = np.unique(train_ids, return_counts=True) | |
| order = np.lexsort((uniq, -cnt)) # count desc, id asc (build order) | |
| sorted_ids = uniq[order] | |
| test_ids_local = np.fromiter( | |
| (fnv1a64(p) for p in model.ept.segment(test)), dtype=np.uint64, | |
| count=len(model.ept.segment(test))) | |
| one_ids_local = np.fromiter( | |
| (fnv1a64(p) for p in model.ept.segment(oneshot[:60_000])), dtype=np.uint64, | |
| count=len(model.ept.segment(oneshot[:60_000]))) | |
| curve = [] | |
| for V in (4096, 6144, 8192, 12288, 16384, 24576, 32768): | |
| top = sorted_ids[:V] | |
| cov_t = float(np.isin(test_ids_local, top).mean()) | |
| cov_o = float(np.isin(one_ids_local, top).mean()) | |
| curve.append({"vocab": int(V), "cov_test": round(cov_t, 4), | |
| "cov_oneshot": round(cov_o, 4)}) | |
| print(f" V={V:>6}: held-out test coverage={cov_t:.4f} " | |
| f"oneshot coverage={cov_o:.4f}") | |
| metrics["vocab_coverage_curve"] = curve | |
| metrics["vocab_coverage"] = { | |
| "distinct_train_patches": int(len(uniq)), | |
| "note": "fraction of held-out target patches inside the top-V train vocabulary"} | |
| dump(metrics, args.out) | |
| # ================================================== 1. compression ===== | |
| if "compression" not in metrics: | |
| print("== [1/7] compression: EPT vs BPE vs gzip ==") | |
| comp = {} | |
| gz_bpb, gz_s = gzip_bits_per_byte(test) | |
| comp["gzip"] = {"bits_per_byte": gz_bpb, "seconds": gz_s} | |
| for vs in (4096, 16384): | |
| r = bpe_encode(train[:4_000_000], test, vs) | |
| h0 = stream_entropy_bits(r["ids"], order=0) | |
| h1 = stream_entropy_bits(r["ids"], order=1) | |
| db = dict_bits(r["vocab_size"], r["vocab_bytes"]) | |
| comp[f"bpe_{vs}"] = { | |
| "vocab_size": r["vocab_size"], | |
| "n_tokens": int(len(r["ids"])), | |
| "avg_token_bytes": len(test) / len(r["ids"]), | |
| "stream_bits_unigram": h0, | |
| "stream_bits_bigram": h1, | |
| "dict_bits": db, | |
| "bits_per_byte_unigram": (h0 * len(r["ids"]) + db) / len(test), | |
| "bits_per_byte_bigram": (h1 * len(r["ids"]) + db) / len(test), | |
| "train_s": r["train_s"], | |
| } | |
| t0 = time.perf_counter() | |
| ept_patches = model.ept.segment(test) | |
| ept_seg_s = time.perf_counter() - t0 | |
| ept_ids = np.fromiter((fnv1a64(p) for p in ept_patches), dtype=np.uint64, | |
| count=len(ept_patches)) | |
| h0 = stream_entropy_bits(ept_ids, order=0) | |
| h1 = stream_entropy_bits(ept_ids, order=1) | |
| reg_ids = model.registry.ids | |
| in_train = model.registry.has(ept_ids) | |
| known_ids = ept_ids[in_train] | |
| pos = np.searchsorted(reg_ids, np.unique(known_ids)) | |
| lens = np.diff(model.registry.offsets)[pos].astype(np.float64) | |
| ept_dict_bits = float(lens.sum() * 8 + 32 * len(lens)) | |
| novelty = float((~in_train).mean()) | |
| comp["ept"] = { | |
| "n_tokens": int(len(ept_ids)), | |
| "avg_token_bytes": float(len(test) / len(ept_ids)), | |
| "distinct_in_test": int(len(np.unique(ept_ids))), | |
| "stream_bits_unigram": h0, | |
| "stream_bits_bigram": h1, | |
| "dict_bits_test_local": ept_dict_bits, | |
| "bits_per_byte_unigram": (h0 * len(ept_ids) + ept_dict_bits) / len(test), | |
| "bits_per_byte_bigram": (h1 * len(ept_ids) + ept_dict_bits) / len(test), | |
| "novelty_rate": novelty, | |
| "bytes_in_novel_patches": float(sum(len(p) for p, k in zip(ept_patches, in_train) if not k)), | |
| "segment_s": ept_seg_s, | |
| "MB_per_s": len(test) / 1e6 / max(ept_seg_s, 1e-9), | |
| } | |
| metrics["compression"] = comp | |
| dump(metrics, args.out) | |
| for k, v in comp.items(): | |
| print(f" {k}: {json.dumps({kk: round(vv, 4) if isinstance(vv, float) else vv for kk, vv in v.items() if 'bits_per_byte' in kk or 'novelty' in kk or 'avg_token' in kk})}") | |
| # ================================================== 2. prediction ===== | |
| if "prediction" not in metrics: | |
| print("== [2/7] next-patch prediction: held-out test + trained region ==") | |
| ev = model.eval_stream(test, max_positions=args.eval_positions) | |
| print(f" TEST: fused={ev['fused_acc']:.4f} (hier={ev['hier_acc']:.4f}, " | |
| f"fast={ev['snn_acc']:.4f}, am={ev['am_acc']:.4f}, coverage={ev['vocab_coverage']:.3f})") | |
| # memorization reference: the FINAL training region (tokens ~550k-600k, | |
| # where the model's operating point was warm-matched; SNN region 0-600k, | |
| # AM region 0-1.2M). Earlier regions are legitimately drifted under | |
| # continual plasticity — that is the forgetting curve, not a bug. | |
| ev_tr = model.eval_stream(train[2_600_000:2_800_000], max_positions=8_000) | |
| print(f" TRAINED-REGION: fused={ev_tr['fused_acc']:.4f} " | |
| f"(hier={ev_tr['hier_acc']:.4f}, am={ev_tr['am_acc']:.4f})") | |
| # baselines over the same test stream, in-vocab positions | |
| tp = model.ept.segment(test) | |
| tids = np.fromiter((fnv1a64(p) for p in tp), dtype=np.uint64, count=len(tp)) | |
| tcols = model.cols_of(tids)[: args.eval_positions + 1] | |
| V = model.cfg.vocab_size | |
| inv_mask = tcols != V | |
| ok_pos = inv_mask[:-1] & inv_mask[1:] | |
| prev_in = tcols[:-1][ok_pos].astype(np.int64) | |
| cur_in = tcols[1:][ok_pos].astype(np.int64) | |
| train_cols = model.cols_of(train_ids[:600_000]) | |
| uni = UnigramBaseline(train_cols, V) | |
| bi = BigramBaseline(train_cols, V) | |
| pred = { | |
| "spira_test": ev, | |
| "spira_trained_region": ev_tr, | |
| "unigram_acc": uni.accuracy(cur_in), | |
| "bigram_acc": bi.accuracy(prev_in, cur_in), | |
| "n_baseline_positions": int(len(cur_in)), | |
| "note": "baselines evaluated on in-vocab target positions of the same stream; " | |
| "compare with spira_test.*_invocab", | |
| } | |
| metrics["prediction"] = pred | |
| dump(metrics, args.out) | |
| print(f" baselines (in-vocab): unigram={pred['unigram_acc']:.4f} " | |
| f"bigram={pred['bigram_acc']:.4f} vs spira fast_invocab={ev['snn_acc_invocab']:.4f} " | |
| f"hier_invocab={ev['hier_acc_invocab']:.4f} fused_invocab={ev['fused_acc_invocab']:.4f}") | |
| # ==================================================== 3. one-shot ===== | |
| if "one_shot" not in metrics: | |
| print("== [3/7] one-shot learning: teach the unseen book ONCE ==") | |
| book = oneshot[:60_000] | |
| generalization = model.eval_stream(book, max_positions=8_000) | |
| t0 = time.perf_counter() | |
| taught = model.teach_bytes(book) | |
| teach_s = time.perf_counter() - t0 | |
| recall = model.eval_stream(book, max_positions=8_000) | |
| oneshot_metrics = { | |
| "book_bytes": len(book), | |
| "before_teaching": generalization, | |
| "teaching": {**taught, "seconds": teach_s}, | |
| "after_one_pass": recall, | |
| } | |
| metrics["one_shot"] = oneshot_metrics | |
| dump(metrics, args.out) | |
| print(f" before: fused={generalization['fused_acc']:.4f} am={generalization['am_acc']:.4f}") | |
| print(f" teach: {taught['tokens']} tokens, {taught['new_patches']} new patches in {teach_s:.1f}s") | |
| print(f" after ONE pass: am recall={recall['am_acc']:.4f} fused={recall['fused_acc']:.4f}") | |
| # ===================================== 3b. replay -> consolidation ===== | |
| if "replay" not in metrics: | |
| print("== [4/7] hippocampal replay -> cortical consolidation ==") | |
| book = oneshot[:60_000] | |
| if not model.episodes: | |
| taught = model.teach_bytes(book) | |
| print(f" (re-taught for this section: {taught['tokens']} tokens)") | |
| before = model.eval_stream(book, use_am=False, max_positions=8_000) | |
| print(f" cortex-only BEFORE replay: hier={before['hier_acc']:.4f} " | |
| f"(in-vocab {before['hier_acc_invocab']:.4f}, fast={before['snn_acc']:.4f})") | |
| curve = [] | |
| stage = max(1, args.replay_tokens // args.replay_stages) | |
| for i in range(args.replay_stages): | |
| r = model.replay(stage, log=lambda s: print(f" {s}")) | |
| probe = model.eval_stream(book, use_am=False, max_positions=2_000) | |
| curve.append({ | |
| "replayed_tokens": (i + 1) * stage, | |
| "hier_acc": probe["hier_acc"], | |
| "hier_acc_invocab": probe["hier_acc_invocab"], | |
| "snn_acc": probe["snn_acc"], | |
| "replay_acc": r.get("accuracy"), | |
| "beta": r.get("beta"), | |
| }) | |
| print(f" after {(i+1)*stage} replayed tokens: hier={probe['hier_acc']:.4f} " | |
| f"(in-vocab {probe['hier_acc_invocab']:.4f}) beta={r.get('beta'):.3f}") | |
| after = model.eval_stream(book, use_am=False, max_positions=8_000) | |
| am_after = model.eval_stream(book, max_positions=8_000) | |
| interference_core = model.eval_stream(test, use_am=False, max_positions=8_000) | |
| interference_fused = model.eval_stream(test, max_positions=8_000) | |
| metrics["replay"] = { | |
| "replay_tokens": int(args.replay_tokens), | |
| "stage_tokens": int(stage), | |
| "before": before, | |
| "curve": curve, | |
| "after": after, | |
| "am_recall_after_replay": am_after, | |
| "interference_test_core_only": interference_core, | |
| "interference_test_fused": interference_fused, | |
| "pre_replay_test": metrics.get("prediction", {}).get("spira_test"), | |
| "note": "the shipped checkpoint is PRE-replay; this section runs in-memory", | |
| } | |
| dump(metrics, args.out) | |
| print(f" cortex-only AFTER replay: hier={after['hier_acc']:.4f} " | |
| f"(in-vocab {after['hier_acc_invocab']:.4f})") | |
| print(f" AM recall after replay: {am_after['am_acc']:.4f} (must be intact)") | |
| print(f" interference on test: core {interference_core['hier_acc']:.4f} / " | |
| f"fused {interference_fused['fused_acc']:.4f}") | |
| # ================================================== 4. correction ===== | |
| if "correction" not in metrics: | |
| print("== [5/7] correction: recall vs corrupted context lanes ==") | |
| corr = {} | |
| for k in (0, 1, 2, 3): | |
| r = model.eval_stream(oneshot[:60_000], corrupt_lanes=k, max_positions=8_000) | |
| corr[f"corrupt_{k}_of_4"] = {kk: r[kk] for kk in ("am_acc", "am_acc_invocab", "fused_acc")} | |
| print(f" corrupt {k}/4 lanes: am={r['am_acc']:.4f} fused={r['fused_acc']:.4f}") | |
| metrics["correction"] = corr | |
| dump(metrics, args.out) | |
| # ================================================ 5. generation ====== | |
| if "generation" not in metrics: | |
| print("== [6/7] generation ==") | |
| gens = {} | |
| prompts = { | |
| "english": b"The old man looked at the ", | |
| "code": b"def process(data):\n ", | |
| "french": b"Il etait une fois un ", | |
| } | |
| for name, p in prompts.items(): | |
| t0 = time.perf_counter() | |
| g = model.generate(p, n_tokens=150, temperature=0.7, top_k=8, seed=42) | |
| dt = time.perf_counter() - t0 | |
| text = g.decode("utf-8", "replace") | |
| gens[name] = { | |
| "prompt": p.decode("utf-8", "replace"), | |
| "output": text, | |
| "chars_per_s": len(text) / max(dt, 1e-9), | |
| "distinct1": distinct_n(text, 1), | |
| "distinct2": distinct_n(text, 2), | |
| } | |
| print(f" [{name}] {len(text)} chars @ {len(text)/max(dt,1e-9):.0f} chars/s") | |
| print(f" {text[:180]!r}") | |
| metrics["generation"] = gens | |
| dump(metrics, args.out) | |
| # ================================================ 6. efficiency ====== | |
| if "efficiency" not in metrics: | |
| print("== [7/7] efficiency accounting ==") | |
| d_in, k, n = model.cfg.d_in, model.cfg.k_in, model.cfg.n_res | |
| V = model.cfg.vocab_size | |
| n_g = model.cfg.n_g | |
| win = model.cfg.win | |
| slow_amortized = (n_g * (V + 1) + n * n_g + n_g * n_g) / win | |
| macs = { | |
| "input_gather": k * n, | |
| "recurrent": n * n, | |
| "readout_fast": n * V, | |
| "readout_slow_amortized": slow_amortized, | |
| "total_per_token": k * n + n * n + n * V + slow_amortized, | |
| } | |
| metrics["efficiency"] = { | |
| "synaptic_macs_per_token": macs, | |
| "vs_transformer_reference": { | |
| "gpt2_small_params": 124e6, | |
| "gpt2_train_flops_per_token": 6 * 124e6, | |
| "spira_vs_gpt2_ops_ratio": macs["total_per_token"] / (6 * 124e6), | |
| }, | |
| } | |
| dump(metrics, args.out) | |
| print(f" {macs['total_per_token']/1e6:.2f}M MACs/token " | |
| f"(fast readout {macs['readout_fast']/1e6:.2f}M, slow amortized " | |
| f"{slow_amortized/1e6:.2f}M) -> " | |
| f"{(6*124e6)/macs['total_per_token']:.0f}x fewer ops/token than GPT-2-small training") | |
| print(f"metrics -> {args.out}") | |
| if __name__ == "__main__": | |
| main() | |