File size: 5,983 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 178 179 | """Focused invariant tests for the PNS experiment (pragmatic set).
Run: python3 -m pytest tests/ -q
"""
import sys
from pathlib import Path
import numpy as np
import pytest
import torch
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from pns.common import EVAL_METADATA_KEYS, MODEL_INPUT_KEYS, SUPERVISION_KEYS
from pns.world.reducer import CAP, S_GOAL, S_OBS, S_SELF, ExactReducer, N_PTR_SLOTS
from pns.world.schema import ET, Mode
from pns.world.sim import generate_lifetime
SEED = 4242
@pytest.fixture(scope="module")
def lt():
return generate_lifetime(SEED, 256)
def test_determinism(lt):
assert generate_lifetime(SEED, 256)["digest"] == lt["digest"]
def test_gold_certified_live(lt):
"""Every pointer gold rid was live (in the bounded store) at its question."""
red = ExactReducer()
for ev in lt["events"]:
red.step(ev)
if ev.ptr_gold_rid >= 0:
assert ev.ptr_gold_rid in red.records
for rid in ev.op_arg_rids:
assert rid in red.records
def test_bounded_stores(lt):
red = ExactReducer()
for ev in lt["events"]:
red.step(ev)
for s in (S_OBS, S_SELF, S_GOAL):
assert len(red.by_store[s]) <= CAP[s]
# questions never write records
if ev.etype == ET.QUESTION:
assert not red.last_adds
def test_slot_stability(lt):
"""A record keeps its slot for its whole live span."""
red = ExactReducer()
seen: dict[int, int] = {}
for ev in lt["events"]:
red.step(ev)
for rid in red.records:
slot = red.slot_index(rid)
if rid in seen:
assert seen[rid] == slot
seen[rid] = slot
def test_pack_replay_parity(tmp_path, lt):
"""The packed per-event deltas reproduce the reducer's live sets and the
packed ptr_gold_slot equals the reducer's slot index - the single-source-
of-truth gate between generation and every training/eval consumer."""
from pns.data import tokenize
from pns.data.pack import pack_lifetimes
tok = tokenize.load_tokenizer()
out = tmp_path / "t.npz"
pack_lifetimes([generate_lifetime(SEED, 256)], tok, out)
z = np.load(out)
red = ExactReducer()
live = {}
rlo = 0
for e, ev in enumerate(lt["events"]):
red.step(ev)
for rid in z["evict_rid"][z["evict_off"][e]:z["evict_off"][e + 1]]:
live.pop(int(z["rec_slot"][rlo + int(rid)]), None)
for rid in z["add_rid"][z["add_off"][e]:z["add_off"][e + 1]]:
live[int(z["rec_slot"][rlo + int(rid)])] = int(rid)
true_live = {red.slot_index(r) for r in red.records}
assert set(live) == true_live, f"live set mismatch at event {e}"
if ev.ptr_gold_rid >= 0:
assert int(z["ptr_gold_slot"][e]) == red.slot_index(ev.ptr_gold_rid)
def test_firewall_disjoint():
assert not (MODEL_INPUT_KEYS & SUPERVISION_KEYS)
assert not (MODEL_INPUT_KEYS & EVAL_METADATA_KEYS)
assert not (SUPERVISION_KEYS & EVAL_METADATA_KEYS)
def test_deadline_balance():
from collections import Counter
g = Counter()
for s in range(5000, 5120):
for ev in generate_lifetime(s, 256)["events"]:
if ev.etype == ET.QUESTION and ev.family == 7:
g[ev.gold_value] += 1
frac = max(g.values()) / sum(g.values())
assert frac < 0.60, dict(g)
# ------------------------------------------------------------------- model
@pytest.fixture(scope="module")
def model():
from pns.model.pnsr import PNSR, PNSRConfig
torch.manual_seed(0)
return PNSR(PNSRConfig()).eval()
def _dummy(B=2, T=48):
tok = torch.randint(1, 8000, (B, T))
tok[:, 40:] = 0
rec = torch.randn(B, N_PTR_SLOTS, 384)
live = torch.zeros(B, N_PTR_SLOTS, dtype=torch.bool)
live[:, :17] = True
return tok, rec, live
def test_noop_bitwise(model):
tok, rec, live = _dummy()
st = model.initial_state(2, "cpu") + 0.31
s2, _ = model.step(st, tok, torch.tensor([int(ET.NOOP)] * 2),
torch.tensor([4, 4]), rec, live)
assert torch.equal(st, s2)
def test_typed_eligibility(model):
"""A QUESTION may write discourse/workspace/temporal but must leave the
entity/belief/goal banks bitwise untouched."""
tok, rec, live = _dummy()
st = model.initial_state(2, "cpu") + 0.11
s2, _ = model.step(st, tok, torch.tensor([int(ET.QUESTION)] * 2),
torch.tensor([4, 4]), rec, live)
ent_belief_goal = slice(0, 40 + 28 + 12)
assert torch.equal(st[:, ent_belief_goal], s2[:, ent_belief_goal])
assert not torch.equal(st, s2) # something else did change
def test_pointer_dead_masked(model):
tok, rec, live = _dummy()
st = model.initial_state(2, "cpu")
_, out = model.step(st, tok, torch.tensor([0, 0]), torch.tensor([4, 4]), rec, live)
assert torch.isinf(out["ptr"][:, 17:]).all()
assert torch.isfinite(out["ptr"][:, :17]).all()
def test_param_match():
from pns.model.pnsr import PNSR, PNSRConfig
from pns.model.tx import TX, TXConfig
np_ = sum(p.numel() for p in PNSR(PNSRConfig()).parameters())
nt = sum(p.numel() for p in TX(TXConfig()).parameters())
assert abs(np_ - nt) / max(np_, nt) < 0.03
def test_loader_live_matches_pack(tmp_path):
from pns.data import tokenize
from pns.data.pack import pack_lifetimes
from pns.data.view import Shard
tok = tokenize.load_tokenizer()
out = tmp_path / "t.npz"
pack_lifetimes([generate_lifetime(SEED + 1, 256)], tok, out)
sh = Shard(out)
lv = sh.lifetime(0)
red = ExactReducer()
for ev_true, ev_view in zip(generate_lifetime(SEED + 1, 256)["events"], lv):
red.step(ev_true)
want = {red.slot_index(r) for r in red.records}
got = {s for s in range(N_PTR_SLOTS) if ev_view.live_slots[s] >= 0}
assert want == got
if ev_view.mode_gold == int(Mode.ANSWER_POINTER):
assert ev_view.ptr_gold_slot in got
|