iomgaa/Video-Tree-TRM5
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(pools): 三池切分以 unit 为原子，孪生对同池不被拆散

Browse Source 
build_pools/_sample_excluding 与 PerCategoryPoolStrategy._split_one_category/
build_incremental 两条切分路径均改为以 QuestionUnit 为采样原子：progressive
exclusion 互斥集合与 train/val 分层划分都按 unit_id 计数（pair 计 1 个 unit），
命中单元整体展开，AR 孪生对两题永不落入不同池/split。

复用 app.harness.question_units 的 build_units/flatten_units，不重写分组逻辑。
single-only 输入下 unit 与 question 一一对应、rng 消耗量不变，采样与划分结果
与逐题口径完全一致；抽出 _unit_correct/_assert_correctness_complete 两个 helper
将 _split_one_category 复杂度压回基线以下。

新增 tests/unit/test_pools_pair_atomic.py 覆盖两条路径的 pair 原子性回归。
This commit is contained in:
iomgaa
2026-07-15 06:09:14 -04:00
parent 5ef5f2b8b7
commit ddb9a44f75
2 changed files with 214 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/harness/pools.py
+98 -48
View File
@@ -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
tests/unit/test_pools_pair_atomic.py
+116
View File
@@ -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}"