| """Filler augmentation for long-context anchor training (SPEC.md §12.4), on a tiny random |
| Qwen3.5 (CPU): a short record embedded in unrelated filler text keeps every gold span on |
| its own text and every target of the plain example; negative (filler only) examples |
| train top-level anchored arrays as empty; the corpus uses train-split documents only |
| and refuses ExtractBench; micro-batches materialize filler draws before packing (the |
| memory planner sees real examples) within the context limit. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import random |
| import zlib |
| from pathlib import Path |
|
|
| import pytest |
|
|
| torch = pytest.importorskip("torch") |
| pytest.importorskip("transformers") |
|
|
| from test_anchors import NESTED, TEAM_CONTEXT |
| from test_array_decode import _team_record |
| from test_model import SNAPSHOT, tiny_lm |
|
|
| from deg.model import load_model, load_tokenizer |
|
|
| FILLER_WORDS = ("The quarterly filing discusses liquidity, covenants and segment revenue " |
| "across fiscal periods; management notes risks in supply and demand. ") |
|
|
|
|
| @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 _val(doc_id: str) -> bool: |
| from deg.datagen.v3 import VAL_MODULUS |
|
|
| return zlib.crc32(doc_id.encode()) % VAL_MODULUS == 0 |
|
|
|
|
| @pytest.fixture() |
| def corpus(tmp_path: Path): |
| from deg.train import FillerCorpus |
|
|
| (tmp_path / "text").mkdir() |
| rows = [] |
| val_id = next(f"doc-{i}" for i in range(10_000) if _val(f"doc-{i}")) |
| for i, doc_id in enumerate(["doc-a", "doc-b", val_id]): |
| body = (f"[{doc_id}] " if doc_id == val_id else "") + FILLER_WORDS * (40 + 10 * i) |
| tp = tmp_path / "text" / f"{doc_id}.txt" |
| tp.write_text(body) |
| rows.append({"id": doc_id, "text_path": str(tp)}) |
| (tmp_path / "docs.jsonl").write_text("".join(json.dumps(r) + "\n" for r in rows)) |
| c = FillerCorpus(tmp_path) |
| assert {d for d, _ in c.docs} == {"doc-a", "doc-b"} |
| return c |
|
|
|
|
| def _targets(ex, text: str): |
| """Every training target as (kind, ptr, gold text), independent of token positions.""" |
| def at(i: int, j: int) -> str: |
| return text[ex.ctx.offsets[i][0]:ex.ctx.offsets[j][1]] |
|
|
| out = [] |
| for it in ex.items: |
| if it.kind == "anchor": |
| out.append(("anchor", it.ptr, tuple(at(t, t) for t in it.anchors))) |
| elif it.kind == "extract": |
| out.append(("extract", it.ptr, tuple( |
| at(a, b) if a >= 0 else None for a, b in zip(it.starts, it.ends)))) |
| else: |
| out.append((it.kind, it.ptr, tuple(it.target or ()), it.aux)) |
| return sorted(out, key=repr) |
|
|
|
|
| |
|
|
|
|
| def test_corpus_refuses_extractbench(tmp_path): |
| from deg.train import FillerCorpus |
|
|
| d = tmp_path / "ExtractBench" |
| d.mkdir() |
| (d / "docs.jsonl").write_text("") |
| with pytest.raises(ValueError): |
| FillerCorpus(d) |
|
|
|
|
| def test_filler_record_keeps_every_target_on_its_text(model, corpus): |
| from deg.train import deg_example, filler_record, record_text |
|
|
| rec = _team_record() |
| plain = deg_example(model, rec, random.Random(0), shuffle_p=0.0, anchor_count=True) |
| want = _targets(plain, record_text(rec)) |
| assert any(t[0] == "anchor" for t in want) and any(t[0] == "count" for t in want) |
| seen_offsets = set() |
| for seed in range(6): |
| rng = random.Random(seed) |
| aug = filler_record(model, rec, corpus, rng, 200) |
| text = aug["context"] |
| assert TEAM_CONTEXT in text and len(text) > len(TEAM_CONTEXT) + 200 |
| shift = text.index(TEAM_CONTEXT) |
| seen_offsets.add(shift) |
| for k, (a, b) in aug["spans"].items(): |
| ra, rb = rec["spans"][k] |
| assert (a - shift, b - shift) == (ra, rb) and text[a:b] == TEAM_CONTEXT[ra:rb] |
| ex = deg_example(model, aug, random.Random(0), shuffle_p=0.0, anchor_count=True) |
| assert len(ex.ctx.ids) > len(plain.ctx.ids) + 150 |
| assert _targets(ex, text) == want |
| assert len(seen_offsets) > 2 |
|
|
|
|
| def test_negative_example_trains_top_level_anchors_empty(model, corpus): |
| from deg.train import negative_example |
|
|
| ex = negative_example(model, _team_record(), corpus, random.Random(0), 150, |
| anchor_count=True) |
| assert ex is not None and ex.aug == "negative" |
| assert 140 <= len(ex.ctx.offsets) <= 151 |
| anchors = [it for it in ex.items if it.kind == "anchor"] |
| counts = [it for it in ex.items if it.kind == "count"] |
| assert [it.ptr for it in anchors] == ["/teams"] and anchors[0].anchors == [] |
| assert len(counts) == 1 and counts[0].aux and counts[0].target.index(1.0) == 0 |
| assert not any(it.kind in ("extract", "choice", "score", "noul") for it in ex.items) |
| min1 = {**NESTED, "properties": {"teams": {**NESTED["properties"]["teams"], |
| "minItems": 1}}} |
| assert negative_example(model, {**_team_record(), "schema": min1}, corpus, |
| random.Random(0), 150) is None |
|
|
|
|
| def test_filler_draws_materialize_before_packing_within_the_limit(model, corpus): |
| from deg.train import ( |
| FillerDraw, |
| Source, |
| batches, |
| deg_example, |
| filler_example, |
| ) |
|
|
| rec = _team_record() |
| ex = deg_example(model, rec, random.Random(0), shuffle_p=0.0, anchor_count=True) |
| ex.rec = rec |
| src = Source(Path("x.jsonl")) |
| src.pool = [ex] |
| max_ctx = 300 |
| seen: list[list] = [] |
|
|
| def fits(batch): |
| seen.append(list(batch)) |
| return True |
|
|
| def augment(e, r): |
| return FillerDraw(e, r.randint(len(e.ctx.ids) + 64, max_ctx), r.random() < 0.3) |
|
|
| def materialize(d, r): |
| return filler_example(model, d, corpus, r, max_ctx, anchor_count=True) |
|
|
| gen = batches([src], None, random.Random(0), 10_000, 0.0, 4, draw=8, |
| materialize=materialize, fits=fits, augment=augment) |
| kinds = set() |
| for _ in range(6): |
| for e in next(gen): |
| assert len(e.ctx.ids) <= max_ctx |
| kinds.add(e.aug) |
| assert kinds == {"filler", "negative"} |
| assert seen and all(not isinstance(e, FillerDraw) for b in seen for e in b) |
|
|
|
|
| def test_filler_example_falls_back_to_the_plain_example(model, corpus): |
| from deg.train import FillerDraw, deg_example, filler_example |
|
|
| rec = _team_record() |
| ex = deg_example(model, rec, random.Random(0)) |
| ex.rec = rec |
| |
| got = filler_example(model, FillerDraw(ex, len(ex.ctx.ids) + 500), corpus, |
| random.Random(0), len(ex.ctx.ids) + 3) |
| assert got is ex |
|
|
