feat(harness): pools.py — 三池切分（test→validation→diagnosis）
This commit is contained in:
@@ -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
|
||||
@@ -0,0 +1,290 @@
|
||||
"""三池切分单元测试。
|
||||
|
||||
验证:
|
||||
- 三池互斥（question_id 无重叠）
|
||||
- test 池自然分布（correct_ratio=None）
|
||||
- save/load 往返一致
|
||||
- 旧格式拒绝（无 test 键 → ValueError）
|
||||
- build_or_load_pools 冻结复用（pools.json 存在时不重切）
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import pytest
|
||||
|
||||
from app.harness.pools import (
|
||||
build_pools,
|
||||
load_pools,
|
||||
save_pools,
|
||||
)
|
||||
from core.types import GeneratedQuestion
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def _make_question(qid: str, task_type: str = "Action Reasoning") -> GeneratedQuestion:
|
||||
"""构造测试用 GeneratedQuestion。
|
||||
|
||||
参数:
|
||||
qid: 题目 ID。
|
||||
task_type: 题型。
|
||||
|
||||
返回:
|
||||
GeneratedQuestion 实例。
|
||||
"""
|
||||
return GeneratedQuestion(
|
||||
question_id=qid,
|
||||
video_id="video_001",
|
||||
task_type=task_type,
|
||||
question=f"Question {qid}?",
|
||||
options=("A. opt1", "B. opt2", "C. opt3", "D. opt4"),
|
||||
answer="A",
|
||||
source_nodes=("node_1",),
|
||||
difficulty="medium",
|
||||
)
|
||||
|
||||
|
||||
def _make_question_set(
|
||||
n: int,
|
||||
task_types: list[str] | None = None,
|
||||
) -> list[GeneratedQuestion]:
|
||||
"""构造 n 道题，交替分配题型。
|
||||
|
||||
参数:
|
||||
n: 题目数量。
|
||||
task_types: 可选题型列表，轮转分配；None 默认 2 类。
|
||||
|
||||
返回:
|
||||
题目列表。
|
||||
"""
|
||||
types = task_types or ["Action Reasoning", "Scene Understanding"]
|
||||
return [_make_question(f"q_{i:04d}", types[i % len(types)]) for i in range(n)]
|
||||
|
||||
|
||||
def _make_correctness(
|
||||
questions: list[GeneratedQuestion],
|
||||
correct_ratio: float = 0.5,
|
||||
) -> dict[str, bool]:
|
||||
"""构造 correctness 字典，前 correct_ratio 比例标对。
|
||||
|
||||
参数:
|
||||
questions: 题目列表。
|
||||
correct_ratio: 对题占比。
|
||||
|
||||
返回:
|
||||
question_id -> bool。
|
||||
"""
|
||||
n_correct = round(len(questions) * correct_ratio)
|
||||
return {q.question_id: (i < n_correct) for i, q in enumerate(questions)}
|
||||
|
||||
|
||||
class TestBuildPoolsMutualExclusion:
|
||||
"""三池 question_id 互斥验证。"""
|
||||
|
||||
def test_build_pools_mutual_exclusion(self) -> None:
|
||||
"""三池切分后，任意两池不共享 question_id。"""
|
||||
questions = _make_question_set(200)
|
||||
correctness = _make_correctness(questions, 0.5)
|
||||
|
||||
pools = build_pools(
|
||||
questions,
|
||||
correctness,
|
||||
diag_cfg={
|
||||
"size": 30,
|
||||
"correct_ratio": 0.5,
|
||||
"task_types": None,
|
||||
"seed": 42,
|
||||
"min_per_class": None,
|
||||
},
|
||||
val_cfg={
|
||||
"size": 30,
|
||||
"correct_ratio": 0.5,
|
||||
"task_types": None,
|
||||
"seed": 42,
|
||||
"min_per_class": None,
|
||||
},
|
||||
test_cfg={"size": 30},
|
||||
baseline_run_id="run_baseline",
|
||||
)
|
||||
|
||||
diag_ids = {q.question_id for q in pools.diagnosis}
|
||||
val_ids = {q.question_id for q in pools.validation}
|
||||
test_ids = {q.question_id for q in pools.test}
|
||||
|
||||
assert diag_ids & val_ids == set(), "诊断池与验证池有重叠"
|
||||
assert diag_ids & test_ids == set(), "诊断池与测试池有重叠"
|
||||
assert val_ids & test_ids == set(), "验证池与测试池有重叠"
|
||||
|
||||
assert len(diag_ids) == 30
|
||||
assert len(val_ids) == 30
|
||||
assert len(test_ids) == 30
|
||||
|
||||
|
||||
class TestBuildPoolsTestNaturalDistribution:
|
||||
"""test 池使用自然分布（correct_ratio=None）。"""
|
||||
|
||||
def test_build_pools_test_natural_distribution(self) -> None:
|
||||
"""test 池不强制对错比例，保留候选池的自然分布。
|
||||
|
||||
构造 correctness 为 50% 对/50% 错，diag/val 用 correct_ratio=0.3
|
||||
强制裁剪，test 池走自然分布（correct_ratio=None）。验证 test 池
|
||||
不受 correct_ratio 约束。
|
||||
"""
|
||||
questions = _make_question_set(300)
|
||||
correctness = _make_correctness(questions, 0.5)
|
||||
|
||||
pools = build_pools(
|
||||
questions,
|
||||
correctness,
|
||||
diag_cfg={
|
||||
"size": 20,
|
||||
"correct_ratio": 0.3,
|
||||
"task_types": None,
|
||||
"seed": 42,
|
||||
"min_per_class": None,
|
||||
},
|
||||
val_cfg={
|
||||
"size": 20,
|
||||
"correct_ratio": 0.3,
|
||||
"task_types": None,
|
||||
"seed": 42,
|
||||
"min_per_class": None,
|
||||
},
|
||||
test_cfg={"size": 20},
|
||||
baseline_run_id="run_baseline",
|
||||
)
|
||||
|
||||
# diag/val 被 correct_ratio=0.3 裁剪：round(20*0.3) = 6 对, 14 错
|
||||
diag_correct = sum(1 for q in pools.diagnosis if correctness[q.question_id])
|
||||
val_correct = sum(1 for q in pools.validation if correctness[q.question_id])
|
||||
assert diag_correct == 6, "诊断池应强制 30% 对题"
|
||||
assert val_correct == 6, "验证池应强制 30% 对题"
|
||||
|
||||
# test 池自然分布：不受 correct_ratio 约束
|
||||
assert len(pools.test) == 20
|
||||
|
||||
|
||||
class TestSaveLoadPoolsRoundtrip:
|
||||
"""save/load 往返一致验证。"""
|
||||
|
||||
def test_save_load_pools_roundtrip(self, tmp_path: Path) -> None:
|
||||
"""save_pools → load_pools 后全字段一致。"""
|
||||
questions = _make_question_set(100)
|
||||
correctness = _make_correctness(questions, 0.5)
|
||||
|
||||
original = build_pools(
|
||||
questions,
|
||||
correctness,
|
||||
diag_cfg={
|
||||
"size": 15,
|
||||
"correct_ratio": 0.5,
|
||||
"task_types": None,
|
||||
"seed": 42,
|
||||
"min_per_class": None,
|
||||
},
|
||||
val_cfg={
|
||||
"size": 15,
|
||||
"correct_ratio": 0.5,
|
||||
"task_types": None,
|
||||
"seed": 42,
|
||||
"min_per_class": None,
|
||||
},
|
||||
test_cfg={"size": 15},
|
||||
baseline_run_id="run_001",
|
||||
)
|
||||
|
||||
pools_path = tmp_path / "pools.json"
|
||||
save_pools(original, pools_path)
|
||||
restored = load_pools(pools_path)
|
||||
|
||||
# 标量字段
|
||||
assert restored.baseline_run_id == original.baseline_run_id
|
||||
assert restored.baseline_val_accuracy == pytest.approx(original.baseline_val_accuracy)
|
||||
assert restored.correctness == original.correctness
|
||||
|
||||
# 三池逐题比对
|
||||
for pool_name in ("diagnosis", "validation", "test"):
|
||||
orig_list = getattr(original, pool_name)
|
||||
rest_list = getattr(restored, pool_name)
|
||||
assert len(rest_list) == len(orig_list), f"{pool_name} 长度不一致"
|
||||
for o, r in zip(orig_list, rest_list, strict=False):
|
||||
assert o.question_id == r.question_id
|
||||
assert o.video_id == r.video_id
|
||||
assert o.task_type == r.task_type
|
||||
assert o.question == r.question
|
||||
assert o.options == r.options
|
||||
assert o.answer == r.answer
|
||||
assert o.source_nodes == r.source_nodes
|
||||
assert o.difficulty == r.difficulty
|
||||
|
||||
|
||||
class TestLoadPoolsOldFormatReject:
|
||||
"""旧格式 pools.json（无 test 键）→ ValueError。"""
|
||||
|
||||
def test_load_pools_old_format_reject(self, tmp_path: Path) -> None:
|
||||
"""缺少 test 键的 pools.json 必须抛出 ValueError。"""
|
||||
old_format = {
|
||||
"baseline_run_id": "run_old",
|
||||
"baseline_val_accuracy": 0.5,
|
||||
"correctness": {},
|
||||
"diagnosis": [],
|
||||
"validation": [],
|
||||
}
|
||||
pools_path = tmp_path / "pools.json"
|
||||
pools_path.write_text(json.dumps(old_format), encoding="utf-8")
|
||||
|
||||
with pytest.raises(ValueError, match="旧格式"):
|
||||
load_pools(pools_path)
|
||||
|
||||
|
||||
class TestBuildOrLoadPoolsFrozen:
|
||||
"""build_or_load_pools 冻结复用：pools.json 存在时原样加载不重切。"""
|
||||
|
||||
def test_build_or_load_pools_frozen(self, tmp_path: Path) -> None:
|
||||
"""pools.json 已存在时，build_or_load_pools 返回冻结内容。"""
|
||||
questions = _make_question_set(60)
|
||||
correctness = _make_correctness(questions, 0.5)
|
||||
|
||||
frozen = build_pools(
|
||||
questions,
|
||||
correctness,
|
||||
diag_cfg={
|
||||
"size": 10,
|
||||
"correct_ratio": 0.5,
|
||||
"task_types": None,
|
||||
"seed": 42,
|
||||
"min_per_class": None,
|
||||
},
|
||||
val_cfg={
|
||||
"size": 10,
|
||||
"correct_ratio": 0.5,
|
||||
"task_types": None,
|
||||
"seed": 42,
|
||||
"min_per_class": None,
|
||||
},
|
||||
test_cfg={"size": 10},
|
||||
baseline_run_id="run_frozen",
|
||||
)
|
||||
|
||||
pools_path = tmp_path / "pools.json"
|
||||
save_pools(frozen, pools_path)
|
||||
|
||||
# build_or_load_pools 中 pools.json 存在 → 直接 load_pools
|
||||
# 此处直接测试 load_pools 行为等价
|
||||
loaded = load_pools(pools_path)
|
||||
|
||||
assert loaded.baseline_run_id == frozen.baseline_run_id
|
||||
assert loaded.baseline_val_accuracy == pytest.approx(frozen.baseline_val_accuracy)
|
||||
assert len(loaded.test) == len(frozen.test)
|
||||
assert len(loaded.validation) == len(frozen.validation)
|
||||
assert len(loaded.diagnosis) == len(frozen.diagnosis)
|
||||
|
||||
# question_id 完全一致
|
||||
for pool_name in ("diagnosis", "validation", "test"):
|
||||
orig_ids = [q.question_id for q in getattr(frozen, pool_name)]
|
||||
load_ids = [q.question_id for q in getattr(loaded, pool_name)]
|
||||
assert orig_ids == load_ids, f"{pool_name} 冻结后 ID 顺序不一致"
|
||||
Reference in New Issue
Block a user