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