"""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