feat(tools): generate_questions.py generate 子命令

- VLM 出题 + embedding 去重 + 断点续跑 + 并发控制
- 单线程汇总点保证去重原子性
- 18 个单元测试覆盖 progress/exemplar/pool rebuild/JSON append

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-07-09 05:42:50 -04:00
parent 6e46d184b8
commit 11f3c90200
2 changed files with 959 additions and 0 deletions
+322
View File
@@ -0,0 +1,322 @@
"""tools/generate_questions.py 单元测试。
覆盖断点续跑、exemplar 选取、embedding 池重建、JSON 追加写入等纯函数。
"""
from __future__ import annotations
import json
import random
import sys
from pathlib import Path
import numpy as np
# 确保项目根目录在 sys.path 中
PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
from core.types import GeneratedQuestion
from tools.generate_questions import (
_append_to_json,
_load_or_init_progress,
_rebuild_embedding_pool,
_save_progress,
_select_exemplars,
)
# ---------------------------------------------------------------------------
# 辅助工厂
# ---------------------------------------------------------------------------
def _make_question(
qid: str = "q1",
vid: str = "v1",
task_type: str = "Object Recognition",
question: str = "What is this?",
answer: str = "A",
) -> GeneratedQuestion:
"""构造测试用 GeneratedQuestion。"""
return GeneratedQuestion(
question_id=qid,
video_id=vid,
task_type=task_type,
question=question,
options=("A. X", "B. Y", "C. Z", "D. W"),
answer=answer,
source_nodes=(),
difficulty="medium",
)
# ---------------------------------------------------------------------------
# TestLoadOrInitProgress
# ---------------------------------------------------------------------------
class TestLoadOrInitProgress:
"""_load_or_init_progress 测试。"""
def test_init_fresh(self, tmp_path: Path) -> None:
"""目录为空时返回初始结构。"""
progress = _load_or_init_progress(tmp_path)
assert progress["completed"] == {}
assert progress["output_dir"] == str(tmp_path)
def test_load_existing(self, tmp_path: Path) -> None:
"""已有 progress.json 时正确加载。"""
data = {
"completed": {"Object Recognition": ["gen-x-001"]},
"output_dir": str(tmp_path),
}
(tmp_path / "progress.json").write_text(json.dumps(data))
progress = _load_or_init_progress(tmp_path)
assert "gen-x-001" in progress["completed"]["Object Recognition"]
def test_corrupted_json_reinits(self, tmp_path: Path) -> None:
"""损坏的 JSON 文件导致重新初始化。"""
(tmp_path / "progress.json").write_text("{invalid json")
progress = _load_or_init_progress(tmp_path)
assert progress["completed"] == {}
def test_invalid_completed_type_reinits(self, tmp_path: Path) -> None:
"""completed 字段类型不正确时重新初始化。"""
(tmp_path / "progress.json").write_text(
json.dumps({"completed": "not-a-dict", "output_dir": str(tmp_path)})
)
progress = _load_or_init_progress(tmp_path)
assert progress["completed"] == {}
# ---------------------------------------------------------------------------
# TestSaveProgress
# ---------------------------------------------------------------------------
class TestSaveProgress:
"""_save_progress 测试。"""
def test_atomic_write(self, tmp_path: Path) -> None:
"""原子写入 progress.json。"""
progress = {
"completed": {"Action Reasoning": ["gen-v1-001"]},
"output_dir": str(tmp_path),
}
_save_progress(tmp_path, progress)
written = json.loads((tmp_path / "progress.json").read_text())
assert written["completed"]["Action Reasoning"] == ["gen-v1-001"]
# 临时文件不应残留
assert not (tmp_path / "progress.json.tmp").exists()
# ---------------------------------------------------------------------------
# TestSelectExemplars
# ---------------------------------------------------------------------------
class TestSelectExemplars:
"""_select_exemplars 测试。"""
def test_selects_correct_type(self) -> None:
"""只选取匹配题型的示例。"""
qs = [
_make_question("q1", "v1", "Object Recognition", "Q1?"),
_make_question("q2", "v2", "Object Recognition", "Q2?"),
_make_question("q3", "v1", "Action Reasoning", "Q3?"),
]
result = _select_exemplars(qs, "Object Recognition", 3, random.Random(42))
assert all(q.task_type == "Object Recognition" for q in result)
assert len(result) == 2 # 只有 2 个可用
def test_cross_video_diversity(self) -> None:
"""优先从不同 video_id 选取示例。"""
qs = [_make_question(f"q{i}", f"v{i}", "Object Recognition", f"Q{i}?") for i in range(10)]
result = _select_exemplars(qs, "Object Recognition", 3, random.Random(42))
video_ids = {q.video_id for q in result}
assert len(video_ids) == 3 # 全部来自不同视频
def test_empty_benchmark(self) -> None:
"""benchmark 为空时返回空列表。"""
result = _select_exemplars([], "Object Recognition", 3, random.Random(42))
assert result == []
def test_no_matching_type(self) -> None:
"""无匹配题型时返回空列表。"""
qs = [_make_question("q1", "v1", "Action Reasoning", "Q1?")]
result = _select_exemplars(qs, "Object Recognition", 3, random.Random(42))
assert result == []
def test_request_more_than_available(self) -> None:
"""请求数超过可用数时返回全部。"""
qs = [
_make_question("q1", "v1", "Object Recognition", "Q1?"),
_make_question("q2", "v2", "Object Recognition", "Q2?"),
]
result = _select_exemplars(qs, "Object Recognition", 10, random.Random(42))
assert len(result) == 2
# ---------------------------------------------------------------------------
# TestProgressResume
# ---------------------------------------------------------------------------
class TestProgressResume:
"""断点续跑集成测试。"""
def test_skips_completed_and_rebuilds_pool(self, tmp_path: Path) -> None:
"""已完成的题目从 progress 加载,embedding 池从已生成 JSON 重建。"""
output_dir = tmp_path / "output"
output_dir.mkdir()
(output_dir / "test_vid.json").write_text(
json.dumps(
[
{
"question_id": "gen-test_vid-001",
"task_type": "Object Recognition",
"question": "Existing question?",
"options": ["A. X", "B. Y", "C. Z", "D. W"],
"answer": "A",
"source_nodes": ["L3_001"],
"difficulty": "medium",
}
]
)
)
progress = {
"completed": {"Object Recognition": ["gen-test_vid-001"]},
"output_dir": str(output_dir),
}
(output_dir / "progress.json").write_text(json.dumps(progress))
loaded = _load_or_init_progress(output_dir)
assert "gen-test_vid-001" in loaded["completed"]["Object Recognition"]
# ---------------------------------------------------------------------------
# TestRebuildEmbeddingPool
# ---------------------------------------------------------------------------
class TestRebuildEmbeddingPool:
"""_rebuild_embedding_pool 测试。"""
@staticmethod
def _fake_embed_fn(texts):
"""伪嵌入函数:返回固定维度的随机向量。"""
if isinstance(texts, str):
texts = [texts]
return np.random.RandomState(0).randn(len(texts), 8).astype(np.float32)
def test_empty_dir(self, tmp_path: Path) -> None:
"""空目录时所有池为空。"""
pools = _rebuild_embedding_pool(tmp_path, self._fake_embed_fn, [])
# 所有 12 种题型都应有条目
assert len(pools) >= 12
for v in pools.values():
assert v.shape[0] == 0 or v.ndim == 2
def test_with_generated_json(self, tmp_path: Path) -> None:
"""从已生成的 JSON 文件重建池。"""
(tmp_path / "vid1.json").write_text(
json.dumps(
[
{
"question_id": "gen-vid1-001",
"task_type": "Object Recognition",
"question": "Test question 1?",
"options": ["A. X", "B. Y", "C. Z", "D. W"],
"answer": "A",
"source_nodes": [],
"difficulty": "medium",
},
{
"question_id": "gen-vid1-002",
"task_type": "Object Recognition",
"question": "Test question 2?",
"options": ["A. X", "B. Y", "C. Z", "D. W"],
"answer": "B",
"source_nodes": [],
"difficulty": "medium",
},
]
)
)
pools = _rebuild_embedding_pool(tmp_path, self._fake_embed_fn, [])
assert pools["Object Recognition"].shape[0] == 2
assert pools["Object Recognition"].shape[1] == 8
def test_with_benchmark_questions(self, tmp_path: Path) -> None:
"""benchmark 题目也加入去重池。"""
benchmark = [
_make_question("bm1", "v1", "Action Reasoning", "Benchmark Q1?"),
_make_question("bm2", "v2", "Action Reasoning", "Benchmark Q2?"),
]
pools = _rebuild_embedding_pool(tmp_path, self._fake_embed_fn, benchmark)
assert pools["Action Reasoning"].shape[0] == 2
def test_progress_json_excluded(self, tmp_path: Path) -> None:
"""progress.json 不被当作题目文件。"""
(tmp_path / "progress.json").write_text(
json.dumps({"completed": {}, "output_dir": str(tmp_path)})
)
pools = _rebuild_embedding_pool(tmp_path, self._fake_embed_fn, [])
for v in pools.values():
assert v.shape[0] == 0 or v.ndim == 2
# ---------------------------------------------------------------------------
# TestAppendToJson
# ---------------------------------------------------------------------------
class TestAppendToJson:
"""_append_to_json 测试。"""
def test_create_new_file(self, tmp_path: Path) -> None:
"""文件不存在时创建新文件。"""
q = _make_question("gen-v1-001", "v1", "Object Recognition", "Q?")
_append_to_json(tmp_path, q)
written = json.loads((tmp_path / "v1.json").read_text())
assert len(written) == 1
assert written[0]["question_id"] == "gen-v1-001"
def test_append_to_existing(self, tmp_path: Path) -> None:
"""追加到已有文件。"""
existing = [
{
"question_id": "gen-v1-001",
"task_type": "Object Recognition",
"question": "Existing?",
"options": ["A. X", "B. Y", "C. Z", "D. W"],
"answer": "A",
"source_nodes": [],
"difficulty": "medium",
}
]
(tmp_path / "v1.json").write_text(json.dumps(existing))
q = _make_question("gen-v1-002", "v1", "Object Recognition", "New?")
_append_to_json(tmp_path, q)
written = json.loads((tmp_path / "v1.json").read_text())
assert len(written) == 2
assert written[1]["question_id"] == "gen-v1-002"
def test_different_video_ids(self, tmp_path: Path) -> None:
"""不同 video_id 写入不同文件。"""
q1 = _make_question("gen-v1-001", "v1", "Object Recognition", "Q1?")
q2 = _make_question("gen-v2-001", "v2", "Action Reasoning", "Q2?")
_append_to_json(tmp_path, q1)
_append_to_json(tmp_path, q2)
assert (tmp_path / "v1.json").exists()
assert (tmp_path / "v2.json").exists()
v1_data = json.loads((tmp_path / "v1.json").read_text())
v2_data = json.loads((tmp_path / "v2.json").read_text())
assert len(v1_data) == 1
assert len(v2_data) == 1
+637
View File
@@ -0,0 +1,637 @@
#!/usr/bin/env python3
"""赛题生成工具:generate + calibrate。
用法:
conda activate Video-Tree-TRM
python tools/generate_questions.py generate --store-dir store ...
python tools/generate_questions.py calibrate ...
app/core/adapters 不 import 此脚本。
"""
from __future__ import annotations
import argparse
import asyncio
import json
import os
import random
import sys
from pathlib import Path
PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT_ROOT))
from dotenv import load_dotenv
from loguru import logger
load_dotenv(PROJECT_ROOT / ".env")
import numpy as np
from app.question_gen.loader import load_benchmark
from app.question_gen.synthesizer import (
TASK_TYPE_LEVEL_MAP,
generate_one,
is_duplicate,
)
from core.types import GeneratedQuestion # noqa: TCH001 — runtime use in _append_to_json
# ---------------------------------------------------------------------------
# 日志配置:不缓存,立即输出
# ---------------------------------------------------------------------------
logger.remove()
logger.add(
sys.stderr,
format="{time:HH:mm:ss} | {level:<7} | {message}",
level="DEBUG",
colorize=True,
)
logger.add(
PROJECT_ROOT / "logs" / "generate_questions.log",
format="{time:YYYY-MM-DD HH:mm:ss} | {level:<7} | {message}",
level="DEBUG",
rotation="50 MB",
)
# ---------------------------------------------------------------------------
# 断点续跑 — progress 文件管理
# ---------------------------------------------------------------------------
def _load_or_init_progress(output_dir: Path) -> dict:
"""加载 progress.json,不存在则返回初始结构。
参数:
output_dir: 输出目录路径。
返回:
{"completed": {task_type: [question_id, ...]}, "output_dir": str}。
文件损坏时返回初始结构并记录警告。
"""
progress_path = output_dir / "progress.json"
if progress_path.exists():
try:
data = json.loads(progress_path.read_text(encoding="utf-8"))
if not isinstance(data.get("completed"), dict):
raise ValueError("completed 字段不是 dict")
return data
except (json.JSONDecodeError, ValueError, KeyError, TypeError) as exc:
logger.warning("progress.json 损坏,重新初始化: {}", exc)
return {"completed": {}, "output_dir": str(output_dir)}
def _save_progress(output_dir: Path, progress: dict) -> None:
"""原子写入 progress.json。
参数:
output_dir: 输出目录路径。
progress: 进度数据。
"""
tmp = output_dir / "progress.json.tmp"
tmp.write_text(
json.dumps(progress, ensure_ascii=False, indent=2),
encoding="utf-8",
)
os.replace(str(tmp), str(output_dir / "progress.json"))
# ---------------------------------------------------------------------------
# Embedding 池重建(断点续跑时从已生成 JSON 重建)
# ---------------------------------------------------------------------------
def _rebuild_embedding_pool(
output_dir: Path,
embed_fn,
benchmark_questions: list[GeneratedQuestion],
) -> dict[str, np.ndarray]:
"""从已生成 JSON + benchmark 题目重建每个题型的 embedding 池。
断点续跑时调用,确保去重池包含所有已有题目。
参数:
output_dir: 包含 {video_id}.json 的输出目录。
embed_fn: 文本嵌入函数(str | list[str] → [N, D] ndarray)。
benchmark_questions: benchmark 题目列表(也要加入去重池)。
返回:
{task_type: [N, D] ndarray},空题型的 ndarray 为 shape (0,)。
"""
pools: dict[str, list[str]] = {}
# Phase 1: 收集 benchmark 题目文本
for q in benchmark_questions:
pools.setdefault(q.task_type, []).append(q.question)
# Phase 2: 收集已生成题目文本
for json_path in sorted(output_dir.glob("*.json")):
if json_path.name == "progress.json":
continue
try:
items = json.loads(json_path.read_text(encoding="utf-8"))
except (json.JSONDecodeError, OSError) as exc:
logger.warning("跳过损坏文件 {}: {}", json_path, exc)
continue
if not isinstance(items, list):
continue
for item in items:
task_type = item.get("task_type", "")
question = item.get("question", "")
if task_type and question:
pools.setdefault(task_type, []).append(question)
# Phase 3: 批量嵌入
result: dict[str, np.ndarray] = {}
for task_type, texts in pools.items():
if texts:
result[task_type] = embed_fn(texts)
else:
result[task_type] = np.empty(0)
# Phase 4: 确保所有 12 题型都有条目
for task_type in TASK_TYPE_LEVEL_MAP:
if task_type not in result:
result[task_type] = np.empty(0)
logger.info(
"embedding 池重建完成: {}",
{k: v.shape[0] if v.ndim == 2 else 0 for k, v in result.items()},
)
return result
# ---------------------------------------------------------------------------
# Exemplar 选取
# ---------------------------------------------------------------------------
def _select_exemplars(
benchmark: list[GeneratedQuestion],
task_type: str,
n: int,
rng: random.Random,
) -> list[GeneratedQuestion]:
"""从 benchmark 中选取同题型示例,优先跨视频多样性。
参数:
benchmark: benchmark 题目全集。
task_type: 目标题型。
n: 期望选取数量。
rng: 可控随机数生成器。
返回:
min(n, 可用数) 个示例,尽量来自不同 video_id。
"""
# Phase 1: 过滤同题型
candidates = [q for q in benchmark if q.task_type == task_type]
if not candidates:
return []
take = min(n, len(candidates))
# Phase 2: 按 video_id 分桶,轮询取样保证跨视频多样性
by_video: dict[str, list[GeneratedQuestion]] = {}
for q in candidates:
by_video.setdefault(q.video_id, []).append(q)
# 每桶内部打乱
for bucket in by_video.values():
rng.shuffle(bucket)
# 轮询选取
video_ids = list(by_video.keys())
rng.shuffle(video_ids)
selected: list[GeneratedQuestion] = []
idx = 0
while len(selected) < take:
vid = video_ids[idx % len(video_ids)]
bucket = by_video[vid]
if bucket:
selected.append(bucket.pop(0))
else:
# 桶空了,从 video_ids 中移除
video_ids.remove(vid)
if not video_ids:
break
# 不递增 idx,因为移除后当前位置是下一个
continue
idx += 1
return selected
# ---------------------------------------------------------------------------
# 客户端构建(从 .env
# ---------------------------------------------------------------------------
def _build_vlm_client():
"""构建 GovernedVLMClient,复用 repair_trees.py 的模式。
从 .env 读取 VL_LLM_MODEL / VL_LLM_BASE_URL / VL_LLM_API_KEY
和 LLM 韧性参数,构造治理栈。
返回:
GovernedVLMClient 实例。
"""
from adapters.breaker import CircuitBreaker
from adapters.llm import GovernedLLMClient
from adapters.telemetry import SQLiteTelemetryRecorder
from adapters.vlm import GovernedVLMClient
(PROJECT_ROOT / "logs").mkdir(exist_ok=True)
telemetry = SQLiteTelemetryRecorder(
str(PROJECT_ROOT / "logs" / "generate_questions_telemetry.db")
)
breaker_threshold = int(os.getenv("LLM_CIRCUIT_BREAKER_THRESHOLD", "5"))
breaker_cooldown = int(os.getenv("LLM_CIRCUIT_BREAKER_COOLDOWN", "60"))
timeout_s = float(os.getenv("LLM_TIMEOUT", "120"))
max_retries = int(os.getenv("LLM_MAX_RETRIES", "3"))
base_delay = float(os.getenv("LLM_RETRY_BASE_DELAY", "2.0"))
max_delay = float(os.getenv("LLM_RETRY_MAX_DELAY", "30.0"))
ttft = float(os.getenv("LLM_TTFT_TIMEOUT", "30"))
inter_token = float(os.getenv("LLM_INTER_TOKEN_TIMEOUT", "15"))
vlm_base = GovernedLLMClient(
model=os.environ["VL_LLM_MODEL"],
base_url=os.environ["VL_LLM_BASE_URL"],
api_key=os.environ["VL_LLM_API_KEY"],
provider="qwen",
thinking=False,
breaker=CircuitBreaker(fail_threshold=breaker_threshold, cooldown_s=breaker_cooldown),
cache=None,
telemetry=telemetry,
timeout_s=timeout_s,
ttft_timeout_s=ttft,
inter_token_timeout_s=inter_token,
max_retries=max_retries,
retry_base_delay_s=base_delay,
retry_max_delay_s=max_delay,
)
return GovernedVLMClient(vlm_base)
def _build_embed_provider():
"""构建 EmbeddingProvider,从 .env 决定 local 或 remote。
环境变量:
EMBED_API_KEY + EMBED_API_URL 都非空 → RemoteEmbeddingProvider
否则 → LocalEmbeddingProvider
模型名称和维度通过 EMBED_MODEL / EMBED_DIM 环境变量配置。
返回:
LocalEmbeddingProvider 或 RemoteEmbeddingProvider 实例。
"""
from adapters.embedding import LocalEmbeddingProvider, RemoteEmbeddingProvider
model_name = os.environ.get("EMBED_MODEL", "BAAI/bge-base-zh-v1.5")
embed_dim = int(os.environ.get("EMBED_DIM", "768"))
api_key = os.environ.get("EMBED_API_KEY", "")
api_url = os.environ.get("EMBED_API_URL", "")
if api_key and api_url:
logger.info("使用远程嵌入: model={}, url={}", model_name, api_url)
return RemoteEmbeddingProvider(
model_name=model_name,
embed_dim=embed_dim,
api_key=api_key,
api_url=api_url,
)
logger.info("使用本地嵌入: model={}, dim={}", model_name, embed_dim)
device = os.environ.get("EMBED_DEVICE", "cpu")
return LocalEmbeddingProvider(
model_name=model_name,
embed_dim=embed_dim,
device=device,
)
# ---------------------------------------------------------------------------
# JSON 追加写入
# ---------------------------------------------------------------------------
def _append_to_json(output_dir: Path, question: GeneratedQuestion) -> None:
"""将生成的题目追加到对应 video_id 的 JSON 文件。
文件格式:[{...}, {...}, ...],每个 video_id 一个文件。
参数:
output_dir: 输出目录。
question: 待写入的题目。
"""
json_path = output_dir / f"{question.video_id}.json"
existing: list[dict] = []
if json_path.exists():
try:
existing = json.loads(json_path.read_text(encoding="utf-8"))
except (json.JSONDecodeError, OSError):
logger.warning("读取 {} 失败,覆盖写入", json_path)
existing = []
entry = {
"question_id": question.question_id,
"task_type": question.task_type,
"question": question.question,
"options": list(question.options),
"answer": question.answer,
"source_nodes": list(question.source_nodes),
"difficulty": question.difficulty,
}
existing.append(entry)
# 原子写入
tmp = json_path.with_suffix(".json.tmp")
tmp.write_text(
json.dumps(existing, ensure_ascii=False, indent=2),
encoding="utf-8",
)
os.replace(str(tmp), str(json_path))
# ---------------------------------------------------------------------------
# generate 主流程
# ---------------------------------------------------------------------------
async def _run_generate(args: argparse.Namespace) -> None:
"""generate 子命令主流程。
按题型顺序生成题目,每个 slot 串行生成并去重,
断点续跑通过 progress.json 跳过已完成 slot。
参数:
args: CLI 参数(store_dir, output_dir, per_type, similarity_threshold,
max_retries, concurrency, seed)。
"""
store_dir = Path(args.store_dir)
output_dir = Path(args.output_dir)
per_type = args.per_type
similarity_threshold = args.similarity_threshold
max_retries = args.max_retries
concurrency = args.concurrency
seed = args.seed
output_dir.mkdir(parents=True, exist_ok=True)
# Phase 1: 加载视频列表
videos_dir = store_dir / "videos"
if not videos_dir.exists():
logger.error("视频目录不存在: {}", videos_dir)
sys.exit(1)
video_ids = sorted(
d.name for d in videos_dir.iterdir() if d.is_dir() and (d / "tree.json").exists()
)
if not video_ids:
logger.error("未找到任何有 tree.json 的视频目录")
sys.exit(1)
logger.info("发现 {} 个视频", len(video_ids))
# Phase 2: 加载 benchmark 题目(用于 exemplars + 去重池初始化)
benchmark_dir = store_dir / "questions" / "benchmarks" / "Video-MME"
benchmark: list[GeneratedQuestion] = []
if benchmark_dir.exists():
benchmark = load_benchmark(benchmark_dir)
logger.info("加载 {} 道 benchmark 题目", len(benchmark))
else:
logger.warning("benchmark 目录不存在: {}", benchmark_dir)
# Phase 3: 构建客户端
vlm = _build_vlm_client()
embed_provider = _build_embed_provider()
embed_fn = embed_provider.embed
# Phase 4: 初始化或恢复 progress
progress = _load_or_init_progress(output_dir)
# Phase 5: 重建 embedding 池
pools = _rebuild_embedding_pool(output_dir, embed_fn, benchmark)
# Phase 6: 加载视频树索引(延迟按需加载)
from app.tree.index import TreeIndex
rng = random.Random(seed)
sem = asyncio.Semaphore(concurrency)
task_types = list(TASK_TYPE_LEVEL_MAP.keys())
total_generated = 0
total_failed = 0
async def _generate_with_sem(
vlm_client,
embed_fn_inner,
tree,
video_id,
task_type,
seq,
*,
exemplars,
used_node_ids,
max_retries_inner,
similarity_threshold_inner,
rng_inner,
session_id,
):
"""Semaphore 包装的 generate_one 调用。"""
async with sem:
return await generate_one(
vlm_client,
embed_fn_inner,
tree,
video_id,
task_type,
seq,
exemplars=exemplars,
used_node_ids=used_node_ids,
max_retries=max_retries_inner,
similarity_threshold=similarity_threshold_inner,
rng=rng_inner,
session_id=session_id,
)
# Phase 7: 逐题型、逐 slot 生成
for task_type in task_types:
completed_ids = set(progress["completed"].get(task_type, []))
start_seq = len(completed_ids)
if start_seq >= per_type:
logger.info("题型 {} 已完成 {}/{}", task_type, start_seq, per_type)
continue
logger.info(
"题型 {} 开始生成: 已完成 {}, 目标 {}",
task_type,
start_seq,
per_type,
)
# 选取 exemplars
exemplars = _select_exemplars(benchmark, task_type, 3, rng)
for seq in range(start_seq, per_type):
# 随机选一个视频
video_id = rng.choice(video_ids)
tree_path = videos_dir / video_id / "tree.json"
try:
tree = TreeIndex.load_json(str(tree_path))
except Exception as exc:
logger.warning("加载树 {} 失败: {}", tree_path, exc)
total_failed += 1
continue
used_node_ids: set[str] = set()
session_id = f"gen-{task_type}-{seq}"
generated = False
for _attempt in range(max_retries):
candidate = await _generate_with_sem(
vlm,
embed_fn,
tree,
video_id,
task_type,
seq,
exemplars=exemplars,
used_node_ids=used_node_ids,
max_retries_inner=1,
similarity_threshold_inner=similarity_threshold,
rng_inner=rng,
session_id=session_id,
)
if candidate is None:
continue
# 去重检查(单线程原子操作)
pool = pools.get(task_type, np.empty(0))
if (
pool.ndim == 2
and pool.shape[0] > 0
and is_duplicate(candidate.question, pool, embed_fn, similarity_threshold)
):
logger.warning("去重: {} 与池中题目相似", candidate.question_id)
continue
# 原子操作:更新池 + 写 JSON + 更新 progress
new_emb = embed_fn(candidate.question) # [1, D]
if pool.ndim == 2 and pool.shape[0] > 0:
pools[task_type] = np.vstack([pool, new_emb])
else:
pools[task_type] = new_emb
_append_to_json(output_dir, candidate)
progress["completed"].setdefault(task_type, []).append(candidate.question_id)
_save_progress(output_dir, progress)
total_generated += 1
generated = True
logger.debug(
"生成: {} (题型={}, 序号={})",
candidate.question_id,
task_type,
seq,
)
break
if not generated:
logger.error("题型 {} seq {} 耗尽 {} 次重试", task_type, seq, max_retries)
total_failed += 1
# Phase 8: 汇总
logger.info("=" * 60)
logger.info("生成完成: 成功 {}, 失败 {}", total_generated, total_failed)
logger.info("=" * 60)
if total_failed > 0:
logger.error("{} 个 slot 生成失败", total_failed)
sys.exit(1)
# 全部完成,删除 progress.json
progress_path = output_dir / "progress.json"
if progress_path.exists():
progress_path.unlink()
logger.info("已删除 progress.json(全部完成)")
# ---------------------------------------------------------------------------
# CLI 解析
# ---------------------------------------------------------------------------
def _parse_args() -> argparse.Namespace:
"""解析命令行参数。"""
parser = argparse.ArgumentParser(description="赛题生成工具:generate + calibrate")
subparsers = parser.add_subparsers(dest="command", required=True)
# generate 子命令
gen_parser = subparsers.add_parser("generate", help="生成新题目")
gen_parser.add_argument(
"--store-dir",
type=str,
required=True,
help="store 根目录(包含 videos/ 和 questions/",
)
gen_parser.add_argument(
"--output-dir",
type=str,
required=True,
help="输出目录(生成的 JSON 写入此处)",
)
gen_parser.add_argument(
"--per-type",
type=int,
required=True,
help="每种题型生成数量",
)
gen_parser.add_argument(
"--similarity-threshold",
type=float,
required=True,
help="embedding 去重阈值(余弦相似度)",
)
gen_parser.add_argument(
"--max-retries",
type=int,
required=True,
help="每个 slot 最大重试次数",
)
gen_parser.add_argument(
"--concurrency",
type=int,
required=True,
help="VLM 调用并发数(Semaphore 容量)",
)
gen_parser.add_argument(
"--seed",
type=int,
required=True,
help="随机种子",
)
# calibrate 子命令(占位,后续任务实现)
subparsers.add_parser("calibrate", help="校准题目难度(待实现)")
return parser.parse_args()
def main() -> None:
"""同步入口。"""
args = _parse_args()
(PROJECT_ROOT / "logs").mkdir(exist_ok=True)
if args.command == "generate":
asyncio.run(_run_generate(args))
elif args.command == "calibrate":
logger.error("calibrate 子命令尚未实现")
sys.exit(1)
if __name__ == "__main__":
main()