SPIRA v2: entropy-patch tokenization + hierarchical dopamine-gated spiking cortex + hippocampal one-shot memory with replay consolidation
032a961 verified | #!/usr/bin/env python3 | |
| """Export a trained SPIRA model as a compact JSON artifact for the web demo. | |
| The TypeScript port (src/lib/spira/*) is compatible with the Python | |
| prototype: FNV-1a-64 / splitmix64 / mulberry32 are shared specs, the input | |
| and recurrent weights are REGENERATED from the seed in the browser, and only | |
| the learned matrices (W_out, W_out_G, W_rec after STDP), the Markov table, | |
| the patch registry, the hippocampal entries, and the learned homeostatic | |
| state are shipped (binary packed as base64). | |
| Format v2 adds the hierarchical slow area (G): its fixed projections ship | |
| dense (they are small), its readout ships int8-quantized like the fast one, | |
| and the learned state includes the slow homeostatic scalars, the top-down | |
| gain beta, and the measured variance-matching fusion scale. | |
| Usage: | |
| python3 -m spira.export_deploy --model artifacts/spira_model.npz \ | |
| --out artifacts/spira_deploy.json [--am-cap 100000] | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import base64 | |
| import json | |
| import numpy as np | |
| from .model import SPIRA | |
| def b64(arr: np.ndarray) -> str: | |
| return base64.b64encode(np.ascontiguousarray(arr).tobytes()).decode("ascii") | |
| def main() -> None: | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--model", required=True) | |
| ap.add_argument("--out", required=True) | |
| ap.add_argument("--am-cap", type=int, default=80_000) | |
| ap.add_argument("--am-min-count", type=int, default=2) | |
| args = ap.parse_args() | |
| model = SPIRA.load(args.model) | |
| cfg = model.cfg | |
| core = model.core | |
| # ---- W_out: per-column int8 quantization ----------------------------- | |
| W = core.W_out.astype(np.float32) # (n_res, n_out) | |
| n_res, n_out = W.shape | |
| col_max = np.abs(W).max(axis=0) | |
| col_max[col_max == 0] = 1.0 | |
| scale = (col_max / 127.0).astype(np.float32) | |
| Wq = np.round(W / scale[None, :]).astype(np.int8) | |
| # ---- W_rec after STDP: sparse triplets (row=pre, col=post) ----------- | |
| rec = core.W_rec | |
| rows, cols = np.nonzero(rec) | |
| vals = rec[rows, cols].astype(np.float32) | |
| rec_entries = np.stack([rows.astype(np.int16), cols.astype(np.int16)], axis=1).astype(np.int16) | |
| # ---- hippocampus: deduplicated entries, binary packed ---------------- | |
| entries = model.am.top_entries(min_count=args.am_min_count, cap=args.am_cap) | |
| n_e = len(entries) | |
| # pack: 4 x uint32 lanes + uint64 patch id + uint8 count = 25 bytes/entry | |
| packed = np.zeros((n_e, 25), dtype=np.uint8) | |
| for i, (k, v, c) in enumerate(entries): | |
| kl = k.astype("<u4").tobytes() | |
| packed[i, :16] = np.frombuffer(kl, dtype=np.uint8) | |
| packed[i, 16:24] = np.frombuffer(np.uint64(v).astype("<u8").tobytes(), dtype=np.uint8) | |
| packed[i, 24] = min(c, 255) | |
| flat = packed.reshape(-1) | |
| # ---- registry (vocab + all registered patches) ------------------------ | |
| reg = model.registry | |
| # patch lens as uint8/16 | |
| lens = np.diff(reg.offsets).astype(np.uint16) | |
| # ---- hierarchical slow area (v2) -------------------------------------- | |
| hier = None | |
| if core.hier: | |
| WG = core.W_outG.astype(np.float32) # (n_g, n_out) | |
| g_max = np.abs(WG).max(axis=0) | |
| g_max[g_max == 0] = 1.0 | |
| g_scale = (g_max / 127.0).astype(np.float32) | |
| WGq = np.round(WG / g_scale[None, :]).astype(np.int8) | |
| hier = { | |
| "n_g": core.n_g, | |
| "win": core.win, | |
| "lam_g": core.lam_g, | |
| "tau_trace_g": core.tau_trace_g, | |
| "target_rate_g": core.target_rate_g, | |
| "eta_homeo_g": core.eta_homeo_g, | |
| "beta": float(core.beta), | |
| "fusion_scale": float(model.fusion["scale"]), | |
| # fixed projections ship dense (small): L->G and G recurrent | |
| "w_ing_b64": b64(core.W_inG), # (n_res, n_g) float32 | |
| "w_recg_b64": b64(core.W_recG), # (n_g, n_g) float32 | |
| "w_outg_q_b64": b64(WGq), | |
| "w_outg_scale_b64": b64(g_scale), | |
| } | |
| deploy = { | |
| "format": "spira-deploy-v2", | |
| "config": { | |
| "d_in": cfg.d_in, "k_in": cfg.k_in, "n_res": cfg.n_res, | |
| "vocab_size": cfg.vocab_size, "seed": cfg.seed, | |
| "in_fanout": core.in_fanout, "rec_fanout": core.rec_fanout, | |
| "inh_frac": core.inh_frac, "lam": core.lam, "base_thr": core.base_thr, | |
| "thr_min": core.thr_min, "thr_max": core.thr_max, "refrac": core.refrac, | |
| "tau_trace": core.tau_trace, "w_in_scale": core.w_in_scale, | |
| "w_rec_scale": core.w_rec_scale, "target_rate": core.target_rate, | |
| "eta_homeo": core.eta_homeo, | |
| "hierarchy": bool(cfg.hierarchy), "n_g": cfg.n_g, "win": cfg.win, | |
| "am_m": cfg.am_m, "am_min_lanes": cfg.am_min_lanes, | |
| "am_size": 1 << 20, # browser-side table is smaller | |
| "beta_am": cfg.beta_am, "am_override": cfg.am_override, | |
| "temperature": cfg.temperature, "top_k": cfg.top_k, | |
| }, | |
| "ept": { | |
| "budget_bits": model.ept.budget_bits, "min_len": model.ept.min_len, | |
| "max_len": model.ept.max_len, "alpha": model.ept.alpha, | |
| # 256x256 float32 log2-probs (row = prev byte); TS negates for surprise | |
| "log2p_b64": b64(model.ept.log2p.astype(np.float32)), | |
| }, | |
| "reservoir": { | |
| "w_rec_rows": b64(rows.astype(np.int32)), | |
| "w_rec_cols": b64(cols.astype(np.int32)), | |
| "w_rec_vals_b64": b64(vals), | |
| "w_rec_n": int(len(rows)), | |
| }, | |
| "readout": { | |
| "shape": [int(n_res), int(n_out)], | |
| "vocab_ids_hex": [hex(int(x)) for x in model.vocab_ids], | |
| "w_out_q_b64": b64(Wq), | |
| "w_out_scale_b64": b64(scale), | |
| }, | |
| "hier": hier, | |
| "am": { | |
| "n_entries": n_e, | |
| "packed_b64": b64(flat), | |
| "entry_bytes": 25, | |
| }, | |
| "registry": { | |
| "n": int(len(reg.ids)), | |
| "ids_b64": b64(reg.ids.astype("<u8")), | |
| "lens_b64": b64(lens), | |
| "buffer_b64": b64(np.frombuffer(reg.buffer, dtype=np.uint8)), | |
| }, | |
| "core_state": { | |
| # learned homeostatic scalars needed for faithful inference | |
| "x_mean_b64": b64(model.core.x_mean.astype(np.float32)), | |
| "thr_b64": b64(model.core.thr.astype(np.float32)), | |
| "rhat_b64": b64(model.core.rhat.astype(np.float32)), | |
| **({ | |
| "x_mean_g_b64": b64(model.core.x_mean_G.astype(np.float32)), | |
| "thr_g_b64": b64(model.core.thr_G.astype(np.float32)), | |
| "rhat_g_b64": b64(model.core.rhat_G.astype(np.float32)), | |
| } if core.hier else {}), | |
| }, | |
| "fusion": model.fusion, | |
| "meta": model.meta, | |
| } | |
| with open(args.out, "w") as f: | |
| json.dump(deploy, f) | |
| import os | |
| sz = os.path.getsize(args.out) | |
| print(f"deploy artifact -> {args.out} ({sz/1e6:.1f} MB)") | |
| print(f" vocab {len(model.vocab_ids)} | registry {len(reg.ids)} | am entries {n_e} " | |
| f"| W_out {Wq.shape} int8 | W_rec triplets {len(rows)}" | |
| + (f" | hier n_g={core.n_g} win={core.win} beta={core.beta:.4f}" if core.hier else "")) | |
| if __name__ == "__main__": | |
| main() | |