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

chore: lint and format per-category pool strategy implementation

Browse Source 
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
iomgaa
2026-07-12 22:58:10 -04:00
parent 671db2f88c
commit 37d4519905
2 changed files with 71 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/harness/pools.py
+28 -37
View File
@@ -254,9 +254,7 @@ def save_pools(
ValueError: split_mode 为 "per_category" 但未提供 config。
"""
if split_mode == "per_category" and config is None:
raise ValueError(
"per_category 模式下 save_pools 必须提供 config 参数以写入元数据。"
)
raise ValueError("per_category 模式下 save_pools 必须提供 config 参数以写入元数据。")
data: dict = {
"split_mode": split_mode,
@@ -285,9 +283,7 @@ def save_pools(
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
)
data["test_source"] = str(config.test_questions_dir) if config.test_questions_dir else None
path.write_text(
json.dumps(data, ensure_ascii=False, indent=2),
@@ -402,23 +398,17 @@ def _validate_per_category_consistency(
"""
mismatches: list[str] = []
if frozen_data.get("seed") != pool_config.seed:
mismatches.append(
f"seed: 冻结={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}"
f"train_ratio: 冻结={frozen_data.get('train_ratio')}, 当前={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}"
f"baseline_run_id: 冻结={frozen_data.get('baseline_run_id')}, 当前={baseline_run_id}"
)
if frozen_data.get("split_mode") != "per_category":
mismatches.append(
f"split_mode: 冻结={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"
@@ -484,13 +474,13 @@ def build_or_load_pools(
"FROM predictions WHERE run_id=?",
(baseline_run_id,),
)
correctness = {
r["question_id"]: r["prediction"] == r["answer"]
for r in rows
}
correctness = {r["question_id"]: r["prediction"] == r["answer"] for r in rows}
new_cats = strategy.build_incremental(
sorted(new_types), questions, correctness, pool_config,
sorted(new_types),
questions,
correctness,
pool_config,
)
# 合并新类别到 categories
frozen_categories.update(new_cats)
@@ -525,7 +515,8 @@ def build_or_load_pools(
)
logger.info(
"per_category 增量追加 {} 个新类别: {}",
len(new_types), sorted(new_types),
len(new_types),
sorted(new_types),
)
return load_pools(pools_path)
@@ -596,7 +587,10 @@ class PerCategoryPoolStrategy:
for task_type in sorted(groups.keys()):
train, val = self._split_one_category(
groups[task_type], correctness, config.train_ratio, rng,
groups[task_type],
correctness,
config.train_ratio,
rng,
)
all_train.extend(train)
all_val.extend(val)
@@ -609,9 +603,7 @@ class PerCategoryPoolStrategy:
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)
)
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(
@@ -654,9 +646,7 @@ class PerCategoryPoolStrategy:
# 校验 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]}"
)
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]]
@@ -667,7 +657,9 @@ class PerCategoryPoolStrategy:
label = "全部正确" if n_correct == n_total else "全部错误"
logger.warning(
"类别 {} {} ({} 题)，退化为非分层随机划分",
questions[0].task_type, label, n_total,
questions[0].task_type,
label,
n_total,
)
shuffled = list(questions)
rng.shuffle(shuffled)
@@ -683,12 +675,8 @@ class PerCategoryPoolStrategy:
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}"
)
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
@@ -721,7 +709,10 @@ class PerCategoryPoolStrategy:
for task_type in sorted(groups.keys()):
train, val = self._split_one_category(
groups[task_type], correctness, config.train_ratio, rng,
groups[task_type],
correctness,
config.train_ratio,
rng,
)
result[task_type] = {
"train": [q.question_id for q in train],