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:
@@ -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
|
||||||
@@ -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()
|
||||||
Reference in New Issue
Block a user