| """Per-draw fork cap (cap_forks): a record with ~1k field rows and generate forks that |
| repeat its whole known-value block costs at most one trunk pass plus about the room in |
| fork tokens per step, a fresh random subset each draw, whole-array forks always kept.""" |
|
|
| from __future__ import annotations |
|
|
| import random |
| from collections import Counter |
| from pathlib import Path |
| from types import SimpleNamespace |
|
|
| from deg.train import ( |
| Example, |
| Item, |
| Source, |
| batches, |
| cap_decide_forks, |
| cap_forks, |
| field_of, |
| fork_room, |
| ) |
|
|
|
|
| def _ex(ctx: int, extract: int = 0, anchors: int = 0, counts: int = 0, gen: int = 0, |
| gen_len: int = 16_000, uid: str = "r") -> Example: |
| c = SimpleNamespace(ids=tuple(range(ctx)), offsets=tuple((i, i + 1) for i in range(ctx))) |
| items = ([Item("extract", tuple(range(80)), ptr=f"/e{i}") for i in range(extract)] |
| + [Item("anchor", tuple(range(120)), ptr=f"/a{i}") for i in range(anchors)] |
| + [Item("count", tuple(range(110)), ptr=f"/a{i}", aux=True) |
| for i in range(counts)] |
| + [Item("generate", tuple(range(gen_len)), prompt_len=gen_len - 8, ptr=f"/g{i}") |
| for i in range(gen)]) |
| return Example("deg", uid, c, items) |
|
|
|
|
| def _tok(ex: Example) -> int: |
| return sum(len(it.seq) for it in ex.items) |
|
|
|
|
| def test_small_record_untouched(): |
| ex = _ex(500, extract=20) |
| assert cap_forks(ex, random.Random(0), 12000) is ex |
|
|
|
|
| def test_big_record_bounded_whole_array_forks_kept_pool_untouched(): |
| |
| ex = _ex(7590, extract=900, anchors=61, counts=0, gen=176) |
| n_items = len(ex.items) |
| room = fork_room(ex, 12000) |
| assert room == 6000 |
| whole = sum(len(it.seq) for it in ex.items if it.kind == "anchor") |
| biggest = max(len(it.seq) for it in ex.items) |
| seen = set() |
| for seed in range(50): |
| got = cap_forks(ex, random.Random(seed), 12000) |
| assert got is not ex and got.ctx is ex.ctx and got.uid == ex.uid |
| assert sum(it.kind == "anchor" for it in got.items) == 61 |
| assert _tok(got) <= room + whole + biggest |
| assert got.subsampled == n_items - len(got.items) > 0 |
| order = [ex.items.index(it) for it in got.items] |
| assert order == sorted(order) |
| seen.add(tuple(order)) |
| assert len(ex.items) == n_items and ex.subsampled == 0 |
| assert len(seen) > 40 |
|
|
|
|
| def test_every_fork_kind_trains_and_short_forks_uniformly(): |
| ex = _ex(4000, extract=200, anchors=5, counts=5, gen=10) |
| hits: Counter = Counter() |
| draws = 4000 |
| for seed in range(draws): |
| for it in cap_forks(ex, random.Random(seed), 12000).items: |
| hits[it.ptr if it.kind != "count" else "c" + it.ptr] += 1 |
| assert all(hits[f"/a{i}"] == draws and hits[f"c/a{i}"] == draws for i in range(5)) |
| gen = [hits[f"/g{i}"] for i in range(10)] |
| assert all(g > 0 for g in gen) |
| ext = [hits[f"/e{i}"] for i in range(200)] |
| mean = sum(ext) / len(ext) |
| assert 0 < mean < draws |
| assert max(ext) < 1.25 * mean and min(ext) > 0.75 * mean |
|
|
|
|
| def test_record_fork_tokens_override(): |
| ex = _ex(1000, extract=300) |
| got = cap_forks(ex, random.Random(0), 12000, record_fork_tokens=800) |
| assert _tok(got) <= 800 and len(got.items) == 10 |
|
|
|
|
| def test_batches_cap_plain_examples_per_draw(): |
| big = _ex(7590, extract=900, anchors=10, gen=50, uid="big") |
| src = Source(Path("/v3/records.jsonl")) |
| src.pool = [big] |
| gen = batches([src], None, random.Random(0), budget=12000, jev_frac=0.0, max_records=4, |
| draw=2, cap=lambda e, r: cap_forks(e, r, 12000)) |
| outs = [next(gen) for _ in range(5)] |
| for batch in outs: |
| assert all(e.uid == "big" and e.subsampled > 0 for e in batch) |
| assert all(_tok(e) <= 6000 + 1200 + 16000 for e in batch) |
| assert len(big.items) == 960 and big.subsampled == 0 |
|
|
|
|
| def _decide_ex() -> Example: |
| c = SimpleNamespace(ids=tuple(range(300)), offsets=tuple((i, i + 1) for i in range(300))) |
| items = ([Item("choice", tuple(range(40)), ptr=f"/rows/{i}/unit") for i in range(50)] |
| + [Item("noul", tuple(range(30)), ptr=f"/rows/{i}/audited") for i in range(20)] |
| + [Item("choice", tuple(range(40)), ptr="/kind"), |
| Item("score", tuple(range(40)), ptr="/risk")] |
| + [Item("extract", tuple(range(30)), ptr=f"/rows/{i}/value") for i in range(50)]) |
| return Example("deg", "d", c, items) |
|
|
|
|
| def test_field_of(): |
| assert field_of("/rows/12/cells/3/unit") == "/rows/*/cells/*/unit" |
| assert field_of("/kind") == "/kind" |
|
|
|
|
| def test_decide_forks_capped_per_field_fresh_subset_others_untouched(): |
| ex = _decide_ex() |
| assert cap_decide_forks(ex, random.Random(0), 0) is ex |
| seen = set() |
| for seed in range(30): |
| got = cap_decide_forks(ex, random.Random(seed), 2) |
| per = Counter(field_of(it.ptr) for it in got.items if it.kind != "extract") |
| assert per == {"/rows/*/unit": 2, "/rows/*/audited": 2, "/kind": 1, "/risk": 1} |
| assert sum(it.kind == "extract" for it in got.items) == 50 |
| assert got.subsampled == 48 + 18 |
| order = [ex.items.index(it) for it in got.items] |
| assert order == sorted(order) |
| seen.add(tuple(order)) |
| assert len(ex.items) == 122 and len(seen) > 20 |
|
|
|
|
| def test_cap_forks_applies_the_decide_cap_before_the_room(): |
| ex = _decide_ex() |
| got = cap_forks(ex, random.Random(0), 12000, decide_per_field=1) |
| assert sum(it.kind in ("choice", "noul", "score") for it in got.items) == 4 |
| assert cap_forks(ex, random.Random(0), 12000) is ex |
|
|
