SPIRA v2: entropy-patch tokenization + hierarchical dopamine-gated spiking cortex + hippocampal one-shot memory with replay consolidation
032a961 verified | #!/usr/bin/env python3 | |
| """SPIRA training CLI. | |
| Usage: | |
| python3 -m spira.train --data data/train.bin \ | |
| --out artifacts/spira_model.npz \ | |
| [--snn-tokens 600000] [--am-tokens 1200000] [--epochs 2] [--quick] | |
| Writes: | |
| <out> : model checkpoint (npz) | |
| <out>.train_log.jsonl : tick-level progress log | |
| <out>.history.json : final training history | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import sys | |
| import time | |
| from .model import SPIRA, SPIRAConfig | |
| def main() -> None: | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--data", required=True) | |
| ap.add_argument("--out", required=True) | |
| ap.add_argument("--snn-tokens", type=int, default=600_000) | |
| ap.add_argument("--am-tokens", type=int, default=1_200_000) | |
| ap.add_argument("--epochs", type=int, default=2) | |
| ap.add_argument("--vocab", type=int, default=4096) | |
| ap.add_argument("--n-res", type=int, default=512) | |
| ap.add_argument("--target-avg-len", type=float, default=4.5) | |
| ap.add_argument("--quick", action="store_true", help="tiny smoke-test run") | |
| args = ap.parse_args() | |
| cfg = SPIRAConfig( | |
| vocab_size=args.vocab, | |
| n_res=args.n_res, | |
| snn_tokens=args.snn_tokens, | |
| am_tokens=args.am_tokens, | |
| epochs=args.epochs, | |
| target_avg_len=args.target_avg_len, | |
| ) | |
| if args.quick: | |
| cfg.snn_tokens, cfg.am_tokens, cfg.epochs = 20_000, 20_000, 1 | |
| data = open(args.data, "rb").read() | |
| t0 = time.perf_counter() | |
| log_path = args.out + ".train_log.jsonl" | |
| log_f = open(log_path, "w") | |
| def log(msg: str) -> None: | |
| print(msg, flush=True) | |
| log_f.write(json.dumps({"t": round(time.perf_counter() - t0, 1), "msg": msg}) + "\n") | |
| log_f.flush() | |
| def tick_log(d: dict) -> None: | |
| log_f.write(json.dumps({"t": round(time.perf_counter() - t0, 1), **d}) + "\n") | |
| log_f.flush() | |
| model = SPIRA.build(data, cfg, log=log) | |
| model.prepare() | |
| # free the patch list — the registry carries everything we need | |
| model._patches_train = None | |
| log("[spira] training starts") | |
| history = model.train( | |
| log_fn=lambda msg: tick_log({"msg": msg}) or print(msg, flush=True), | |
| ckpt_path=args.out, | |
| ) | |
| history["build_meta"] = model.meta | |
| history["config"] = {k: v for k, v in vars(cfg).items()} | |
| history["wall_clock_s"] = time.perf_counter() - t0 | |
| with open(args.out + ".history.json", "w") as f: | |
| json.dump(history, f, indent=2) | |
| model.save(args.out) | |
| log(f"[spira] DONE in {history['wall_clock_s']:.0f}s -> {args.out}") | |
| log_f.close() | |
| if __name__ == "__main__": | |
| main() | |