iomgaa
e5a28e8e77
data.py（对应 docs/03 §5 U3）: - 加 prompt_only 开关：True 时输出 prompts/prompt_attention_mask（不产 labels， 由 U4 生成后重建）；False 时 SFT 双预算路径逐字不变 - max_length 改可选：prompt-only 无总预算；SFT 模式缺它构造即报错 - 兑现 T3 为 on-policy 生成预留的口子；生成用左 padding（右边界对齐） test_data.py: - 新增 prompt_only 模式：返回 prompt 张量/不报错、左 padding、截断、剥末轮 assistant - 回归守卫：SFT 模式仍拒绝 prompt-only 行（"SFT 路径行为不变"） Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
319 lines
11 KiB
Python
319 lines
11 KiB
Python
"""层 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 模式下注入 "<T></T>"（模仿
|
||
Qwen3 的空 <think>\\n\\n</think>）；
|
||
- 完整渲染 == 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 "<T></T>"
|
||
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]<T></T>"
|
||
|
||
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]<T></T>"
|
||
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("<T></T>")
|
||
|
||
|
||
def test_sft模式仍拒绝prompt_only行():
|
||
# 回归守卫（docs/03 §5 U3："SFT 路径行为不变"）：SFT 模式下 prompt-only 行
|
||
# 仍是静默空训练闸门，必须报错——放开只发生在显式 prompt_only=True 模式
|
||
with pytest.raises(ValueError, match="prompt-only"):
|
||
make_collator()([row("没有答案的题")])
|
||
|
||
|
||
def test_sft模式缺max_length构造即报错():
|
||
with pytest.raises(ValueError, match="max_length"):
|
||
SFTCollator(
|
||
ToyTokenizer(), max_prompt_length=50
|
||
) # 非 prompt_only 却无 max_length
|
||
|
||
|
||
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
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# SFTCollator：prompt_only 模式（层 2 / U3）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def make_prompt_collator(**kw):
|
||
defaults = dict(max_prompt_length=100, enable_thinking=False, prompt_only=True)
|
||
defaults.update(kw)
|
||
return SFTCollator(ToyTokenizer(), **defaults)
|
||
|
||
|
||
def test_prompt_only模式返回prompt张量且不报错():
|
||
# 层 2：prompt-only 行是常态，不再报错；渲染带生成引导符供 generate 续写
|
||
collator = make_prompt_collator()
|
||
batch = collator([row("ab")])
|
||
expected = "[user]ab[assistant]<T></T>"
|
||
|
||
assert set(batch.keys()) == {"prompts", "prompt_attention_mask"}
|
||
assert batch["prompts"][0].tolist() == ids_of(expected)
|
||
assert batch["prompt_attention_mask"][0].tolist() == [1] * len(expected)
|
||
|
||
|
||
def test_prompt_only模式左padding对齐():
|
||
# 生成要求左 padding：短 prompt 在左侧补 pad，右边界对齐
|
||
collator = make_prompt_collator()
|
||
batch = collator([row("abc"), row("a")])
|
||
t = batch["prompts"].shape[1]
|
||
short_mask = batch["prompt_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["prompts"][1].tolist()[:n_pad] == [0] * n_pad # pad_token_id=0
|
||
|
||
|
||
def test_prompt_only模式截断到prompt预算():
|
||
collator = make_prompt_collator(max_prompt_length=8)
|
||
batch = collator([row("很长的题目" * 20)])
|
||
assert batch["prompts"].shape[1] == 8 # 恰截到 max_prompt_length
|
||
|
||
|
||
def test_prompt_only模式剥掉末轮assistant():
|
||
# 若数据碰巧带了 assistant 轮，取生成前上下文（剥掉它再加生成引导符）
|
||
collator = make_prompt_collator()
|
||
batch = collator([row("q", "已有答案")])
|
||
expected = "[user]q[assistant]<T></T>" # 不含"已有答案"
|
||
assert batch["prompts"][0].tolist() == ids_of(expected)
|