diff --git a/ars_opd/data.py b/ars_opd/data.py new file mode 100644 index 0000000..a38bc30 --- /dev/null +++ b/ars_opd/data.py @@ -0,0 +1,328 @@ +"""数据管线(IO 边缘,无论文锚点):加载 → messages 归一 → 挂接 teacher 解答 → collator。 + +层 1 的数据流(对应 docs/02 §1 的基线定义:SFT = 在 teacher rollout 上的离线蒸馏): + + DAPO parquet(prompt-only) + → to_messages 归一成 [{"role","content"}] 列表 + → 按 seed 抽子集 + → attach_teacher_completions 从 JSONL 缓存挂上 teacher 解答(assistant 轮) + → SFTCollator 分词、双预算截断、-100 掩码、左 padding + +本模块与 teacher.py 的缓存契约由 `prompt_key` 单点定义:teacher.py 生成缓存、 +本模块消费缓存,双方必须用同一个函数算键。 +""" + +from __future__ import annotations + +import ast +import hashlib +import json +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch +from datasets import Dataset, load_dataset + +if TYPE_CHECKING: + from ars_opd.configs import SFTConfig + +# F.cross_entropy 的 ignore_index 默认值;标了它的位置不产生 loss +IGNORE_INDEX = -100 + + +# --------------------------------------------------------------------------- +# messages 归一 +# --------------------------------------------------------------------------- + + +def _parse_stringified_list(value: str, column: str) -> list: + """parquet 有时把 list 存成其字符串形态,用 ast 还原。 + + 差异标注:参考实现(train_distillation.py:299-305)在这里 `except: pass` 静默吞错, + 坏行会以原始字符串流进 collator,在 apply_chat_template 处以难懂的方式炸; + 我们显式报错,错误信息直接指向坏数据本身。 + """ + try: + parsed = ast.literal_eval(value) + except (ValueError, SyntaxError) as e: + raise ValueError( + f"列 {column!r} 是字符串但无法解析为 Python 字面量(坏数据行):" + f"{value[:200]!r}" + ) from e + if not isinstance(parsed, (list, tuple)): + raise ValueError(f"列 {column!r} 解析结果不是列表:{type(parsed).__name__}") + return list(parsed) + + +def to_messages(example: dict[str, Any]) -> dict[str, list[dict[str, str]]]: + """把三种来源格式归一成 messages 列:[{"role": ..., "content": ...}, ...]。 + + 支持(与参考实现 train_distillation.py:295-321 相同的三分支): + - ``messages`` 列:直取; + - ``prompt`` 列(DAPO parquet,列名不副实——装的是完整 chat 列表):改名; + - ``question`` 列(gsm8k 风格纯文本):包成单 user 轮。 + + 差异标注:参考实现对不认识的行 `return x` 静默放行,我们显式报错。 + """ + if "messages" in example: + msgs = example["messages"] + column = "messages" + elif "prompt" in example: + msgs = example["prompt"] + column = "prompt" + elif "question" in example: + return {"messages": [{"role": "user", "content": example["question"]}]} + else: + raise ValueError( + f"无法识别的数据行:既无 messages/prompt 也无 question 列," + f"实有列 {sorted(example.keys())}" + ) + + if isinstance(msgs, str): + msgs = _parse_stringified_list(msgs, column) + msgs = list(msgs) + + if not msgs: + raise ValueError(f"列 {column!r} 是空列表(坏数据行)") + for m in msgs: + if not isinstance(m, dict) or "role" not in m or "content" not in m: + raise ValueError( + f"列 {column!r} 中存在非 {{role, content}} 结构的元素:{m!r}" + ) + + return {"messages": [{"role": m["role"], "content": m["content"]} for m in msgs]} + + +# --------------------------------------------------------------------------- +# teacher 解答缓存(与 teacher.py 的契约) +# --------------------------------------------------------------------------- + + +def prompt_key(messages: list[dict[str, str]]) -> str: + """teacher 缓存的键:对 messages 的规范化 JSON 取 sha256。 + + 差异标注:参考实现(distillation_trainer.py:984)用 `str(hash(prompt))`—— + Python 对 str 的 hash 默认加盐,跨进程/跨次运行不稳定,缓存必然失效重生成。 + sha256 内容寻址:同一道题永远同一个键。 + + 只取 role/content 两个字段参与哈希:DAPO 行里其余元数据(data_source 等) + 变了不应导致缓存失效。 + """ + canon = [{"role": m["role"], "content": m["content"]} for m in messages] + return hashlib.sha256(json.dumps(canon, ensure_ascii=False).encode()).hexdigest() + + +def attach_teacher_completions(dataset: Dataset, jsonl_path: str) -> Dataset: + """把 teacher 解答缓存(JSONL,每行 {"key", "completion"})挂到数据集上。 + + - 末轮已是 assistant 的行保持原样(数据自带解答,不覆盖); + - 任何 prompt-only 行在缓存中查不到键 → 收集齐所有缺失后一次性报错, + 提示先运行 teacher 生成——绝不静默跳过(跳过 = 悄悄改变训练集组成)。 + """ + path = Path(jsonl_path) + if not path.exists(): + raise FileNotFoundError( + f"teacher 解答缓存不存在:{jsonl_path}。先运行 teacher.py 的批量生成。" + ) + + cache: dict[str, str] = {} + 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) # 坏行直接炸,带行号 + if "key" not in rec or "completion" not in rec: + raise ValueError(f"{jsonl_path}:{line_no} 缺少 key/completion 字段") + cache[rec["key"]] = rec["completion"] + + # 先整体扫描缺失,一次性报全——比在 .map 里炸第一条更省来回 + missing = [ + i + for i, ex in enumerate(dataset) + if ex["messages"][-1]["role"] != "assistant" + and prompt_key(ex["messages"]) not in cache + ] + if missing: + raise KeyError( + f"{len(missing)}/{len(dataset)} 行在 teacher 缓存中查不到解答" + f"(首个缺失行 index={missing[0]})。检查:teacher 生成是否用了同一" + f"子集与同一 seed?(子集抽取在 load_sft_dataset 中先于挂接发生," + f"两侧 seed 不同则键集合不同)" + ) + + def _attach(ex: dict[str, Any]) -> dict[str, Any]: + msgs = ex["messages"] + if msgs[-1]["role"] == "assistant": + return ex + completion = cache[prompt_key(msgs)] + return {"messages": msgs + [{"role": "assistant", "content": completion}]} + + return dataset.map(_attach) + + +# --------------------------------------------------------------------------- +# 数据集加载(入口) +# --------------------------------------------------------------------------- + + +def load_sft_dataset(cfg: "SFTConfig") -> Dataset: + """层 1 数据管线入口:加载 → 归一 → 抽子集 → 挂 teacher 解答。 + + 返回只含 ``messages`` 一列的 Dataset,每行末轮是 assistant(可直接喂 SFTCollator)。 + """ + ds = _load_raw(cfg.dataset_path, cfg.dataset_split) + ds = ds.map( + to_messages, + remove_columns=[c for c in ds.column_names if c != "messages"], + ) + if cfg.subset_size is not None and cfg.subset_size < len(ds): + # 非显然约束:抽子集必须在挂接 teacher 解答之前、且由 seed 完全确定—— + # teacher.py 生成缓存时会走完全相同的"加载→归一→抽子集"路径,两侧 seed + # 一致才能得到同一批题;否则 attach 处大面积缓存 miss 报错。 + ds = ds.shuffle(seed=cfg.seed).select(range(cfg.subset_size)) + if cfg.teacher_completions_path is not None: + ds = attach_teacher_completions(ds, cfg.teacher_completions_path) + return ds + + +def _load_raw(dataset_path: str, split: str) -> Dataset: + """三分支加载:parquet 目录 / 单 parquet 文件 / HF Hub 数据集名。 + + 差异标注:参考实现(train_distillation.py:292)对 Hub 分支硬编码 config 名 + "main"(gsm8k 专用);我们不硬编码——需要特定 config 的数据集请下载成 + parquet 本地加载。 + """ + p = Path(dataset_path) + if p.is_dir(): + return load_dataset("parquet", data_dir=dataset_path, split=split) + if dataset_path.endswith(".parquet"): + return load_dataset("parquet", data_files=dataset_path, split=split) + return load_dataset(dataset_path, split=split) + + +# --------------------------------------------------------------------------- +# Collator:本层最核心的一段(对拍 distillation_trainer.py:210-343) +# --------------------------------------------------------------------------- + + +class SFTCollator: + """把一个 batch 的 messages 变成 input_ids/attention_mask/labels。 + + 核心设计(继承参考实现的双预算方案,docs/02 §2.3):prompt 与 completion + 各自独立预算——prompt 用 max_prompt_length 截断,completion 上限是 + max_length - len(截断后 prompt)。若只用一个总预算从右截断,超长解答会把 + prompt 挤空,模型在"没有题目"的样本上学解答。 + + 与参考实现的差异: + - 不支持 prompt-only 行(直接报错)。参考实现支持是为 on-policy 生成留口, + 纯 SFT 下 prompt-only 行只会静默产生零 loss;层 2 接 on-policy 时再放开。 + - 不返回 prompts/prompt_attention_mask(参考实现留给 vLLM 生成用,层 1 用不到)。 + - 空 的一次性诊断打印改为单元测试断言(契约进测试,不进运行时日志)。 + """ + + def __init__( + self, + tokenizer: "Any", + max_length: int, + max_prompt_length: int, + enable_thinking: bool = False, + ) -> None: + """tokenizer 需实现 HF 接口:apply_chat_template / __call__ / pad_token_id。""" + self.tokenizer = tokenizer + self.max_length = max_length + self.max_prompt_length = max_prompt_length + self.enable_thinking = enable_thinking + # pad→eos 回退:左 padding 位置的 attention_mask 恒为 0,pad 值不参与 + # 任何计算,只需要一个合法 token id 占位,借用 eos 即可 + if tokenizer.pad_token_id is not None: + self.pad_token_id: int = tokenizer.pad_token_id + elif tokenizer.eos_token_id is not None: + self.pad_token_id = tokenizer.eos_token_id + else: + raise ValueError("tokenizer 既无 pad_token 也无 eos_token,无法 padding") + + def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]: + """batch 的 messages → 定长张量。 + + 返回(B = batch 大小,T = batch 内最长序列长度): + - input_ids: (B, T) 左 padding + - attention_mask: (B, T) padding 位置为 0 + - labels: (B, T) padding 与 prompt 位置为 -100,completion 位置为 token id + """ + all_input_ids: list[list[int]] = [] + all_labels: list[list[int]] = [] + + for example in examples: + messages = example["messages"] + if len(messages) < 2 or messages[-1]["role"] != "assistant": + raise ValueError( + "SFTCollator 收到 prompt-only 行(末轮不是 assistant)。纯 SFT 下" + "它只会产生全 -100 的零 loss 样本——静默空训练。检查 teacher " + "解答是否挂接成功。" + ) + + # prompt = 末轮 assistant 之前的全部轮次,渲染时带生成引导符 + # ("<|im_start|>assistant\n..."),这样 completion 是纯解答文本的分词 + formatted_prompt = self.tokenizer.apply_chat_template( + messages[:-1], + tokenize=False, + add_generation_prompt=True, + enable_thinking=self.enable_thinking, + ) + # prompt 自己的预算内截断。沿用 tokenizer 默认右截断(与参考实现一致): + # 超预算的题目被截掉尾部(含生成引导符)——1024 预算下 DAPO 极少触发, + # 触发时该样本退化但不会污染边界(边界用未截断长度算,见下) + prompt_ids: list[int] = self.tokenizer( + formatted_prompt, + truncation=True, + max_length=self.max_prompt_length, + add_special_tokens=False, + )["input_ids"] + + # 非显然约束(docs/02 坑一/坑二):completion 边界必须用"未截断 prompt + # 的分词长度"从整段渲染中切出。BPE 分词不满足拼接稳定性,分开渲染 + # prompt 和 completion 再拼接 ≠ 整段渲染后分词;而若用截断后长度当切分 + # 点,会把 prompt 尾部的 token 误标成 completion——静默的语义错误。 + formatted_full = self.tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=False, + enable_thinking=self.enable_thinking, + ) + full_ids: list[int] = self.tokenizer( + formatted_full, truncation=False, add_special_tokens=False + )["input_ids"] + untruncated_prompt_len = len( + self.tokenizer( + formatted_prompt, truncation=False, add_special_tokens=False + )["input_ids"] + ) + completion_ids = full_ids[untruncated_prompt_len:] + + # completion 预算 = 总预算 - 截断后 prompt 实长。配置校验 + # (max_prompt_length < max_length)保证它恒 > 0 + completion_budget = self.max_length - len(prompt_ids) + completion_ids = completion_ids[:completion_budget] + + all_input_ids.append(prompt_ids + completion_ids) + # prompt 位置标 -100:题目不产生 loss,只学解答 + all_labels.append([IGNORE_INDEX] * len(prompt_ids) + completion_ids) + + # 左 padding:batch 内所有序列右对齐。纯 SFT 用右 padding 也行,但左 padding + # 让 trainer 能用一个标量 prompt_length 切 batch(docs/02 §2.4),且与 + # 层 2+ 的生成场景(生成必须左 padding)统一,全项目只有一种 padding 约定 + return { + "input_ids": _left_pad(all_input_ids, self.pad_token_id), # (B, T) + "attention_mask": _left_pad( + [[1] * len(ids) for ids in all_input_ids], 0 + ), # (B, T) + "labels": _left_pad(all_labels, IGNORE_INDEX), # (B, T) + } + + +def _left_pad(seqs: list[list[int]], pad_value: int) -> torch.Tensor: + """把变长序列在左侧补齐成 (B, T) 张量,T = batch 内最大长度。""" + t_max = max(len(s) for s in seqs) + return torch.tensor( + [[pad_value] * (t_max - len(s)) + s for s in seqs], dtype=torch.long + ) diff --git a/tests/test_data.py b/tests/test_data.py new file mode 100644 index 0000000..e22289c --- /dev/null +++ b/tests/test_data.py @@ -0,0 +1,260 @@ +"""层 1 / T3:数据管线单测(docs/02 §5.1 规定的验证项)。 + +用字符级玩具 tokenizer 在 CPU 上对拍 collator 行为,不依赖网络下载真模型。 +玩具模板刻意模仿 Qwen3 的关键结构:生成引导符 + no-think 时注入空思考块。 +""" + +import json + +import pytest +from datasets import Dataset + +from ars_opd.data import ( + IGNORE_INDEX, + SFTCollator, + attach_teacher_completions, + prompt_key, + to_messages, +) + + +class ToyTokenizer: + """字符级 tokenizer:一个字符一个 token(id = 码点)。 + + 模板契约与真 chat 模板同构: + - 每轮渲染成 "[role]content"; + - assistant 轮(或生成引导符后)在 no-think 模式下注入 ""(模仿 + Qwen3 的空 \\n\\n); + - 完整渲染 == prompt 渲染 + 解答文本,保证边界可精确断言。 + """ + + def __init__(self, pad_token_id=0, eos_token_id=1): + self.pad_token_id = pad_token_id + self.eos_token_id = eos_token_id + + def apply_chat_template( + self, + messages, + tokenize=False, + add_generation_prompt=False, + enable_thinking=False, + ): + think = "" if enable_thinking else "" + parts = [] + for m in messages: + prefix = think if m["role"] == "assistant" else "" + parts.append(f"[{m['role']}]{prefix}{m['content']}") + text = "".join(parts) + if add_generation_prompt: + text += f"[assistant]{think}" + return text + + def __call__( + self, + text, + truncation=False, + max_length=None, + padding=False, + add_special_tokens=False, + ): + ids = [ord(c) for c in text] + if truncation and max_length is not None: + ids = ids[:max_length] + return {"input_ids": ids} + + +def ids_of(text): + return [ord(c) for c in text] + + +def row(question, answer=None): + msgs = [{"role": "user", "content": question}] + if answer is not None: + msgs.append({"role": "assistant", "content": answer}) + return {"messages": msgs} + + +# --------------------------------------------------------------------------- +# to_messages +# --------------------------------------------------------------------------- + + +def test_dapo_prompt列直接归一(): + ex = {"prompt": [{"role": "user", "content": "1+1=?"}], "data_source": "dapo"} + assert to_messages(ex) == {"messages": [{"role": "user", "content": "1+1=?"}]} + + +def test_字符串化的列表被还原(): + ex = {"prompt": "[{'role': 'user', 'content': 'hi'}]"} + assert to_messages(ex)["messages"] == [{"role": "user", "content": "hi"}] + + +def test_坏字符串显式报错而非静默放行(): + # 参考实现 except:pass 会让这行以字符串形态流进 collator + with pytest.raises(ValueError, match="无法解析"): + to_messages({"prompt": "[{'role': broken"}) + + +def test_question列包成单user轮(): + assert to_messages({"question": "2+2=?"}) == { + "messages": [{"role": "user", "content": "2+2=?"}] + } + + +def test_无法识别的行报错(): + with pytest.raises(ValueError, match="无法识别"): + to_messages({"foo": 1}) + + +# --------------------------------------------------------------------------- +# prompt_key(与 teacher.py 的缓存契约) +# --------------------------------------------------------------------------- + + +def test_同题同键_不同题不同键(): + m1 = [{"role": "user", "content": "q"}] + m2 = [{"role": "user", "content": "q'"}] + assert prompt_key(m1) == prompt_key(m1) + assert prompt_key(m1) != prompt_key(m2) + + +def test_额外元数据字段不影响键(): + plain = [{"role": "user", "content": "q"}] + noisy = [{"role": "user", "content": "q", "source": "dapo"}] + assert prompt_key(plain) == prompt_key(noisy) + + +# --------------------------------------------------------------------------- +# attach_teacher_completions +# --------------------------------------------------------------------------- + + +def write_cache(path, entries): + with open(path, "w", encoding="utf-8") as f: + for msgs, completion in entries: + f.write( + json.dumps({"key": prompt_key(msgs), "completion": completion}) + "\n" + ) + + +def test_挂接teacher解答(tmp_path): + q = [{"role": "user", "content": "1+1=?"}] + cache = tmp_path / "cache.jsonl" + write_cache(cache, [(q, "答案是 2")]) + ds = Dataset.from_list([{"messages": q}]) + out = attach_teacher_completions(ds, str(cache)) + assert out[0]["messages"][-1] == {"role": "assistant", "content": "答案是 2"} + + +def test_缓存缺键一次性报全部缺失(tmp_path): + cache = tmp_path / "cache.jsonl" + write_cache(cache, []) + ds = Dataset.from_list([row("q1"), row("q2")]) + with pytest.raises(KeyError, match="2/2"): + attach_teacher_completions(ds, str(cache)) + + +def test_自带解答的行不被覆盖(tmp_path): + cache = tmp_path / "cache.jsonl" + write_cache(cache, []) + ds = Dataset.from_list([row("q", "人写的答案")]) + out = attach_teacher_completions(ds, str(cache)) + assert out[0]["messages"][-1]["content"] == "人写的答案" + + +def test_缓存文件不存在报错(): + ds = Dataset.from_list([row("q")]) + with pytest.raises(FileNotFoundError): + attach_teacher_completions(ds, "/不存在/cache.jsonl") + + +# --------------------------------------------------------------------------- +# SFTCollator +# --------------------------------------------------------------------------- + + +def make_collator(**kw): + defaults = dict(max_length=1000, max_prompt_length=100, enable_thinking=False) + defaults.update(kw) + return SFTCollator(ToyTokenizer(), **defaults) + + +def test_基本形态_掩码与边界(): + collator = make_collator() + batch = collator([row("ab", "cd")]) + prompt_text = "[user]ab[assistant]" + + labels = batch["labels"][0].tolist() + # prompt 全 -100,completion 位置是解答的 token + assert labels[: len(prompt_text)] == [IGNORE_INDEX] * len(prompt_text) + assert labels[len(prompt_text) :] == ids_of("cd") + assert batch["input_ids"][0].tolist() == ids_of(prompt_text + "cd") + assert batch["attention_mask"][0].tolist() == [1] * (len(prompt_text) + 2) + + +def test_超长解答不挤占prompt(): + # 头号正确性卖点:completion 被截,prompt 一个 token 不少 + prompt_text = "[user]ab[assistant]" + collator = make_collator(max_length=len(prompt_text) + 3) + batch = collator([row("ab", "x" * 50)]) + + input_ids = batch["input_ids"][0].tolist() + assert input_ids[: len(prompt_text)] == ids_of(prompt_text) # prompt 完整 + assert len(input_ids) == len(prompt_text) + 3 # completion 只剩预算内 3 个 + + +def test_超长题目截断但边界不错位(): + # 坑二场景:prompt 超预算被截断,completion 的 token 必须仍然精确 + # (切分点用未截断长度,而非截断后长度) + collator = make_collator(max_prompt_length=10) + batch = collator([row("很长的题目" * 20, "答案")]) + + labels = batch["labels"][0].tolist() + non_masked = [t for t in labels if t != IGNORE_INDEX] + assert non_masked == ids_of("答案") # 解答 token 一个不错 + assert sum(t == IGNORE_INDEX for t in labels) == 10 # prompt 恰被截到预算 + + +def test_enable_thinking两种取值边界都正确(): + for thinking in (False, True): + collator = make_collator(enable_thinking=thinking) + batch = collator([row("q", "ans")]) + non_masked = [t for t in batch["labels"][0].tolist() if t != IGNORE_INDEX] + assert non_masked == ids_of("ans"), f"enable_thinking={thinking} 时边界错位" + + +def test_nothink模板注入空思考块(): + # 参考实现的一次性诊断打印,在这里变成永久契约 + text = ToyTokenizer().apply_chat_template( + [{"role": "user", "content": "q"}], + add_generation_prompt=True, + enable_thinking=False, + ) + assert text.endswith("") + + +def test_prompt_only行显式报错(): + with pytest.raises(ValueError, match="prompt-only"): + make_collator()([row("没有答案的题")]) + + +def test_左padding对齐(): + collator = make_collator() + batch = collator([row("ab", "cd"), row("a", "c")]) + t = batch["input_ids"].shape[1] + short_mask = batch["attention_mask"][1].tolist() + n_pad = t - short_mask.count(1) + + assert n_pad > 0 + assert short_mask[:n_pad] == [0] * n_pad # padding 在左 + assert batch["labels"][1].tolist()[:n_pad] == [IGNORE_INDEX] * n_pad + assert batch["input_ids"][1].tolist()[:n_pad] == [0] * n_pad # pad_token_id=0 + + +def test_pad回退到eos(): + collator = SFTCollator( + ToyTokenizer(pad_token_id=None, eos_token_id=7), + max_length=100, + max_prompt_length=50, + ) + assert collator.pad_token_id == 7