diff --git a/app/harness/pools.py b/app/harness/pools.py index bffed72..4e9b5b4 100644 --- a/app/harness/pools.py +++ b/app/harness/pools.py @@ -2,8 +2,8 @@ 三池切分对应训练循环中的 DataLoader 阶段——从题目全集中按 test -> validation -> diagnosis 的顺序 progressive exclusion, -保证 question_id 互斥。test 池用自然分布(correct_ratio=None), -验证池/诊断池按对错比例分层采样。 +以 unit 为原子保证 unit_id 互斥(AR 孪生对两题永不被劈到不同池)。 +test 池用自然分布(correct_ratio=None),验证池/诊断池按对错比例分层采样。 """ from __future__ import annotations @@ -17,6 +17,7 @@ from typing import TYPE_CHECKING from loguru import logger +from app.harness.question_units import build_units, flatten_units from app.question_gen import stratified_sample from core.types import GeneratedQuestion, PoolConfig @@ -25,6 +26,7 @@ if TYPE_CHECKING: from app.harness.config import RunConfig from app.ports import PoolStrategy + from core.types import QuestionUnit @dataclass @@ -70,11 +72,15 @@ def build_pools( 冻结的三池 Pools。 关键实现细节: - 切分顺序 test -> validation -> diagnosis;后两步从剩余题中采样以保证 - question_id 互斥。test 池用 correct_ratio=None 的自然分布采样。 + 切分顺序 test -> validation -> diagnosis;后两步从剩余单元中采样以保证 + unit_id 互斥。test 池用 correct_ratio=None 的自然分布采样。以 unit 为采样 + 原子(pair 计 1 个 unit),孪生对两题永不被劈到不同池;size/correct_ratio + 按 unit 计数,single-only 输入下 unit 与 question 一一对应,行为完全不变。 """ + units = build_units(questions) + test = _sample_excluding( - questions, + units, set(), correctness, size=test_cfg["size"], @@ -83,12 +89,12 @@ def build_pools( seed=test_cfg.get("seed", 0), min_per_class=None, ) - selected_ids = {q.question_id for q in test} + selected_units = {q.unit_id for q in test} - validation = _sample_excluding(questions, selected_ids, correctness, **val_cfg) - selected_ids |= {q.question_id for q in validation} + validation = _sample_excluding(units, selected_units, correctness, **val_cfg) + selected_units |= {q.unit_id for q in validation} - diagnosis = _sample_excluding(questions, selected_ids, correctness, **diag_cfg) + diagnosis = _sample_excluding(units, selected_units, correctness, **diag_cfg) val_correct = sum(1 for q in validation if correctness.get(q.question_id)) baseline_val_accuracy = val_correct / len(validation) if validation else 0.0 @@ -167,26 +173,70 @@ class GlobalPoolStrategy: ) +def _unit_correct(unit: QuestionUnit, correctness: dict[str, bool]) -> bool: + """单元级正确性:成员全部答对才算对(缺失按 False,宽松口径)。 + + 参数: + unit: 目标单元(single 1 题,pair 2 题)。 + correctness: question_id -> 基线是否答对。 + + 返回: + pair 走双向 AND、single 即单题正确性;任一成员缺失或答错即 False。 + """ + return all(correctness.get(q.question_id, False) for q in unit.questions) + + +def _assert_correctness_complete( + units: list[QuestionUnit], + correctness: dict[str, bool], +) -> None: + """校验 correctness 覆盖所有单元成员题(含 pair 两题),缺失即 fail-fast。 + + 参数: + units: 待校验单元列表。 + correctness: question_id -> 基线是否答对。 + + 异常: + ValueError: correctness 中缺少某些 question_id。 + """ + missing = [ + q.question_id for u in units for q in u.questions if q.question_id not in correctness + ] + if missing: + raise ValueError(f"correctness 缺失 {len(missing)} 题: {missing[:5]}") + + def _sample_excluding( - questions: list[GeneratedQuestion], - exclude_ids: set[str], + units: list[QuestionUnit], + exclude_unit_ids: set[str], correctness: dict[str, bool], **cfg: object, ) -> list[GeneratedQuestion]: - """排除已选 question_id 后,按 cfg 对剩余题做分层采样。 + """排除已选 unit 后,以 unit 为原子按 cfg 分层采样,返回展开后的逐题列表。 + + 每个单元以其首题作为分层采样的代表参与 stratified_sample,correct_ratio / + size 因此按 unit 计数(pair 计 1 个 unit);命中的单元整体展开,孪生对两题 + 永远同进同出。single-only 输入下 unit 与 question 一一对应、顺序不变,采样 + 结果与逐题采样完全一致。 参数: - questions: 题目全集。 - exclude_ids: 已被其他池选走的 question_id,从候选中剔除以保证三池互斥。 - correctness: question_id -> 基线是否答对。 + units: 单元全集(single 单封、pair 成对聚合)。 + exclude_unit_ids: 已被其他池选走的 unit_id,从候选中剔除以保证三池互斥。 + correctness: question_id -> 基线是否答对;单元级正确性取成员的 AND + (缺失按 False,与 stratified_sample 的宽松口径一致)。 cfg: 透传给 stratified_sample 的采样配置 (size/correct_ratio/task_types[/seed/min_per_class])。 返回: - 采样后的题目列表。 + 采样命中单元展开后的题目列表。 """ - pool = [q for q in questions if q.question_id not in exclude_ids] - return stratified_sample(pool, correctness, **cfg) + candidates = [u for u in units if u.unit_id not in exclude_unit_ids] + rep_to_unit = {u.questions[0].question_id: u for u in candidates} + reps = [u.questions[0] for u in candidates] + unit_correct = {u.questions[0].question_id: _unit_correct(u, correctness) for u in candidates} + sampled_reps = stratified_sample(reps, unit_correct, **cfg) + sampled_units = [rep_to_unit[rep.question_id] for rep in sampled_reps] + return flatten_units(sampled_units) def _q_to_dict(q: GeneratedQuestion) -> dict: @@ -604,14 +654,14 @@ class PerCategoryPoolStrategy: rng = random.Random(config.seed) for task_type in sorted(groups.keys()): - train, val = self._split_one_category( - groups[task_type], + train_units, val_units = self._split_one_category( + build_units(groups[task_type]), correctness, config.train_ratio, rng, ) - all_train.extend(train) - all_val.extend(val) + all_train.extend(flatten_units(train_units)) + all_val.extend(flatten_units(val_units)) # Phase 4: test 池(从外部目录加载,无则空;按 task_types 过滤) test: list[GeneratedQuestion] = [] @@ -759,48 +809,48 @@ class PerCategoryPoolStrategy: def _split_one_category( self, - questions: list[GeneratedQuestion], + units: list[QuestionUnit], correctness: dict[str, bool], train_ratio: float, rng: random.Random, - ) -> tuple[list[GeneratedQuestion], list[GeneratedQuestion]]: - """单类别 correctness 分层划分。 + ) -> tuple[list[QuestionUnit], list[QuestionUnit]]: + """单类别 correctness 分层划分,以 unit 为原子(pair 计 1 个 unit)。 + + 孪生对两题作为一个整体落入 train 或 val,绝不被拆散;single-only 输入下 + unit 与 question 一一对应、rng 消耗量不变,划分结果与逐题划分完全一致。 参数: - questions: 单类别全部题目。 - correctness: question_id -> 基线是否答对。 - train_ratio: train 占总量的比例。 + units: 单类别全部单元。 + correctness: question_id -> 基线是否答对;单元级正确性取成员的 AND。 + train_ratio: train 占单元总量的比例。 rng: 随机数生成器(保证跨类别可复现)。 返回: - (train, val) 题目列表元组,两池互斥且总量 == len(questions)。 + (train_units, val_units) 单元列表元组,两侧互斥且总量 == len(units)。 异常: ValueError: correctness 中缺少某些 question_id。 """ - n_total = len(questions) + n_total = len(units) n_train = round(n_total * train_ratio) n_val = n_total - n_train - # 校验 correctness 完整性 - missing = [q.question_id for q in questions if q.question_id not in correctness] - if missing: - raise ValueError(f"correctness 缺失 {len(missing)} 题: {missing[:5]}") + _assert_correctness_complete(units, correctness) - correct_qs = [q for q in questions if correctness[q.question_id]] - wrong_qs = [q for q in questions if not correctness[q.question_id]] - n_correct = len(correct_qs) + correct_units = [u for u in units if _unit_correct(u, correctness)] + wrong_units = [u for u in units if not _unit_correct(u, correctness)] + n_correct = len(correct_units) # 全 correct 或全 wrong -> 退化为非分层随机划分 if n_correct == 0 or n_correct == n_total: label = "全部正确" if n_correct == n_total else "全部错误" logger.warning( - "类别 {} {} ({} 题),退化为非分层随机划分", - questions[0].task_type, + "类别 {} {} ({} 单元),退化为非分层随机划分", + units[0].task_type, label, n_total, ) - shuffled = list(questions) + shuffled = list(units) rng.shuffle(shuffled) return shuffled[:n_train], shuffled[n_train:] @@ -808,11 +858,11 @@ class PerCategoryPoolStrategy: train_correct = math.floor(n_correct * n_train / n_total) train_wrong = n_train - train_correct - rng.shuffle(correct_qs) - rng.shuffle(wrong_qs) + rng.shuffle(correct_units) + rng.shuffle(wrong_units) - train = correct_qs[:train_correct] + wrong_qs[:train_wrong] - val = correct_qs[train_correct:] + wrong_qs[train_wrong:] + train = correct_units[:train_correct] + wrong_units[:train_wrong] + val = correct_units[train_correct:] + wrong_units[train_wrong:] assert len(train) == n_train, f"train 数量不匹配: {len(train)} != {n_train}" assert len(val) == n_val, f"val 数量不匹配: {len(val)} != {n_val}" @@ -847,15 +897,15 @@ class PerCategoryPoolStrategy: rng = random.Random(config.seed) for task_type in sorted(groups.keys()): - train, val = self._split_one_category( - groups[task_type], + train_units, val_units = self._split_one_category( + build_units(groups[task_type]), correctness, config.train_ratio, rng, ) result[task_type] = { - "train": [q.question_id for q in train], - "val": [q.question_id for q in val], + "train": [q.question_id for q in flatten_units(train_units)], + "val": [q.question_id for q in flatten_units(val_units)], } return result diff --git a/tests/unit/test_pools_pair_atomic.py b/tests/unit/test_pools_pair_atomic.py new file mode 100644 index 0000000..8567a9f --- /dev/null +++ b/tests/unit/test_pools_pair_atomic.py @@ -0,0 +1,116 @@ +"""pair 原子性回归测试:三池切分不得把孪生对劈到不同池。 + +覆盖 pools.py 的两条切分路径: +- GlobalPoolStrategy 路径(build_pools 的 test->val->diag progressive exclusion); +- PerCategoryPoolStrategy 路径(_split_one_category 的 train/val 分层划分)。 + +核心不变量:同一 pair_id 的两题必落在同一个池(或同一 split),绝不被拆散。 +""" + +from __future__ import annotations + +from app.harness.pools import PerCategoryPoolStrategy, build_pools +from core.types import GeneratedQuestion, PoolConfig + + +def _pair(pid: str) -> list[GeneratedQuestion]: + """构造一个合法孪生对(original + mirror),共享 pair_id / unit_id / flip_axis。 + + 参数: + pid: 该孪生对的共享标识。 + + 返回: + 含 original 与 mirror 两条题目的列表。 + """ + base = { + "video_id": "v", + "task_type": "Action Reasoning", + "question": "?", + "options": ("A. a", "B. b", "C. c", "D. d"), + "answer": "A", + "source_nodes": ("n",), + "difficulty": "hard", + "pair_id": pid, + "flip_axis": "before_after", + } + return [ + GeneratedQuestion(question_id=f"{pid}_o", question_role="pair_original", **base), + GeneratedQuestion(question_id=f"{pid}_m", question_role="pair_mirror", **base), + ] + + +def test_pair_never_split_across_pools() -> None: + """build_pools(Global 路径)三池切分后,任一 pair_id 只出现在一个池。""" + qs = [q for pid in [f"p{i:02d}" for i in range(12)] for q in _pair(pid)] + cfg = { + "size": 4, + "correct_ratio": None, + "task_types": None, + "seed": 7, + "min_per_class": None, + } + pools = build_pools( + qs, + correctness={}, + diag_cfg=cfg, + val_cfg=cfg, + test_cfg={"size": 4, "seed": 7}, + baseline_run_id="b", + ) + + by_name = { + "diagnosis": pools.diagnosis, + "validation": pools.validation, + "test": pools.test, + } + loc: dict[str, set[str]] = {} + for name, pool in by_name.items(): + for q in pool: + loc.setdefault(q.pair_id, set()).add(name) + + split = {pid: names for pid, names in loc.items() if len(names) > 1} + assert not split, f"pair 被劈到多个池: {split}" + + # 每个被选中的 pair 必须两题齐全(同池内成对),不得只落单题 + per_pool_pair_count: dict[tuple[str, str], int] = {} + for name, pool in by_name.items(): + for q in pool: + key = (name, q.pair_id) + per_pool_pair_count[key] = per_pool_pair_count.get(key, 0) + 1 + assert all(c == 2 for c in per_pool_pair_count.values()), ( + f"存在池内落单的 pair 成员: {per_pool_pair_count}" + ) + + +def test_pair_never_split_across_train_val() -> None: + """PerCategoryPoolStrategy 路径:孪生对不得被 train/val 划分劈开。""" + qs = [q for pid in [f"p{i:02d}" for i in range(15)] for q in _pair(pid)] + # 让部分 pair 正确、部分错误,触发分层划分(非退化随机) + correctness: dict[str, bool] = {} + for i, pid in enumerate(f"p{i:02d}" for i in range(15)): + val = i < 9 + correctness[f"{pid}_o"] = val + correctness[f"{pid}_m"] = val + + config = PoolConfig( + task_types=None, + seed=42, + baseline_run_id="b", + diag_size=0, + diag_correct_ratio=0.0, + val_size=0, + val_correct_ratio=0.0, + test_size=0, + eval_min_per_class=0, + train_ratio=2 / 3, + test_questions_dir=None, + ) + pools = PerCategoryPoolStrategy().build(qs, correctness, config) + + loc: dict[str, set[str]] = {} + for name, pool in (("diagnosis", pools.diagnosis), ("validation", pools.validation)): + for q in pool: + loc.setdefault(q.pair_id, set()).add(name) + + split = {pid: names for pid, names in loc.items() if len(names) > 1} + assert not split, f"pair 被 train/val 划分劈开: {split}"