diff --git a/app/question_gen/loader.py b/app/question_gen/loader.py index 6febe36..d7f11b8 100644 --- a/app/question_gen/loader.py +++ b/app/question_gen/loader.py @@ -7,6 +7,7 @@ from __future__ import annotations import json +import random from typing import TYPE_CHECKING from core.types import GeneratedQuestion @@ -48,3 +49,119 @@ def load_benchmark(questions_dir: Path) -> list[GeneratedQuestion]: ) ) 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: 候选题目全集。 + correctness: question_id -> 基线是否答对。 + size: 采样总量。 + correct_ratio: 采样中"基线答对"题的占比;None 表示自然分布。 + task_types: 限定题型;None 表示不限。 + seed: 随机种子,保证可复现。 + min_per_class: 每个题型补足到的下限;None 表示不补足。 + + 返回: + 采样后的题目列表。 + + 异常: + ValueError: 自然分布时池不足 size,或分层时某层题目不足。 + """ + rng = random.Random(seed) + pool = [q for q in questions if task_types is None or q.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 sampled + + +def _ratio_stratified_sample( + pool: list[GeneratedQuestion], + correctness: dict[str, bool], + size: int, + correct_ratio: float, + rng: random.Random, +) -> list[GeneratedQuestion]: + """按对错比例分层采样:对题占 correct_ratio,其余为错题。 + + 参数: + pool: 题型过滤后的候选题。 + correctness: question_id -> 基线是否答对。 + size: 采样总量。 + correct_ratio: 对题占比。 + rng: 随机数发生器。 + + 返回: + 采样后的题目列表(对题在前、错题在后)。 + + 异常: + ValueError: 对题或错题层不足。 + """ + correct = [q for q in pool if correctness.get(q.question_id, False)] + wrong = [q for q in pool if not correctness.get(q.question_id, 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[GeneratedQuestion], + pool: list[GeneratedQuestion], + min_per_class: int, + rng: random.Random, +) -> list[GeneratedQuestion]: + """对候选池中出现的每个题型,将采样结果补足到 min_per_class 道。 + + 遍历对象是候选池 pool 里出现的全部题型(非仅 sampled 命中的), + 保证任意稀疏题型都能拿到足额样本。 + + 参数: + sampled: 主采样结果(不修改,返回新列表)。 + pool: 候选题全集(补足来源 + 题型枚举来源)。 + min_per_class: 每个题型的下限。 + rng: 随机数发生器。 + + 返回: + 补足后的题目列表。 + """ + selected_ids = {q.question_id for q in sampled} + result = list(sampled) + counts: dict[str, int] = {} + for q in sampled: + counts[q.task_type] = counts.get(q.task_type, 0) + 1 + ordered_task_types: dict[str, None] = {} + for q in pool: + ordered_task_types.setdefault(q.task_type, None) + for task_type in ordered_task_types: + deficit = min_per_class - counts.get(task_type, 0) + if deficit <= 0: + continue + candidates = [ + q for q in pool if q.task_type == task_type and q.question_id not in selected_ids + ] + take = rng.sample(candidates, min(deficit, len(candidates))) + for q in take: + selected_ids.add(q.question_id) + result.append(q) + return result diff --git a/tests/unit/test_question_loader.py b/tests/unit/test_question_loader.py index 7df8963..ebc645d 100644 --- a/tests/unit/test_question_loader.py +++ b/tests/unit/test_question_loader.py @@ -7,7 +7,7 @@ from pathlib import Path import pytest -from app.question_gen.loader import load_benchmark +from app.question_gen.loader import load_benchmark, stratified_sample from core.types import GeneratedQuestion @@ -127,3 +127,235 @@ class TestLoadBenchmark: (tmp_path / "vid.json").write_text(json.dumps(data), encoding="utf-8") with pytest.raises(KeyError): load_benchmark(tmp_path) + + +def _make_questions(n: int, task_type: str = "T") -> list[GeneratedQuestion]: + """辅助函数:批量构造题目。""" + return [ + GeneratedQuestion( + question_id=f"{task_type}-{i}", + video_id="v1", + task_type=task_type, + question=f"Q{i}?", + options=("A", "B", "C", "D"), + answer="A", + source_nodes=(), + difficulty="medium", + ) + for i in range(n) + ] + + +class TestStratifiedSample: + def test_natural_distribution(self) -> None: + questions = _make_questions(20) + result = stratified_sample( + questions=questions, + correctness={}, + size=10, + correct_ratio=None, + task_types=None, + seed=42, + min_per_class=None, + ) + assert len(result) == 10 + + def test_natural_distribution_pool_insufficient(self) -> None: + questions = _make_questions(5) + with pytest.raises(ValueError, match="自然分布采样不足"): + stratified_sample( + questions=questions, + correctness={}, + size=10, + correct_ratio=None, + task_types=None, + seed=42, + min_per_class=None, + ) + + def test_ratio_stratified(self) -> None: + questions = _make_questions(20) + correctness = {f"T-{i}": i < 10 for i in range(20)} + result = stratified_sample( + questions=questions, + correctness=correctness, + size=10, + correct_ratio=0.6, + task_types=None, + seed=42, + min_per_class=None, + ) + assert len(result) == 10 + correct_count = sum(1 for q in result if correctness.get(q.question_id, False)) + assert correct_count == 6 + + def test_ratio_stratified_correct_first(self) -> None: + questions = _make_questions(20) + correctness = {f"T-{i}": i < 10 for i in range(20)} + result = stratified_sample( + questions=questions, + correctness=correctness, + size=10, + correct_ratio=0.5, + task_types=None, + seed=42, + min_per_class=None, + ) + n_correct = round(10 * 0.5) + for q in result[:n_correct]: + assert correctness.get(q.question_id, False) is True + for q in result[n_correct:]: + assert correctness.get(q.question_id, False) is False + + def test_ratio_stratified_pool_insufficient(self) -> None: + questions = _make_questions(10) + correctness = {f"T-{i}": True for i in range(10)} + with pytest.raises(ValueError, match="分层不足"): + stratified_sample( + questions=questions, + correctness=correctness, + size=10, + correct_ratio=0.5, + task_types=None, + seed=42, + min_per_class=None, + ) + + def test_task_types_filter(self) -> None: + q_a = _make_questions(10, task_type="TypeA") + q_b = _make_questions(10, task_type="TypeB") + result = stratified_sample( + questions=q_a + q_b, + correctness={}, + size=5, + correct_ratio=None, + task_types=["TypeA"], + seed=42, + min_per_class=None, + ) + assert all(q.task_type == "TypeA" for q in result) + + def test_unknown_correctness_treated_as_wrong(self) -> None: + questions = _make_questions(20) + correctness = {f"T-{i}": True for i in range(10)} + result = stratified_sample( + questions=questions, + correctness=correctness, + size=10, + correct_ratio=0.5, + task_types=None, + seed=42, + min_per_class=None, + ) + n_correct = round(10 * 0.5) + for q in result[:n_correct]: + assert q.question_id in correctness + + def test_seed_reproducibility(self) -> None: + questions = _make_questions(20) + r1 = stratified_sample( + questions=questions, + correctness={}, + size=10, + correct_ratio=None, + task_types=None, + seed=123, + min_per_class=None, + ) + r2 = stratified_sample( + questions=questions, + correctness={}, + size=10, + correct_ratio=None, + task_types=None, + seed=123, + min_per_class=None, + ) + assert [q.question_id for q in r1] == [q.question_id for q in r2] + + def test_different_seeds_differ(self) -> None: + questions = _make_questions(20) + r1 = stratified_sample( + questions=questions, + correctness={}, + size=10, + correct_ratio=None, + task_types=None, + seed=1, + min_per_class=None, + ) + r2 = stratified_sample( + questions=questions, + correctness={}, + size=10, + correct_ratio=None, + task_types=None, + seed=2, + min_per_class=None, + ) + assert [q.question_id for q in r1] != [q.question_id for q in r2] + + def test_min_per_class_backfill(self) -> None: + q_a = _make_questions(10, task_type="TypeA") + q_b = _make_questions(10, task_type="TypeB") + all_q = q_a + q_b + correctness = {q.question_id: True for q in q_a[:5]} + result = stratified_sample( + questions=all_q, + correctness=correctness, + size=3, + correct_ratio=None, + task_types=None, + seed=42, + min_per_class=2, + ) + type_counts: dict[str, int] = {} + for q in result: + type_counts[q.task_type] = type_counts.get(q.task_type, 0) + 1 + assert type_counts.get("TypeA", 0) >= 2 + assert type_counts.get("TypeB", 0) >= 2 + + def test_min_per_class_partial_backfill(self) -> None: + q_sparse = _make_questions(1, task_type="Sparse") + q_main = _make_questions(10, task_type="Main") + result = stratified_sample( + questions=q_sparse + q_main, + correctness={}, + size=5, + correct_ratio=None, + task_types=None, + seed=42, + min_per_class=3, + ) + sparse_in_result = [q for q in result if q.task_type == "Sparse"] + assert len(sparse_in_result) == 1 + + def test_min_per_class_no_duplicates(self) -> None: + q_a = _make_questions(5, task_type="TypeA") + q_b = _make_questions(5, task_type="TypeB") + result = stratified_sample( + questions=q_a + q_b, + correctness={}, + size=3, + correct_ratio=None, + task_types=None, + seed=42, + min_per_class=2, + ) + ids = [q.question_id for q in result] + assert len(ids) == len(set(ids)) + + def test_backfill_enumerates_all_pool_types(self) -> None: + q_main = _make_questions(10, task_type="Main") + q_rare = _make_questions(3, task_type="Rare") + result = stratified_sample( + questions=q_main + q_rare, + correctness={}, + size=2, + correct_ratio=None, + task_types=None, + seed=0, + min_per_class=1, + ) + types_in_result = {q.task_type for q in result} + assert "Rare" in types_in_result