File size: 7,234 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 166 167 168 169 170 171 172 173 174 175 | #!/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()
|