SPIRA v2: entropy-patch tokenization + hierarchical dopamine-gated spiking cortex + hippocampal one-shot memory with replay consolidation
032a961 verified | #!/usr/bin/env python3 | |
| """Quick smoke test of every SPIRA component on a tiny slice of the corpus. | |
| Validates: hashing determinism, EPT segmentation, reservoir dynamics, | |
| three-factor learning signal, hippocampal write/query/corruption, generation. | |
| """ | |
| import sys | |
| import time | |
| import numpy as np | |
| sys.path.insert(0, "/home/z/my-project/download/spira") | |
| from spira import EPT, Hippocampus, SpikingCore, SPIRA, SPIRAConfig | |
| from spira.hashing import Mulberry32, fnv1a64, input_neurons, lane_key32, splitmix64 | |
| data = open("/home/z/my-project/data/train.bin", "rb").read() | |
| print(f"corpus: {len(data)/1e6:.2f} MB") | |
| small = data[:120_000] | |
| # ---------------------------------------------------------------- hashing | |
| h1 = fnv1a64(b"hello world") | |
| h2 = fnv1a64(b"hello world") | |
| assert h1 == h2 and h1 != fnv1a64(b"hello worle") | |
| assert splitmix64(42) == splitmix64(42) | |
| r = Mulberry32(7) | |
| vals = [r.next_float() for _ in range(5)] | |
| r2 = Mulberry32(7) | |
| assert vals == [r2.next_float() for _ in range(5)] | |
| neurons = input_neurons(h1, 1024, 12) | |
| assert len(set(neurons)) >= 10 and all(0 <= n < 1024 for n in neurons) | |
| print(f"[ok] hashing: fnv1a64('hello world') = {h1:#x}, 12 neurons, deterministic PRNG") | |
| # -------------------------------------------------------------------- EPT | |
| ept = EPT() | |
| ept.fit(small[:80_000]) | |
| budget = ept.calibrate(small[80_000:100_000], target_len=4.5) | |
| patches = ept.segment(small[:40_000]) | |
| st = ept.stats(patches) | |
| print(f"[ok] EPT: budget={budget:.2f} bits -> avg patch len {st['avg_len']:.2f} B " | |
| f"(p50={st['p50_len']:.0f} p90={st['p90_len']:.0f} max={st['max_len']})") | |
| assert 3.0 < st["avg_len"] < 7.0 | |
| # show a sample segmentation | |
| sample = small[:220].decode("utf-8", "replace") | |
| seg = ept.segment(small[:220]) | |
| lens = [len(p) for p in seg] | |
| print(f"[ok] EPT sample: {' | '.join(repr(p.decode('utf-8','replace')) for p in seg[:24])}") | |
| ids = [fnv1a64(p) for p in patches] | |
| assert len(set(ids)) > 100 | |
| # ------------------------------------------------------------- hippocampus | |
| am = Hippocampus(m=4, size=1 << 16) | |
| ids_small = ids[:3000] | |
| lane = am.lane_keys_stream(np.array(ids_small, dtype=np.uint64)) | |
| # scalar vs vectorized lane keys must agree | |
| for i in (0, 17, 999): | |
| for l in range(4): | |
| assert int(lane[i, l]) == lane_key32(ids_small[i], l), (i, l) | |
| print("[ok] hippocampus: scalar and vectorized lane keys agree") | |
| ctx_lane_list = [] | |
| m = 4 | |
| n_ctx = len(ids_small) - m + 1 | |
| ctx = np.empty((n_ctx, m), dtype=np.uint32) | |
| for li in range(m): | |
| ctx[:, li] = lane[li : li + n_ctx, li] | |
| vals_idx = np.arange(m, m + n_ctx) # continuation token index for each context | |
| am.write_stream(ctx, vals_idx) | |
| # context tokens 10..13 -> continuation = token 14 | |
| q, w = am.query(ctx[10]) | |
| assert q == 14 and w > 0, (q, w) | |
| print(f"[ok] hippocampus: one-shot write + exact recall (query->next={q}, weight={w:.2f})") | |
| # corruption: replace 1 of 4 context patches (slot 2 = token 12) | |
| corrupt = ctx[10].copy() | |
| corrupt[2] = lane_key32(ids_small[500], 2) | |
| q2, w2 = am.query(corrupt) | |
| print(f"[ok] hippocampus: corrupted-context query -> {q2} (true=14), weight={w2:.2f}") | |
| # ------------------------------------------------------------- spiking core | |
| core = SpikingCore(n_out=128, seed=1) | |
| acc = [] | |
| T = 4000 | |
| rng = np.random.default_rng(0) | |
| toy_ids = np.array(ids_small[:T], dtype=np.uint64) | |
| neurons_stream = SPIRA.input_neurons_stream(toy_ids, 1024, 12) | |
| # toy target: next id mod 128 -> learnable map? no: use the actual next patch (hard) | |
| targets = (toy_ids % np.uint64(127)).astype(np.int64)[1:] | |
| core.reset() | |
| t0 = time.perf_counter() | |
| res = core.train_pass(neurons_stream[: T - 1], targets, log_every=2000, | |
| log_fn=lambda d: print(f" {d}")) | |
| dt = time.perf_counter() - t0 | |
| print(f"[ok] spiking core: {res['ticks']} ticks in {dt:.1f}s " | |
| f"({res['ticks_per_sec']:.0f} t/s), final acc_ema={res['acc_ema_final']:.3f}") | |
| assert res["ticks_per_sec"] > 200 | |
| # ------------------------------------------------------- end-to-end (tiny) | |
| cfg = SPIRAConfig(snn_tokens=6_000, am_tokens=6_000, epochs=1, vocab_size=1024) | |
| model = SPIRA.build(small, cfg, log=lambda m: print(" ", m)) | |
| model.prepare() | |
| model._patches_train = None | |
| hist = model.train(log_fn=lambda m: print(" ", m)) | |
| print(f"[ok] end-to-end train: {hist['epochs'][0]['accuracy']:.3f} acc, am={hist['am']}") | |
| # generation | |
| gen = model.generate(b"The ", n_tokens=40, seed=1) | |
| print(f"[ok] generation sample: {gen[:200]!r}") | |
| # eval incl. one-shot teaching on held-out text | |
| oneshot = open("/home/z/my-project/data/oneshot.bin", "rb").read()[:30_000] | |
| ev0 = model.eval_stream(oneshot, max_positions=3000) | |
| ev1 = model.eval_stream(oneshot, teach_am_first=True, max_positions=3000) | |
| print(f"[ok] eval on held-out book: before-teach fused={ev0['fused_acc']:.3f}; " | |
| f"after ONE-SHOT teach fused={ev1['fused_acc']:.3f} (am={ev1['am_acc']:.3f})") | |
| ev2 = model.eval_stream(oneshot, teach_am_first=True, corrupt_lanes=1, max_positions=3000) | |
| print(f"[ok] eval with 1/4 corrupted context lanes: fused={ev2['fused_acc']:.3f}") | |
| print("\nALL SMOKE TESTS PASSED") | |