diff --git a/app/harness/pools.py b/app/harness/pools.py index 9c0cda5..6dc0533 100644 --- a/app/harness/pools.py +++ b/app/harness/pools.py @@ -24,6 +24,7 @@ if TYPE_CHECKING: from pathlib import Path from app.harness.config import RunConfig + from app.ports import PoolStrategy @dataclass @@ -232,26 +233,64 @@ def _dict_to_q(d: dict) -> GeneratedQuestion: ) -def save_pools(pools: Pools, path: Path) -> None: +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 + ) + 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, - ), + json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8", ) @@ -259,6 +298,9 @@ def save_pools(pools: Pools, path: Path) -> None: def load_pools(path: Path) -> Pools: """从 JSON 恢复冻结的三池。 + 兼容新旧格式:有无 split_mode 字段都能加载。新格式(含 split_mode / + categories)的额外元数据在加载时忽略——Pools 对象只关心三池列表和标量。 + 参数: path: 冻结的 pools.json 路径。 @@ -289,10 +331,105 @@ def load_pools(path: Path) -> Pools: ) +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, + ) + + +def _read_baseline_run_id(config: RunConfig) -> str: + """从 workspace 的 seed.json 读取 baseline_run_id。 + + workspace 由 init_workspace_from_seed 从种子创建,seed.json 保存在 + store/seeds//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')}, " + f"当前={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')}, " + f"当前={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, - run_id: str, - task_types: list[str] | None = None, + strategy: PoolStrategy, + db_path: Path, ) -> Pools: """train 模式的三池获取入口:pools.json 已存在则加载,否则从基线 db 切分并冻结。 @@ -302,53 +439,114 @@ def build_or_load_pools( 参数: config: 运行配置,提供 workspace_dir 与三池采样旋钮(diag/val/test 各项)。 - run_id: 基线全量记录的 run_id(fresh 时来自 seed.json,决定从哪个 run 读对错)。 - task_types: 可选题型过滤,限定诊断/验证池只采样这些题型;None 表示不过滤。 + strategy: 池构建策略(GlobalPoolStrategy / PerCategoryPoolStrategy)。 + db_path: harness.db 路径,用于读取基线推理对错。 返回: 冻结的三池 Pools。 关键实现: - 切分前从基线 db 的 predictions 表读该 run_id 的逐题对错,作为分层采样依据。 - pools.json 落在 config.workspace_dir 下,存在即视为已冻结,原样加载不重切。 + 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"] + }, + } + # 重新冻结 + pools_path.write_text( + json.dumps(raw, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + 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(paths.db_path), run_id) as log: - rows = log.query( + with HarnessLog(str(db_path), baseline_run_id) as hlog: + rows = hlog.query( "SELECT question_id, prediction, answer FROM predictions WHERE run_id=?", - (run_id,), + (baseline_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, + + pools = strategy.build(questions, correctness, pool_config) + save_pools( + pools, + pools_path, + split_mode=config.pool_split_mode, + config=pool_config, ) - save_pools(pools, pools_path) return pools diff --git a/tests/unit/test_harness_pools.py b/tests/unit/test_harness_pools.py index 3a2da7f..bc4b905 100644 --- a/tests/unit/test_harness_pools.py +++ b/tests/unit/test_harness_pools.py @@ -460,3 +460,122 @@ class TestPerCategoryPoolStrategy: assert len(pools.validation) == 20 types_in_diag = {q.task_type for q in pools.diagnosis} assert types_in_diag == {"Action Reasoning", "Scene Understanding"} + + +class TestPerCategorySaveLoad: + """per_category 格式的 pools.json 冻结/加载。""" + + def test_save_load_per_category_roundtrip(self, tmp_path: Path) -> None: + """per_category 模式 save -> load 往返一致。""" + questions = _make_per_category_questions() + correctness = {q.question_id: True for q in questions} + config = PoolConfig( + task_types=("Action Reasoning",), seed=42, baseline_run_id="baseline_v2", + diag_size=0, diag_correct_ratio=0.0, val_size=0, val_correct_ratio=0.0, + test_size=0, eval_min_per_class=0, train_ratio=20 / 30, + test_questions_dir=None, + ) + strategy = PerCategoryPoolStrategy() + pools = strategy.build(questions, correctness, config) + pools_path = tmp_path / "pools.json" + save_pools(pools, pools_path, split_mode="per_category", config=config) + loaded = load_pools(pools_path) + assert loaded.baseline_run_id == pools.baseline_run_id + assert len(loaded.diagnosis) == len(pools.diagnosis) + assert len(loaded.validation) == len(pools.validation) + + # 验证 per_category 格式内容 + data = json.loads(pools_path.read_text()) + assert data["split_mode"] == "per_category" + assert "categories" in data + assert data["seed"] == 42 + assert data["train_ratio"] == pytest.approx(20 / 30) + assert data["test_source"] is None + + # categories 内容校验 + cats = data["categories"] + assert "Action Reasoning" in cats + assert len(cats["Action Reasoning"]["train"]) == 20 + assert len(cats["Action Reasoning"]["val"]) == 10 + + def test_save_per_category_without_config_raises(self, tmp_path: Path) -> None: + """per_category 模式未提供 config 时报 ValueError。""" + from app.harness.pools import Pools + + pools = Pools( + diagnosis=[], validation=[], test=[], + baseline_run_id="r", baseline_val_accuracy=0.0, + ) + with pytest.raises(ValueError, match="per_category 模式下.*必须提供 config"): + save_pools(pools, tmp_path / "pools.json", split_mode="per_category") + + def test_save_global_mode_has_split_mode_field(self, tmp_path: Path) -> None: + """global 模式 save 也写入 split_mode 字段。""" + questions = _make_question_set(60) + correctness = _make_correctness(questions, 0.5) + original = 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_001", + ) + pools_path = tmp_path / "pools.json" + save_pools(original, pools_path, split_mode="global") + data = json.loads(pools_path.read_text()) + assert data["split_mode"] == "global" + assert "categories" not in data + + # 仍能正常 load + loaded = load_pools(pools_path) + assert loaded.baseline_run_id == "run_001" + assert len(loaded.diagnosis) == 10 + + def test_load_legacy_format_without_split_mode(self, tmp_path: Path) -> None: + """旧格式(无 split_mode 字段)仍可加载。""" + legacy = { + "baseline_run_id": "run_legacy", + "baseline_val_accuracy": 0.75, + "correctness": {"q1": True}, + "diagnosis": [{ + "question_id": "q1", "video_id": "v1", "task_type": "AR", + "question": "Q?", "options": ["A", "B", "C", "D"], + "answer": "A", "source_nodes": [], "difficulty": "medium", + "skill_target": None, "difficulty_steps": None, + }], + "validation": [], + "test": [], + } + pools_path = tmp_path / "pools.json" + pools_path.write_text(json.dumps(legacy), encoding="utf-8") + loaded = load_pools(pools_path) + assert loaded.baseline_run_id == "run_legacy" + assert len(loaded.diagnosis) == 1 + + def test_per_category_categories_multi_type(self, tmp_path: Path) -> None: + """多类别 per_category save 后 categories 包含所有类别。""" + questions = _make_per_category_questions() + correctness = {q.question_id: True for q in questions} + config = PoolConfig( + task_types=("Action Reasoning", "Scene Understanding"), seed=0, + baseline_run_id="b", diag_size=0, diag_correct_ratio=0.0, + val_size=0, val_correct_ratio=0.0, test_size=0, + eval_min_per_class=0, train_ratio=20 / 30, test_questions_dir=None, + ) + strategy = PerCategoryPoolStrategy() + pools = strategy.build(questions, correctness, config) + pools_path = tmp_path / "pools.json" + save_pools(pools, pools_path, split_mode="per_category", config=config) + + data = json.loads(pools_path.read_text()) + assert set(data["categories"].keys()) == { + "Action Reasoning", "Scene Understanding", + } + for tt in data["categories"]: + cat = data["categories"][tt] + assert len(cat["train"]) == 20 + assert len(cat["val"]) == 10 + # train + val 的 qid 互斥 + assert set(cat["train"]) & set(cat["val"]) == set()