File size: 5,857 Bytes
f8dbde1 | 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 | """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) # type: ignore[arg-type]
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():
# the run-2b record that stalled a step for 20+ min: 7590 context tokens, 1153 rows
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) # item order kept
seen.add(tuple(order))
assert len(ex.items) == n_items and ex.subsampled == 0 # pooled example untouched
assert len(seen) > 40 # a fresh subset per draw
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) # a fork longer than the room still trains
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 # uniform among equals
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) # type: ignore[arg-type]
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() # fits the room: only the decide cap applies
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
|