File size: 4,602 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 | """Read-side view of packed shards: per-lifetime iteration with incremental
live-cache materialisation. Single implementation shared by the null attacks,
the TX window sampler, the PNSR streaming batcher and the evaluators, so every
consumer sees byte-identical cache states.
"""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
import numpy as np
from ..world.reducer import N_PTR_SLOTS
# Split identifier -> directory inside the published dataset repo. The
# identifiers are the ones used throughout the study and in every result file;
# the directory names are the published layout.
SPLIT_DIRS = {
"train": "e1_original/train",
"val": "e1_original/val",
"val_1k": "e1_original/val_1k",
"val_4k": "e1_original/val_4k",
"test": "e1_original/test",
"test_1k": "e1_original/test_1k",
"test_4k": "e1_original/test_4k",
"train_relabel": "e1_corrected_eval/train_relabel",
"val_relabel": "e1_corrected_eval/val_relabel",
"e3_train": "e3_corrected/train",
"e3_dev": "e3_corrected/dev",
"e3_dev_1k": "e3_corrected/dev_1k",
"e3_conf": "e3_corrected/confirmation",
"e3_conf_1k": "e3_corrected/confirmation_1k",
}
def shard_paths(split: str, root: Path) -> list[Path]:
"""Shards of `split` under `root` (= <dataset repo>/data).
Accepts either a study split identifier (see SPLIT_DIRS) or a literal
subdirectory name, so a locally regenerated corpus also works.
"""
d = root / SPLIT_DIRS.get(split, split)
if not d.is_dir():
d = root / split
return sorted(d.glob("shard_*.npz"))
class Shard:
def __init__(self, path: Path):
self.z = np.load(path, mmap_mode=None) # npz members load lazily on access
self.path = path
for k in self.z.files:
setattr(self, k, self.z[k])
self.n_lifetimes = len(self.lt_off) - 1
def lifetime(self, i: int) -> "LifetimeView":
return LifetimeView(self, i)
@dataclass
class EventView:
idx: int # event index within lifetime
tokens: np.ndarray
etype: int
dt: int
live_slots: np.ndarray # int32 [N_PTR_SLOTS] -> record row (global) or -1
mode_gold: int
family: int
enum_gold: int
enum_legal: np.ndarray # uint8 [8] bitmask
ptr_gold_slot: int
op_gold: int
op_arg_slots: np.ndarray
meta: dict
class LifetimeView:
def __init__(self, sh: Shard, i: int):
self.sh = sh
self.lo = int(sh.lt_off[i])
self.hi = int(sh.lt_off[i + 1])
self.rec_lo = int(sh.rec_lt_off[i])
self.rec_hi = int(sh.rec_lt_off[i + 1])
self.seed = int(sh.lt_seed[i])
self.length = self.hi - self.lo
def records(self) -> dict[str, np.ndarray]:
s = slice(self.rec_lo, self.rec_hi)
sh = self.sh
return dict(store=sh.rec_store[s], kind=sh.rec_kind[s], ent=sh.rec_ent[s],
key=sh.rec_key[s], birth_ev=sh.rec_birth_ev[s], slot=sh.rec_slot[s],
val_hash=sh.rec_val_hash[s], val_toks=sh.rec_val_toks[s],
key_toks=sh.rec_key_toks[s])
def __iter__(self):
sh = self.sh
live = np.full(N_PTR_SLOTS, -1, np.int32)
for e in range(self.lo, self.hi):
for rid in sh.evict_rid[sh.evict_off[e]:sh.evict_off[e + 1]]:
slot = sh.rec_slot[self.rec_lo + int(rid)]
live[slot] = -1
for rid in sh.add_rid[sh.add_off[e]:sh.add_off[e + 1]]:
slot = sh.rec_slot[self.rec_lo + int(rid)]
live[slot] = self.rec_lo + int(rid) # global record row
yield EventView(
idx=e - self.lo,
tokens=sh.tokens[sh.ev_tok_off[e]:sh.ev_tok_off[e + 1]],
etype=int(sh.etype[e]), dt=int(sh.dt[e]),
live_slots=live.copy(),
mode_gold=int(sh.mode_gold[e]), family=int(sh.family[e]),
enum_gold=int(sh.enum_gold[e]), enum_legal=sh.enum_legal[e],
ptr_gold_slot=int(sh.ptr_gold_slot[e]), op_gold=int(sh.op_gold[e]),
op_arg_slots=sh.op_arg_slots[e],
meta=dict(delay=int(sh.delay[e]), kappa=int(sh.kappa[e]),
gold_age_rank=int(sh.gold_age_rank[e]),
corrected=int(sh.corrected[e]), reverted=int(sh.reverted[e]),
chain_len=int(sh.chain_len[e]),
gold_val_hash=int(sh.gold_val_hash[e])),
)
|