iomgaa
6c6fb576ee
核心算法保真 #3（断点续跑）：checkpoint 从逐题 question_id 改为存 unit_id 序列（孪生对折叠为单个 unit_id），恢复时 build_units + 按完整 unit 展开， 续跑后 pair 两成员同进同出、绝不被劈开。 - _batch_unit_ids/_batch_from_ids 对称折叠/展开，保序去重，纯非 AR 下 unit_id==question_id、与旧逐题序列逐字节一致。 - momentum 采样抽取为 _sample_momentum_candidates 纯函数，docstring 显式 记录 Phase 1 设计偏差：仅保证纯非 AR byte-identical，混格 momentum 不保证。 - 新增 test_checkpoint_pair（unit_id 落盘往返、pair 不拆）与 test_non_ar_byte_identical（pools→batching→checkpoint→momentum 端到端黄金）。
298 lines
12 KiB
Python
298 lines
12 KiB
Python
"""纯非 AR 题库端到端 byte-identical 黄金测试（Task 10）。
|
||
|
||
Phase 1 验收铁律：引入 QuestionUnit 后，纯 single 题库过 pools 抽样 → batching 分批
|
||
→ checkpoint 折叠/恢复 → momentum 采样，端到端结果与"引入 QuestionUnit 前"的逐题
|
||
旧算法逐字节一致。golden 用固定 seed 的确定性对照（忠实重实现旧逐题逻辑），非空断言。
|
||
|
||
momentum 设计偏差：momentum 采样在逐题粒度进行、不折叠 unit，故仅保证纯非 AR 的抽样
|
||
序列 byte-identical；Phase 1 不保证混格 momentum byte-identical（见 runner
|
||
_sample_momentum_candidates docstring）。本测试保守证明纯非 AR momentum 不漂。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
import random
|
||
|
||
from app.harness.batching import build_batches
|
||
from app.harness.pools import build_pools
|
||
from app.harness.question_units import build_units
|
||
from app.harness.runner import _batch_from_ids, _batch_unit_ids, _sample_momentum_candidates
|
||
from core.types import GeneratedQuestion
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 辅助构造
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _single(qid: str, task_type: str = "RETRIEVAL", video_id: str = "v1") -> GeneratedQuestion:
|
||
"""构造最小 single 题目。"""
|
||
return GeneratedQuestion(
|
||
question_id=qid,
|
||
video_id=video_id,
|
||
task_type=task_type,
|
||
question=f"q_{qid}",
|
||
options=("A. a", "B. b", "C. c", "D. d"),
|
||
answer="A",
|
||
source_nodes=("n1",),
|
||
difficulty="medium",
|
||
question_role="single",
|
||
)
|
||
|
||
|
||
def _pair(
|
||
pair_id: str, task_type: str = "AR", video_id: str = "v1"
|
||
) -> tuple[GeneratedQuestion, GeneratedQuestion]:
|
||
"""构造一个孪生对（original + mirror）。"""
|
||
common = {
|
||
"video_id": video_id,
|
||
"task_type": task_type,
|
||
"question": "?",
|
||
"options": ("A. a", "B. b", "C. c", "D. d"),
|
||
"answer": "A",
|
||
"source_nodes": ("n1",),
|
||
"difficulty": "medium",
|
||
"pair_id": pair_id,
|
||
"flip_axis": "before_after",
|
||
}
|
||
original = GeneratedQuestion(
|
||
question_id=f"{pair_id}_o", question_role="pair_original", **common
|
||
)
|
||
mirror = GeneratedQuestion(question_id=f"{pair_id}_m", question_role="pair_mirror", **common)
|
||
return original, mirror
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 引入 QuestionUnit 前的旧逐题 build_batches 忠实副本（单一 rng，无 unit 折叠）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _reference_select_mixed(
|
||
items: list[GeneratedQuestion],
|
||
correctness: dict[str, bool],
|
||
correct_ratio: float,
|
||
rng: random.Random,
|
||
) -> dict[str, list[GeneratedQuestion]]:
|
||
"""旧版 _select_mixed_by_task_type 的忠实副本（逐题、单一 rng）。"""
|
||
errors_by_type: dict[str, list[GeneratedQuestion]] = {}
|
||
correct_by_type: dict[str, list[GeneratedQuestion]] = {}
|
||
for q in items:
|
||
qid = q.question_id
|
||
if correctness.get(qid) is False:
|
||
errors_by_type.setdefault(q.task_type, []).append(q)
|
||
elif correctness.get(qid, False):
|
||
correct_by_type.setdefault(q.task_type, []).append(q)
|
||
if correct_ratio <= 0:
|
||
return errors_by_type
|
||
grouped: dict[str, list[GeneratedQuestion]] = {}
|
||
for task_type in sorted(errors_by_type):
|
||
errs = errors_by_type[task_type]
|
||
n_correct = round(len(errs) * correct_ratio / (1 - correct_ratio))
|
||
available = correct_by_type.get(task_type, [])
|
||
sampled = (
|
||
list(available) if len(available) <= n_correct else rng.sample(available, n_correct)
|
||
)
|
||
grouped[task_type] = errs + sampled
|
||
return grouped
|
||
|
||
|
||
def _reference_build_batches(
|
||
items: list[GeneratedQuestion],
|
||
correctness: dict[str, bool],
|
||
batch_size: int,
|
||
min_class_per_batch: int,
|
||
seed: int,
|
||
correct_ratio: float = 0.0,
|
||
) -> list[list[GeneratedQuestion]]:
|
||
"""引入 QuestionUnit 前的旧版 build_batches 忠实副本（逐题、单一 rng）。"""
|
||
rng = random.Random(seed)
|
||
grouped = _reference_select_mixed(items, correctness, correct_ratio, rng)
|
||
total = sum(len(g) for g in grouped.values())
|
||
if total == 0:
|
||
return []
|
||
nb = max(1, math.ceil(total / batch_size))
|
||
batches: list[list[GeneratedQuestion]] = [[] for _ in range(nb)]
|
||
small = {t: g for t, g in grouped.items() if len(g) <= min_class_per_batch}
|
||
large = {t: g for t, g in grouped.items() if len(g) > min_class_per_batch}
|
||
for t in sorted(small, key=lambda t: (-len(small[t]), t)):
|
||
group = small[t]
|
||
placed = False
|
||
for b in batches:
|
||
if len(b) + len(group) <= batch_size:
|
||
b.extend(group)
|
||
placed = True
|
||
break
|
||
if not placed:
|
||
batches.append(list(group))
|
||
nb_live = len(batches)
|
||
pointer = 0
|
||
for task_type in sorted(large):
|
||
group = list(large[task_type])
|
||
rng.shuffle(group)
|
||
for q in group:
|
||
for offset in range(nb_live):
|
||
idx = (pointer + offset) % nb_live
|
||
if len(batches[idx]) < batch_size:
|
||
batches[idx].append(q)
|
||
pointer = (idx + 1) % nb_live
|
||
break
|
||
return [b for b in batches if b]
|
||
|
||
|
||
def _ids(batches: list[list[GeneratedQuestion]]) -> list[list[str]]:
|
||
"""提取 batch 的 question_id 序列，便于逐字节对比。"""
|
||
return [[q.question_id for q in b] for b in batches]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# (1) 纯非 AR 走 size=1 unit：unit_id == question_id
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestPureSingleAreSizeOneUnits:
|
||
"""纯 single 输入折叠出的每个 unit 都是 size=1、unit_id 等于 question_id。"""
|
||
|
||
def test_all_units_size_one(self) -> None:
|
||
items = [_single(f"s{i}", task_type=f"t{i % 3}") for i in range(15)]
|
||
units = build_units(items)
|
||
assert len(units) == len(items)
|
||
assert all(u.kind == "single" and u.size == 1 for u in units)
|
||
assert [u.unit_id for u in units] == [q.question_id for q in items]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# (2) batching 逐字节一致（黄金对照）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestBatchingByteIdentical:
|
||
"""纯 single build_batches 与旧逐题算法逐字节一致。"""
|
||
|
||
def _assert(self, items, correctness, batch_size, min_cls, seed, ratio) -> None:
|
||
got, _ = build_batches(items, correctness, batch_size, min_cls, seed, ratio)
|
||
ref = _reference_build_batches(items, correctness, batch_size, min_cls, seed, ratio)
|
||
assert _ids(got) == _ids(ref)
|
||
|
||
def test_pure_errors(self) -> None:
|
||
items = [_single(f"q{i}", task_type=f"t{i % 3}") for i in range(20)]
|
||
correctness = {f"q{i}": False for i in range(20)}
|
||
self._assert(items, correctness, 5, 2, 42, 0.0)
|
||
|
||
def test_mixed_ratio(self) -> None:
|
||
items = [_single(f"q{i}", task_type=f"t{i % 4}") for i in range(40)]
|
||
correctness = {f"q{i}": (i % 3 == 0) for i in range(40)}
|
||
self._assert(items, correctness, 8, 3, 7, 0.5)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# (3) 端到端：pools 抽样 → batching → checkpoint 折叠/恢复 逐字节一致
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestEndToEndByteIdentical:
|
||
"""pools → batching → checkpoint unit_id round-trip 全链纯非 AR 逐字节稳定。"""
|
||
|
||
def _make_pool_questions(self) -> tuple[list[GeneratedQuestion], dict[str, bool]]:
|
||
items = [_single(f"s{i}", task_type=f"t{i % 3}") for i in range(60)]
|
||
correctness = {f"s{i}": (i % 2 == 0) for i in range(60)}
|
||
return items, correctness
|
||
|
||
def test_pools_deterministic_and_pure_single(self) -> None:
|
||
items, correctness = self._make_pool_questions()
|
||
cfg = {
|
||
"diag_cfg": {
|
||
"size": 18,
|
||
"correct_ratio": 0.5,
|
||
"task_types": None,
|
||
"seed": 5,
|
||
"min_per_class": None,
|
||
},
|
||
"val_cfg": {
|
||
"size": 12,
|
||
"correct_ratio": 0.5,
|
||
"task_types": None,
|
||
"seed": 5,
|
||
"min_per_class": None,
|
||
},
|
||
"test_cfg": {"size": 10, "seed": 5},
|
||
"baseline_run_id": "baseline",
|
||
}
|
||
p1 = build_pools(items, correctness, **cfg)
|
||
p2 = build_pools(items, correctness, **cfg)
|
||
# 确定性：同 seed 同输入池划分逐字节一致
|
||
assert [q.question_id for q in p1.diagnosis] == [q.question_id for q in p2.diagnosis]
|
||
# 诊断池纯 single：折叠出的 unit 与逐题一一对应
|
||
units = build_units(p1.diagnosis)
|
||
assert [u.unit_id for u in units] == [q.question_id for q in p1.diagnosis]
|
||
|
||
def test_batching_then_checkpoint_roundtrip(self) -> None:
|
||
items, correctness = self._make_pool_questions()
|
||
p = build_pools(
|
||
items,
|
||
correctness,
|
||
diag_cfg={
|
||
"size": 18,
|
||
"correct_ratio": 0.5,
|
||
"task_types": None,
|
||
"seed": 5,
|
||
"min_per_class": None,
|
||
},
|
||
val_cfg={
|
||
"size": 12,
|
||
"correct_ratio": 0.5,
|
||
"task_types": None,
|
||
"seed": 5,
|
||
"min_per_class": None,
|
||
},
|
||
test_cfg={"size": 10, "seed": 5},
|
||
baseline_run_id="baseline",
|
||
)
|
||
batches, _ = build_batches(p.diagnosis, p.correctness, 6, 2, seed=3, correct_ratio=0.5)
|
||
|
||
# checkpoint 折叠为 unit_id（纯 single 即 question_id）后恢复
|
||
epoch_batches = [_batch_unit_ids(b) for b in batches]
|
||
assert epoch_batches == _ids(batches) # 纯非 AR：unit_id 序列 == question_id 序列
|
||
rebuilt = [_batch_from_ids(p, ids) for ids in epoch_batches]
|
||
# 恢复的 batch 与原 batch 逐字节一致（inference 消费的题序不变）
|
||
assert _ids(rebuilt) == _ids(batches)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# (4) momentum 纯非 AR byte-identical（混格偏差已在 runner docstring 记录）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestMomentumPureNonARByteIdentical:
|
||
"""momentum 采样纯非 AR 与旧逐题 random.Random(epoch).sample 逐字节一致。"""
|
||
|
||
def test_matches_legacy_rng(self) -> None:
|
||
pool = [_single(f"s{i}", task_type="RETRIEVAL") for i in range(30)]
|
||
allowed = {"RETRIEVAL"}
|
||
epoch = 11
|
||
samples = 8
|
||
got = _sample_momentum_candidates(pool, allowed, samples, epoch)
|
||
|
||
candidates = [q for q in pool if q.task_type in allowed]
|
||
ref = random.Random(epoch).sample(candidates, samples)
|
||
assert [q.question_id for q in got] == [q.question_id for q in ref]
|
||
|
||
def test_ar_pairs_of_other_type_do_not_shift(self) -> None:
|
||
"""向池中加入其它题型的 AR pair，不改变非 AR momentum 的抽样序列。"""
|
||
singles = [_single(f"s{i}", task_type="RETRIEVAL") for i in range(30)]
|
||
allowed = {"RETRIEVAL"}
|
||
epoch = 11
|
||
samples = 8
|
||
base = _sample_momentum_candidates(singles, allowed, samples, epoch)
|
||
|
||
mixed = list(singles)
|
||
for k in range(4):
|
||
po, pm = _pair(f"p{k}", task_type="AR")
|
||
mixed.extend([po, pm])
|
||
after = _sample_momentum_candidates(mixed, allowed, samples, epoch)
|
||
assert [q.question_id for q in base] == [q.question_id for q in after]
|
||
|
||
def test_fewer_candidates_than_samples_returns_all(self) -> None:
|
||
pool = [_single(f"s{i}", task_type="RETRIEVAL") for i in range(3)]
|
||
got = _sample_momentum_candidates(pool, {"RETRIEVAL"}, 10, epoch=1)
|
||
assert {q.question_id for q in got} == {"s0", "s1", "s2"}
|