198 lines
7.8 KiB
Python
198 lines
7.8 KiB
Python
"""teacher rollout 采样（IO 边缘，论文 §3.2.1）。
|
||
|
||
层 1 起步能力：给一批 prompt 批量生成解答，落盘 sha256 键的 JSONL 缓存
|
||
（键契约在 data.prompt_key 单点定义，本模块与 data.attach_teacher_completions
|
||
共用）。层 5 在此长出 chunk 前缀续写的 MC rollout 能力。
|
||
|
||
连接信息从 `.env` 读取（TEACHER_API_BASE / TEACHER_API_KEY / TEACHER_MODEL），
|
||
密钥永不出现在代码与配置类里。
|
||
|
||
缓存即断点：生成过程逐条追加写盘，任何中断（网络、Ctrl-C、单条失败）后
|
||
重跑同一命令，已完成的条目自动跳过——API 花的钱不会白花。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import re
|
||
import time
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from pathlib import Path
|
||
|
||
from dotenv import load_dotenv
|
||
from openai import OpenAI
|
||
|
||
from ars_opd.configs import TeacherGenConfig
|
||
from ars_opd.data import prompt_key
|
||
|
||
Messages = list[dict[str, str]]
|
||
|
||
|
||
def _load_teacher_env(env_file: str | None = None) -> tuple[str, str, str]:
|
||
"""从 .env（及进程环境）读取 API 连接三元组，缺一项都显式报错。"""
|
||
load_dotenv(env_file)
|
||
values = {}
|
||
for name in ("TEACHER_API_BASE", "TEACHER_API_KEY", "TEACHER_MODEL"):
|
||
value = os.environ.get(name, "").strip()
|
||
if not value:
|
||
raise ValueError(
|
||
f"环境变量 {name} 未设置。复制 .env.example 为 .env 并填入真实值。"
|
||
)
|
||
values[name] = value
|
||
return (
|
||
values["TEACHER_API_BASE"],
|
||
values["TEACHER_API_KEY"],
|
||
values["TEACHER_MODEL"],
|
||
)
|
||
|
||
|
||
def _strip_leading_think(text: str) -> str:
|
||
"""剥离 content 开头的 <think>...</think> 段（M3 等 reasoning 模型会内联思考）。
|
||
|
||
只剥开头一段：解答正文里若出现字面 "<think>" 字样（例如题目在讨论标签本身），
|
||
不应被误删。
|
||
"""
|
||
return re.sub(r"^\s*<think>.*?</think>\s*", "", text, count=1, flags=re.DOTALL)
|
||
|
||
|
||
class TeacherClient:
|
||
"""OpenAI 兼容的 teacher 客户端：单条生成 + 采样参数收口。
|
||
|
||
差异标注：参考实现是 OpenRouter 专用客户端（带其私有请求头与站点字段）；
|
||
我们用通用 OpenAI 客户端 + base_url 配置驱动，任何兼容网关（new-api、
|
||
vLLM serve、官方 API）都无需改代码。
|
||
|
||
测试注入口：传入 client/model 可绕过 .env 与真实网络（见 tests/test_teacher.py）。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
gen_config: TeacherGenConfig,
|
||
client: OpenAI | None = None,
|
||
model: str | None = None,
|
||
) -> None:
|
||
self.cfg = gen_config
|
||
if client is None:
|
||
base, key, env_model = _load_teacher_env()
|
||
client = OpenAI(
|
||
base_url=base, api_key=key, max_retries=gen_config.max_retries
|
||
)
|
||
model = model or env_model
|
||
if model is None:
|
||
raise ValueError("注入 client 时必须同时指定 model")
|
||
self.client = client
|
||
self.model = model
|
||
|
||
def generate(self, messages: Messages) -> str:
|
||
"""对单条 prompt（messages 列表，末轮为 user）生成解答文本。
|
||
|
||
返回剥离思考段、去首尾空白后的解答。空解答直接报错——空字符串写进
|
||
缓存会在训练时变成全 -100 的空样本（trainer 会炸，但应在这里更早炸）。
|
||
"""
|
||
if self.cfg.system_prompt is not None:
|
||
messages = [
|
||
{"role": "system", "content": self.cfg.system_prompt}
|
||
] + messages
|
||
resp = self.client.chat.completions.create(
|
||
model=self.model,
|
||
messages=messages,
|
||
temperature=self.cfg.temperature,
|
||
top_p=self.cfg.top_p,
|
||
max_tokens=self.cfg.max_tokens,
|
||
)
|
||
content = resp.choices[0].message.content or ""
|
||
if self.cfg.strip_think:
|
||
content = _strip_leading_think(content)
|
||
content = content.strip()
|
||
if not content:
|
||
raise ValueError(
|
||
"teacher 返回空解答（可能：max_tokens 太小把思考截断在半途，"
|
||
"或模型拒答）。该条不会入缓存。"
|
||
)
|
||
return content
|
||
|
||
|
||
def generate_completions(
|
||
prompts: list[Messages],
|
||
cache_path: str,
|
||
teacher: TeacherClient,
|
||
) -> None:
|
||
"""批量生成解答并追加写入 JSONL 缓存（每行 {"key", "completion", "preview"}）。
|
||
|
||
- 已在缓存中的键直接跳过（断点续传）；
|
||
- 并发线程池执行，每完成一条立即写盘并 flush（中断不丢已完成的结果）；
|
||
- 单条失败不中断其余任务（并发中的兄弟请求已经花了钱，先让它们落盘），
|
||
全部结束后若有失败则汇总显式报错——重跑即续传，绝不静默缺数据。
|
||
"""
|
||
path = Path(cache_path)
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
done_keys = _cached_keys(path)
|
||
todo = [(prompt_key(p), p) for p in prompts]
|
||
todo = [(k, p) for k, p in todo if k not in done_keys]
|
||
print(
|
||
f"[teacher] 共 {len(prompts)} 条：缓存命中 {len(prompts) - len(todo)}，"
|
||
f"待生成 {len(todo)}，并发 {teacher.cfg.concurrency}",
|
||
flush=True,
|
||
)
|
||
if not todo:
|
||
return
|
||
|
||
failures: list[tuple[str, str]] = []
|
||
finished = 0
|
||
start = time.monotonic()
|
||
# 写盘收口在主线程（as_completed 消费端），工作线程只跑网络请求——
|
||
# 多线程同写一个文件句柄会交错损坏 JSONL
|
||
with open(path, "a", encoding="utf-8") as f:
|
||
with ThreadPoolExecutor(max_workers=teacher.cfg.concurrency) as pool:
|
||
futures = {pool.submit(teacher.generate, p): (k, p) for k, p in todo}
|
||
for fut in as_completed(futures):
|
||
key, p = futures[fut]
|
||
try:
|
||
completion = fut.result()
|
||
except Exception as e: # noqa: BLE001 —— 收集后统一显式报错，非静默吞错
|
||
failures.append((key, repr(e)))
|
||
continue
|
||
finally:
|
||
finished += 1
|
||
if finished % 20 == 0 or finished == len(todo):
|
||
elapsed = time.monotonic() - start
|
||
rate = finished / elapsed * 60 # 条/分
|
||
eta = (len(todo) - finished) / rate if rate > 0 else 0
|
||
print(
|
||
f"[teacher] {finished}/{len(todo)} 完成 | "
|
||
f"{rate:.1f} 条/分 | 已用 {elapsed / 60:.1f} 分 | "
|
||
f"预计剩余 {eta:.0f} 分",
|
||
flush=True,
|
||
)
|
||
record = {
|
||
"key": key,
|
||
"completion": completion,
|
||
# preview 仅供人工抽查缓存文件，消费端(attach)只认 key/completion
|
||
"preview": p[-1]["content"][:80],
|
||
}
|
||
f.write(json.dumps(record, ensure_ascii=False) + "\n")
|
||
f.flush()
|
||
|
||
if failures:
|
||
examples = "; ".join(f"{k[:12]}…: {err}" for k, err in failures[:3])
|
||
raise RuntimeError(
|
||
f"{len(failures)}/{len(todo)} 条生成失败（成功的已入缓存，重跑本命令"
|
||
f"即断点续传）。前几条错误：{examples}"
|
||
)
|
||
|
||
|
||
def _cached_keys(path: Path) -> set[str]:
|
||
"""读取缓存中已有的键集合；文件不存在视为空缓存（首跑）。"""
|
||
if not path.exists():
|
||
return set()
|
||
keys = set()
|
||
with open(path, encoding="utf-8") as f:
|
||
for line_no, line in enumerate(f, 1):
|
||
if not line.strip():
|
||
continue
|
||
rec = json.loads(line) # 坏行直接炸：缓存损坏必须暴露，不能悄悄重新生成
|
||
keys.add(rec["key"])
|
||
return keys
|