File size: 8,973 Bytes
aff03ff 16d760d aff03ff 16d760d aff03ff 16d760d aff03ff 16d760d aff03ff 16d760d aff03ff 16d760d aff03ff 16d760d aff03ff 16d760d aff03ff 16d760d aff03ff 16d760d aff03ff | 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 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 | """Count-then-peaks decoding of anchored arrays (SPEC.md §12.4, §12.6), on a tiny random
Qwen3.5 (CPU):
* count aux: anchored arrays of whole contexts get an auxiliary count fork (target
capped at the head's capacity) that trains the ``anchor_count`` loss, never in cropped
windows, never without the flag;
* array_decode: threshold is unchanged (no extra fork), count_peaks takes the count
head's N then the N best non-adjacent anchor peaks (the default), hybrid falls back
only when nothing fires; outputs stay bitwise identical alone and batched with other
requests.
"""
from __future__ import annotations
import random
import pytest
torch = pytest.importorskip("torch")
pytest.importorskip("transformers")
from test_anchors import NESTED, TEAM_CONTEXT # noqa: E402
from test_model import SNAPSHOT, tiny_lm # noqa: E402
from deg.infer import anchor_peaks, fill, fill_many, resolve_anchors # noqa: E402
from deg.model import MAX_COUNT, load_model, load_tokenizer # noqa: E402
from deg.request import FillOptions, RequestError, parse_fill_request # noqa: E402
@pytest.fixture(scope="module")
def model():
if SNAPSHOT is None:
pytest.skip("Qwen3.5-0.8B tokenizer not cached")
m = load_model(None, device="cpu", dtype=torch.float32, lm=tiny_lm(),
tokenizer=load_tokenizer(SNAPSHOT))
torch.manual_seed(1)
with torch.no_grad():
for p in m.heads.parameters():
p.normal_(0.0, 0.5)
return m
def _team_record() -> dict:
text = TEAM_CONTEXT
def sp(s):
p = text.index(s)
return [p, p + len(s)]
return {"uid": "t1", "schema": NESTED, "context": text,
"value": {"teams": [{"team": "Red", "members": ["Ann", "Bob"]},
{"team": "Blue", "members": ["Cy", "Dee", "Eve"]},
{"team": "Gold", "members": ["Fay"]}]},
"spans": {"/teams/0/team": sp("Red"), "/teams/1/team": sp("Blue"),
"/teams/2/team": sp("Gold"), "/teams/0/members/0": sp("Ann"),
"/teams/0/members/1": sp("Bob"), "/teams/1/members/0": sp("Cy"),
"/teams/1/members/1": sp("Dee"), "/teams/1/members/2": sp("Eve"),
"/teams/2/members/0": sp("Fay")},
"span_alternatives": {"/teams/0/team": [sp("Red")]}, "span_probs": {},
"decide_probs": {}}
# -------------------------------------------------------------------- count aux
def test_anchor_count_items_only_with_the_flag_and_capped(model):
from deg.train import deg_example, forward_losses, loss_counts
rec = _team_record()
off = deg_example(model, rec, random.Random(0), shuffle_p=0.0)
assert not any(it.kind == "count" for it in off.items)
on = deg_example(model, rec, random.Random(0), shuffle_p=0.0, anchor_count=True)
aux = {it.ptr: it for it in on.items if it.kind == "count"}
assert set(aux) == {"/teams", "/teams/0/members", "/teams/1/members",
"/teams/2/members"}
assert all(it.aux for it in aux.values())
assert aux["/teams"].target.index(1.0) == 3 and aux["/teams/1/members"].target.index(1.0) == 3
assert aux["/teams"].bounds == (0, MAX_COUNT)
assert loss_counts([on])["anchor_count"] == 4 and loss_counts([on])["count"] == 0
model.backbone.train()
try:
losses, _ = forward_losses(model, [on], 4096, 16)
finally:
model.backbone.eval()
assert "anchor_count" in losses and "count" not in losses
losses["anchor_count"].backward()
assert model.heads.count.weight.grad is not None
for p in [*model.heads.parameters(), *model.backbone.parameters()]:
p.grad = None
# more items than the head holds: the target saturates at its top class
many = "Results:\n" + "".join(f"- swimmer {k}\n" for k in range(70))
schema = {"type": "object", "properties": {"results": {"type": "array", "items": {
"type": "object", "properties": {"name": {"type": ["string", "null"]}},
"required": ["name"]}}}, "required": ["results"]}
spans = {}
pos = -1
for k in range(70):
pos = many.index(f"swimmer {k}\n", pos + 1)
spans[f"/results/{k}/name"] = [pos, pos + len(f"swimmer {k}")]
big = {"uid": "b", "schema": schema, "context": many,
"value": {"results": [{"name": f"swimmer {k}"} for k in range(70)]},
"spans": spans, "span_alternatives": {}, "span_probs": {}, "decide_probs": {}}
ex = deg_example(model, big, random.Random(0), anchor_count=True)
(c,) = [it for it in ex.items if it.kind == "count"]
assert c.target.index(1.0) == MAX_COUNT
def test_cropped_windows_never_train_the_count(model):
from deg.train import deg_example
rec = _team_record()
a = TEAM_CONTEXT.index("Team Blue")
ex = deg_example(model, rec, random.Random(0), window=(a, len(TEAM_CONTEXT)),
anchor_count=True)
assert any(it.kind == "anchor" for it in ex.items)
assert not any(it.kind == "count" for it in ex.items)
# ------------------------------------------------------------------- decoding
def test_anchor_peaks_nms_and_order():
ninf = float("-inf")
lg = torch.tensor([-1.0, 3.0, 2.9, 2.8, -2.0, -3.0, 1.0, -4.0, 0.5, -1.0, -1.0, ninf])
# local maxima best first: 1, 6, 8 (2 and 3 are the slope of the peak at 1)
assert anchor_peaks(lg, 1, radius=2) == [1]
assert anchor_peaks(lg, 3, radius=1) == [1, 6, 8]
# fewer peaks than asked: best remaining tokens outside every radius, never off-region
assert anchor_peaks(lg, 5, radius=1) == [1, 3, 6, 8, 10]
assert anchor_peaks(lg, 5, radius=2) == [1, 6, 9] # 8 is within 2 of 6; then 9
assert anchor_peaks(lg, 0) == []
assert anchor_peaks(torch.tensor([ninf, ninf]), 2) == []
def test_resolve_anchors_modes():
lg = torch.full((40,), -3.0)
lg[[5, 15, 25]] = torch.tensor([-0.5, -0.2, -1.0]) # nothing fires (all p < 0.5)
cnt = torch.full((MAX_COUNT + 1,), -5.0)
cnt[2] = 5.0
assert resolve_anchors(lg, cnt, 0, None, MAX_COUNT, "threshold") == []
assert resolve_anchors(lg, cnt, 0, None, MAX_COUNT, "count_peaks") == [5, 15]
assert resolve_anchors(lg, cnt, 0, None, MAX_COUNT, "hybrid") == [5, 15]
fired = lg.clone()
fired[25] = 2.0
assert resolve_anchors(fired, cnt, 0, None, MAX_COUNT, "hybrid") == [25] # threshold
assert resolve_anchors(fired, cnt, 0, None, MAX_COUNT, "count_peaks") == [15, 25]
assert resolve_anchors(lg, cnt, 0, 1, MAX_COUNT, "count_peaks") == [15] # maxItems
assert resolve_anchors(lg, cnt, 3, None, MAX_COUNT, "count_peaks") == [5, 15, 25]
zero = torch.full((MAX_COUNT + 1,), -5.0)
zero[0] = 5.0
assert resolve_anchors(lg, zero, 0, None, MAX_COUNT, "hybrid") == []
# saturated count head: the threshold count when it is larger
sat = torch.full((MAX_COUNT + 1,), -5.0)
sat[MAX_COUNT] = 5.0
many = torch.full((400,), -3.0)
many[::5] = 2.0
assert len(resolve_anchors(many, sat, 0, None, MAX_COUNT, "count_peaks")) == 80
def test_array_decode_option_parses():
base = {"context": "x", "schema": {"type": "object", "properties": {
"a": {"type": ["string", "null"]}}, "required": ["a"]}}
assert parse_fill_request(base).options.array_decode == "count_peaks"
for mode in ("threshold", "count_peaks", "hybrid"):
req = parse_fill_request({**base, "options": {"array_decode": mode}})
assert req.options.array_decode == mode
with pytest.raises(RequestError):
parse_fill_request({**base, "options": {"array_decode": "peaks"}})
def test_fill_array_decode_modes_and_independence(model, monkeypatch):
from deg.model import DegModel
seen: list[int] = []
real = DegModel.answer_many
def spy(self, trunks, queries):
seen.append(sum(q.kind == "count" for _, q in queries))
return real(self, trunks, queries)
monkeypatch.setattr(DegModel, "answer_many", spy)
text = TEAM_CONTEXT
calls = {}
for mode in ("threshold", "count_peaks", "hybrid"):
seen.clear()
r = fill(model, text, NESTED, FillOptions(array_decode=mode))
calls[mode] = (r, sum(seen))
# threshold issues no count fork; the others add one per anchored array instance
assert calls["threshold"][1] == 0 and calls["count_peaks"][1] >= 1
assert calls["hybrid"][1] >= 1
other = ("Invoice 7 for Dana: 3 lamps.", {"type": "object", "properties": {
"who": {"type": ["string", "null"]}}, "required": ["who"]}, None)
for mode in ("count_peaks", "hybrid"):
alone = calls[mode][0]
batched = fill_many(model, [other, (text, NESTED, FillOptions(array_decode=mode)),
other])
assert batched[1] == alone
again = fill(model, text, NESTED, FillOptions(array_decode=mode))
assert again == alone
assert fill(model, text, NESTED) == calls["count_peaks"][0] # the default
|