206 lines
8.1 KiB
Python
206 lines
8.1 KiB
Python
"""题目加载与分层采样。
|
||
|
||
从 benchmark JSON 目录加载题目，提供按对错比例的分层采样。
|
||
对应训练循环中的 DataLoader 角色。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import random
|
||
from typing import TYPE_CHECKING
|
||
|
||
from core.types import GeneratedQuestion
|
||
|
||
if TYPE_CHECKING:
|
||
from pathlib import Path
|
||
|
||
from core.types import QuestionUnit
|
||
|
||
_LEGACY_DEFAULT_DIFFICULTY = "medium"
|
||
|
||
|
||
def load_benchmark(questions_dir: Path) -> list[GeneratedQuestion]:
|
||
"""从 benchmark JSON 目录加载题目列表。
|
||
|
||
video_id 优先使用题目 JSON 中的 ``video_id`` 字段；若缺失则回退到
|
||
文件名（不含扩展名）。Video-MME benchmark 按视频拆文件（文件名即
|
||
video_id），v2 生成题把多视频题目合并在单个 JSON 中（每条记录自带
|
||
``video_id``），两种格式均兼容。
|
||
|
||
pair 契约字段（``pair_id`` / ``question_role`` / ``flip_axis`` / ``unit_id``）
|
||
用 ``.get`` 读取：旧 benchmark 无这些键时退化为 single（``question_role``
|
||
默认 "single"，``unit_id`` 留空由 __post_init__ 回填为 question_id），
|
||
保证历史题库可无缝加载。
|
||
|
||
参数:
|
||
questions_dir: 包含 *.json 文件的目录路径。
|
||
|
||
返回:
|
||
按文件名排序加载的题目列表。
|
||
"""
|
||
results: list[GeneratedQuestion] = []
|
||
for path in sorted(questions_dir.glob("*.json")):
|
||
fallback_video_id = path.stem
|
||
with open(path, encoding="utf-8") as f:
|
||
qa_list: list[dict] = json.load(f)
|
||
for qa in qa_list:
|
||
results.append(
|
||
GeneratedQuestion(
|
||
question_id=qa["question_id"],
|
||
video_id=qa.get("video_id", fallback_video_id),
|
||
task_type=qa["task_type"],
|
||
question=qa["question"],
|
||
options=tuple(qa["options"]),
|
||
answer=qa["answer"],
|
||
source_nodes=tuple(qa.get("source_nodes", ())),
|
||
difficulty=qa.get("difficulty", _LEGACY_DEFAULT_DIFFICULTY),
|
||
family=qa.get("family"),
|
||
skill_target=qa.get("skill_target"),
|
||
difficulty_steps=qa.get("difficulty_steps"),
|
||
sub_pattern=qa.get("sub_pattern"),
|
||
# pair 契约字段：旧 benchmark 无这些键时按 single 默认兜底，
|
||
# unit_id 留空交由 GeneratedQuestion.__post_init__ 回填。
|
||
pair_id=qa.get("pair_id"),
|
||
question_role=qa.get("question_role", "single"),
|
||
flip_axis=qa.get("flip_axis"),
|
||
unit_id=qa.get("unit_id", ""),
|
||
)
|
||
)
|
||
return results
|
||
|
||
|
||
def stratified_sample(
|
||
questions: list[GeneratedQuestion],
|
||
correctness: dict[str, bool],
|
||
size: int,
|
||
correct_ratio: float | None,
|
||
task_types: list[str] | None,
|
||
seed: int,
|
||
min_per_class: int | None,
|
||
) -> list[GeneratedQuestion]:
|
||
"""按题型过滤后采样 size 个单元，可选按对错比例分层并按题型保底。
|
||
|
||
参数:
|
||
questions: 候选题目全集（single 与孪生对成员可混含）。
|
||
correctness: question_id -> 基线是否答对（单元级正确性取成员 AND）。
|
||
size: 采样单元总量（single 计 1、pair 计 1）。
|
||
correct_ratio: 采样中"基线答对"单元的占比；None 表示自然分布。
|
||
task_types: 限定题型；None 表示不限。
|
||
seed: 随机种子，保证可复现。
|
||
min_per_class: 每个题型补足到的单元下限；None 表示不补足。
|
||
|
||
返回:
|
||
采样后的题目列表（pair 单元展开为原始的两道题）。
|
||
|
||
异常:
|
||
ValueError: 自然分布时单元池不足 size，或分层时某层单元不足。
|
||
|
||
关键实现:
|
||
以 **QuestionUnit 为采样原子**（single 计 1、pair 计 1），size /
|
||
correct_ratio / min_per_class 均按 unit 计数，孪生对两题永不被劈开。
|
||
采样完成后 flatten_units 展开回逐题列表。纯 single 输入时 build_units
|
||
与题目一一对应、顺序不变，rng 消耗与旧逐题实现完全一致（字节级回归）。
|
||
|
||
build_units / flatten_units 采用函数内延迟导入：loader 属 question_gen，
|
||
question_units 属 harness，而 harness 包初始化会反向 import question_gen，
|
||
模块级导入将触发循环依赖（沿用 adversarial_filter 的既有做法）。
|
||
"""
|
||
from app.harness.question_units import build_units, flatten_units
|
||
|
||
rng = random.Random(seed)
|
||
units = build_units(questions)
|
||
pool = [u for u in units if task_types is None or u.task_type in task_types]
|
||
|
||
if correct_ratio is None:
|
||
if len(pool) < size:
|
||
raise ValueError(f"自然分布采样不足: 需 {size} 个单元, 实有 {len(pool)} 个")
|
||
sampled = rng.sample(pool, size)
|
||
else:
|
||
sampled = _ratio_stratified_sample(pool, correctness, size, correct_ratio, rng)
|
||
|
||
if min_per_class is not None:
|
||
sampled = _backfill_per_class(sampled, pool, min_per_class, rng)
|
||
return flatten_units(sampled)
|
||
|
||
|
||
def _ratio_stratified_sample(
|
||
pool: list[QuestionUnit],
|
||
correctness: dict[str, bool],
|
||
size: int,
|
||
correct_ratio: float,
|
||
rng: random.Random,
|
||
) -> list[QuestionUnit]:
|
||
"""按对错比例分层采样：对单元占 correct_ratio，其余为错单元。
|
||
|
||
参数:
|
||
pool: 题型过滤后的候选单元。
|
||
correctness: question_id -> 基线是否答对。
|
||
size: 采样单元总量。
|
||
correct_ratio: 对单元占比。
|
||
rng: 随机数发生器。
|
||
|
||
返回:
|
||
采样后的单元列表（对单元在前、错单元在后）。
|
||
|
||
异常:
|
||
ValueError: 对单元或错单元层不足。
|
||
|
||
关键实现:
|
||
unit_correctness 采用函数内延迟导入：loader 属 question_gen，
|
||
question_units 属 harness，模块级导入将触发循环依赖（沿用 build_units /
|
||
flatten_units 的既有做法）。以 strict=False 保持"缺基线对错即视为未答对"的
|
||
原 loose 语义不变。
|
||
"""
|
||
from app.harness.question_units import unit_correctness
|
||
|
||
correct = [u for u in pool if unit_correctness(u, correctness, strict=False)]
|
||
wrong = [u for u in pool if not unit_correctness(u, correctness, strict=False)]
|
||
n_correct = round(size * correct_ratio)
|
||
n_wrong = size - n_correct
|
||
if len(correct) < n_correct or len(wrong) < n_wrong:
|
||
raise ValueError(
|
||
f"分层不足: 需对{n_correct}/错{n_wrong}, 实有对{len(correct)}/错{len(wrong)}"
|
||
)
|
||
return rng.sample(correct, n_correct) + rng.sample(wrong, n_wrong)
|
||
|
||
|
||
def _backfill_per_class(
|
||
sampled: list[QuestionUnit],
|
||
pool: list[QuestionUnit],
|
||
min_per_class: int,
|
||
rng: random.Random,
|
||
) -> list[QuestionUnit]:
|
||
"""对候选池中出现的每个题型，将采样单元补足到 min_per_class 个。
|
||
|
||
遍历对象是候选池 pool 里出现的全部题型（非仅 sampled 命中的），
|
||
保证任意稀疏题型都能拿到足额样本。补足以 unit 为原子，孪生对整进整出。
|
||
|
||
参数:
|
||
sampled: 主采样结果单元（不修改，返回新列表）。
|
||
pool: 候选单元全集（补足来源 + 题型枚举来源）。
|
||
min_per_class: 每个题型的单元下限。
|
||
rng: 随机数发生器。
|
||
|
||
返回:
|
||
补足后的单元列表。
|
||
"""
|
||
selected_ids = {u.unit_id for u in sampled}
|
||
result = list(sampled)
|
||
counts: dict[str, int] = {}
|
||
for u in sampled:
|
||
counts[u.task_type] = counts.get(u.task_type, 0) + 1
|
||
ordered_task_types: dict[str, None] = {}
|
||
for u in pool:
|
||
ordered_task_types.setdefault(u.task_type, None)
|
||
for task_type in ordered_task_types:
|
||
deficit = min_per_class - counts.get(task_type, 0)
|
||
if deficit <= 0:
|
||
continue
|
||
candidates = [u for u in pool if u.task_type == task_type and u.unit_id not in selected_ids]
|
||
take = rng.sample(candidates, min(deficit, len(candidates)))
|
||
for u in take:
|
||
selected_ids.add(u.unit_id)
|
||
result.append(u)
|
||
return result
|