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

refactor(sampler): replace family_spec param with level+constraint

Browse Source 
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
iomgaa
2026-07-14 05:38:14 -04:00
parent b6b6a48503
commit e2325b6535
3 changed files with 100 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/question_gen/pipeline_v2.py
+70 -31
View File
@@ -38,7 +38,7 @@ from app.question_gen.families import QuestionFamilySpec, get_family_for_slot
from app.question_gen.gates import GateReport, GateResult, GateVerdict, run_gates
from app.question_gen.generator_v2 import CandidateQuestion, generate_one_v2
from app.question_gen.postprocess import run_postprocess
from app.question_gen.sampler_v2 import sample_material_v2
from app.question_gen.sampler_v2 import _TASK_TYPE_TO_LEVEL, sample_material_v2
from core.types import GeneratedQuestion
if TYPE_CHECKING:
@@ -101,8 +101,6 @@ class PipelineConfig:
concurrency: 并发 slot 数上限。
seed: 随机种子。
output_dir: 输出目录。
gate_models: 门控模型配置字典。
heavy_agent_model: 重量抽检使用的模型名。
"""
family_ratios: dict[str, float]
@@ -113,8 +111,6 @@ class PipelineConfig:
concurrency: int
seed: int
output_dir: Path
gate_models: dict[str, str]
heavy_agent_model: str
# ---------------------------------------------------------------------------
@@ -163,8 +159,6 @@ def load_pipeline_config(yaml_path: Path) -> PipelineConfig:
concurrency=int(section["concurrency"]),
seed=int(section["seed"]),
output_dir=Path(section["output_dir"]),
gate_models=section["gate"],
heavy_agent_model=str(section["heavy_agent_model"]),
)
@@ -269,6 +263,7 @@ def _is_duplicate(
def _to_generated_question(
candidate: CandidateQuestion,
*,
family: str,
options: tuple[str, ...] | None = None,
answer: str | None = None,
) -> GeneratedQuestion:
@@ -276,6 +271,7 @@ def _to_generated_question(
参数:
candidate: 门控通过的候选题目。
family: 问题家族名称（如 "RETRIEVAL"）。
options: 洗牌后的选项元组（若为 None 则使用 candidate 原始选项）。
answer: 重映射后的答案字母（若为 None 则使用 candidate 原始答案）。
@@ -291,6 +287,7 @@ def _to_generated_question(
answer=answer if answer is not None else candidate.answer,
source_nodes=candidate.source_nodes,
difficulty=candidate.difficulty,
family=family,
skill_target=candidate.skill_target,
difficulty_steps=None,
)
@@ -316,6 +313,7 @@ async def _process_one_slot(
*,
session_id: str,
run_id: str,
all_trees: dict[str, TreeIndex] | None = None,
) -> GeneratedQuestion | None:
"""处理单个 slot 的完整重出循环。
@@ -347,18 +345,36 @@ async def _process_one_slot(
返回:
GeneratedQuestion（通过全部检查）或 None（重出耗尽）。
"""
_RESAMPLE_VIDEO_INTERVAL = 1
async with sem:
prev_reason: str | None = None
current_tree = tree
current_video_id = slot.video_id
for attempt in range(1, config.retry_limit + 1):
# 连续失败 _RESAMPLE_VIDEO_INTERVAL 次后换视频
if attempt > 1 and (attempt - 1) % _RESAMPLE_VIDEO_INTERVAL == 0 and all_trees:
alt_ids = [v for v in all_trees if v != current_video_id]
if alt_ids:
current_video_id = rng.choice(alt_ids)
current_tree = all_trees[current_video_id]
logger.info(
"slot {} 连续 {} 次失败，换视频 {} 重试",
slot.slot_id,
attempt - 1,
current_video_id,
)
# Phase 1: 采样素材
try:
material = sample_material_v2(
tree=tree,
family_spec=slot.family,
tree=current_tree,
task_type=slot.task_type,
used_node_ids=used_node_ids,
rng=rng,
level=_TASK_TYPE_TO_LEVEL[slot.task_type],
constraint=slot.family.sampling,
)
except (RuntimeError, KeyError) as e:
logger.warning(
@@ -374,16 +390,16 @@ async def _process_one_slot(
try:
candidate = await generate_one_v2(
vlm=vlm,
tree=tree,
tree=current_tree,
material=material,
family_spec=slot.family,
task_type=slot.task_type,
seq=slot.seq,
video_id=slot.video_id,
video_id=current_video_id,
reject_reason=prev_reason,
session_id=session_id,
)
except (ValueError, FileNotFoundError) as e:
except (ValueError, FileNotFoundError, OSError, Exception) as e:
logger.warning(
"slot {} 生成失败 (attempt {}/{}): {}",
slot.slot_id,
@@ -399,7 +415,7 @@ async def _process_one_slot(
item_id=item_id,
run_id=run_id,
slot_id=slot.slot_id,
video_id=slot.video_id,
video_id=current_video_id,
family=slot.family.name,
task_type=slot.task_type,
skill_target=slot.family.skill_target,
@@ -446,15 +462,27 @@ async def _process_one_slot(
store.update_gates(item_id, verbatim_report)
continue
# Phase 6: 四门质量检查
report = await run_gates(
candidate=candidate,
tree=tree,
llm=llm,
family_spec=slot.family,
postprocess=pp,
session_id=session_id,
)
# Phase 6: 四门质量检查（key_verify 使用 VLM 看帧+文本）
try:
report = await run_gates(
candidate=candidate,
tree=tree,
llm=llm,
family_spec=slot.family,
postprocess=pp,
vlm=vlm,
session_id=session_id,
)
except Exception as e:
logger.warning(
"slot {} 门控调用异常 (attempt {}/{}): {}",
slot.slot_id,
attempt,
config.retry_limit,
e,
)
prev_reason = f"gate_error: {e}"
continue
store.update_gates(item_id, report)
if not report.passed:
@@ -481,7 +509,12 @@ async def _process_one_slot(
continue
# Phase 8: 通过全部检查 → 接受（使用洗牌后的选项和答案）
result = _to_generated_question(candidate, options=pp.options, answer=pp.answer)
result = _to_generated_question(
candidate,
family=slot.family.name,
options=pp.options,
answer=pp.answer,
)
# 将题目 embedding 加入池（flatten 确保 1D）
embed_pool.append(embed_fn(candidate.question).flatten())
# 标记使用的节点
@@ -571,16 +604,16 @@ async def _heavy_check_one(
_DEFAULT_TASK_TYPES: list[str] = [
"Action Recognition",
"Action Reasoning",
"Action Prediction",
"Action Sequence",
"Attribute Perception",
"Counting Problem",
"Information Synopsis",
"Object Recognition",
"Object Reasoning",
"Object Interaction",
"Scene Understanding",
"Event Reasoning",
"Causal Reasoning",
"Temporal Reasoning",
"OCR Problems",
"Spatial Perception",
"Spatial Reasoning",
"Temporal Perception",
"Temporal Reasoning",
]
@@ -697,6 +730,7 @@ async def run_pipeline_v2(
*,
task_types: list[str] | None = None,
progress: dict[str, str] | None = None,
on_accept: Callable[[GeneratedQuestion], None] | None = None,
) -> PipelineResult:
"""v2 出题管线主入口 — 编排全部 slot 的生成、检查与抽检。
@@ -718,6 +752,7 @@ async def run_pipeline_v2(
config: 管线配置。
task_types: 任务类型列表（默认使用 12 类标准集）。
progress: 已完成 slot 映射 {slot_id → "accepted"|"rejected"}。
on_accept: 每接受一题时的回调（用于实时持久化，防崩溃丢数据）。
返回:
PipelineResult 实例。
@@ -749,7 +784,7 @@ async def run_pipeline_v2(
if tree is None:
logger.warning("slot {} 对应视频 {} 的树不存在，跳过", slot.slot_id, slot.video_id)
return None
return await _process_one_slot(
result = await _process_one_slot(
slot=slot,
tree=tree,
vlm=vlm,
@@ -763,7 +798,11 @@ async def run_pipeline_v2(
sem=sem,
session_id=session_id,
run_id=run_id,
all_trees=trees,
)
if result is not None and on_accept is not None:
on_accept(result)
return result
results = await asyncio.gather(*[_process_wrapper(s) for s in pending_slots])