534 lines
18 KiB
Python
534 lines
18 KiB
Python
"""三池：held-out test + 验证 + 诊断，分层采样 + 冻结持久化。
|
||
|
||
三池切分对应训练循环中的 DataLoader 阶段——从题目全集中按
|
||
test -> validation -> diagnosis 的顺序 progressive exclusion，
|
||
保证 question_id 互斥。test 池用自然分布（correct_ratio=None），
|
||
验证池/诊断池按对错比例分层采样。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import math
|
||
import random
|
||
from collections import defaultdict
|
||
from dataclasses import dataclass, field
|
||
from typing import TYPE_CHECKING
|
||
|
||
from loguru import logger
|
||
|
||
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
|
||
|
||
|
||
@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；后两步从剩余题中采样以保证
|
||
question_id 互斥。test 池用 correct_ratio=None 的自然分布采样。
|
||
"""
|
||
test = _sample_excluding(
|
||
questions,
|
||
set(),
|
||
correctness,
|
||
size=test_cfg["size"],
|
||
correct_ratio=None,
|
||
task_types=None,
|
||
seed=test_cfg.get("seed", 0),
|
||
min_per_class=None,
|
||
)
|
||
selected_ids = {q.question_id for q in test}
|
||
|
||
validation = _sample_excluding(questions, selected_ids, correctness, **val_cfg)
|
||
selected_ids |= {q.question_id for q in validation}
|
||
|
||
diagnosis = _sample_excluding(questions, selected_ids, 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
|
||
},
|
||
)
|
||
|
||
|
||
class GlobalPoolStrategy:
|
||
"""全局三分策略：test -> val -> diag progressive exclusion。
|
||
|
||
封装现有 build_pools 逻辑为 PoolStrategy 接口。
|
||
"""
|
||
|
||
def build(
|
||
self,
|
||
questions: list[GeneratedQuestion],
|
||
correctness: dict[str, bool],
|
||
config: PoolConfig,
|
||
) -> 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 _sample_excluding(
|
||
questions: list[GeneratedQuestion],
|
||
exclude_ids: set[str],
|
||
correctness: dict[str, bool],
|
||
**cfg: object,
|
||
) -> list[GeneratedQuestion]:
|
||
"""排除已选 question_id 后，按 cfg 对剩余题做分层采样。
|
||
|
||
参数:
|
||
questions: 题目全集。
|
||
exclude_ids: 已被其他池选走的 question_id，从候选中剔除以保证三池互斥。
|
||
correctness: question_id -> 基线是否答对。
|
||
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)
|
||
|
||
|
||
def _q_to_dict(q: GeneratedQuestion) -> dict:
|
||
"""将 GeneratedQuestion 转为可序列化字典。
|
||
|
||
参数:
|
||
q: 题目对象。
|
||
|
||
返回:
|
||
包含全部字段的字典（options/source_nodes 从 tuple 转为 list）。
|
||
"""
|
||
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,
|
||
"skill_target": q.skill_target,
|
||
"difficulty_steps": q.difficulty_steps,
|
||
}
|
||
|
||
|
||
def _dict_to_q(d: dict) -> GeneratedQuestion:
|
||
"""从字典恢复 GeneratedQuestion。
|
||
|
||
参数:
|
||
d: 由 _q_to_dict 产出的字典。
|
||
|
||
返回:
|
||
恢复的 GeneratedQuestion 实例（options/source_nodes 恢复为 tuple）。
|
||
"""
|
||
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"),
|
||
skill_target=d.get("skill_target"),
|
||
difficulty_steps=d.get("difficulty_steps"),
|
||
)
|
||
|
||
|
||
def save_pools(pools: Pools, path: Path) -> None:
|
||
"""将三池及基线指标冻结为 JSON。
|
||
|
||
参数:
|
||
pools: 待冻结的三池。
|
||
path: 目标 JSON 文件路径。
|
||
"""
|
||
path.write_text(
|
||
json.dumps(
|
||
{
|
||
"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],
|
||
},
|
||
ensure_ascii=False,
|
||
indent=2,
|
||
),
|
||
encoding="utf-8",
|
||
)
|
||
|
||
|
||
def load_pools(path: Path) -> Pools:
|
||
"""从 JSON 恢复冻结的三池。
|
||
|
||
参数:
|
||
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 build_or_load_pools(
|
||
config: RunConfig,
|
||
run_id: str,
|
||
task_types: list[str] | None = None,
|
||
) -> 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 各项）。
|
||
run_id: 基线全量记录的 run_id（fresh 时来自 seed.json，决定从哪个 run 读对错）。
|
||
task_types: 可选题型过滤，限定诊断/验证池只采样这些题型；None 表示不过滤。
|
||
|
||
返回:
|
||
冻结的三池 Pools。
|
||
|
||
关键实现:
|
||
切分前从基线 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
|
||
|
||
pools_path = config.workspace_dir / "pools.json"
|
||
if pools_path.exists():
|
||
return load_pools(pools_path)
|
||
|
||
paths = resolve_paths(config.workspace_dir)
|
||
questions = load_benchmark(paths.questions_dir)
|
||
with HarnessLog(str(paths.db_path), run_id) as log:
|
||
rows = log.query(
|
||
"SELECT question_id, prediction, answer FROM predictions WHERE run_id=?",
|
||
(run_id,),
|
||
)
|
||
correctness = {r["question_id"]: r["prediction"] == r["answer"] for r in rows}
|
||
pools = build_pools(
|
||
questions,
|
||
correctness,
|
||
diag_cfg={
|
||
"size": config.diag_size,
|
||
"correct_ratio": config.diag_correct_ratio,
|
||
"task_types": task_types,
|
||
"seed": 0,
|
||
"min_per_class": None,
|
||
},
|
||
val_cfg={
|
||
"size": config.val_size,
|
||
"correct_ratio": config.val_correct_ratio,
|
||
"task_types": task_types,
|
||
"seed": 0,
|
||
"min_per_class": config.eval_min_per_class,
|
||
},
|
||
test_cfg={"size": config.test_size},
|
||
baseline_run_id=run_id,
|
||
)
|
||
save_pools(pools, pools_path)
|
||
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,
|
||
) -> Pools:
|
||
"""按题型分组后，每组做 correctness 分层的 train/val 划分。
|
||
|
||
参数:
|
||
questions: 题目全集。
|
||
correctness: question_id -> 基线是否答对。
|
||
config: 池构建配置（使用 train_ratio, task_types, seed,
|
||
baseline_run_id, test_questions_dir）。
|
||
|
||
返回:
|
||
冻结的 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 3: 每组分层划分
|
||
all_train: list[GeneratedQuestion] = []
|
||
all_val: list[GeneratedQuestion] = []
|
||
rng = random.Random(config.seed)
|
||
|
||
for task_type in sorted(groups.keys()):
|
||
train, val = self._split_one_category(
|
||
groups[task_type], correctness, config.train_ratio, rng,
|
||
)
|
||
all_train.extend(train)
|
||
all_val.extend(val)
|
||
|
||
# Phase 4: test 池（从外部目录加载，无则空）
|
||
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)
|
||
|
||
# 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 _split_one_category(
|
||
self,
|
||
questions: list[GeneratedQuestion],
|
||
correctness: dict[str, bool],
|
||
train_ratio: float,
|
||
rng: random.Random,
|
||
) -> tuple[list[GeneratedQuestion], list[GeneratedQuestion]]:
|
||
"""单类别 correctness 分层划分。
|
||
|
||
参数:
|
||
questions: 单类别全部题目。
|
||
correctness: question_id -> 基线是否答对。
|
||
train_ratio: train 占总量的比例。
|
||
rng: 随机数生成器（保证跨类别可复现）。
|
||
|
||
返回:
|
||
(train, val) 题目列表元组，两池互斥且总量 == len(questions)。
|
||
|
||
异常:
|
||
ValueError: correctness 中缺少某些 question_id。
|
||
"""
|
||
n_total = len(questions)
|
||
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]}"
|
||
)
|
||
|
||
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 或全 wrong -> 退化为非分层随机划分
|
||
if n_correct == 0 or n_correct == n_total:
|
||
label = "全部正确" if n_correct == n_total else "全部错误"
|
||
logger.warning(
|
||
"类别 {} {} ({} 题)，退化为非分层随机划分",
|
||
questions[0].task_type, label, n_total,
|
||
)
|
||
shuffled = list(questions)
|
||
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_qs)
|
||
rng.shuffle(wrong_qs)
|
||
|
||
train = correct_qs[:train_correct] + wrong_qs[:train_wrong]
|
||
val = correct_qs[train_correct:] + wrong_qs[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, val = self._split_one_category(
|
||
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],
|
||
}
|
||
|
||
return result
|