SPIRA v2: entropy-patch tokenization + hierarchical dopamine-gated spiking cortex + hippocampal one-shot memory with replay consolidation
032a961 verified | #!/usr/bin/env python3 | |
| """Chunked foreground training driver for SPIRA. | |
| Background processes do not survive between tool calls in this sandbox, so | |
| training runs in foreground chunks (each fits within a single tool-call | |
| timeout). Each chunk checkpoints the model; later chunks resume from it. | |
| Usage: | |
| python3 train_chunk.py --start 0 --count 250000 # chunk 1 (fresh build) | |
| python3 train_chunk.py --start 250000 --count 250000 # chunk 2 (resume) | |
| python3 train_chunk.py --start 500000 --count 100000 --am # chunk 3 + AM pass | |
| python3 train_chunk.py --start 0 --count 300000 --am --raster # extra epoch + AM | |
| python3 train_chunk.py --replay 120000 # hippocampal replay pass | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| import sys | |
| import time | |
| import numpy as np | |
| _ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) | |
| sys.path.insert(0, _ROOT) | |
| from spira import SPIRA, SPIRAConfig | |
| from spira.hashing import fnv1a64 | |
| BASE = _ROOT # repo python/ dir | |
| DATA = f"{BASE}/data/train.bin" | |
| CKPT = f"{BASE}/artifacts/spira_model.npz" | |
| LOG = f"{BASE}/artifacts/chunks.jsonl" | |
| def log(msg: str) -> None: | |
| print(msg, flush=True) | |
| with open(LOG, "a") as f: | |
| f.write(json.dumps({"t": time.time(), "msg": msg}) + "\n") | |
| def main() -> None: | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--start", type=int, default=0) | |
| ap.add_argument("--count", type=int, default=250_000) | |
| ap.add_argument("--am", action="store_true", help="run the hippocampal write pass") | |
| ap.add_argument("--warmup", type=int, default=0, | |
| help="step-only reservoir warmup (no learning) to re-adapt the " | |
| "homeostatic state, then save — restores the operating point") | |
| ap.add_argument("--am-tokens", type=int, default=1_200_000) | |
| ap.add_argument("--raster", action="store_true", help="record the spike raster") | |
| ap.add_argument("--vocab", type=int, default=16384) | |
| ap.add_argument("--replay", type=int, default=0, | |
| help="hippocampal replay pass: re-run this many tokens of " | |
| "episodic traces through the cortex (consolidation)") | |
| ap.add_argument("--am-only", action="store_true", | |
| help="skip the SNN pass; run only the hippocampal write pass " | |
| "(+ optional replay) and save") | |
| args = ap.parse_args() | |
| t0 = time.perf_counter() | |
| data = open(DATA, "rb").read() | |
| import os | |
| if args.start == 0 and not os.path.exists(CKPT): | |
| cfg = SPIRAConfig(vocab_size=args.vocab) | |
| model = SPIRA.build(data, cfg, log=log) | |
| model._patches_train = None | |
| log("[chunk] fresh build done") | |
| else: | |
| model = SPIRA.load(CKPT) | |
| # rebuild the token stream for this chunk (registry/vocab already inside) | |
| patches = model.ept.segment(data) | |
| model._ids_train = np.fromiter( | |
| (fnv1a64(p) for p in patches), dtype=np.uint64, count=len(patches)) | |
| del patches | |
| log("[chunk] resumed from checkpoint") | |
| ids = model._ids_train | |
| n = len(ids) | |
| start = min(args.start, n - 2) | |
| count = min(args.count, n - 1 - start) | |
| if args.warmup > 0: | |
| w = min(args.warmup, n - 1 - start) | |
| log(f"[chunk] step-only warmup over tokens[{start}:{start+w}] " | |
| f"(matching final training-region statistics)") | |
| wn = SPIRA.input_neurons_stream(ids[start : start + w], model.cfg.d_in, model.cfg.k_in) | |
| wc = model.cols_of(ids[start : start + w + 1]) | |
| beta_saved = model.core.beta if model.core.hier else None | |
| model.core.reset() | |
| if beta_saved is not None: | |
| model.core.beta = beta_saved # beta is LEARNED state — never wipe it | |
| t1 = time.perf_counter() | |
| for t in range(w): | |
| x = model.core.step(wn[t]) | |
| u_fast = (x - model.core.x_mean) @ model.core.W_out | |
| if model.core.hier: | |
| u = u_fast + np.float32(model.core.beta) * model.core.u_slow | |
| else: | |
| u = u_fast | |
| pred = int(np.argmax(u[: model.cfg.vocab_size])) | |
| model.core.learn_reservoir(pred, int(wc[t + 1]), x, u=u_fast) | |
| if model.core.hier and (t + 1) % model.core.win == 0: | |
| model.core.step_G() | |
| log(f"[chunk] warmup done in {time.perf_counter()-t1:.0f}s " | |
| f"(firing rate {float(model.core.s.sum())/model.core.n_res:.3f})") | |
| model.save(CKPT) | |
| log(f"[chunk] checkpoint saved with core state -> {CKPT}") | |
| return | |
| if args.am_only: | |
| t1 = time.perf_counter() | |
| n_am = min(args.am_tokens, n - 1) | |
| log(f"[chunk] hippocampus one-shot write pass over {n_am} tokens") | |
| taught = model.teach_ids(ids[:n_am], episodic=False) | |
| log(f"[chunk] AM: {taught['associations']} associations in " | |
| f"{time.perf_counter()-t1:.0f}s; {model.am.stats()}") | |
| if args.replay > 0: | |
| res_r = model.replay(args.replay, log=log) | |
| log(f"[chunk] replay: {res_r}") | |
| model.save(CKPT) | |
| log(f"[chunk] checkpoint saved -> {CKPT} (total {time.perf_counter()-t0:.0f}s)") | |
| return | |
| neurons = SPIRA.input_neurons_stream(ids[start : start + count], model.cfg.d_in, model.cfg.k_in) | |
| cols = model.cols_of(ids[start + 1 : start + count + 1]) | |
| log(f"[chunk] SNN pass tokens[{start}:{start+count}]") | |
| res = model.core.train_pass( | |
| neurons, cols, log_every=25_000, | |
| log_fn=lambda d: log(f" tick {d['tick']}/{d['total']} acc={d['acc_ema']:.3f} " | |
| f"dop={d['dopamine_ema']:+.3f} {d['ticks_per_sec']:.0f} t/s"), | |
| record_raster_from=(count - 1500 if args.raster else None), | |
| ) | |
| raster = res.pop("raster", None) | |
| if raster: | |
| ticks = np.concatenate([np.full(len(sp), t, dtype=np.int32) for t, sp in raster]) | |
| neurons_r = np.concatenate([sp for _, sp in raster]).astype(np.int32) | |
| np.savez_compressed(CKPT + ".raster.npz", ticks=ticks, neurons=neurons_r) | |
| log(f"[chunk] raster saved ({len(ticks)} spikes over {len(raster)} ticks)") | |
| log(f"[chunk] SNN done: acc={res['accuracy']:.4f} " | |
| f"fast={res.get('accuracy_fast_only')} beta={res.get('beta')} " | |
| f"ema={res['acc_ema_final'] if res['acc_ema_final'] is not None else 0:.4f} " | |
| f"({res['elapsed_s']:.0f}s, {res['ticks_per_sec']:.0f} t/s)") | |
| del neurons, cols | |
| if args.am: | |
| n_am = min(args.am_tokens, n - 1) | |
| log(f"[chunk] hippocampus one-shot write pass over {n_am} tokens") | |
| t1 = time.perf_counter() | |
| taught = model.teach_ids(ids[:n_am], episodic=False) | |
| log(f"[chunk] AM: {taught['associations']} associations in " | |
| f"{time.perf_counter()-t1:.0f}s; {model.am.stats()}") | |
| if args.replay > 0: | |
| res_r = model.replay(args.replay, log=log) | |
| log(f"[chunk] replay: {res_r}") | |
| model.save(CKPT) | |
| log(f"[chunk] checkpoint saved -> {CKPT} (total {time.perf_counter()-t0:.0f}s)") | |
| if __name__ == "__main__": | |
| main() | |