File size: 8,087 Bytes
f930dac | 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 176 177 | """Encode lifetimes into npz shards.
One shard = many lifetimes, flat arrays + offset indexes. Cache contents are
stored as per-event add/evict deltas against the per-lifetime record table, so
every consumer (PNSR trainer, TX window sampler, nulls, eval) materialises the
exact same bounded stores the generation-time reducer certified.
"""
from __future__ import annotations
import hashlib
import os
import tempfile
from pathlib import Path
import numpy as np
from ..world.reducer import N_PTR_SLOTS
from .binding import N_BIND_SLOTS, binding_plan
from ..world.schema import ET, Event
from .tokenize import encode_event
VAL_T, KEY_T = 12, 14 # token pieces kept per record value/key
MAX_EV_TOKENS = 112
N_ENT_BUCKETS = 64
def val_hash(s: str) -> np.uint64:
return np.uint64(int.from_bytes(hashlib.sha256(s.encode()).digest()[:8], "little"))
def pack_lifetimes(lifetimes: list[dict], tok, out_path: Path) -> dict:
A = {k: [] for k in ("tokens", "etype", "dt", "mode_gold", "family", "enum_gold",
"ptr_gold_slot", "op_gold", "gold_val_hash", "delay", "kappa",
"gold_age_rank", "corrected", "reverted", "chain_len",
"bind_write", "bind_read", "rec_store", "rec_kind",
"rec_ent", "rec_key", "rec_birth_ev", "rec_slot", "rec_val_hash",
"add_rid", "evict_rid")}
ev_tok_off, lt_off, rec_lt_off = [0], [0], [0]
add_off, evict_off = [0], [0]
op_arg_slots, enum_legal = [], []
rec_val_toks, rec_key_toks = [], []
lt_seed, lt_len, lt_digest = [], [], []
n_trunc = 0
slot_ent_all, slot_attr_all = [], []
for lt in lifetimes:
red = lt["reducer"]
slot_hist: dict[int, int] = {} # rid -> stable slot (recorded at add time)
events: list[Event] = lt["events"]
bp = binding_plan(lt)
A["bind_write"].extend(bp["write_slot"].tolist())
A["bind_read"].extend(bp["read_slot"].tolist())
slot_ent_all.append(bp["slot_ent"])
slot_attr_all.append(bp["slot_attr"])
for ev in events:
ids = encode_event(tok, ev)
if len(ids) > MAX_EV_TOKENS:
assert ev.etype != ET.QUESTION, f"question truncated: {ev.text!r}"
ids = ids[:MAX_EV_TOKENS]
n_trunc += 1
A["tokens"].extend(ids)
ev_tok_off.append(len(A["tokens"]))
A["etype"].append(int(ev.etype))
A["dt"].append(ev.dt)
A["mode_gold"].append(int(ev.mode_gold))
A["family"].append(ev.family)
A["enum_gold"].append(ev.enum_gold)
mask = np.zeros(8, np.uint8)
for e in ev.enum_legal:
mask[e >> 3] |= 1 << (e & 7)
enum_legal.append(mask)
# pointer gold -> stable slot id
if ev.ptr_gold_rid >= 0:
A["ptr_gold_slot"].append(slot_hist[ev.ptr_gold_rid])
else:
A["ptr_gold_slot"].append(-1)
A["op_gold"].append(ev.op_gold)
args = [slot_hist[r] for r in ev.op_arg_rids[:2]]
op_arg_slots.append(args + [-1] * (2 - len(args)))
A["gold_val_hash"].append(val_hash(ev.gold_value) if ev.gold_value else np.uint64(0))
m = ev.meta
A["delay"].append(m.get("delay", -1))
A["kappa"].append(m.get("kappa", -1))
A["gold_age_rank"].append(m.get("gold_age_rank", -1))
A["corrected"].append(m.get("corrected", -1))
A["reverted"].append(m.get("reverted", -1))
A["chain_len"].append(m.get("chain_len", -1))
# cache deltas; slots were assigned by the reducer at add time
for rid in m["evicts"]:
A["evict_rid"].append(rid)
for rid, gslot in zip(m["adds"], m["add_slots"]):
slot_hist[rid] = gslot
A["add_rid"].append(rid)
add_off.append(len(A["add_rid"]))
evict_off.append(len(A["evict_rid"]))
lt_off.append(len(A["etype"]))
# record table
for r in red.archive:
A["rec_store"].append(r.store)
A["rec_kind"].append(r.kind)
A["rec_ent"].append(r.ent % N_ENT_BUCKETS)
A["rec_key"].append(r.key_id)
A["rec_birth_ev"].append(r.birth_ev)
A["rec_slot"].append(slot_hist[r.rid])
A["rec_val_hash"].append(val_hash(r.value))
vt = tok.encode(r.value).ids[:VAL_T]
kt = tok.encode(r.key_text).ids[:KEY_T]
rec_val_toks.append(vt + [0] * (VAL_T - len(vt)))
rec_key_toks.append(kt + [0] * (KEY_T - len(kt)))
rec_lt_off.append(len(A["rec_store"]))
lt_seed.append(lt["seed"])
lt_len.append(lt["length"])
lt_digest.append(val_hash(lt["digest"]))
arrays = dict(
tokens=np.asarray(A["tokens"], np.uint16),
ev_tok_off=np.asarray(ev_tok_off, np.int64),
lt_off=np.asarray(lt_off, np.int64),
rec_lt_off=np.asarray(rec_lt_off, np.int64),
add_off=np.asarray(add_off, np.int64),
evict_off=np.asarray(evict_off, np.int64),
etype=np.asarray(A["etype"], np.uint8),
dt=np.asarray(A["dt"], np.uint16),
mode_gold=np.asarray(A["mode_gold"], np.uint8),
family=np.asarray(A["family"], np.int8),
enum_gold=np.asarray(A["enum_gold"], np.int16),
enum_legal=np.stack(enum_legal) if enum_legal else np.zeros((0, 8), np.uint8),
ptr_gold_slot=np.asarray(A["ptr_gold_slot"], np.int16),
op_gold=np.asarray(A["op_gold"], np.int8),
op_arg_slots=np.asarray(op_arg_slots, np.int16),
gold_val_hash=np.asarray(A["gold_val_hash"], np.uint64),
delay=np.asarray(A["delay"], np.int32),
kappa=np.asarray(A["kappa"], np.int16),
gold_age_rank=np.asarray(A["gold_age_rank"], np.int16),
corrected=np.asarray(A["corrected"], np.int8),
bind_write=np.asarray(A["bind_write"], np.int16),
bind_read=np.asarray(A["bind_read"], np.int16),
bind_slot_ent=np.stack(slot_ent_all) if slot_ent_all
else np.zeros((0, N_BIND_SLOTS), np.int16),
bind_slot_attr=np.stack(slot_attr_all) if slot_attr_all
else np.zeros((0, N_BIND_SLOTS), np.int16),
reverted=np.asarray(A["reverted"], np.int8),
chain_len=np.asarray(A["chain_len"], np.int16),
rec_store=np.asarray(A["rec_store"], np.uint8),
rec_kind=np.asarray(A["rec_kind"], np.uint8),
rec_ent=np.asarray(A["rec_ent"], np.int16),
rec_key=np.asarray(A["rec_key"], np.uint8),
rec_birth_ev=np.asarray(A["rec_birth_ev"], np.int32),
rec_slot=np.asarray(A["rec_slot"], np.int16),
rec_val_hash=np.asarray(A["rec_val_hash"], np.uint64),
rec_val_toks=np.asarray(rec_val_toks, np.uint16) if rec_val_toks
else np.zeros((0, VAL_T), np.uint16),
rec_key_toks=np.asarray(rec_key_toks, np.uint16) if rec_key_toks
else np.zeros((0, KEY_T), np.uint16),
add_rid=np.asarray(A["add_rid"], np.int32),
evict_rid=np.asarray(A["evict_rid"], np.int32),
lt_seed=np.asarray(lt_seed, np.int64),
lt_len=np.asarray(lt_len, np.int32),
lt_digest=np.asarray(lt_digest, np.uint64),
)
assert arrays["ev_tok_off"][-1] == len(arrays["tokens"])
out_path.parent.mkdir(parents=True, exist_ok=True)
fd, tmp = tempfile.mkstemp(dir=out_path.parent, suffix=".tmp")
os.close(fd)
np.savez(tmp, **arrays) # uncompressed: memmap-friendly at load
saved = tmp if tmp.endswith(".npz") else tmp + ".npz" # numpy appends .npz
os.replace(saved, out_path)
if os.path.exists(tmp):
os.unlink(tmp)
return dict(n_lifetimes=len(lifetimes), n_events=int(arrays["etype"].shape[0]),
n_tokens=int(arrays["tokens"].shape[0]),
n_records=int(arrays["rec_store"].shape[0]), n_trunc=n_trunc,
sha=hashlib.sha256(out_path.read_bytes()).hexdigest()[:16])
assert N_PTR_SLOTS == 100
|