feat(question_gen): stratified_sample — 分层采样 + 题型保底
算法 100% 保真 TRM4: task_types 过滤、correctness.get(id, False) 语义、 对题在前返回顺序、min_per_class 遍历 pool 全部题型(含稀疏类)。 所有参数显式传入,无默认值。 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -7,6 +7,7 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import json
|
import json
|
||||||
|
import random
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
from core.types import GeneratedQuestion
|
from core.types import GeneratedQuestion
|
||||||
@@ -48,3 +49,119 @@ def load_benchmark(questions_dir: Path) -> list[GeneratedQuestion]:
|
|||||||
)
|
)
|
||||||
)
|
)
|
||||||
return results
|
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
|
||||||
|
|||||||
@@ -7,7 +7,7 @@ from pathlib import Path
|
|||||||
|
|
||||||
import pytest
|
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
|
from core.types import GeneratedQuestion
|
||||||
|
|
||||||
|
|
||||||
@@ -127,3 +127,235 @@ class TestLoadBenchmark:
|
|||||||
(tmp_path / "vid.json").write_text(json.dumps(data), encoding="utf-8")
|
(tmp_path / "vid.json").write_text(json.dumps(data), encoding="utf-8")
|
||||||
with pytest.raises(KeyError):
|
with pytest.raises(KeyError):
|
||||||
load_benchmark(tmp_path)
|
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
|
||||||
|
|||||||
Reference in New Issue
Block a user