405 lines
16 KiB
Python
405 lines
16 KiB
Python
"""CE-Gate 信息量阶梯与基线缓存（unit 粒度，核心算法保真 #5）。
|
||
|
||
阶梯（每题型一条）：gate 的出题顺序表，键为 **unit_id**（single 题 unit_id
|
||
等于 question_id，AR pair 折叠为一个单元、unit_id 等于共享 pair_id）。冷启动
|
||
（FRESH）用种子基线的**单元级**对错两档粗排（错 unit 高优先 2:1 交错 + 全错
|
||
unit 的 probe_quota 探针插尾）；epoch >=1 用非 gate run 观测**折叠成单元观测**后做
|
||
gamma-EMA 更新 p_hat，按信息量 p_hat(1-p_hat) 降序、剔 p_hat 不在 [p_low, p_high]。
|
||
单元错 = 该单元任一成员错（AR pair 双向 AND）。防泄露铁律：gate 内 rollout 永不
|
||
回流 p_hat（调用方以 run_id 含 "_gate_" 过滤观测源），本迁移不改此过滤。
|
||
|
||
持久化门控：gate_pools.json 带 schema_version（当前 = 2，unit 键）。旧版无
|
||
schema_version（v1、qid 键）加载时**直接报错**，拒绝静默混用 qid/unit 键。
|
||
|
||
BaselineCache：基线侧单元级对错缓存，键 = (task_type, skill_hash,
|
||
prompts_version, unit_id) 内容寻址、无显式失效。JSON 持久化到 workspace，
|
||
供 resume 后合法复用已冻结阶梯上的新鲜 draw。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import json
|
||
import os
|
||
import random
|
||
from dataclasses import dataclass
|
||
from typing import TYPE_CHECKING
|
||
|
||
from loguru import logger
|
||
|
||
from app.harness.question_units import build_units, unit_correctness
|
||
|
||
if TYPE_CHECKING:
|
||
from pathlib import Path
|
||
|
||
from core.types import GeneratedQuestion, QuestionUnit
|
||
|
||
# gate_pools.json 结构版本。v1（隐式、无此字段）为逐题 qid 键的存量格式；
|
||
# v2 起改为 unit_id 键。load 时严格校验，不匹配即报错（不静默迁移/混用）。
|
||
SCHEMA_VERSION = 2
|
||
|
||
|
||
def skill_hash(content: str) -> str:
|
||
"""对 skill 正文取 sha1 摘要，作缓存键的内容维度。
|
||
|
||
参数:
|
||
content: skill 文件全文（基线侧为解析后生效文件的正文）。
|
||
|
||
返回:
|
||
sha1 十六进制摘要。
|
||
"""
|
||
return hashlib.sha1(content.encode("utf-8")).hexdigest()
|
||
|
||
|
||
@dataclass
|
||
class LadderEntry:
|
||
"""阶梯单元：题目单元与其估计答对率。
|
||
|
||
字段:
|
||
unit_id: 单元唯一标识（single 等于 question_id，AR pair 等于共享 pair_id）。
|
||
p_hat: 估计答对率。冷启动为 Beta(1,1) 平滑的单次观测后验均值
|
||
（错=1/3、对=2/3），此后经 gamma-EMA 更新。
|
||
"""
|
||
|
||
unit_id: str
|
||
p_hat: float
|
||
|
||
|
||
def build_cold_entries(
|
||
units: list[QuestionUnit],
|
||
correctness: dict[str, bool],
|
||
probe_quota: float,
|
||
seed: int,
|
||
) -> list[LadderEntry]:
|
||
"""冷启动排序（unit 粒度）：错 unit 高优先 2:1 交错 + 全错 unit 探针插尾。
|
||
|
||
参数:
|
||
units: 该题型的全部候选单元（已排除 test 池；AR pair 已折叠成单元）。
|
||
correctness: question_id -> 种子基线是否答对（900 题全量逐题对错）。
|
||
单元级对错由 unit_correctness(strict=False) 折叠（任一成员错 → 单元错）。
|
||
probe_quota: 从错 unit 中随机抽出插到梯尾的探针比例（防"解锁新能力"盲区）。
|
||
seed: 洗牌种子，保证确定性重建。
|
||
|
||
返回:
|
||
排序后的 LadderEntry 列表（键=unit_id；p_hat 用 Beta(1,1) 平滑：错=1/3、
|
||
对=2/3，与 warm 阶段 gamma-EMA / 信息量排序自然衔接）。
|
||
|
||
关键实现细节:
|
||
与逐题版**同公式、同比例、同顺序**，仅把调度粒度从题换成单元：错 unit、
|
||
对 unit 各自固定种子洗牌 -> 按 probe_quota 从错 unit 抽探针 -> 剩余按
|
||
错错对 2:1 交错（一方耗尽后顺排另一方）-> 探针追加尾部。
|
||
"""
|
||
rng = random.Random(seed)
|
||
wrong = [u for u in units if not unit_correctness(u, correctness, strict=False)]
|
||
right = [u for u in units if unit_correctness(u, correctness, strict=False)]
|
||
rng.shuffle(wrong)
|
||
rng.shuffle(right)
|
||
|
||
n_probe = int(len(wrong) * probe_quota)
|
||
probes, wrong_main = wrong[:n_probe], wrong[n_probe:]
|
||
|
||
interleaved: list[QuestionUnit] = []
|
||
wi, ri = 0, 0
|
||
while wi < len(wrong_main) or ri < len(right):
|
||
for _ in range(2):
|
||
if wi < len(wrong_main):
|
||
interleaved.append(wrong_main[wi])
|
||
wi += 1
|
||
if ri < len(right):
|
||
interleaved.append(right[ri])
|
||
ri += 1
|
||
interleaved.extend(probes)
|
||
|
||
def _p0(u: QuestionUnit) -> float:
|
||
return 2 / 3 if unit_correctness(u, correctness, strict=False) else 1 / 3
|
||
|
||
return [LadderEntry(u.unit_id, _p0(u)) for u in interleaved]
|
||
|
||
|
||
def order_ladder(entries: list[LadderEntry], p_low: float, p_high: float) -> list[LadderEntry]:
|
||
"""warm 排序：剔 p_hat 不在 [p_low, p_high] 的零信息题，按信息量 p_hat(1-p_hat) 降序。
|
||
|
||
参数:
|
||
entries: 待排序的阶梯单元。
|
||
p_low / p_high: p_hat 保留区间。
|
||
|
||
返回:
|
||
过滤并排序后的新列表（稳定排序，同信息量保持原相对序）。
|
||
"""
|
||
kept = [e for e in entries if p_low <= e.p_hat <= p_high]
|
||
return sorted(kept, key=lambda e: e.p_hat * (1 - e.p_hat), reverse=True)
|
||
|
||
|
||
@dataclass
|
||
class GatePools:
|
||
"""全部题型的阶梯容器，含构建种子与数据指纹（确定性重建凭据）。
|
||
|
||
字段:
|
||
entries: task_type -> 冷启动序 LadderEntry 列表（warm 排序在取用时做，
|
||
保持存储序稳定、避免每次更新重写全表顺序）。
|
||
seed: 冷启动洗牌种子。
|
||
fingerprint: 构建输入指纹（基线 run_id + 题集 hash 等），resume 校验用。
|
||
"""
|
||
|
||
entries: dict[str, list[LadderEntry]]
|
||
seed: int
|
||
fingerprint: str
|
||
|
||
def ladder_for(
|
||
self,
|
||
task_type: str,
|
||
exclude_units: set[str],
|
||
p_low: float,
|
||
p_high: float,
|
||
cold: bool,
|
||
) -> list[str]:
|
||
"""取该题型的 gate 出题序（unit_id 列表），排除本 step 进化案例包所在单元。
|
||
|
||
参数:
|
||
task_type: 目标题型。
|
||
exclude_units: 本 step 案例包（failure/success cases）所在单元的
|
||
unit_id，防止在"刚学的那道题"上自测。按 **unit** 排除：命中单元
|
||
整体剔除，避免只排 AR pair 半个成员而向 gate 池灌入半个 pair。
|
||
p_low / p_high: warm 阶段的 p_hat 保留区间。
|
||
cold: True 表示尚无 epoch 级观测（epoch 1），用冷启动存储序；
|
||
False 走 order_ladder 信息量排序。
|
||
|
||
返回:
|
||
排除后的有序 unit_id 列表。
|
||
|
||
异常:
|
||
ValueError: 该题型无阶梯（冷启动构建缺失），或该题型阶梯为空。
|
||
"""
|
||
if task_type not in self.entries:
|
||
raise ValueError(f"task_type={task_type} 无阶梯，冷启动构建缺失该题型")
|
||
pool = self.entries[task_type]
|
||
if not pool:
|
||
raise ValueError(f"task_type={task_type} 阶梯为空，无可出题目")
|
||
ordered = pool if cold else order_ladder(pool, p_low, p_high)
|
||
return [e.unit_id for e in ordered if e.unit_id not in exclude_units]
|
||
|
||
def update_probs(
|
||
self,
|
||
per_q_observations: dict[str, bool],
|
||
units_by_id: dict[str, QuestionUnit],
|
||
gamma: float,
|
||
) -> None:
|
||
"""gamma-EMA 更新 p_hat：先把逐题观测折叠成单元观测，再按 unit_id 匹配更新。
|
||
|
||
p_hat <- gamma * p_hat + (1-gamma) * unit_obs。只更新"整个单元都被观测到"
|
||
的单元；单元观测 = 成员逐题对错的 AND（任一成员错 → 单元错）。折叠是必需的：
|
||
AR pair 的 unit_id 是 pair_id，若直接按 unit_id 去逐题观测里匹配将永不命中、
|
||
导致 gamma-EMA 停摆（核心算法保真 #5）。
|
||
|
||
参数:
|
||
per_q_observations: question_id -> 本 epoch 非 gate run 的最新逐题对错。
|
||
调用方必须已按 run_id 过滤掉 gate 内 rollout（防泄露铁律）。
|
||
units_by_id: unit_id -> QuestionUnit，用于把逐题观测折叠成单元观测。
|
||
gamma: EMA 衰减系数。
|
||
|
||
关键实现细节:
|
||
单元只有在其**全部**成员都出现在 per_q_observations 时才更新；半观测
|
||
（AR pair 只见一半）跳过，避免用不完整证据污染 p_hat。
|
||
"""
|
||
for entries in self.entries.values():
|
||
for e in entries:
|
||
unit = units_by_id.get(e.unit_id)
|
||
if unit is None:
|
||
continue
|
||
if not all(q.question_id in per_q_observations for q in unit.questions):
|
||
continue
|
||
unit_correct = all(per_q_observations[q.question_id] for q in unit.questions)
|
||
obs = 1.0 if unit_correct else 0.0
|
||
e.p_hat = gamma * e.p_hat + (1 - gamma) * obs
|
||
|
||
def save(self, path: Path) -> None:
|
||
"""原子写 gate_pools.json（.tmp 再 replace），落 schema_version + unit_id 键。
|
||
|
||
参数:
|
||
path: 目标 JSON 路径。
|
||
"""
|
||
payload = {
|
||
"schema_version": SCHEMA_VERSION,
|
||
"seed": self.seed,
|
||
"fingerprint": self.fingerprint,
|
||
"entries": {
|
||
t: [{"unit_id": e.unit_id, "p_hat": e.p_hat} for e in es]
|
||
for t, es in self.entries.items()
|
||
},
|
||
}
|
||
tmp = path.with_suffix(".json.tmp")
|
||
tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
os.replace(tmp, path)
|
||
|
||
@classmethod
|
||
def load(cls, path: Path) -> GatePools:
|
||
"""从 gate_pools.json 恢复；schema_version 不匹配直接报错（不静默混用）。
|
||
|
||
参数:
|
||
path: gate_pools.json 路径。
|
||
|
||
返回:
|
||
复活的 GatePools。
|
||
|
||
异常:
|
||
RuntimeError: 缺 schema_version（存量 v1、qid 键）或版本不等于
|
||
SCHEMA_VERSION——拒绝把 qid 键当 unit 键静默复用，须 FRESH 重建。
|
||
"""
|
||
d = json.loads(path.read_text(encoding="utf-8"))
|
||
version = d.get("schema_version")
|
||
if version != SCHEMA_VERSION:
|
||
raise RuntimeError(
|
||
f"gate_pools.json schema_version={version!r} 与当前 {SCHEMA_VERSION} 不符"
|
||
f"（存量 qid 键池不可当 unit 键复用），请删除后 FRESH 重建: {path}"
|
||
)
|
||
return cls(
|
||
entries={
|
||
t: [LadderEntry(x["unit_id"], x["p_hat"]) for x in es]
|
||
for t, es in d["entries"].items()
|
||
},
|
||
seed=d["seed"],
|
||
fingerprint=d["fingerprint"],
|
||
)
|
||
|
||
|
||
def build_or_load_gate_pools(
|
||
workspace_dir: Path,
|
||
questions: list[GeneratedQuestion],
|
||
test_qids: set[str],
|
||
baseline_correctness: dict[str, bool],
|
||
task_types: list[str],
|
||
probe_quota: float,
|
||
seed: int,
|
||
baseline_run_id: str,
|
||
) -> GatePools:
|
||
"""gate 阶梯获取入口：gate_pools.json 存在且指纹一致则加载，否则冷启动构建。
|
||
|
||
参数:
|
||
workspace_dir: workspace 根目录（gate_pools.json 落其下）。
|
||
questions: benchmark 全量题（900 题）。
|
||
test_qids: held-out test 池题目 id（阶梯题源必须排除）。
|
||
baseline_correctness: 种子基线 900 题全量对错（从基线 run 的 db 读）。
|
||
task_types: 参与进化的题型列表。
|
||
probe_quota: 冷启动探针比例。
|
||
seed: 冷启动洗牌种子。
|
||
baseline_run_id: 指纹成分。
|
||
|
||
返回:
|
||
GatePools。
|
||
|
||
关键实现细节:
|
||
指纹 = sha1(baseline_run_id|全 qid|seed|probe_quota|task_types|test_qids)。
|
||
指纹不一致（题集/基线/参数变了）直接报错——FRESH 语义下不该发生，
|
||
防御性拒绝而非静默重建。
|
||
"""
|
||
joined = ",".join(sorted(q.question_id for q in questions))
|
||
fp_src = (
|
||
f"{baseline_run_id}|{joined}|{seed}|{probe_quota}"
|
||
f"|{','.join(sorted(task_types))}|{','.join(sorted(test_qids))}"
|
||
)
|
||
fingerprint = hashlib.sha1(fp_src.encode()).hexdigest()
|
||
path = workspace_dir / "gate_pools.json"
|
||
if path.exists():
|
||
pools = GatePools.load(path)
|
||
if pools.fingerprint != fingerprint:
|
||
raise RuntimeError(
|
||
f"gate_pools.json 指纹不一致（题集或基线变更），拒绝静默重建: {path}"
|
||
)
|
||
return pools
|
||
|
||
entries: dict[str, list[LadderEntry]] = {}
|
||
for t in task_types:
|
||
units = _task_units_excluding_test(questions, t, test_qids)
|
||
if not units:
|
||
raise ValueError(f"task_type={t} 无非 test 单元，无法建阶梯")
|
||
entries[t] = build_cold_entries(units, baseline_correctness, probe_quota, seed)
|
||
logger.info("gate 阶梯[{}]: {} 单元（冷启动）", t, len(entries[t]))
|
||
pools = GatePools(entries=entries, seed=seed, fingerprint=fingerprint)
|
||
pools.save(path)
|
||
return pools
|
||
|
||
|
||
def _task_units_excluding_test(
|
||
questions: list[GeneratedQuestion], task_type: str, test_qids: set[str]
|
||
) -> list[QuestionUnit]:
|
||
"""取某题型的非 test 候选单元：先按 unit 折叠，再整体排除含 test 成员的单元。
|
||
|
||
先折叠后排除保证 AR pair 不被拆半（否则半个 pair 交给下游会触发 build_units 的
|
||
孤儿 fail-fast）；single 单元等价于逐题排除（核心算法保真 #5）。
|
||
|
||
参数:
|
||
questions: benchmark 全量题。
|
||
task_type: 目标题型。
|
||
test_qids: held-out test 池题目 id。
|
||
|
||
返回:
|
||
该题型下不含任何 test 成员的候选单元列表。
|
||
"""
|
||
pool = [q for q in questions if q.task_type == task_type]
|
||
return [
|
||
u for u in build_units(pool) if all(q.question_id not in test_qids for q in u.questions)
|
||
]
|
||
|
||
|
||
class BaselineCache:
|
||
"""基线侧单元级对错缓存（内容寻址，JSON 持久化）。
|
||
|
||
键 = (task_type, skill_hash, prompts_version, unit_id)：任何影响该题型
|
||
有效 skill 的变化（含共享 default-strategy.md 被他类 accept 改写）
|
||
都使 skill_hash 变化、缓存自然 miss；prompts 版本变化同理。unit_id 维度
|
||
使 single 题以自身 question_id、AR pair 以共享 pair_id 寻址，缓存单元级
|
||
对错（pair 双向 AND 折叠后一个布尔）。
|
||
"""
|
||
|
||
def __init__(self, path: Path) -> None:
|
||
"""加载或初始化缓存文件。
|
||
|
||
参数:
|
||
path: 缓存 JSON 路径（workspace/baseline_cache.json）。
|
||
"""
|
||
self._path = path
|
||
self._store: dict[str, bool] = {}
|
||
if path.exists():
|
||
self._store = json.loads(path.read_text(encoding="utf-8"))
|
||
|
||
@staticmethod
|
||
def _key(task_type: str, s_hash: str, prompts_version: str, unit_id: str) -> str:
|
||
"""拼缓存键（四维内容寻址，第四维为 unit_id）。"""
|
||
return f"{task_type}|{s_hash}|{prompts_version}|{unit_id}"
|
||
|
||
def get(self, task_type: str, s_hash: str, prompts_version: str, unit_id: str) -> bool | None:
|
||
"""读缓存；未命中返回 None。
|
||
|
||
参数:
|
||
task_type: 题型。
|
||
s_hash: 基线侧生效 skill 文件的内容哈希。
|
||
prompts_version: 当前 prompts 版本。
|
||
unit_id: 单元 id（single=question_id，AR pair=pair_id）。
|
||
|
||
返回:
|
||
缓存的单元级对错；未命中 None。
|
||
"""
|
||
return self._store.get(self._key(task_type, s_hash, prompts_version, unit_id))
|
||
|
||
def put(
|
||
self, task_type: str, s_hash: str, prompts_version: str, unit_id: str, correct: bool
|
||
) -> None:
|
||
"""写缓存并落盘（原子写，gate 频度低、全量重写成本可忽略）。
|
||
|
||
参数:
|
||
task_type / s_hash / prompts_version / unit_id: 缓存键四维。
|
||
correct: 基线侧该单元对错（AR pair 双向 AND 折叠后一个布尔）。
|
||
|
||
关键实现细节:
|
||
先盘后存：新条目先原子落盘（tmp 写 + os.replace）成功后才更新
|
||
内存，磁盘写失败时内存与磁盘一致（均无新条目），无分裂窗口。
|
||
"""
|
||
updated = {
|
||
**self._store,
|
||
self._key(task_type, s_hash, prompts_version, unit_id): correct,
|
||
}
|
||
tmp = self._path.with_suffix(".json.tmp")
|
||
tmp.write_text(json.dumps(updated, ensure_ascii=False), encoding="utf-8")
|
||
os.replace(tmp, self._path)
|
||
self._store = updated
|