feat : preprocessing 支持 DPO/GRPO 多输出格式
- InputConfig 新增 sources 字段驱动多输出映射 - SectionedMaskBuilder 提取 _process_sections/_build_multi 模板方法 - Pipeline 泛化 accumulate 逻辑处理多 key 结果 - 测试拆分为 config/builder/pipeline 三文件，纯函数风格
This commit is contained in:
@@ -0,0 +1,396 @@
|
||||
from astrai.config.preprocess_config import (
|
||||
InputConfig,
|
||||
OutputConfig,
|
||||
PipelineConfig,
|
||||
ProcessingConfig,
|
||||
)
|
||||
from astrai.preprocessing.builder import (
|
||||
MaskBuilderFactory,
|
||||
SectionedMaskBuilder,
|
||||
)
|
||||
from tests.data.conftest import (
|
||||
_CHAT_SECTIONS,
|
||||
_INSTRUCTION_SECTIONS,
|
||||
_TEXT_SECTIONS,
|
||||
make_chat_config,
|
||||
make_dpo_chat_config,
|
||||
make_grpo_config,
|
||||
make_instruction_config,
|
||||
make_text_config,
|
||||
)
|
||||
|
||||
|
||||
def test_chat_simple(chat_tokenizer):
|
||||
config = make_chat_config()
|
||||
builder = SectionedMaskBuilder()
|
||||
item = {
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are helpful."},
|
||||
{"role": "user", "content": "Hello."},
|
||||
{"role": "assistant", "content": "Hi there!"},
|
||||
]
|
||||
}
|
||||
result = builder.build(item, config, chat_tokenizer)
|
||||
assert result is not None
|
||||
assert "sequence" in result
|
||||
assert "loss_mask" in result
|
||||
assert len(result["sequence"]) == len(result["loss_mask"])
|
||||
|
||||
ids = chat_tokenizer.decode(result["sequence"], skip_special_tokens=False)
|
||||
assert "system" in ids.lower() or "<|im_start|>system" in ids
|
||||
assert "assistant" in ids.lower() or "<|im_start|>assistant" in ids
|
||||
|
||||
total = len(result["sequence"])
|
||||
trained = sum(result["loss_mask"])
|
||||
assert trained > 0
|
||||
assert trained < total
|
||||
|
||||
|
||||
def test_chat_mask_only_assistant(chat_tokenizer):
|
||||
config = make_chat_config()
|
||||
builder = SectionedMaskBuilder()
|
||||
item = {
|
||||
"messages": [
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
{"role": "assistant", "content": "4"},
|
||||
]
|
||||
}
|
||||
result = builder.build(item, config, chat_tokenizer)
|
||||
mask = result["loss_mask"]
|
||||
ids = result["sequence"]
|
||||
assert len(ids) == len(mask)
|
||||
|
||||
trained = [i for i, m in enumerate(mask) if m == 1]
|
||||
masked = [i for i, m in enumerate(mask) if m == 0]
|
||||
assert len(trained) > 0
|
||||
assert len(masked) > 0
|
||||
|
||||
|
||||
def test_chat_all_masked(chat_tokenizer):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(sections=_CHAT_SECTIONS),
|
||||
mask={"system": "mask", "user": "mask", "assistant": "mask"},
|
||||
mask_default="mask",
|
||||
preprocessing=ProcessingConfig(max_seq_len=2048),
|
||||
)
|
||||
builder = SectionedMaskBuilder()
|
||||
item = {
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are helpful."},
|
||||
{"role": "assistant", "content": "Hi there!"},
|
||||
]
|
||||
}
|
||||
result = builder.build(item, config, chat_tokenizer)
|
||||
assert sum(result["loss_mask"]) == 0
|
||||
|
||||
|
||||
def test_chat_all_trained(chat_tokenizer):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(sections=_CHAT_SECTIONS),
|
||||
mask={},
|
||||
mask_default="train",
|
||||
preprocessing=ProcessingConfig(max_seq_len=2048),
|
||||
)
|
||||
builder = SectionedMaskBuilder()
|
||||
item = {
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are helpful."},
|
||||
{"role": "assistant", "content": "Hi there!"},
|
||||
]
|
||||
}
|
||||
result = builder.build(item, config, chat_tokenizer)
|
||||
assert sum(result["loss_mask"]) == len(result["sequence"]) - 1
|
||||
|
||||
|
||||
def test_chat_empty_messages(chat_tokenizer):
|
||||
config = make_chat_config()
|
||||
builder = SectionedMaskBuilder()
|
||||
assert builder.build({"messages": []}, config, chat_tokenizer) is None
|
||||
assert builder.build({}, config, chat_tokenizer) is None
|
||||
|
||||
|
||||
def test_chat_domain_extraction(chat_tokenizer):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(sections=_CHAT_SECTIONS),
|
||||
mask={"assistant": "train"},
|
||||
mask_default="mask",
|
||||
preprocessing=ProcessingConfig(max_seq_len=2048),
|
||||
output=OutputConfig(domain_key="source"),
|
||||
)
|
||||
builder = SectionedMaskBuilder()
|
||||
item = {
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hi"},
|
||||
{"role": "assistant", "content": "Hello"},
|
||||
],
|
||||
"source": "wiki",
|
||||
}
|
||||
result = builder.build(item, config, chat_tokenizer)
|
||||
assert result["domain"] == "wiki"
|
||||
|
||||
|
||||
def test_chat_truncation(chat_tokenizer):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(sections=_CHAT_SECTIONS),
|
||||
mask={"assistant": "train"},
|
||||
mask_default="mask",
|
||||
preprocessing=ProcessingConfig(max_seq_len=10),
|
||||
)
|
||||
builder = SectionedMaskBuilder()
|
||||
item = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Tell me a very long story about dragons and knights and magic.",
|
||||
},
|
||||
{"role": "assistant", "content": "Sure! Here is a tale..."},
|
||||
]
|
||||
}
|
||||
result = builder.build(item, config, chat_tokenizer)
|
||||
assert len(result["sequence"]) <= 10
|
||||
assert len(result["loss_mask"]) == len(result["sequence"])
|
||||
|
||||
|
||||
def test_instruction_basic(test_tokenizer):
|
||||
config = make_instruction_config()
|
||||
builder = SectionedMaskBuilder()
|
||||
item = {"prompt": "Translate to French: Hello", "response": "Bonjour"}
|
||||
result = builder.build(item, config, test_tokenizer)
|
||||
assert result is not None
|
||||
assert len(result["sequence"]) == len(result["loss_mask"])
|
||||
|
||||
|
||||
def test_instruction_prompt_masked(test_tokenizer):
|
||||
config = make_instruction_config()
|
||||
builder = SectionedMaskBuilder()
|
||||
item = {"prompt": "hello", "response": "world"}
|
||||
result = builder.build(item, config, test_tokenizer)
|
||||
mask = result["loss_mask"]
|
||||
ids = result["sequence"]
|
||||
|
||||
prompt_ids = test_tokenizer.encode("hello", add_special_tokens=True)
|
||||
p_len = min(len(prompt_ids), len(ids))
|
||||
assert all(m == 0 for m in mask[:p_len])
|
||||
if p_len < len(ids):
|
||||
assert all(m == 1 for m in mask[p_len:])
|
||||
|
||||
|
||||
def test_instruction_train_on_prompt(test_tokenizer):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(
|
||||
sections=[
|
||||
{"field": "prompt", "action": "train", "add_special_tokens": True},
|
||||
{"field": "response", "action": "mask"},
|
||||
]
|
||||
),
|
||||
preprocessing=ProcessingConfig(max_seq_len=2048),
|
||||
)
|
||||
builder = SectionedMaskBuilder()
|
||||
item = {"prompt": "hello", "response": "world"}
|
||||
result = builder.build(item, config, test_tokenizer)
|
||||
mask = result["loss_mask"]
|
||||
ids = result["sequence"]
|
||||
|
||||
prompt_ids = test_tokenizer.encode("hello", add_special_tokens=True)
|
||||
p_len = min(len(prompt_ids), len(ids))
|
||||
assert all(m == 1 for m in mask[:p_len])
|
||||
|
||||
|
||||
def test_text_basic(test_tokenizer):
|
||||
config = make_text_config()
|
||||
builder = SectionedMaskBuilder()
|
||||
item = {"text": "Hello world. This is a test document."}
|
||||
result = builder.build(item, config, test_tokenizer)
|
||||
assert result is not None
|
||||
assert "sequence" in result
|
||||
assert len(result["sequence"]) > 0
|
||||
assert "loss_mask" not in result
|
||||
|
||||
|
||||
def test_text_empty(test_tokenizer):
|
||||
config = make_text_config()
|
||||
builder = SectionedMaskBuilder()
|
||||
assert builder.build({"text": ""}, config, test_tokenizer) is None
|
||||
assert builder.build({"text": " "}, config, test_tokenizer) is None
|
||||
|
||||
|
||||
def test_text_too_short(test_tokenizer):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(sections=_TEXT_SECTIONS),
|
||||
preprocessing=ProcessingConfig(min_chars=100),
|
||||
)
|
||||
builder = SectionedMaskBuilder()
|
||||
assert builder.build({"text": "short"}, config, test_tokenizer) is None
|
||||
|
||||
|
||||
def test_text_truncation(test_tokenizer):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(sections=_TEXT_SECTIONS),
|
||||
preprocessing=ProcessingConfig(max_seq_len=3, min_chars=1),
|
||||
)
|
||||
builder = SectionedMaskBuilder()
|
||||
item = {"text": "This is a very long text that should be truncated"}
|
||||
result = builder.build(item, config, test_tokenizer)
|
||||
assert len(result["sequence"]) <= 3
|
||||
|
||||
|
||||
def test_sectioned_chat(chat_tokenizer):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(sections=_CHAT_SECTIONS),
|
||||
mask={"system": "mask", "user": "mask", "assistant": "train"},
|
||||
mask_default="mask",
|
||||
preprocessing=ProcessingConfig(max_seq_len=2048),
|
||||
)
|
||||
builder = SectionedMaskBuilder()
|
||||
item = {
|
||||
"messages": [
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
{"role": "assistant", "content": "4"},
|
||||
]
|
||||
}
|
||||
result = builder.build(item, config, chat_tokenizer)
|
||||
assert result is not None
|
||||
assert len(result["sequence"]) == len(result["loss_mask"])
|
||||
assert sum(result["loss_mask"]) > 0
|
||||
assert 0 in result["loss_mask"]
|
||||
|
||||
|
||||
def test_sectioned_instruction(test_tokenizer):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(sections=_INSTRUCTION_SECTIONS),
|
||||
preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=0),
|
||||
)
|
||||
builder = SectionedMaskBuilder()
|
||||
item = {"prompt": "Q: Why?", "response": "A: Because."}
|
||||
result = builder.build(item, config, test_tokenizer)
|
||||
assert result is not None
|
||||
mask = result["loss_mask"]
|
||||
assert mask[0] == 0
|
||||
assert mask[-1] == 1
|
||||
|
||||
|
||||
def test_sectioned_text(test_tokenizer):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(sections=_TEXT_SECTIONS),
|
||||
preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=1),
|
||||
)
|
||||
builder = SectionedMaskBuilder()
|
||||
item = {"text": "Hello world, this is a test."}
|
||||
result = builder.build(item, config, test_tokenizer)
|
||||
assert result is not None
|
||||
assert "loss_mask" not in result
|
||||
|
||||
|
||||
def test_sectioned_text_too_short(test_tokenizer):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(sections=_TEXT_SECTIONS),
|
||||
preprocessing=ProcessingConfig(max_seq_len=2048, min_chars=100),
|
||||
)
|
||||
builder = SectionedMaskBuilder()
|
||||
assert builder.build({"text": "short"}, config, test_tokenizer) is None
|
||||
|
||||
|
||||
def test_factory_registered():
|
||||
names = MaskBuilderFactory._registry.list_names()
|
||||
assert "sectioned" in names
|
||||
|
||||
|
||||
def test_factory_create():
|
||||
builder = MaskBuilderFactory.create("sectioned")
|
||||
assert isinstance(builder, SectionedMaskBuilder)
|
||||
|
||||
|
||||
def test_dpo_chat_basic(chat_tokenizer):
|
||||
config = make_dpo_chat_config()
|
||||
builder = SectionedMaskBuilder()
|
||||
item = {
|
||||
"chosen": [
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
{"role": "assistant", "content": "4"},
|
||||
],
|
||||
"rejected": [
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
{"role": "assistant", "content": "5"},
|
||||
],
|
||||
}
|
||||
result = builder.build(item, config, chat_tokenizer)
|
||||
assert result is not None
|
||||
assert "chosen" in result
|
||||
assert "rejected" in result
|
||||
assert "chosen_mask" in result
|
||||
assert "rejected_mask" in result
|
||||
assert "domain" in result
|
||||
assert len(result["chosen"]) == len(result["chosen_mask"])
|
||||
assert len(result["rejected"]) == len(result["rejected_mask"])
|
||||
assert sum(result["chosen_mask"]) > 0
|
||||
assert sum(result["rejected_mask"]) > 0
|
||||
|
||||
|
||||
def test_dpo_chosen_only_trained(chat_tokenizer):
|
||||
config = make_dpo_chat_config()
|
||||
builder = SectionedMaskBuilder()
|
||||
item = {
|
||||
"chosen": [
|
||||
{"role": "user", "content": "Hi"},
|
||||
{"role": "assistant", "content": "Hello"},
|
||||
],
|
||||
"rejected": [
|
||||
{"role": "user", "content": "Hi"},
|
||||
{"role": "assistant", "content": "Go away"},
|
||||
],
|
||||
}
|
||||
result = builder.build(item, config, chat_tokenizer)
|
||||
assert 0 in result["chosen_mask"]
|
||||
assert 1 in result["chosen_mask"]
|
||||
assert 0 in result["rejected_mask"]
|
||||
assert 1 in result["rejected_mask"]
|
||||
|
||||
|
||||
def test_dpo_missing_field_is_none(chat_tokenizer):
|
||||
config = make_dpo_chat_config()
|
||||
builder = SectionedMaskBuilder()
|
||||
assert builder.build({"chosen": [], "rejected": []}, config, chat_tokenizer) is None
|
||||
|
||||
|
||||
def test_grpo_basic(chat_tokenizer):
|
||||
config = make_grpo_config()
|
||||
builder = SectionedMaskBuilder()
|
||||
item = {
|
||||
"prompt": [{"role": "user", "content": "What is 2+2?"}],
|
||||
"responses": ["4", "The answer is four", "Four", "2+2=4"],
|
||||
"rewards": [1.0, 0.5, 0.8, 0.2],
|
||||
}
|
||||
result = builder.build(item, config, chat_tokenizer)
|
||||
assert result is not None
|
||||
assert "prompts" in result
|
||||
assert "responses" in result
|
||||
assert "masks" in result
|
||||
assert "rewards" in result
|
||||
assert len(result["responses"]) == len(result["masks"])
|
||||
assert result["rewards"] == [1.0, 0.5, 0.8, 0.2]
|
||||
|
||||
|
||||
def test_grpo_response_tokens_all_trained(chat_tokenizer):
|
||||
config = make_grpo_config()
|
||||
builder = SectionedMaskBuilder()
|
||||
item = {
|
||||
"prompt": [{"role": "user", "content": "Q"}],
|
||||
"responses": ["A", "B"],
|
||||
"rewards": [0.8, 0.2],
|
||||
}
|
||||
result = builder.build(item, config, chat_tokenizer)
|
||||
masks = result["masks"]
|
||||
assert all(m == 1 for m in masks)
|
||||
assert len(masks) == len(result["responses"])
|
||||
|
||||
|
||||
def test_grpo_single_reward(chat_tokenizer):
|
||||
config = make_grpo_config()
|
||||
builder = SectionedMaskBuilder()
|
||||
item = {
|
||||
"prompt": [{"role": "user", "content": "Q"}],
|
||||
"responses": ["A"],
|
||||
"rewards": 0.9,
|
||||
}
|
||||
result = builder.build(item, config, chat_tokenizer)
|
||||
assert result["rewards"] == [0.9]
|
||||
Reference in New Issue
Block a user