#!/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()