feat(harness): pools.py — 三池切分（test→validation→diagnosis）

2026-07-07 12:48:54 -04:00
parent 9800fef37a
commit e0f3ee10ec
2 changed files with 573 additions and 0 deletions
@@ -0,0 +1,283 @@
+"""三池：held-out test + 验证 + 诊断，分层采样 + 冻结持久化。
+
+三池切分对应训练循环中的 DataLoader 阶段——从题目全集中按
+test -> validation -> diagnosis 的顺序 progressive exclusion，
+保证 question_id 互斥。test 池用自然分布（correct_ratio=None），
+验证池/诊断池按对错比例分层采样。
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from app.question_gen import stratified_sample
+from core.types import GeneratedQuestion
+
+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
+        },
+    )
+
+
+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,
+    }
+
+
+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"),
+    )
+
+
+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