From 11f3c9020083f77d8930b4713d7db18696389621 Mon Sep 17 00:00:00 2001 From: iomgaa Date: Thu, 9 Jul 2026 05:42:50 -0400 Subject: [PATCH] =?UTF-8?q?feat(tools):=20generate=5Fquestions.py=20genera?= =?UTF-8?q?te=20=E5=AD=90=E5=91=BD=E4=BB=A4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - VLM 出题 + embedding 去重 + 断点续跑 + 并发控制 - 单线程汇总点保证去重原子性 - 18 个单元测试覆盖 progress/exemplar/pool rebuild/JSON append Co-Authored-By: Claude Opus 4.6 (1M context) --- tests/unit/test_generate_questions.py | 322 +++++++++++++ tools/generate_questions.py | 637 ++++++++++++++++++++++++++ 2 files changed, 959 insertions(+) create mode 100644 tests/unit/test_generate_questions.py create mode 100644 tools/generate_questions.py diff --git a/tests/unit/test_generate_questions.py b/tests/unit/test_generate_questions.py new file mode 100644 index 0000000..3d608af --- /dev/null +++ b/tests/unit/test_generate_questions.py @@ -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 diff --git a/tools/generate_questions.py b/tools/generate_questions.py new file mode 100644 index 0000000..741de37 --- /dev/null +++ b/tools/generate_questions.py @@ -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()