iomgaa/Video-Tree-TRM5
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(harness): refactor build_or_load_pools to accept PoolStrategy + per_category freeze format

Browse Source 
- save_pools: extended with split_mode and config params; per_category
  mode writes categories metadata (seed, train_ratio, test_source) for
  incremental append and consistency validation
- load_pools: compatible with both old format (no split_mode) and new
  format; extra metadata fields ignored during load
- build_or_load_pools: signature changed to (config, strategy, db_path);
  baseline_run_id read from seed.json (not config.run_id); per_category
  mode does consistency check on reload and supports incremental category
  append via strategy.build_incremental
- Added _to_pool_config, _read_baseline_run_id,
  _validate_per_category_consistency helpers
- Tests: TestPerCategorySaveLoad with 5 test cases covering roundtrip,
  missing config error, global split_mode field, legacy format compat,
  multi-type categories

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
iomgaa
2026-07-12 22:53:19 -04:00
parent e5b07ac974
commit c66a00c924
2 changed files with 359 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/harness/pools.py
+240 -42
View File
@@ -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/<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')}, "
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