feat(harness): checkpoint 存 unit_id 序列,断点续跑孪生对不拆
核心算法保真 #3(断点续跑):checkpoint 从逐题 question_id 改为存 unit_id 序列(孪生对折叠为单个 unit_id),恢复时 build_units + 按完整 unit 展开, 续跑后 pair 两成员同进同出、绝不被劈开。 - _batch_unit_ids/_batch_from_ids 对称折叠/展开,保序去重,纯非 AR 下 unit_id==question_id、与旧逐题序列逐字节一致。 - momentum 采样抽取为 _sample_momentum_candidates 纯函数,docstring 显式 记录 Phase 1 设计偏差:仅保证纯非 AR byte-identical,混格 momentum 不保证。 - 新增 test_checkpoint_pair(unit_id 落盘往返、pair 不拆)与 test_non_ar_byte_identical(pools→batching→checkpoint→momentum 端到端黄金)。
This commit is contained in:
@@ -233,7 +233,8 @@ def write_checkpoint(
|
||||
global_step: 全局 step 序号。
|
||||
total_steps: 全局总 step 数。
|
||||
version_snapshot: skills/prompts 版本快照。
|
||||
epoch_batches: 本 epoch 的 batch 划分(question_id 列表的列表)。
|
||||
epoch_batches: 本 epoch 的 batch 划分(unit_id 列表的列表,孪生对折叠为
|
||||
单个 unit_id;纯非 AR 下 unit_id==question_id)。
|
||||
config: 训练配置对象,用于计算 config_fingerprint。
|
||||
|
||||
关键实现细节:
|
||||
|
||||
+74
-15
@@ -199,18 +199,78 @@ def _accumulate_slow_packs(diagnosis: DiagnosisResult, state: _TrainState) -> No
|
||||
state.tool_packs.extend(diagnosis.tool_case_packs.values())
|
||||
|
||||
|
||||
def _batch_from_ids(pools: Pools, ids: list[str]) -> list[GeneratedQuestion]:
|
||||
"""按 question_id 从诊断池重建一个 batch(保持原 epoch 划分)。
|
||||
def _batch_unit_ids(batch: list[GeneratedQuestion]) -> list[str]:
|
||||
"""把一个 batch 的扁平题目折叠为 unit_id 序列(孪生对成员去重为单个 unit_id)。
|
||||
|
||||
checkpoint 存 unit_id 序列而非逐题 question_id:断点续跑恢复时按完整 unit 展开,
|
||||
保证孪生对整体重建、绝不被劈开(核心算法保真 #3 断点续跑)。
|
||||
|
||||
参数:
|
||||
batch: 一个 mini-batch 的扁平题目列表(pair 两成员相邻)。
|
||||
|
||||
返回:
|
||||
unit_id 列表,按题目在 batch 中的首次出现顺序去重;single 的 unit_id 即
|
||||
question_id,故纯非 AR 输入下与旧逐题 question_id 序列逐字节一致。
|
||||
|
||||
关键实现:
|
||||
用 dict 保序去重(pair 两成员共享 unit_id,仅记一次),无需额外集合。
|
||||
"""
|
||||
ordered: dict[str, None] = {}
|
||||
for q in batch:
|
||||
ordered[q.unit_id] = None
|
||||
return list(ordered)
|
||||
|
||||
|
||||
def _batch_from_ids(pools: Pools, unit_ids: list[str]) -> list[GeneratedQuestion]:
|
||||
"""按 unit_id 序列从诊断池重建一个 batch,按完整 unit 展开成题目列表。
|
||||
|
||||
与 _batch_unit_ids 对称:恢复时以完整 unit 为单位展开(pair 两成员同进同出),
|
||||
断点续跑后孪生对绝不被拆开(核心算法保真 #3)。
|
||||
|
||||
参数:
|
||||
pools: 三池容器。
|
||||
ids: 一个 batch 的 question_id 列表。
|
||||
unit_ids: 一个 batch 的 unit_id 序列(checkpoint 存的粒度)。
|
||||
|
||||
返回:
|
||||
按 ids 顺序取出的 GeneratedQuestion 列表。
|
||||
按 unit_ids 顺序展开的 GeneratedQuestion 列表;每个 unit_id 展开为其全部
|
||||
成员题(single 1 题、pair 2 题),顺序与原 batch 一致。
|
||||
|
||||
关键实现:
|
||||
直接以 units_by_id[uid] 取值,unit_id 缺失触发 KeyError(P5 防静默兜底),
|
||||
强制 checkpoint 与当前诊断池一致;纯非 AR 下 unit_id==question_id、单元即
|
||||
单题,与旧逐题重建逐字节一致。
|
||||
"""
|
||||
by_id = {q.question_id: q for q in pools.diagnosis}
|
||||
return [by_id[i] for i in ids]
|
||||
units_by_id = {u.unit_id: u for u in build_units(pools.diagnosis)}
|
||||
return [q for uid in unit_ids for q in units_by_id[uid].questions]
|
||||
|
||||
|
||||
def _sample_momentum_candidates(
|
||||
pool: list[GeneratedQuestion],
|
||||
allowed_task_types: set[str],
|
||||
momentum_samples: int,
|
||||
epoch: int,
|
||||
) -> list[GeneratedQuestion]:
|
||||
"""为 momentum 从诊断池按题型过滤后做确定性抽样(逐题粒度)。
|
||||
|
||||
设计偏差(Phase 1 显式记录):momentum 采样在逐题粒度进行、不折叠 QuestionUnit,
|
||||
故仅保证「纯非 AR 题库」的抽样序列与引入 QuestionUnit 前逐字节一致;混格题库下
|
||||
孪生对可能被半采样、且候选集长度/顺序随 AR 成员增减而漂移,**Phase 1 不保证混格
|
||||
momentum 的 byte-identical**(属可接受偏差,纯非 AR 必须不漂)。
|
||||
|
||||
参数:
|
||||
pool: 诊断池扁平题目列表。
|
||||
allowed_task_types: 允许参与的题型集合。
|
||||
momentum_samples: 目标采样数上限。
|
||||
epoch: 采样种子(同 epoch 可复现)。
|
||||
|
||||
返回:
|
||||
采样到的题目列表;候选不足则全取,候选为空返回空列表。
|
||||
"""
|
||||
candidates = [q for q in pool if q.task_type in allowed_task_types]
|
||||
n = min(momentum_samples, len(candidates))
|
||||
if n <= 0:
|
||||
return []
|
||||
return random.Random(epoch).sample(candidates, n)
|
||||
|
||||
|
||||
def _snapshot_current_skills(skills_dir: Path) -> dict[str, str]:
|
||||
@@ -753,7 +813,8 @@ class Runner:
|
||||
correct_ratio=self._config.batch_correct_ratio,
|
||||
)
|
||||
step_from = 0
|
||||
batch_ids = [[q.question_id for q in b] for b in batches]
|
||||
# checkpoint 存 unit_id 序列:断点续跑按完整 unit 展开,孪生对不拆
|
||||
batch_unit_ids = [_batch_unit_ids(b) for b in batches]
|
||||
for step in range(step_from, len(batches)):
|
||||
await self._run_step(epoch, step, total_steps, batches[step], pools, state)
|
||||
state.global_step += 1
|
||||
@@ -766,7 +827,7 @@ class Runner:
|
||||
global_step=state.global_step,
|
||||
total_steps=total_steps,
|
||||
version_snapshot=self._current_version_snapshot(),
|
||||
epoch_batches=batch_ids,
|
||||
epoch_batches=batch_unit_ids,
|
||||
config=self._config,
|
||||
)
|
||||
await self._slow_update_cycle(epoch, pools, state)
|
||||
@@ -782,7 +843,7 @@ class Runner:
|
||||
global_step=state.global_step,
|
||||
total_steps=total_steps,
|
||||
version_snapshot=self._current_version_snapshot(),
|
||||
epoch_batches=batch_ids,
|
||||
epoch_batches=batch_unit_ids,
|
||||
config=self._config,
|
||||
)
|
||||
if _should_early_stop(
|
||||
@@ -1610,12 +1671,10 @@ class Runner:
|
||||
prev_skill = state.epoch_start_skills.get(target_file, skill_content)
|
||||
prev_guidance = momentum_inner(skill_content)
|
||||
|
||||
# 采样
|
||||
allowed = set(task_types)
|
||||
candidates = [q for q in pools.diagnosis if q.task_type in allowed]
|
||||
rng = random.Random(epoch)
|
||||
n = min(self._config.momentum_samples, len(candidates))
|
||||
sampled = rng.sample(candidates, n) if n > 0 else []
|
||||
# 采样(逐题粒度,不折叠 unit;混格偏差见 _sample_momentum_candidates docstring)
|
||||
sampled = _sample_momentum_candidates(
|
||||
pools.diagnosis, set(task_types), self._config.momentum_samples, epoch
|
||||
)
|
||||
|
||||
if not sampled:
|
||||
skill_path.write_text(
|
||||
|
||||
@@ -0,0 +1,254 @@
|
||||
"""断点续跑 checkpoint 的 unit_id 粒度契约集成测试(Task 10)。
|
||||
|
||||
覆盖三条铁律:
|
||||
- checkpoint 存 **unit_id 序列**而非逐题 question_id(孪生对折叠为单个 unit_id);
|
||||
- 续跑恢复以 **完整 unit** 展开(pair 两成员同进同出,绝不被劈开);
|
||||
- 真实经 write_checkpoint → load_checkpoint 落盘往返后,unit_id 序列无损存活、
|
||||
重建的 batch 与原 batch 逐字节一致(core 算法保真 #3 断点续跑)。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from app.harness.batching import build_batches
|
||||
from app.harness.checkpoint import load_checkpoint, write_checkpoint
|
||||
from app.harness.pools import Pools
|
||||
from app.harness.question_units import build_units, validate_units
|
||||
from app.harness.runner import _batch_from_ids, _batch_unit_ids
|
||||
from core.types import GeneratedQuestion
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pathlib import Path
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 辅助构造
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_single(qid: str, task_type: str = "RETRIEVAL", video_id: str = "v1") -> GeneratedQuestion:
|
||||
"""构造最小 single 题目。"""
|
||||
return GeneratedQuestion(
|
||||
question_id=qid,
|
||||
video_id=video_id,
|
||||
task_type=task_type,
|
||||
question=f"q_{qid}",
|
||||
options=("A. a", "B. b", "C. c", "D. d"),
|
||||
answer="A",
|
||||
source_nodes=("n1",),
|
||||
difficulty="medium",
|
||||
question_role="single",
|
||||
)
|
||||
|
||||
|
||||
def _make_pair(
|
||||
pair_id: str, task_type: str = "AR", video_id: str = "v1"
|
||||
) -> tuple[GeneratedQuestion, GeneratedQuestion]:
|
||||
"""构造一个孪生对(original + mirror),共享 pair_id / flip_axis。"""
|
||||
common = {
|
||||
"video_id": video_id,
|
||||
"task_type": task_type,
|
||||
"question": "?",
|
||||
"options": ("A. a", "B. b", "C. c", "D. d"),
|
||||
"answer": "A",
|
||||
"source_nodes": ("n1",),
|
||||
"difficulty": "medium",
|
||||
"pair_id": pair_id,
|
||||
"flip_axis": "before_after",
|
||||
}
|
||||
original = GeneratedQuestion(
|
||||
question_id=f"{pair_id}_o", question_role="pair_original", **common
|
||||
)
|
||||
mirror = GeneratedQuestion(question_id=f"{pair_id}_m", question_role="pair_mirror", **common)
|
||||
return original, mirror
|
||||
|
||||
|
||||
def _pools_with(diagnosis: list[GeneratedQuestion]) -> Pools:
|
||||
"""用给定诊断池构造最小 Pools(其余池置空,_batch_from_ids 只读 diagnosis)。"""
|
||||
return Pools(
|
||||
diagnosis=diagnosis,
|
||||
validation=[],
|
||||
test=[],
|
||||
baseline_run_id="baseline",
|
||||
baseline_val_accuracy=0.0,
|
||||
correctness={q.question_id: False for q in diagnosis},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# checkpoint 落盘所需的最小 _TrainState / RunConfig 替身
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeState:
|
||||
"""serialize_state 只读的可持久化字段(本测试用空累加包即可)。"""
|
||||
|
||||
correctness: dict[str, bool] = field(default_factory=dict)
|
||||
eval_prev_acc: float = 0.0
|
||||
eval_prev_run_id: str = "run-0"
|
||||
baseline_skills_version: str = "v1"
|
||||
baseline_prompts_version: str = "v1"
|
||||
steps_since_best_improved: int = 0
|
||||
epoch_start_skills: str = "v1"
|
||||
changed_task_types_this_epoch: set[str] = field(default_factory=set)
|
||||
rejected_buffer: dict = field(default_factory=dict)
|
||||
system_packs: list = field(default_factory=list)
|
||||
tool_packs: list = field(default_factory=list)
|
||||
probations: dict = field(default_factory=dict)
|
||||
gate_cooldown: dict = field(default_factory=dict)
|
||||
gate_epoch_observed: dict = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _FakeConfig:
|
||||
"""compute_fingerprint 读取的结构性 + 决策性字段。"""
|
||||
|
||||
batch_size: int = 4
|
||||
min_class_per_batch: int = 2
|
||||
epochs: int = 3
|
||||
diag_size: int = 30
|
||||
val_size: int = 50
|
||||
batch_correct_ratio: float = 0.0
|
||||
edit_budget_start: int = 6
|
||||
edit_budget_end: int = 3
|
||||
early_stop_patience: int = 3
|
||||
use_slow_momentum: bool = True
|
||||
skill_update_mode: str = "patch"
|
||||
appendix_consolidate_threshold: int = 10
|
||||
momentum_samples: int = 20
|
||||
gate_e_confirm: float = 20.0
|
||||
gate_e_provisional: float = 6.0
|
||||
gate_w_net_min: int = 2
|
||||
gate_delta_min: float = 0.02
|
||||
gate_lambda_dir: float = -3.0
|
||||
gate_e_rollback: float = 10.0
|
||||
gate_block: int = 4
|
||||
gate_n_max: int = 40
|
||||
gate_p_low: float = 0.1
|
||||
gate_p_high: float = 0.9
|
||||
gate_probe_quota: float = 0.2
|
||||
gate_gamma_decay: float = 0.9
|
||||
gate_cooldown_steps: int = 2
|
||||
gate_guard_err: float = 0.3
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# checkpoint 存 unit_id 序列(孪生对折叠为单个 unit_id)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCheckpointStoresUnitIds:
|
||||
"""一个含孪生对的 batch,checkpoint 序列应折叠为 unit_id(pair 只记一次)。"""
|
||||
|
||||
def test_pair_folds_to_single_unit_id(self) -> None:
|
||||
po, pm = _make_pair("pA")
|
||||
s1, s2 = _make_single("s1"), _make_single("s2")
|
||||
batch = [s1, po, pm, s2] # pair 两成员相邻
|
||||
|
||||
unit_ids = _batch_unit_ids(batch)
|
||||
|
||||
# 4 题折叠为 3 个 unit:s1 / pA(pair) / s2
|
||||
assert unit_ids == ["s1", "pA", "s2"]
|
||||
# 逐题 question_id 序列(旧行为)会含 4 项且把 pair 拆成两条
|
||||
assert [q.question_id for q in batch] == ["s1", "pA_o", "pA_m", "s2"]
|
||||
|
||||
def test_pure_single_unit_ids_equal_question_ids(self) -> None:
|
||||
"""纯 single batch 的 unit_id 序列与逐题 question_id 序列逐字节一致。"""
|
||||
batch = [_make_single(f"s{i}") for i in range(5)]
|
||||
assert _batch_unit_ids(batch) == [q.question_id for q in batch]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 恢复以完整 unit 展开:pair 不拆
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestResumeExpandsFullUnit:
|
||||
"""按 unit_id 序列恢复时,孪生对整体展开、绝不被劈开。"""
|
||||
|
||||
def test_pair_not_split_on_resume(self) -> None:
|
||||
po, pm = _make_pair("pA")
|
||||
s1, s2 = _make_single("s1"), _make_single("s2")
|
||||
batch = [s1, po, pm, s2]
|
||||
# 诊断池顺序故意与 batch 不同,验证恢复以池折叠的 unit 为准
|
||||
pools = _pools_with([s2, pm, s1, po])
|
||||
|
||||
unit_ids = _batch_unit_ids(batch)
|
||||
restored = _batch_from_ids(pools, unit_ids)
|
||||
|
||||
# 逐字节一致(顺序保留)
|
||||
assert [q.question_id for q in restored] == ["s1", "pA_o", "pA_m", "s2"]
|
||||
# 恢复后可无损重建为完整孪生对单元
|
||||
units = validate_units(build_units(restored))
|
||||
pair_units = [u for u in units if u.kind == "pair"]
|
||||
assert len(pair_units) == 1
|
||||
assert pair_units[0].unit_id == "pA"
|
||||
assert pair_units[0].size == 2
|
||||
|
||||
def test_roundtrip_preserves_order_mixed(self) -> None:
|
||||
"""混格 batch 经 unit_id 折叠 → 展开 round-trip 逐字节还原原顺序。"""
|
||||
po, pm = _make_pair("pA")
|
||||
qo, qm = _make_pair("pB")
|
||||
singles = [_make_single(f"s{i}") for i in range(3)]
|
||||
batch = [singles[0], po, pm, singles[1], qo, qm, singles[2]]
|
||||
pools = _pools_with(list(batch))
|
||||
|
||||
restored = _batch_from_ids(pools, _batch_unit_ids(batch))
|
||||
assert [q.question_id for q in restored] == [q.question_id for q in batch]
|
||||
|
||||
def test_missing_unit_id_raises(self) -> None:
|
||||
"""checkpoint 引用了诊断池不存在的 unit_id 时硬失败(防静默兜底)。"""
|
||||
pools = _pools_with([_make_single("s1")])
|
||||
try:
|
||||
_batch_from_ids(pools, ["s1", "ghost"])
|
||||
except KeyError:
|
||||
return
|
||||
raise AssertionError("引用缺失 unit_id 应触发 KeyError")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 真实落盘往返:write_checkpoint → load_checkpoint → 重建 batch
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCheckpointPersistenceRoundtrip:
|
||||
"""unit_id 序列经真实 checkpoint.json 原子写/读回后无损、重建逐字节一致。"""
|
||||
|
||||
def test_persisted_unit_ids_rebuild_batches(self, tmp_path: Path) -> None:
|
||||
po, pm = _make_pair("pA")
|
||||
singles = [_make_single(f"s{i}", task_type="RETRIEVAL") for i in range(6)]
|
||||
items = [*singles, po, pm]
|
||||
correctness = {q.question_id: False for q in items}
|
||||
batches, _ = build_batches(items, correctness, batch_size=4, min_class_per_batch=2, seed=7)
|
||||
pools = _pools_with(list(items))
|
||||
|
||||
# 存 unit_id 序列(train() 落盘路径)
|
||||
epoch_batches = [_batch_unit_ids(b) for b in batches]
|
||||
write_checkpoint(
|
||||
tmp_path,
|
||||
state=_FakeState(correctness=correctness),
|
||||
epoch=1,
|
||||
step_completed=0,
|
||||
phase="in_epoch",
|
||||
global_step=1,
|
||||
total_steps=len(batches),
|
||||
version_snapshot={"skills": "skills/v1", "prompts": "prompts/v1"},
|
||||
epoch_batches=epoch_batches,
|
||||
config=_FakeConfig(),
|
||||
)
|
||||
|
||||
ckpt = load_checkpoint(tmp_path)
|
||||
assert ckpt is not None
|
||||
rebuilt = [_batch_from_ids(pools, ids) for ids in ckpt["epoch_batches"]]
|
||||
|
||||
# 落盘往返后逐批逐字节还原
|
||||
assert [[q.question_id for q in b] for b in rebuilt] == [
|
||||
[q.question_id for q in b] for b in batches
|
||||
]
|
||||
# 含孪生对的批次恢复后 pair 仍成对
|
||||
for b in rebuilt:
|
||||
for u in validate_units(build_units(b)):
|
||||
if u.kind == "pair":
|
||||
assert u.size == 2
|
||||
@@ -0,0 +1,297 @@
|
||||
"""纯非 AR 题库端到端 byte-identical 黄金测试(Task 10)。
|
||||
|
||||
Phase 1 验收铁律:引入 QuestionUnit 后,纯 single 题库过 pools 抽样 → batching 分批
|
||||
→ checkpoint 折叠/恢复 → momentum 采样,端到端结果与"引入 QuestionUnit 前"的逐题
|
||||
旧算法逐字节一致。golden 用固定 seed 的确定性对照(忠实重实现旧逐题逻辑),非空断言。
|
||||
|
||||
momentum 设计偏差:momentum 采样在逐题粒度进行、不折叠 unit,故仅保证纯非 AR 的抽样
|
||||
序列 byte-identical;Phase 1 不保证混格 momentum byte-identical(见 runner
|
||||
_sample_momentum_candidates docstring)。本测试保守证明纯非 AR momentum 不漂。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import random
|
||||
|
||||
from app.harness.batching import build_batches
|
||||
from app.harness.pools import build_pools
|
||||
from app.harness.question_units import build_units
|
||||
from app.harness.runner import _batch_from_ids, _batch_unit_ids, _sample_momentum_candidates
|
||||
from core.types import GeneratedQuestion
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 辅助构造
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _single(qid: str, task_type: str = "RETRIEVAL", video_id: str = "v1") -> GeneratedQuestion:
|
||||
"""构造最小 single 题目。"""
|
||||
return GeneratedQuestion(
|
||||
question_id=qid,
|
||||
video_id=video_id,
|
||||
task_type=task_type,
|
||||
question=f"q_{qid}",
|
||||
options=("A. a", "B. b", "C. c", "D. d"),
|
||||
answer="A",
|
||||
source_nodes=("n1",),
|
||||
difficulty="medium",
|
||||
question_role="single",
|
||||
)
|
||||
|
||||
|
||||
def _pair(
|
||||
pair_id: str, task_type: str = "AR", video_id: str = "v1"
|
||||
) -> tuple[GeneratedQuestion, GeneratedQuestion]:
|
||||
"""构造一个孪生对(original + mirror)。"""
|
||||
common = {
|
||||
"video_id": video_id,
|
||||
"task_type": task_type,
|
||||
"question": "?",
|
||||
"options": ("A. a", "B. b", "C. c", "D. d"),
|
||||
"answer": "A",
|
||||
"source_nodes": ("n1",),
|
||||
"difficulty": "medium",
|
||||
"pair_id": pair_id,
|
||||
"flip_axis": "before_after",
|
||||
}
|
||||
original = GeneratedQuestion(
|
||||
question_id=f"{pair_id}_o", question_role="pair_original", **common
|
||||
)
|
||||
mirror = GeneratedQuestion(question_id=f"{pair_id}_m", question_role="pair_mirror", **common)
|
||||
return original, mirror
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 引入 QuestionUnit 前的旧逐题 build_batches 忠实副本(单一 rng,无 unit 折叠)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _reference_select_mixed(
|
||||
items: list[GeneratedQuestion],
|
||||
correctness: dict[str, bool],
|
||||
correct_ratio: float,
|
||||
rng: random.Random,
|
||||
) -> dict[str, list[GeneratedQuestion]]:
|
||||
"""旧版 _select_mixed_by_task_type 的忠实副本(逐题、单一 rng)。"""
|
||||
errors_by_type: dict[str, list[GeneratedQuestion]] = {}
|
||||
correct_by_type: dict[str, list[GeneratedQuestion]] = {}
|
||||
for q in items:
|
||||
qid = q.question_id
|
||||
if correctness.get(qid) is False:
|
||||
errors_by_type.setdefault(q.task_type, []).append(q)
|
||||
elif correctness.get(qid, False):
|
||||
correct_by_type.setdefault(q.task_type, []).append(q)
|
||||
if correct_ratio <= 0:
|
||||
return errors_by_type
|
||||
grouped: dict[str, list[GeneratedQuestion]] = {}
|
||||
for task_type in sorted(errors_by_type):
|
||||
errs = errors_by_type[task_type]
|
||||
n_correct = round(len(errs) * correct_ratio / (1 - correct_ratio))
|
||||
available = correct_by_type.get(task_type, [])
|
||||
sampled = (
|
||||
list(available) if len(available) <= n_correct else rng.sample(available, n_correct)
|
||||
)
|
||||
grouped[task_type] = errs + sampled
|
||||
return grouped
|
||||
|
||||
|
||||
def _reference_build_batches(
|
||||
items: list[GeneratedQuestion],
|
||||
correctness: dict[str, bool],
|
||||
batch_size: int,
|
||||
min_class_per_batch: int,
|
||||
seed: int,
|
||||
correct_ratio: float = 0.0,
|
||||
) -> list[list[GeneratedQuestion]]:
|
||||
"""引入 QuestionUnit 前的旧版 build_batches 忠实副本(逐题、单一 rng)。"""
|
||||
rng = random.Random(seed)
|
||||
grouped = _reference_select_mixed(items, correctness, correct_ratio, rng)
|
||||
total = sum(len(g) for g in grouped.values())
|
||||
if total == 0:
|
||||
return []
|
||||
nb = max(1, math.ceil(total / batch_size))
|
||||
batches: list[list[GeneratedQuestion]] = [[] for _ in range(nb)]
|
||||
small = {t: g for t, g in grouped.items() if len(g) <= min_class_per_batch}
|
||||
large = {t: g for t, g in grouped.items() if len(g) > min_class_per_batch}
|
||||
for t in sorted(small, key=lambda t: (-len(small[t]), t)):
|
||||
group = small[t]
|
||||
placed = False
|
||||
for b in batches:
|
||||
if len(b) + len(group) <= batch_size:
|
||||
b.extend(group)
|
||||
placed = True
|
||||
break
|
||||
if not placed:
|
||||
batches.append(list(group))
|
||||
nb_live = len(batches)
|
||||
pointer = 0
|
||||
for task_type in sorted(large):
|
||||
group = list(large[task_type])
|
||||
rng.shuffle(group)
|
||||
for q in group:
|
||||
for offset in range(nb_live):
|
||||
idx = (pointer + offset) % nb_live
|
||||
if len(batches[idx]) < batch_size:
|
||||
batches[idx].append(q)
|
||||
pointer = (idx + 1) % nb_live
|
||||
break
|
||||
return [b for b in batches if b]
|
||||
|
||||
|
||||
def _ids(batches: list[list[GeneratedQuestion]]) -> list[list[str]]:
|
||||
"""提取 batch 的 question_id 序列,便于逐字节对比。"""
|
||||
return [[q.question_id for q in b] for b in batches]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# (1) 纯非 AR 走 size=1 unit:unit_id == question_id
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPureSingleAreSizeOneUnits:
|
||||
"""纯 single 输入折叠出的每个 unit 都是 size=1、unit_id 等于 question_id。"""
|
||||
|
||||
def test_all_units_size_one(self) -> None:
|
||||
items = [_single(f"s{i}", task_type=f"t{i % 3}") for i in range(15)]
|
||||
units = build_units(items)
|
||||
assert len(units) == len(items)
|
||||
assert all(u.kind == "single" and u.size == 1 for u in units)
|
||||
assert [u.unit_id for u in units] == [q.question_id for q in items]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# (2) batching 逐字节一致(黄金对照)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBatchingByteIdentical:
|
||||
"""纯 single build_batches 与旧逐题算法逐字节一致。"""
|
||||
|
||||
def _assert(self, items, correctness, batch_size, min_cls, seed, ratio) -> None:
|
||||
got, _ = build_batches(items, correctness, batch_size, min_cls, seed, ratio)
|
||||
ref = _reference_build_batches(items, correctness, batch_size, min_cls, seed, ratio)
|
||||
assert _ids(got) == _ids(ref)
|
||||
|
||||
def test_pure_errors(self) -> None:
|
||||
items = [_single(f"q{i}", task_type=f"t{i % 3}") for i in range(20)]
|
||||
correctness = {f"q{i}": False for i in range(20)}
|
||||
self._assert(items, correctness, 5, 2, 42, 0.0)
|
||||
|
||||
def test_mixed_ratio(self) -> None:
|
||||
items = [_single(f"q{i}", task_type=f"t{i % 4}") for i in range(40)]
|
||||
correctness = {f"q{i}": (i % 3 == 0) for i in range(40)}
|
||||
self._assert(items, correctness, 8, 3, 7, 0.5)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# (3) 端到端:pools 抽样 → batching → checkpoint 折叠/恢复 逐字节一致
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEndToEndByteIdentical:
|
||||
"""pools → batching → checkpoint unit_id round-trip 全链纯非 AR 逐字节稳定。"""
|
||||
|
||||
def _make_pool_questions(self) -> tuple[list[GeneratedQuestion], dict[str, bool]]:
|
||||
items = [_single(f"s{i}", task_type=f"t{i % 3}") for i in range(60)]
|
||||
correctness = {f"s{i}": (i % 2 == 0) for i in range(60)}
|
||||
return items, correctness
|
||||
|
||||
def test_pools_deterministic_and_pure_single(self) -> None:
|
||||
items, correctness = self._make_pool_questions()
|
||||
cfg = {
|
||||
"diag_cfg": {
|
||||
"size": 18,
|
||||
"correct_ratio": 0.5,
|
||||
"task_types": None,
|
||||
"seed": 5,
|
||||
"min_per_class": None,
|
||||
},
|
||||
"val_cfg": {
|
||||
"size": 12,
|
||||
"correct_ratio": 0.5,
|
||||
"task_types": None,
|
||||
"seed": 5,
|
||||
"min_per_class": None,
|
||||
},
|
||||
"test_cfg": {"size": 10, "seed": 5},
|
||||
"baseline_run_id": "baseline",
|
||||
}
|
||||
p1 = build_pools(items, correctness, **cfg)
|
||||
p2 = build_pools(items, correctness, **cfg)
|
||||
# 确定性:同 seed 同输入池划分逐字节一致
|
||||
assert [q.question_id for q in p1.diagnosis] == [q.question_id for q in p2.diagnosis]
|
||||
# 诊断池纯 single:折叠出的 unit 与逐题一一对应
|
||||
units = build_units(p1.diagnosis)
|
||||
assert [u.unit_id for u in units] == [q.question_id for q in p1.diagnosis]
|
||||
|
||||
def test_batching_then_checkpoint_roundtrip(self) -> None:
|
||||
items, correctness = self._make_pool_questions()
|
||||
p = build_pools(
|
||||
items,
|
||||
correctness,
|
||||
diag_cfg={
|
||||
"size": 18,
|
||||
"correct_ratio": 0.5,
|
||||
"task_types": None,
|
||||
"seed": 5,
|
||||
"min_per_class": None,
|
||||
},
|
||||
val_cfg={
|
||||
"size": 12,
|
||||
"correct_ratio": 0.5,
|
||||
"task_types": None,
|
||||
"seed": 5,
|
||||
"min_per_class": None,
|
||||
},
|
||||
test_cfg={"size": 10, "seed": 5},
|
||||
baseline_run_id="baseline",
|
||||
)
|
||||
batches, _ = build_batches(p.diagnosis, p.correctness, 6, 2, seed=3, correct_ratio=0.5)
|
||||
|
||||
# checkpoint 折叠为 unit_id(纯 single 即 question_id)后恢复
|
||||
epoch_batches = [_batch_unit_ids(b) for b in batches]
|
||||
assert epoch_batches == _ids(batches) # 纯非 AR:unit_id 序列 == question_id 序列
|
||||
rebuilt = [_batch_from_ids(p, ids) for ids in epoch_batches]
|
||||
# 恢复的 batch 与原 batch 逐字节一致(inference 消费的题序不变)
|
||||
assert _ids(rebuilt) == _ids(batches)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# (4) momentum 纯非 AR byte-identical(混格偏差已在 runner docstring 记录)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestMomentumPureNonARByteIdentical:
|
||||
"""momentum 采样纯非 AR 与旧逐题 random.Random(epoch).sample 逐字节一致。"""
|
||||
|
||||
def test_matches_legacy_rng(self) -> None:
|
||||
pool = [_single(f"s{i}", task_type="RETRIEVAL") for i in range(30)]
|
||||
allowed = {"RETRIEVAL"}
|
||||
epoch = 11
|
||||
samples = 8
|
||||
got = _sample_momentum_candidates(pool, allowed, samples, epoch)
|
||||
|
||||
candidates = [q for q in pool if q.task_type in allowed]
|
||||
ref = random.Random(epoch).sample(candidates, samples)
|
||||
assert [q.question_id for q in got] == [q.question_id for q in ref]
|
||||
|
||||
def test_ar_pairs_of_other_type_do_not_shift(self) -> None:
|
||||
"""向池中加入其它题型的 AR pair,不改变非 AR momentum 的抽样序列。"""
|
||||
singles = [_single(f"s{i}", task_type="RETRIEVAL") for i in range(30)]
|
||||
allowed = {"RETRIEVAL"}
|
||||
epoch = 11
|
||||
samples = 8
|
||||
base = _sample_momentum_candidates(singles, allowed, samples, epoch)
|
||||
|
||||
mixed = list(singles)
|
||||
for k in range(4):
|
||||
po, pm = _pair(f"p{k}", task_type="AR")
|
||||
mixed.extend([po, pm])
|
||||
after = _sample_momentum_candidates(mixed, allowed, samples, epoch)
|
||||
assert [q.question_id for q in base] == [q.question_id for q in after]
|
||||
|
||||
def test_fewer_candidates_than_samples_returns_all(self) -> None:
|
||||
pool = [_single(f"s{i}", task_type="RETRIEVAL") for i in range(3)]
|
||||
got = _sample_momentum_candidates(pool, {"RETRIEVAL"}, 10, epoch=1)
|
||||
assert {q.question_id for q in got} == {"s0", "s1", "s2"}
|
||||
Reference in New Issue
Block a user