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],
tests/integration/test_pool_strategy.py
+43 -18
View File
@@ -7,9 +7,10 @@ from __future__ import annotations
import json
from collections import Counter
from pathlib import Path
from typing import TYPE_CHECKING
import pytest
if TYPE_CHECKING:
from pathlib import Path
from app.harness.pools import (
PerCategoryPoolStrategy,
@@ -22,10 +23,14 @@ from core.types import GeneratedQuestion, PoolConfig
def _make_question(qid: str, task_type: str) -> GeneratedQuestion:
"""构造测试用 GeneratedQuestion。"""
return GeneratedQuestion(
question_id=qid, video_id="v1", task_type=task_type,
question_id=qid,
video_id="v1",
task_type=task_type,
question=f"Q {qid}?",
options=("A. a", "B. b", "C. c", "D. d"),
answer="A", source_nodes=("n1",), difficulty="medium",
answer="A",
source_nodes=("n1",),
difficulty="medium",
)
@@ -35,10 +40,18 @@ class TestPerCategoryE2E:
def test_full_flow(self, tmp_path: Path) -> None:
"""完整流程：12 类各 30 题 → 策略构建 → 冻结 → 加载 → 三池校验。"""
task_types = [
"Action Prediction", "Action Reasoning", "Action Recognition",
"Action Sequence", "Causal Reasoning", "Event Reasoning",
"Object Interaction", "Object Reasoning", "Object Recognition",
"Scene Understanding", "Spatial Reasoning", "Temporal Reasoning",
"Action Prediction",
"Action Reasoning",
"Action Recognition",
"Action Sequence",
"Causal Reasoning",
"Event Reasoning",
"Object Interaction",
"Object Reasoning",
"Object Recognition",
"Scene Understanding",
"Spatial Reasoning",
"Temporal Reasoning",
]
questions = []
for tt in task_types:
@@ -52,11 +65,17 @@ class TestPerCategoryE2E:
correctness[q.question_id] = idx < 18
config = PoolConfig(
task_types=None, 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,
task_types=None,
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()
@@ -107,11 +126,17 @@ class TestPerCategoryE2E:
correctness = {q.question_id: (i < 20) for i, q in enumerate(questions)}
config = PoolConfig(
task_types=("Object Recognition",), seed=42, baseline_run_id="bl",
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,
task_types=("Object Recognition",),
seed=42,
baseline_run_id="bl",
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()