| """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 |
| from test_model import SNAPSHOT, tiny_lm |
|
|
| from deg.infer import anchor_peaks, fill, fill_many, resolve_anchors |
| from deg.model import MAX_COUNT, load_model, load_tokenizer |
| from deg.request import FillOptions, RequestError, parse_fill_request |
|
|
|
|
| @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": {}} |
|
|
|
|
| |
|
|
|
|
| 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 |
| |
| 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) |
|
|
|
|
| |
|
|
|
|
| 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]) |
| |
| assert anchor_peaks(lg, 1, radius=2) == [1] |
| assert anchor_peaks(lg, 3, radius=1) == [1, 6, 8] |
| |
| assert anchor_peaks(lg, 5, radius=1) == [1, 3, 6, 8, 10] |
| assert anchor_peaks(lg, 5, radius=2) == [1, 6, 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]) |
| 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] |
| 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] |
| 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") == [] |
| |
| 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)) |
| |
| 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] |
|
|
|
|
