File size: 7,107 Bytes
032a961 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | #!/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()
|