1183 lines
46 KiB
Python
1183 lines
46 KiB
Python
"""三池：held-out test + 验证 + 诊断，分层采样 + 冻结持久化。
|
||
|
||
三池切分对应训练循环中的 DataLoader 阶段——从题目全集中按
|
||
test -> validation -> diagnosis 的顺序 progressive exclusion，
|
||
以 unit 为原子保证 unit_id 互斥（AR 孪生对两题永不被劈到不同池）。
|
||
test 池用自然分布（correct_ratio=None），验证池/诊断池按对错比例分层采样。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import math
|
||
import os
|
||
import random
|
||
from collections import defaultdict
|
||
from dataclasses import dataclass, field
|
||
from typing import TYPE_CHECKING
|
||
|
||
from loguru import logger
|
||
|
||
from app.harness.question_units import build_units, flatten_units, unit_correctness
|
||
from app.question_gen import stratified_sample
|
||
from core.types import GeneratedQuestion, PoolConfig
|
||
|
||
if TYPE_CHECKING:
|
||
from pathlib import Path
|
||
|
||
from app.harness.config import RunConfig
|
||
from app.ports import PoolStrategy
|
||
from core.types import QuestionUnit
|
||
|
||
|
||
@dataclass
|
||
class Pools:
|
||
"""冻结的三池及其基线指标。
|
||
|
||
字段:
|
||
diagnosis: 诊断池（用于错误归因，对应 loss.backward）。
|
||
validation: 验证池（按类局部验证，每题型有保底样本）。
|
||
test: held-out 测试池（自然分布，用于最终无偏评估）。
|
||
baseline_run_id: 基线 run 标识。
|
||
baseline_val_accuracy: 基线在验证池上的准确率。
|
||
correctness: 三池所有题的 question_id -> 基线是否答对。
|
||
"""
|
||
|
||
diagnosis: list[GeneratedQuestion]
|
||
validation: list[GeneratedQuestion]
|
||
test: list[GeneratedQuestion]
|
||
baseline_run_id: str
|
||
baseline_val_accuracy: float
|
||
correctness: dict[str, bool] = field(default_factory=dict)
|
||
|
||
|
||
def build_pools(
|
||
questions: list[GeneratedQuestion],
|
||
correctness: dict[str, bool],
|
||
diag_cfg: dict,
|
||
val_cfg: dict,
|
||
test_cfg: dict,
|
||
baseline_run_id: str,
|
||
) -> Pools:
|
||
"""先抽 held-out test，再抽验证集，最后抽诊断池，三池互斥。
|
||
|
||
参数:
|
||
questions: 题目全集。
|
||
correctness: question_id -> 基线是否答对。
|
||
diag_cfg: 诊断池采样配置（size/correct_ratio/task_types[/seed]）。
|
||
val_cfg: 验证池采样配置，可含 min_per_class 做按类保底。
|
||
test_cfg: 测试池配置（size[/seed]）；走自然分布，不强制对错比与题型。
|
||
baseline_run_id: 基线 run 标识。
|
||
|
||
返回:
|
||
冻结的三池 Pools。
|
||
|
||
关键实现细节:
|
||
切分顺序 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(
|
||
units,
|
||
set(),
|
||
correctness,
|
||
size=test_cfg["size"],
|
||
correct_ratio=None,
|
||
task_types=None,
|
||
seed=test_cfg.get("seed", 0),
|
||
min_per_class=None,
|
||
)
|
||
selected_units = {q.unit_id for q in test}
|
||
|
||
validation = _sample_excluding(units, selected_units, correctness, **val_cfg)
|
||
selected_units |= {q.unit_id for q in validation}
|
||
|
||
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
|
||
return Pools(
|
||
diagnosis=diagnosis,
|
||
validation=validation,
|
||
test=test,
|
||
baseline_run_id=baseline_run_id,
|
||
baseline_val_accuracy=baseline_val_accuracy,
|
||
correctness={
|
||
q.question_id: correctness.get(q.question_id, False)
|
||
for q in test + validation + diagnosis
|
||
},
|
||
)
|
||
|
||
|
||
_VIDEO_ASSIGNMENT_LABELS = ("trainval", "test")
|
||
|
||
|
||
class InsufficientValSignal(Exception): # noqa: N818 领域名「验证信号不足」，非通用错误后缀更贴切
|
||
"""validation 池错题数不足以支撑可靠验证信号（如 McNemar 检验功效）时抛出。
|
||
|
||
fail loud（P5）：不静默兜底、不放宽阈值，直接暴露 val 池错题数与所需下限，
|
||
由调用方决定放大 val_ratio / 换 trainval 归属或调低 val_wrong_min。
|
||
"""
|
||
|
||
|
||
def split_by_video_assignment(
|
||
questions: list[GeneratedQuestion],
|
||
assignment: dict[str, str],
|
||
correctness: dict[str, bool],
|
||
val_ratio: float,
|
||
seed: int,
|
||
baseline_run_id: str = "",
|
||
val_wrong_min: int = 0,
|
||
wrong_tier_by_video: dict[str, int] | None = None,
|
||
) -> Pools:
|
||
"""按视频归属做原子切分：同一视频所有题绝不跨 trainval/test 池。
|
||
|
||
切分原子从 unit 提升为 **视频组**（同 video 的全部题同进同出），彻底杜绝
|
||
同视频多题散落不同池造成的内容泄漏。trainval 题集内部再以视频组为原子做
|
||
correctness 分层，切出 validation（占 val_ratio）与 diagnosis（其余）。
|
||
|
||
参数:
|
||
questions: 题目全集。
|
||
assignment: video_id -> "trainval" | "test" 归属字典（由选择器上游产出）。
|
||
correctness: question_id -> 基线是否答对；trainval 分层与验证池准确率均依赖它。
|
||
val_ratio: validation 占 trainval 视频组总数的比例，[0.0, 1.0]。
|
||
seed: 随机种子，保证视频组 shuffle 可复现。
|
||
baseline_run_id: 基线 run 标识；离线切分阶段可留空，由调用方回填。
|
||
val_wrong_min: validation 池最少错题数（默认 0 = 不检查，保持既有调用契约）。
|
||
> 0 时切分时保证（不足则从 diag 换入低 T2 错题组补足，耗尽 fail-loud，
|
||
见 InsufficientValSignal）。
|
||
wrong_tier_by_video: video_id -> 该视频错题中 T2(defect) 数量；透传给
|
||
_split_trainval_by_video_group 做 tier 感知 diag/val 分配，None 时退化为
|
||
原随机 shuffle。
|
||
|
||
返回:
|
||
冻结的三池 Pools：diagnosis/validation 仍是逐题 GeneratedQuestion 列表
|
||
（元素粒度不变，仅改变"哪些视频进哪个池"），test 为全部 test 题。
|
||
baseline_val_accuracy = validation 池正确率。
|
||
|
||
异常:
|
||
ValueError: assignment 缺失某题 video_id（fail-fast 不静默丢题）、
|
||
assignment 取值非法、correctness 缺失任一参与 Pools 的题（trainval 或
|
||
test）、或 val_ratio 越界。
|
||
InsufficientValSignal: val_wrong_min > 0 且 validation 池错题数 < val_wrong_min
|
||
（P5，验证信号不足以支撑可靠比较，直接报错而非静默放行）。
|
||
|
||
关键实现细节:
|
||
视频组 correctness 取组内全部题的 AND（组内均答对才记为 correct 组），
|
||
据此在 trainval 内做与 _split_one_category 同构的比例分层，但原子是视频组。
|
||
val_ratio 决定 validation 组数：val_correct = floor(n_correct * n_val / n_total)，
|
||
余额补 wrong 组，全 correct / 全 wrong 时退化为非分层随机划分。下游
|
||
gate_ladder（信息阶梯冷启动 2:1，核心算法保真 #5）消费的 unit 结构不变。
|
||
"""
|
||
if not 0.0 <= val_ratio <= 1.0:
|
||
raise ValueError(f"val_ratio 必须在 [0.0, 1.0]，实际 {val_ratio}")
|
||
|
||
trainval_qs, test_qs = _partition_by_video_assignment(questions, assignment, correctness)
|
||
|
||
diagnosis, validation = _split_trainval_by_video_group(
|
||
trainval_qs, correctness, val_ratio, random.Random(seed),
|
||
wrong_tier_by_video=wrong_tier_by_video,
|
||
val_wrong_min=val_wrong_min,
|
||
)
|
||
|
||
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
|
||
return Pools(
|
||
diagnosis=diagnosis,
|
||
validation=validation,
|
||
test=test_qs,
|
||
baseline_run_id=baseline_run_id,
|
||
baseline_val_accuracy=baseline_val_accuracy,
|
||
correctness={
|
||
q.question_id: correctness[q.question_id] for q in test_qs + validation + diagnosis
|
||
},
|
||
)
|
||
|
||
|
||
def _partition_by_video_assignment(
|
||
questions: list[GeneratedQuestion],
|
||
assignment: dict[str, str],
|
||
correctness: dict[str, bool],
|
||
) -> tuple[list[GeneratedQuestion], list[GeneratedQuestion]]:
|
||
"""校验归属字典并按 video 归属把题划成 (trainval_qs, test_qs)。
|
||
|
||
参数:
|
||
questions: 题目全集。
|
||
assignment: video_id -> "trainval" | "test" 归属字典。
|
||
correctness: question_id -> 基线是否答对；对全部参与 Pools 的题（trainval
|
||
与 test 双侧）强制完整。
|
||
|
||
返回:
|
||
(trainval_qs, test_qs) 逐题列表元组，划分依据每题的 video_id 归属。
|
||
|
||
异常:
|
||
ValueError: assignment 取值非法、缺失某题 video_id、或 correctness 缺失
|
||
任一参与 Pools 的题（trainval 或 test，fail-fast，不静默丢题）。
|
||
"""
|
||
_assert_valid_assignment(questions, assignment)
|
||
|
||
trainval_qs = [q for q in questions if assignment[q.video_id] == "trainval"]
|
||
test_qs = [q for q in questions if assignment[q.video_id] == "test"]
|
||
|
||
# Pools.correctness 会为 test + validation + diagnosis 全体写入基线对错，
|
||
# 故 test 侧同样必须有 correctness，缺失即报错而非静默兜底 False（P5）。
|
||
missing_correctness = [
|
||
q.question_id for q in trainval_qs + test_qs if q.question_id not in correctness
|
||
]
|
||
if missing_correctness:
|
||
raise ValueError(
|
||
f"correctness 缺失 {len(missing_correctness)} 道题: {missing_correctness[:5]}"
|
||
)
|
||
|
||
return trainval_qs, test_qs
|
||
|
||
|
||
def _assert_valid_assignment(
|
||
questions: list[GeneratedQuestion],
|
||
assignment: dict[str, str],
|
||
) -> None:
|
||
"""校验归属字典取值合法且覆盖全部题的 video_id，否则 fail-fast。
|
||
|
||
参数:
|
||
questions: 题目全集。
|
||
assignment: video_id -> "trainval" | "test" 归属字典。
|
||
|
||
异常:
|
||
ValueError: assignment 含非法取值，或缺失某题的 video_id。
|
||
"""
|
||
bad_labels = {v for v in assignment.values() if v not in _VIDEO_ASSIGNMENT_LABELS}
|
||
if bad_labels:
|
||
raise ValueError(
|
||
f"assignment 含非法归属值 {sorted(bad_labels)}，仅允许 {_VIDEO_ASSIGNMENT_LABELS}"
|
||
)
|
||
|
||
missing_videos = sorted({q.video_id for q in questions if q.video_id not in assignment})
|
||
if missing_videos:
|
||
raise ValueError(f"assignment 缺失 {len(missing_videos)} 个 video_id: {missing_videos[:5]}")
|
||
|
||
|
||
def _partition_video_groups_by_correctness(
|
||
groups: dict[str, list[GeneratedQuestion]],
|
||
correctness: dict[str, bool],
|
||
) -> tuple[list[str], list[str]]:
|
||
"""按视频组正确性把 video_id 分成 (correct_vids, wrong_vids)。
|
||
|
||
组正确性取组内全部题的 AND（组内均答对才记为 correct 组），排序保证确定性。
|
||
|
||
参数:
|
||
groups: video_id -> 该视频全部题列表。
|
||
correctness: question_id -> 基线是否答对（调用方已校验完整）。
|
||
|
||
返回:
|
||
(correct_vids, wrong_vids) 两个 video_id 列表，按 video_id 升序。
|
||
"""
|
||
correct_vids: list[str] = []
|
||
wrong_vids: list[str] = []
|
||
for vid in sorted(groups.keys()):
|
||
if all(correctness[q.question_id] for q in groups[vid]):
|
||
correct_vids.append(vid)
|
||
else:
|
||
wrong_vids.append(vid)
|
||
return correct_vids, wrong_vids
|
||
|
||
|
||
def _split_trainval_by_video_group(
|
||
trainval_qs: list[GeneratedQuestion],
|
||
correctness: dict[str, bool],
|
||
val_ratio: float,
|
||
rng: random.Random,
|
||
wrong_tier_by_video: dict[str, int] | None = None,
|
||
val_wrong_min: int = 0,
|
||
) -> tuple[list[GeneratedQuestion], list[GeneratedQuestion]]:
|
||
"""以视频组为原子对 trainval 题集做 correctness 分层，切出 (diagnosis, validation)。
|
||
|
||
参数:
|
||
trainval_qs: trainval 归属的全部题（correctness 已在调用方校验完整）。
|
||
correctness: question_id -> 基线是否答对；视频组正确性取组内全部题 AND。
|
||
val_ratio: validation 占视频组总数的比例。
|
||
rng: 随机数生成器，保证视频组 shuffle 可复现。
|
||
wrong_tier_by_video: video_id -> 该视频错题中 T2(defect) 的数量。提供时错题
|
||
视频组按 T2 含量升序进 val（T2 高的组保留在 diagnosis，把高价值缺陷信号
|
||
留给诊断），确定性排序取代随机 shuffle；None 时退化为原随机 shuffle。
|
||
val_wrong_min: validation 池最少错题数（切分时保证功效）。> 0 且初分 val 错题
|
||
不足时，从 diag 侧的错题组按 T2 升序换入 val 直到满足（每组至多移动一次），
|
||
耗尽仍不足则抛 InsufficientValSignal（fail loud，P5）。
|
||
|
||
返回:
|
||
(diagnosis, validation) 逐题列表元组；同一 video 的全部题整组落在同一侧，
|
||
两侧互斥且并集 == trainval_qs。
|
||
|
||
关键实现细节:
|
||
与 _split_one_category 同构：先按视频组 correctness 分正确组/错误组，按比例
|
||
把 n_val 个组分层落入 validation（全正确退化为非分层随机划分；全错误时若有
|
||
wrong_tier_by_video 仍按 T2 升序分配，否则随机划分），
|
||
再把选中组内所有题展开。视频组按 video_id 排序后再 shuffle，保证确定性。
|
||
"""
|
||
groups: dict[str, list[GeneratedQuestion]] = defaultdict(list)
|
||
for q in trainval_qs:
|
||
groups[q.video_id].append(q)
|
||
|
||
video_ids = sorted(groups.keys())
|
||
n_total = len(video_ids)
|
||
n_val = round(n_total * val_ratio)
|
||
|
||
correct_vids, wrong_vids = _partition_video_groups_by_correctness(groups, correctness)
|
||
n_correct = len(correct_vids)
|
||
|
||
if n_correct == 0 and wrong_tier_by_video is not None:
|
||
# 全部错误 + 有 tier 信号：按 T2 升序，低 T2 组优先进 val（保留高 T2 在 diag）
|
||
wrong_vids.sort(key=lambda v: (wrong_tier_by_video.get(v, 0), v))
|
||
val_vids = set(wrong_vids[:n_val])
|
||
elif n_correct == 0 or n_correct == n_total:
|
||
label = "全部正确" if n_correct == n_total else "全部错误"
|
||
logger.warning("trainval 视频组 {} ({} 组)，退化为非分层随机划分", label, n_total)
|
||
shuffled = list(video_ids)
|
||
rng.shuffle(shuffled)
|
||
val_vids = set(shuffled[:n_val])
|
||
else:
|
||
val_correct = math.floor(n_correct * n_val / n_total)
|
||
val_wrong = n_val - val_correct
|
||
rng.shuffle(correct_vids)
|
||
if wrong_tier_by_video is None:
|
||
rng.shuffle(wrong_vids)
|
||
else:
|
||
# T2 少的错题组优先进 val（保留 T2 高的组在 diag），确定性排序
|
||
wrong_vids.sort(key=lambda v: (wrong_tier_by_video.get(v, 0), v))
|
||
val_vids = set(correct_vids[:val_correct] + wrong_vids[:val_wrong])
|
||
|
||
if val_wrong_min > 0:
|
||
val_wrong_now = sum(
|
||
1 for v in val_vids for q in groups[v] if not correctness[q.question_id]
|
||
)
|
||
# diag 侧仍在的错题组，按 T2 升序（低价值优先移交 val）
|
||
diag_wrong_pool = sorted(
|
||
(v for v in wrong_vids if v not in val_vids),
|
||
key=lambda v: ((wrong_tier_by_video or {}).get(v, 0), v),
|
||
)
|
||
for v in diag_wrong_pool:
|
||
if val_wrong_now >= val_wrong_min:
|
||
break
|
||
val_vids.add(v)
|
||
val_wrong_now += sum(1 for q in groups[v] if not correctness[q.question_id])
|
||
if val_wrong_now < val_wrong_min:
|
||
raise InsufficientValSignal(
|
||
f"trainval 错题不足以让 val 达到 val_wrong_min={val_wrong_min}"
|
||
f"（修复后仅 {val_wrong_now}），请放大 val_ratio 或调整 trainval 归属。"
|
||
)
|
||
|
||
diagnosis = [q for q in trainval_qs if q.video_id not in val_vids]
|
||
validation = [q for q in trainval_qs if q.video_id in val_vids]
|
||
return diagnosis, validation
|
||
|
||
|
||
class GlobalPoolStrategy:
|
||
"""全局三分策略：test -> val -> diag progressive exclusion。
|
||
|
||
封装现有 build_pools 逻辑为 PoolStrategy 接口。
|
||
"""
|
||
|
||
def build(
|
||
self,
|
||
questions: list[GeneratedQuestion],
|
||
correctness: dict[str, bool],
|
||
config: PoolConfig,
|
||
*,
|
||
db_path: Path | None = None,
|
||
) -> Pools:
|
||
"""委托给现有 build_pools 函数。
|
||
|
||
参数:
|
||
questions: 题目全集。
|
||
correctness: question_id -> 基线是否答对。
|
||
config: 池构建统一配置。
|
||
|
||
返回:
|
||
冻结的三池 Pools。
|
||
"""
|
||
return build_pools(
|
||
questions,
|
||
correctness,
|
||
diag_cfg={
|
||
"size": config.diag_size,
|
||
"correct_ratio": config.diag_correct_ratio,
|
||
"task_types": list(config.task_types) if config.task_types else None,
|
||
"seed": config.seed,
|
||
"min_per_class": None,
|
||
},
|
||
val_cfg={
|
||
"size": config.val_size,
|
||
"correct_ratio": config.val_correct_ratio,
|
||
"task_types": list(config.task_types) if config.task_types else None,
|
||
"seed": config.seed,
|
||
"min_per_class": config.eval_min_per_class,
|
||
},
|
||
test_cfg={"size": config.test_size, "seed": config.seed},
|
||
baseline_run_id=config.baseline_run_id,
|
||
)
|
||
|
||
def build_incremental(
|
||
self,
|
||
new_task_types: list[str],
|
||
questions: list[GeneratedQuestion],
|
||
correctness: dict[str, bool],
|
||
config: PoolConfig,
|
||
) -> dict[str, dict[str, list[str]]]:
|
||
"""全局策略不支持增量。
|
||
|
||
异常:
|
||
NotImplementedError: 始终抛出。
|
||
"""
|
||
raise NotImplementedError(
|
||
"GlobalPoolStrategy 不支持增量构建，请使用 PerCategoryPoolStrategy。"
|
||
)
|
||
|
||
|
||
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(
|
||
units: list[QuestionUnit],
|
||
exclude_unit_ids: set[str],
|
||
correctness: dict[str, bool],
|
||
**cfg: object,
|
||
) -> list[GeneratedQuestion]:
|
||
"""排除已选 unit 后，以 unit 为原子按 cfg 分层采样，返回展开后的逐题列表。
|
||
|
||
候选单元展开为逐题列表后透传给 stratified_sample，后者内部重新 build_units
|
||
做单元原子采样：correct_ratio / size 按 unit 计数（pair 计 1 个 unit），单元级
|
||
正确性由 stratified_sample 内部对成员取 AND，命中的孪生对两题永远同进同出。
|
||
single-only 输入下 unit 与 question 一一对应、顺序不变，采样结果与逐题采样一致。
|
||
|
||
参数:
|
||
units: 单元全集（single 单封、pair 成对聚合）。
|
||
exclude_unit_ids: 已被其他池选走的 unit_id，从候选中剔除以保证三池互斥。
|
||
correctness: question_id -> 基线是否答对；单元级正确性由 stratified_sample
|
||
对成员取 AND（缺失按 False，宽松口径）。
|
||
cfg: 透传给 stratified_sample 的采样配置
|
||
（size/correct_ratio/task_types[/seed/min_per_class]）。
|
||
|
||
返回:
|
||
采样命中单元展开后的题目列表。
|
||
"""
|
||
candidates = [u for u in units if u.unit_id not in exclude_unit_ids]
|
||
return stratified_sample(flatten_units(candidates), correctness, **cfg)
|
||
|
||
|
||
def _q_to_dict(q: GeneratedQuestion) -> dict:
|
||
"""将 GeneratedQuestion 转为可序列化字典。
|
||
|
||
参数:
|
||
q: 题目对象。
|
||
|
||
返回:
|
||
包含全部字段的字典（options/source_nodes 从 tuple 转为 list）。
|
||
|
||
关键实现细节:
|
||
pair 四字段（pair_id/question_role/flip_axis/unit_id）必须写出——pools.json
|
||
是训练主回路读回题目的地方，漏写会让孪生对解冻后退化成孤儿 single。
|
||
"""
|
||
return {
|
||
"question_id": q.question_id,
|
||
"video_id": q.video_id,
|
||
"task_type": q.task_type,
|
||
"question": q.question,
|
||
"options": list(q.options),
|
||
"answer": q.answer,
|
||
"source_nodes": list(q.source_nodes),
|
||
"difficulty": q.difficulty,
|
||
"family": q.family,
|
||
"skill_target": q.skill_target,
|
||
"difficulty_steps": q.difficulty_steps,
|
||
"pair_id": q.pair_id,
|
||
"question_role": q.question_role,
|
||
"flip_axis": q.flip_axis,
|
||
"unit_id": q.unit_id,
|
||
}
|
||
|
||
|
||
def _dict_to_q(d: dict) -> GeneratedQuestion:
|
||
"""从字典恢复 GeneratedQuestion。
|
||
|
||
参数:
|
||
d: 由 _q_to_dict 产出的字典。
|
||
|
||
返回:
|
||
恢复的 GeneratedQuestion 实例（options/source_nodes 恢复为 tuple）。
|
||
|
||
关键实现细节:
|
||
pair 四字段用 .get 兼容旧 workspace 的 pools.json（无这些字段不崩，默认退化
|
||
为 single）——断点续跑铁律。unit_id 缺省时传 ""，交给 GeneratedQuestion
|
||
的 __post_init__ 回填为 pair_id 或 question_id，避免孤儿 single。
|
||
"""
|
||
return GeneratedQuestion(
|
||
question_id=d["question_id"],
|
||
video_id=d["video_id"],
|
||
task_type=d["task_type"],
|
||
question=d["question"],
|
||
options=tuple(d["options"]),
|
||
answer=d["answer"],
|
||
source_nodes=tuple(d.get("source_nodes", ())),
|
||
difficulty=d.get("difficulty", "medium"),
|
||
family=d.get("family"),
|
||
skill_target=d.get("skill_target"),
|
||
difficulty_steps=d.get("difficulty_steps"),
|
||
pair_id=d.get("pair_id"),
|
||
question_role=d.get("question_role", "single"),
|
||
flip_axis=d.get("flip_axis"),
|
||
unit_id=d.get("unit_id", ""),
|
||
)
|
||
|
||
|
||
def _atomic_write_json(path: Path, obj: object) -> None:
|
||
"""原子写 JSON：先写 <path>.tmp 再 os.replace，避免半截文件。
|
||
|
||
崩溃或并发写入时，直接 write_text 可能留下被截断的 JSON；本助手先把完整
|
||
内容写入同目录临时文件，再用同一文件系统上的原子 rename 替换目标，
|
||
保证读者只会看到旧完整文件或新完整文件。
|
||
|
||
参数:
|
||
path: 目标 JSON 文件路径。
|
||
obj: 可 json 序列化对象。
|
||
"""
|
||
tmp = path.with_suffix(path.suffix + ".tmp")
|
||
tmp.write_text(json.dumps(obj, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
os.replace(tmp, path)
|
||
|
||
|
||
def save_pools(
|
||
pools: Pools,
|
||
path: Path,
|
||
*,
|
||
split_mode: str = "global",
|
||
config: PoolConfig | None = None,
|
||
) -> None:
|
||
"""将三池及基线指标冻结为 JSON。
|
||
|
||
参数:
|
||
pools: 待冻结的三池。
|
||
path: 目标 JSON 文件路径。
|
||
split_mode: 池划分策略标记（"global" / "per_category"），写入 JSON 用于
|
||
加载时识别格式。
|
||
config: 池构建配置。per_category 模式下必须提供，用于写入 categories
|
||
元数据（seed, train_ratio, test_source）以支持增量追加和一致性校验。
|
||
|
||
异常:
|
||
ValueError: split_mode 为 "per_category" 但未提供 config。
|
||
"""
|
||
if split_mode == "per_category" and config is None:
|
||
raise ValueError("per_category 模式下 save_pools 必须提供 config 参数以写入元数据。")
|
||
|
||
data: dict = {
|
||
"split_mode": split_mode,
|
||
"baseline_run_id": pools.baseline_run_id,
|
||
"baseline_val_accuracy": pools.baseline_val_accuracy,
|
||
"correctness": pools.correctness,
|
||
"diagnosis": [_q_to_dict(q) for q in pools.diagnosis],
|
||
"validation": [_q_to_dict(q) for q in pools.validation],
|
||
"test": [_q_to_dict(q) for q in pools.test],
|
||
}
|
||
|
||
if split_mode == "per_category" and config is not None:
|
||
# 按 task_type 记录 train/val 的 qid 列表，用于增量追加和一致性校验
|
||
categories: dict[str, dict[str, list[str]]] = {}
|
||
diag_by_type: dict[str, list[str]] = defaultdict(list)
|
||
val_by_type: dict[str, list[str]] = defaultdict(list)
|
||
for q in pools.diagnosis:
|
||
diag_by_type[q.task_type].append(q.question_id)
|
||
for q in pools.validation:
|
||
val_by_type[q.task_type].append(q.question_id)
|
||
for task_type in sorted(set(diag_by_type) | set(val_by_type)):
|
||
categories[task_type] = {
|
||
"train": diag_by_type.get(task_type, []),
|
||
"val": val_by_type.get(task_type, []),
|
||
}
|
||
data["categories"] = categories
|
||
data["seed"] = config.seed
|
||
data["train_ratio"] = config.train_ratio
|
||
data["test_source"] = str(config.test_questions_dir) if config.test_questions_dir else None
|
||
|
||
_atomic_write_json(path, data)
|
||
|
||
|
||
def load_pools(path: Path) -> Pools:
|
||
"""从 JSON 恢复冻结的三池。
|
||
|
||
兼容新旧格式：有无 split_mode 字段都能加载。新格式（含 split_mode /
|
||
categories）的额外元数据在加载时忽略——Pools 对象只关心三池列表和标量。
|
||
|
||
参数:
|
||
path: 冻结的 pools.json 路径。
|
||
|
||
返回:
|
||
恢复的三池 Pools。
|
||
|
||
异常:
|
||
ValueError: 旧格式 pools.json（无 test 池）。
|
||
|
||
关键实现细节:
|
||
旧格式 pools.json（无 test 池）会以清晰的 ValueError 中止——本项目不做
|
||
向后兼容，也不为缺失字段填默认值。删除旧文件后 build_pools 会重新采样切分，
|
||
无需重新推理。
|
||
"""
|
||
d = json.loads(path.read_text(encoding="utf-8"))
|
||
if "test" not in d:
|
||
raise ValueError(
|
||
f"{path} 为旧格式 pools.json（缺 test 池），"
|
||
"请删除后重新切分（build_pools 会重新采样，无需重新推理）。"
|
||
)
|
||
return Pools(
|
||
diagnosis=[_dict_to_q(x) for x in d["diagnosis"]],
|
||
validation=[_dict_to_q(x) for x in d["validation"]],
|
||
test=[_dict_to_q(x) for x in d["test"]],
|
||
baseline_run_id=d["baseline_run_id"],
|
||
baseline_val_accuracy=d["baseline_val_accuracy"],
|
||
correctness=d["correctness"],
|
||
)
|
||
|
||
|
||
def _to_pool_config(config: RunConfig, baseline_run_id: str) -> PoolConfig:
|
||
"""从 RunConfig + 外部 baseline_run_id 提取 PoolConfig。
|
||
|
||
baseline_run_id 必须由调用方从 workspace manifest / seed.json 读取，
|
||
绝不能用 config.run_id（那是训练 run ID）。
|
||
|
||
参数:
|
||
config: 运行配置。
|
||
baseline_run_id: 基线 run 标识（来自 workspace manifest 或 seed.json）。
|
||
|
||
返回:
|
||
PoolConfig 实例。
|
||
"""
|
||
test_questions_dir: Path | None = None
|
||
if config.test_questions:
|
||
from app.harness.workspace import resolve_paths
|
||
|
||
paths = resolve_paths(config.workspace_dir)
|
||
test_questions_dir = paths.store_dir / "questions" / config.test_questions
|
||
|
||
return PoolConfig(
|
||
task_types=config.task_types,
|
||
seed=0,
|
||
baseline_run_id=baseline_run_id,
|
||
diag_size=config.diag_size,
|
||
diag_correct_ratio=config.diag_correct_ratio,
|
||
val_size=config.val_size,
|
||
val_correct_ratio=config.val_correct_ratio,
|
||
test_size=config.test_size,
|
||
eval_min_per_class=config.eval_min_per_class,
|
||
train_ratio=config.train_ratio,
|
||
test_questions_dir=test_questions_dir,
|
||
batch_correct_ratio=config.batch_correct_ratio,
|
||
)
|
||
|
||
|
||
def _read_baseline_run_id(config: RunConfig) -> str:
|
||
"""从 workspace 的 seed.json 读取 baseline_run_id。
|
||
|
||
workspace 由 init_workspace_from_seed 从种子创建，seed.json 保存在
|
||
store/seeds/<name>/seed.json 中。manifest.json 中 history 首条或 seed
|
||
配置字段指向对应种子。
|
||
|
||
参数:
|
||
config: 运行配置（提供 workspace_dir, store_dir, seed）。
|
||
|
||
返回:
|
||
baseline_run_id 字符串。
|
||
"""
|
||
from app.harness.store import read_seed
|
||
|
||
meta = read_seed(config.store_dir, config.seed)
|
||
return meta["baseline_run_id"]
|
||
|
||
|
||
def _validate_per_category_consistency(
|
||
frozen_data: dict,
|
||
pool_config: PoolConfig,
|
||
baseline_run_id: str,
|
||
) -> None:
|
||
"""校验已冻结的 per_category pools.json 与当前配置的一致性。
|
||
|
||
参数:
|
||
frozen_data: pools.json 解析后的原始字典。
|
||
pool_config: 当前构建配置。
|
||
baseline_run_id: 当前基线 run 标识。
|
||
|
||
异常:
|
||
ValueError: 任一关键参数与冻结值不一致。
|
||
"""
|
||
mismatches: list[str] = []
|
||
if frozen_data.get("seed") != pool_config.seed:
|
||
mismatches.append(f"seed: 冻结={frozen_data.get('seed')}, 当前={pool_config.seed}")
|
||
if frozen_data.get("train_ratio") != pool_config.train_ratio:
|
||
mismatches.append(
|
||
f"train_ratio: 冻结={frozen_data.get('train_ratio')}, 当前={pool_config.train_ratio}"
|
||
)
|
||
if frozen_data.get("baseline_run_id") != baseline_run_id:
|
||
mismatches.append(
|
||
f"baseline_run_id: 冻结={frozen_data.get('baseline_run_id')}, 当前={baseline_run_id}"
|
||
)
|
||
if frozen_data.get("split_mode") != "per_category":
|
||
mismatches.append(f"split_mode: 冻结={frozen_data.get('split_mode')}, 当前=per_category")
|
||
if mismatches:
|
||
raise ValueError(
|
||
"per_category pools.json 与当前配置不一致:\n"
|
||
+ "\n".join(f" - {m}" for m in mismatches)
|
||
)
|
||
|
||
|
||
def build_or_load_pools(
|
||
config: RunConfig,
|
||
strategy: PoolStrategy,
|
||
db_path: Path,
|
||
) -> Pools:
|
||
"""train 模式的三池获取入口：pools.json 已存在则加载，否则从基线 db 切分并冻结。
|
||
|
||
把 main.py train 分支「pools.json 存在则 load_pools 否则 build_pools 再 save_pools」
|
||
那段抽成纯函数，使 main 与集成测试共用同一切分逻辑、避免重复。pools.json 是
|
||
一次 fresh 训练的冻结切分，resume/重跑同一 workspace 时直接复用以保证三池一致。
|
||
|
||
参数:
|
||
config: 运行配置，提供 workspace_dir 与三池采样旋钮（diag/val/test 各项）。
|
||
strategy: 池构建策略（GlobalPoolStrategy / PerCategoryPoolStrategy）。
|
||
db_path: harness.db 路径，用于读取基线推理对错。
|
||
|
||
返回:
|
||
冻结的三池 Pools。
|
||
|
||
关键实现:
|
||
baseline_run_id 从 seed.json 读取（非 config.run_id）。per_category 模式
|
||
加载时做一致性校验，并支持新类别的增量追加。切分前从基线 db 的 predictions
|
||
表读该 run_id 的逐题对错，作为分层采样依据。pools.json 落在
|
||
config.workspace_dir 下，存在即视为已冻结。
|
||
"""
|
||
from app.harness.log import HarnessLog
|
||
from app.harness.workspace import resolve_paths
|
||
from app.question_gen import load_benchmark
|
||
|
||
baseline_run_id = _read_baseline_run_id(config)
|
||
pool_config = _to_pool_config(config, baseline_run_id)
|
||
pools_path = config.workspace_dir / "pools.json"
|
||
|
||
if pools_path.exists():
|
||
# ── 加载已冻结的 pools ──
|
||
raw = json.loads(pools_path.read_text(encoding="utf-8"))
|
||
frozen_split_mode = raw.get("split_mode", "global")
|
||
|
||
if frozen_split_mode == "per_category":
|
||
_validate_per_category_consistency(raw, pool_config, baseline_run_id)
|
||
|
||
# 检查是否有新类别需要增量追加
|
||
frozen_categories = raw.get("categories", {})
|
||
if pool_config.task_types is not None:
|
||
requested_types = set(pool_config.task_types)
|
||
existing_types = set(frozen_categories.keys())
|
||
new_types = requested_types - existing_types
|
||
|
||
if new_types:
|
||
# 增量构建新类别
|
||
paths = resolve_paths(config.workspace_dir)
|
||
questions = load_benchmark(paths.questions_dir)
|
||
with HarnessLog(str(db_path), baseline_run_id) as hlog:
|
||
rows = hlog.query(
|
||
"SELECT question_id, prediction, answer "
|
||
"FROM predictions WHERE run_id=?",
|
||
(baseline_run_id,),
|
||
)
|
||
correctness = {r["question_id"]: r["prediction"] == r["answer"] for r in rows}
|
||
|
||
new_cats = strategy.build_incremental(
|
||
sorted(new_types),
|
||
questions,
|
||
correctness,
|
||
pool_config,
|
||
)
|
||
# 合并新类别到 categories
|
||
frozen_categories.update(new_cats)
|
||
raw["categories"] = frozen_categories
|
||
|
||
# 从 categories 重建 diagnosis/validation 列表
|
||
qid_map = {q.question_id: q for q in questions}
|
||
new_diag: list[dict] = []
|
||
new_val: list[dict] = []
|
||
for tt in sorted(frozen_categories.keys()):
|
||
cat = frozen_categories[tt]
|
||
for qid in cat["train"]:
|
||
if qid in qid_map:
|
||
new_diag.append(_q_to_dict(qid_map[qid]))
|
||
for qid in cat["val"]:
|
||
if qid in qid_map:
|
||
new_val.append(_q_to_dict(qid_map[qid]))
|
||
raw["diagnosis"] = new_diag
|
||
raw["validation"] = new_val
|
||
raw["correctness"] = {
|
||
**raw.get("correctness", {}),
|
||
**{
|
||
qid: correctness.get(qid, False)
|
||
for cat in new_cats.values()
|
||
for qid in cat["train"] + cat["val"]
|
||
},
|
||
}
|
||
# 重新冻结
|
||
_atomic_write_json(pools_path, raw)
|
||
logger.info(
|
||
"per_category 增量追加 {} 个新类别: {}",
|
||
len(new_types),
|
||
sorted(new_types),
|
||
)
|
||
|
||
return load_pools(pools_path)
|
||
|
||
# ── 全新构建 ──
|
||
paths = resolve_paths(config.workspace_dir)
|
||
questions = load_benchmark(paths.questions_dir)
|
||
with HarnessLog(str(db_path), baseline_run_id) as hlog:
|
||
rows = hlog.query(
|
||
"SELECT question_id, prediction, answer FROM predictions WHERE run_id=?",
|
||
(baseline_run_id,),
|
||
)
|
||
correctness = {r["question_id"]: r["prediction"] == r["answer"] for r in rows}
|
||
|
||
pools = strategy.build(questions, correctness, pool_config, db_path=db_path)
|
||
save_pools(
|
||
pools,
|
||
pools_path,
|
||
split_mode=config.pool_split_mode,
|
||
config=pool_config,
|
||
)
|
||
return pools
|
||
|
||
|
||
class PerCategoryPoolStrategy:
|
||
"""Per-category 分层池构建策略。
|
||
|
||
按题型分组，每个题型内部按 correctness 分层，以 train_ratio 比例
|
||
划分 train（映射到 diagnosis 池）和 val（映射到 validation 池）。
|
||
与 GlobalPoolStrategy 的全局 progressive exclusion 不同，本策略
|
||
保证每个类别内部的 train/val 比例精确对齐。
|
||
"""
|
||
|
||
def build(
|
||
self,
|
||
questions: list[GeneratedQuestion],
|
||
correctness: dict[str, bool],
|
||
config: PoolConfig,
|
||
*,
|
||
db_path: Path | None = None,
|
||
) -> Pools:
|
||
"""按题型分组后，每组做 correctness 分层的 train/val 划分。
|
||
|
||
参数:
|
||
questions: 题目全集。
|
||
correctness: question_id -> 基线是否答对。
|
||
config: 池构建配置（使用 train_ratio, task_types, seed,
|
||
baseline_run_id, test_questions_dir, batch_correct_ratio）。
|
||
db_path: harness.db 路径，用于查询 benchmark 历史推理记录
|
||
（maintenance 补入时需要）。
|
||
|
||
返回:
|
||
冻结的 Pools（diagnosis=train, validation=val,
|
||
test 从 test_questions_dir 加载或为空列表）。
|
||
"""
|
||
# Phase 1: 按 task_types 过滤
|
||
if config.task_types is not None:
|
||
allowed = set(config.task_types)
|
||
filtered = [q for q in questions if q.task_type in allowed]
|
||
else:
|
||
filtered = list(questions)
|
||
|
||
# Phase 2: 按 task_type 分组
|
||
groups: dict[str, list[GeneratedQuestion]] = defaultdict(list)
|
||
for q in filtered:
|
||
groups[q.task_type].append(q)
|
||
|
||
# Phase 2.5: 正确率检查 + maintenance 补入
|
||
if config.batch_correct_ratio is not None:
|
||
self._check_and_supplement_maintenance(
|
||
groups,
|
||
correctness,
|
||
config,
|
||
db_path,
|
||
)
|
||
|
||
# Phase 3: 每组分层划分
|
||
all_train: list[GeneratedQuestion] = []
|
||
all_val: list[GeneratedQuestion] = []
|
||
rng = random.Random(config.seed)
|
||
|
||
for task_type in sorted(groups.keys()):
|
||
train_units, val_units = self._split_one_category(
|
||
build_units(groups[task_type]),
|
||
correctness,
|
||
config.train_ratio,
|
||
rng,
|
||
)
|
||
all_train.extend(flatten_units(train_units))
|
||
all_val.extend(flatten_units(val_units))
|
||
|
||
# Phase 4: test 池（从外部目录加载，无则空；按 task_types 过滤）
|
||
test: list[GeneratedQuestion] = []
|
||
if config.test_questions_dir is not None:
|
||
from app.question_gen import load_benchmark
|
||
|
||
test = load_benchmark(config.test_questions_dir)
|
||
if config.task_types is not None:
|
||
allowed = set(config.task_types)
|
||
test = [q for q in test if q.task_type in allowed]
|
||
|
||
# Phase 5: 计算 baseline_val_accuracy
|
||
val_correct = sum(1 for q in all_val if correctness.get(q.question_id, False))
|
||
baseline_val_accuracy = val_correct / len(all_val) if all_val else 0.0
|
||
|
||
return Pools(
|
||
diagnosis=all_train,
|
||
validation=all_val,
|
||
test=test,
|
||
baseline_run_id=config.baseline_run_id,
|
||
baseline_val_accuracy=baseline_val_accuracy,
|
||
correctness={
|
||
q.question_id: correctness.get(q.question_id, False)
|
||
for q in all_train + all_val + test
|
||
},
|
||
)
|
||
|
||
def _check_and_supplement_maintenance(
|
||
self,
|
||
groups: dict[str, list[GeneratedQuestion]],
|
||
correctness: dict[str, bool],
|
||
config: PoolConfig,
|
||
db_path: Path | None,
|
||
) -> None:
|
||
"""按 task_type 检查正确率，过高警告，过低则从 benchmark 补入正确题。
|
||
|
||
修改 groups 和 correctness（原地更新）。
|
||
|
||
参数:
|
||
groups: task_type -> 题目列表映射（原地追加补入题）。
|
||
correctness: question_id -> 是否正确映射（原地追加补入题标记）。
|
||
config: 含 batch_correct_ratio 和 test_questions_dir。
|
||
db_path: harness.db 路径，用于查询 benchmark 历史推理记录。
|
||
"""
|
||
import sqlite3
|
||
|
||
r = config.batch_correct_ratio
|
||
assert r is not None # 调用方已保证
|
||
|
||
# Phase 2.5a: 正确率检查（不依赖 test_questions_dir）
|
||
for task_type, group in groups.items():
|
||
c = sum(1 for q in group if correctness.get(q.question_id, False))
|
||
n = len(group)
|
||
ratio = c / n if n > 0 else 0.0
|
||
if ratio > 1 - r:
|
||
logger.warning(
|
||
"类别 {} 正确率 {:.1%} 过高（阈值 {:.1%}），出题可能太简单",
|
||
task_type,
|
||
ratio,
|
||
1 - r,
|
||
)
|
||
|
||
# Phase 2.5b: maintenance 补入（需要 test_questions_dir）
|
||
if config.test_questions_dir is None:
|
||
return
|
||
|
||
from app.question_gen import load_benchmark
|
||
|
||
bench_questions = load_benchmark(config.test_questions_dir)
|
||
|
||
# 查询 DB 中 benchmark 题的历史正确性
|
||
bench_correctness: dict[str, bool] = {}
|
||
if db_path is not None and db_path.exists():
|
||
conn = sqlite3.connect(str(db_path))
|
||
bench_qids = [q.question_id for q in bench_questions]
|
||
if bench_qids:
|
||
placeholders = ",".join("?" for _ in bench_qids)
|
||
rows = conn.execute(
|
||
f"SELECT question_id, prediction, answer FROM predictions " # noqa: S608
|
||
f"WHERE question_id IN ({placeholders}) "
|
||
f"ORDER BY timestamp DESC",
|
||
bench_qids,
|
||
).fetchall()
|
||
for qid, pred, ans in rows:
|
||
if qid not in bench_correctness:
|
||
bench_correctness[qid] = pred == ans
|
||
conn.close()
|
||
|
||
# 按 task_type 索引 benchmark 题
|
||
bench_by_type: dict[str, list[GeneratedQuestion]] = defaultdict(list)
|
||
for q in bench_questions:
|
||
bench_by_type[q.task_type].append(q)
|
||
|
||
for task_type, group in groups.items():
|
||
c = sum(1 for q in group if correctness.get(q.question_id, False))
|
||
w = len(group) - c
|
||
n = len(group)
|
||
ratio = c / n if n > 0 else 0.0
|
||
|
||
if ratio >= r:
|
||
continue
|
||
|
||
k = math.ceil((r * w - (1 - r) * c) / (1 - r))
|
||
|
||
existing_ids = {q.question_id for q in group}
|
||
candidates = [
|
||
q
|
||
for q in bench_by_type.get(task_type, [])
|
||
if bench_correctness.get(q.question_id, False) and q.question_id not in existing_ids
|
||
]
|
||
|
||
if not candidates:
|
||
logger.warning(
|
||
"类别 {} 需补入 {} 道正确题，但 benchmark 中无可用候选",
|
||
task_type,
|
||
k,
|
||
)
|
||
continue
|
||
|
||
actual = min(k, len(candidates))
|
||
for q in candidates[:actual]:
|
||
supplemented = GeneratedQuestion(
|
||
question_id=q.question_id,
|
||
video_id=q.video_id,
|
||
task_type=q.task_type,
|
||
question=q.question,
|
||
options=q.options,
|
||
answer=q.answer,
|
||
source_nodes=q.source_nodes,
|
||
difficulty=q.difficulty,
|
||
family="VME_MAINTENANCE",
|
||
skill_target=q.skill_target,
|
||
difficulty_steps=q.difficulty_steps,
|
||
)
|
||
group.append(supplemented)
|
||
correctness[supplemented.question_id] = True
|
||
|
||
logger.info(
|
||
"类别 {} 正确率 {:.1%} < {:.1%}，从 benchmark 补入 {} 道 maintenance 正确题",
|
||
task_type,
|
||
ratio,
|
||
r,
|
||
actual,
|
||
)
|
||
|
||
def _split_one_category(
|
||
self,
|
||
units: list[QuestionUnit],
|
||
correctness: dict[str, bool],
|
||
train_ratio: float,
|
||
rng: random.Random,
|
||
) -> tuple[list[QuestionUnit], list[QuestionUnit]]:
|
||
"""单类别 correctness 分层划分，以 unit 为原子（pair 计 1 个 unit）。
|
||
|
||
孪生对两题作为一个整体落入 train 或 val，绝不被拆散；single-only 输入下
|
||
unit 与 question 一一对应、rng 消耗量不变，划分结果与逐题划分完全一致。
|
||
|
||
参数:
|
||
units: 单类别全部单元。
|
||
correctness: question_id -> 基线是否答对；单元级正确性取成员的 AND。
|
||
train_ratio: train 占单元总量的比例。
|
||
rng: 随机数生成器（保证跨类别可复现）。
|
||
|
||
返回:
|
||
(train_units, val_units) 单元列表元组，两侧互斥且总量 == len(units)。
|
||
|
||
异常:
|
||
ValueError: correctness 中缺少某些 question_id。
|
||
"""
|
||
n_total = len(units)
|
||
n_train = round(n_total * train_ratio)
|
||
n_val = n_total - n_train
|
||
|
||
_assert_correctness_complete(units, correctness)
|
||
|
||
correct_units = [u for u in units if unit_correctness(u, correctness, strict=False)]
|
||
wrong_units = [u for u in units if not unit_correctness(u, correctness, strict=False)]
|
||
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(
|
||
"类别 {} {} ({} 单元)，退化为非分层随机划分",
|
||
units[0].task_type,
|
||
label,
|
||
n_total,
|
||
)
|
||
shuffled = list(units)
|
||
rng.shuffle(shuffled)
|
||
return shuffled[:n_train], shuffled[n_train:]
|
||
|
||
# 分层: 按 correctness 比例分配到 train
|
||
train_correct = math.floor(n_correct * n_train / n_total)
|
||
train_wrong = n_train - train_correct
|
||
|
||
rng.shuffle(correct_units)
|
||
rng.shuffle(wrong_units)
|
||
|
||
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}"
|
||
|
||
return train, val
|
||
|
||
def build_incremental(
|
||
self,
|
||
new_task_types: list[str],
|
||
questions: list[GeneratedQuestion],
|
||
correctness: dict[str, bool],
|
||
config: PoolConfig,
|
||
) -> dict[str, dict[str, list[str]]]:
|
||
"""增量划分：仅处理 new_task_types 中的类别。
|
||
|
||
参数:
|
||
new_task_types: 需要增量划分的类别列表。
|
||
questions: 题目全集（从中筛选指定类别）。
|
||
correctness: question_id -> 基线是否答对。
|
||
config: 池构建配置（使用 train_ratio, seed）。
|
||
|
||
返回:
|
||
{task_type: {"train": [qid, ...], "val": [qid, ...]}}。
|
||
"""
|
||
target_types = set(new_task_types)
|
||
groups: dict[str, list[GeneratedQuestion]] = defaultdict(list)
|
||
for q in questions:
|
||
if q.task_type in target_types:
|
||
groups[q.task_type].append(q)
|
||
|
||
result: dict[str, dict[str, list[str]]] = {}
|
||
rng = random.Random(config.seed)
|
||
|
||
for task_type in sorted(groups.keys()):
|
||
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 flatten_units(train_units)],
|
||
"val": [q.question_id for q in flatten_units(val_units)],
|
||
}
|
||
|
||
return result
|