refactor : 基于声明式 JSON 配置的预处理管线重构
- 用工厂注册的 MaskBuilder（chat/instruction/text）替换硬编码的 _transform_* 方法 - mask 规则以 role-to-action 映射声明在配置中，与 chat_template 完全解耦 - 单次编码 + role-span 追踪替代两次编码 + 长度差计算 mask 的方式 - 支持多轮对话训练：所有 assistant 轮次参与训练，而非仅最后一轮 - 新建 astrai.preprocessing 包（builder.py + pipeline.py），删除 astrai/preprocess.py - CLI 精简为 --config 参数，所有参数通过 PipelineConfig JSON 配置 - 新增 PipelineConfig、InputConfig、ProcessingConfig、OutputConfig dataclass - 文档：assets/docs/preprocessing.md - 27 个测试覆盖 mask builder、pipeline、配置序列化、工厂注册
This commit is contained in:
@@ -0,0 +1,522 @@
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import pytest
|
||||
|
||||
from astrai.config.preprocess_config import (
|
||||
InputConfig,
|
||||
OutputConfig,
|
||||
PipelineConfig,
|
||||
ProcessingConfig,
|
||||
)
|
||||
from astrai.preprocessing.builder import (
|
||||
ChatMaskBuilder,
|
||||
InstructionMaskBuilder,
|
||||
MaskBuilderFactory,
|
||||
TextMaskBuilder,
|
||||
)
|
||||
from astrai.preprocessing.pipeline import Pipeline, dedup_signature, filter_by_length
|
||||
from astrai.tokenize import AutoTokenizer
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def real_tokenizer():
|
||||
return AutoTokenizer.from_pretrained("params")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_dir():
|
||||
d = tempfile.mkdtemp()
|
||||
yield d
|
||||
import shutil
|
||||
|
||||
shutil.rmtree(d, ignore_errors=True)
|
||||
|
||||
|
||||
def make_chat_config():
|
||||
return PipelineConfig(
|
||||
input=InputConfig(type="chat", messages_key="messages"),
|
||||
mask={"system": "mask", "user": "mask", "assistant": "train"},
|
||||
mask_default="mask",
|
||||
preprocessing=ProcessingConfig(max_seq_len=2048),
|
||||
)
|
||||
|
||||
|
||||
def make_instruction_config():
|
||||
return PipelineConfig(
|
||||
input=InputConfig(
|
||||
type="instruction", prompt_key="prompt", response_key="response"
|
||||
),
|
||||
mask={"prompt": "mask", "response": "train"},
|
||||
mask_default="mask",
|
||||
preprocessing=ProcessingConfig(max_seq_len=2048),
|
||||
)
|
||||
|
||||
|
||||
def make_text_config():
|
||||
return PipelineConfig(
|
||||
input=InputConfig(type="text", text_key="text"),
|
||||
preprocessing=ProcessingConfig(
|
||||
max_seq_len=2048, min_chars=1, max_chars=2_000_000
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class TestPipelineConfig:
|
||||
def test_default_values(self):
|
||||
config = PipelineConfig()
|
||||
assert config.version == 1
|
||||
assert config.input.type == "chat"
|
||||
assert config.mask == {}
|
||||
assert config.mask_default == "mask"
|
||||
assert config.preprocessing.max_seq_len == 2048
|
||||
assert config.output.storage_format == "bin"
|
||||
|
||||
def test_from_dict_flat(self):
|
||||
data = {
|
||||
"version": 1,
|
||||
"input": {"type": "chat", "messages_key": "msgs"},
|
||||
"mask": {"system": "mask", "assistant": "train"},
|
||||
"mask_default": "mask",
|
||||
"preprocessing": {"max_seq_len": 1024},
|
||||
"output": {"storage_format": "h5"},
|
||||
}
|
||||
config = PipelineConfig.from_dict(data)
|
||||
assert config.input.type == "chat"
|
||||
assert config.input.messages_key == "msgs"
|
||||
assert config.mask == {"system": "mask", "assistant": "train"}
|
||||
assert config.preprocessing.max_seq_len == 1024
|
||||
assert config.output.storage_format == "h5"
|
||||
|
||||
def test_to_dict_roundtrip(self):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(type="instruction", prompt_key="q", response_key="a"),
|
||||
mask={"prompt": "mask", "response": "train"},
|
||||
mask_default="mask",
|
||||
)
|
||||
d = config.to_dict()
|
||||
config2 = PipelineConfig.from_dict(d)
|
||||
assert config2.input.type == "instruction"
|
||||
assert config2.input.prompt_key == "q"
|
||||
assert config2.mask == {"prompt": "mask", "response": "train"}
|
||||
|
||||
def test_to_json_from_json(self, temp_dir):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(type="text", text_key="body"),
|
||||
mask={"text": "train"},
|
||||
mask_default="mask",
|
||||
)
|
||||
path = os.path.join(temp_dir, "config.json")
|
||||
config.to_json(path)
|
||||
loaded = PipelineConfig.from_json(path)
|
||||
assert loaded.input.type == "text"
|
||||
assert loaded.input.text_key == "body"
|
||||
assert loaded.mask == {"text": "train"}
|
||||
|
||||
|
||||
class TestChatMaskBuilder:
|
||||
def test_simple_chat_mask(self, real_tokenizer):
|
||||
config = make_chat_config()
|
||||
builder = ChatMaskBuilder()
|
||||
item = {
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are helpful."},
|
||||
{"role": "user", "content": "Hello."},
|
||||
{"role": "assistant", "content": "Hi there!"},
|
||||
]
|
||||
}
|
||||
result = builder.build(item, config, real_tokenizer)
|
||||
assert result is not None
|
||||
assert "ids" in result
|
||||
assert "loss_mask" in result
|
||||
assert len(result["ids"]) == len(result["loss_mask"])
|
||||
|
||||
ids = real_tokenizer.decode(result["ids"], 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["ids"])
|
||||
trained = sum(result["loss_mask"])
|
||||
assert trained > 0, "At least assistant tokens should be trained"
|
||||
assert trained < total, "System and user tokens should be masked"
|
||||
|
||||
def test_mask_only_assistant_trained(self, real_tokenizer):
|
||||
config = make_chat_config()
|
||||
builder = ChatMaskBuilder()
|
||||
item = {
|
||||
"messages": [
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
{"role": "assistant", "content": "4"},
|
||||
]
|
||||
}
|
||||
result = builder.build(item, config, real_tokenizer)
|
||||
mask = result["loss_mask"]
|
||||
ids = result["ids"]
|
||||
|
||||
assert len(ids) == len(mask)
|
||||
|
||||
trained_positions = [i for i, m in enumerate(mask) if m == 1]
|
||||
assert len(trained_positions) > 0, "At least some tokens should be trained"
|
||||
|
||||
masked_positions = [i for i, m in enumerate(mask) if m == 0]
|
||||
assert len(masked_positions) > 0, "User tokens should be masked"
|
||||
|
||||
def test_chat_all_masked(self, real_tokenizer):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(type="chat", messages_key="messages"),
|
||||
mask={"system": "mask", "user": "mask", "assistant": "mask"},
|
||||
mask_default="mask",
|
||||
preprocessing=ProcessingConfig(max_seq_len=2048),
|
||||
)
|
||||
builder = ChatMaskBuilder()
|
||||
item = {
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are helpful."},
|
||||
{"role": "assistant", "content": "Hi there!"},
|
||||
]
|
||||
}
|
||||
result = builder.build(item, config, real_tokenizer)
|
||||
assert sum(result["loss_mask"]) == 0
|
||||
|
||||
def test_chat_all_trained(self, real_tokenizer):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(type="chat", messages_key="messages"),
|
||||
mask={},
|
||||
mask_default="train",
|
||||
preprocessing=ProcessingConfig(max_seq_len=2048),
|
||||
)
|
||||
builder = ChatMaskBuilder()
|
||||
item = {
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are helpful."},
|
||||
{"role": "assistant", "content": "Hi there!"},
|
||||
]
|
||||
}
|
||||
result = builder.build(item, config, real_tokenizer)
|
||||
assert sum(result["loss_mask"]) == len(result["ids"])
|
||||
|
||||
def test_empty_messages_returns_none(self, real_tokenizer):
|
||||
config = make_chat_config()
|
||||
builder = ChatMaskBuilder()
|
||||
assert builder.build({"messages": []}, config, real_tokenizer) is None
|
||||
assert builder.build({}, config, real_tokenizer) is None
|
||||
|
||||
def test_domain_extraction(self, real_tokenizer):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(type="chat", messages_key="messages"),
|
||||
mask={"assistant": "train"},
|
||||
mask_default="mask",
|
||||
preprocessing=ProcessingConfig(max_seq_len=2048),
|
||||
output=OutputConfig(domain_key="source"),
|
||||
)
|
||||
builder = ChatMaskBuilder()
|
||||
item = {
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hi"},
|
||||
{"role": "assistant", "content": "Hello"},
|
||||
],
|
||||
"source": "wiki",
|
||||
}
|
||||
result = builder.build(item, config, real_tokenizer)
|
||||
assert result["domain"] == "wiki"
|
||||
|
||||
def test_truncation_to_max_len(self, real_tokenizer):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(type="chat", messages_key="messages"),
|
||||
mask={"assistant": "train"},
|
||||
mask_default="mask",
|
||||
preprocessing=ProcessingConfig(max_seq_len=10),
|
||||
)
|
||||
builder = ChatMaskBuilder()
|
||||
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, real_tokenizer)
|
||||
assert len(result["ids"]) <= 10
|
||||
assert len(result["loss_mask"]) == len(result["ids"])
|
||||
|
||||
|
||||
class TestInstructionMaskBuilder:
|
||||
def test_basic_instruction_mask(self, test_tokenizer):
|
||||
config = make_instruction_config()
|
||||
builder = InstructionMaskBuilder()
|
||||
item = {"prompt": "Translate to French: Hello", "response": "Bonjour"}
|
||||
result = builder.build(item, config, test_tokenizer)
|
||||
assert result is not None
|
||||
assert len(result["ids"]) == len(result["loss_mask"])
|
||||
|
||||
def test_prompt_masked_response_trained(self, test_tokenizer):
|
||||
config = make_instruction_config()
|
||||
builder = InstructionMaskBuilder()
|
||||
item = {"prompt": "hello", "response": "world"}
|
||||
result = builder.build(item, config, test_tokenizer)
|
||||
mask = result["loss_mask"]
|
||||
ids = result["ids"]
|
||||
|
||||
prompt_ids = test_tokenizer.encode("hello", add_special_tokens=True)
|
||||
response_ids = test_tokenizer.encode("world", add_special_tokens=False)
|
||||
|
||||
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_train_on_prompt(self, test_tokenizer):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(
|
||||
type="instruction", prompt_key="prompt", response_key="response"
|
||||
),
|
||||
mask={"prompt": "train", "response": "mask"},
|
||||
mask_default="mask",
|
||||
preprocessing=ProcessingConfig(max_seq_len=2048),
|
||||
)
|
||||
builder = InstructionMaskBuilder()
|
||||
item = {"prompt": "hello", "response": "world"}
|
||||
result = builder.build(item, config, test_tokenizer)
|
||||
mask = result["loss_mask"]
|
||||
ids = result["ids"]
|
||||
|
||||
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])
|
||||
|
||||
|
||||
class TestTextMaskBuilder:
|
||||
def test_basic_text(self, test_tokenizer):
|
||||
config = make_text_config()
|
||||
builder = TextMaskBuilder()
|
||||
item = {"text": "Hello world. This is a test document."}
|
||||
result = builder.build(item, config, test_tokenizer)
|
||||
assert result is not None
|
||||
assert "ids" in result
|
||||
assert len(result["ids"]) > 0
|
||||
assert "loss_mask" not in result
|
||||
|
||||
def test_empty_text_returns_none(self, test_tokenizer):
|
||||
config = make_text_config()
|
||||
builder = TextMaskBuilder()
|
||||
assert builder.build({"text": ""}, config, test_tokenizer) is None
|
||||
assert builder.build({"text": " "}, config, test_tokenizer) is None
|
||||
|
||||
def test_too_short_text(self, test_tokenizer):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(type="text", text_key="text"),
|
||||
preprocessing=ProcessingConfig(min_chars=100),
|
||||
)
|
||||
builder = TextMaskBuilder()
|
||||
assert builder.build({"text": "short"}, config, test_tokenizer) is None
|
||||
|
||||
def test_truncation(self, test_tokenizer):
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(type="text", text_key="text"),
|
||||
preprocessing=ProcessingConfig(max_seq_len=3, min_chars=1),
|
||||
)
|
||||
builder = TextMaskBuilder()
|
||||
item = {"text": "This is a very long text that should be truncated"}
|
||||
result = builder.build(item, config, test_tokenizer)
|
||||
assert len(result["ids"]) <= 3
|
||||
|
||||
|
||||
class TestPipeline:
|
||||
def test_full_chat_pipeline(self, temp_dir, real_tokenizer):
|
||||
jsonl_path = os.path.join(temp_dir, "chat.jsonl")
|
||||
with open(jsonl_path, "w", encoding="utf-8") as f:
|
||||
f.write(
|
||||
json.dumps(
|
||||
{
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are helpful."},
|
||||
{"role": "user", "content": "Hi."},
|
||||
{"role": "assistant", "content": "Hello!"},
|
||||
]
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
f.write(
|
||||
json.dumps(
|
||||
{
|
||||
"messages": [
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
{"role": "assistant", "content": "4"},
|
||||
]
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(type="chat", messages_key="messages"),
|
||||
mask={"system": "mask", "user": "mask", "assistant": "train"},
|
||||
mask_default="mask",
|
||||
preprocessing=ProcessingConfig(max_seq_len=2048, deduplicate=True),
|
||||
output=OutputConfig(storage_format="bin", domain_key=None),
|
||||
)
|
||||
|
||||
out_dir = os.path.join(temp_dir, "output")
|
||||
Pipeline(
|
||||
config=config,
|
||||
input_paths=[jsonl_path],
|
||||
output_dir=out_dir,
|
||||
tokenizer_path="params",
|
||||
).run()
|
||||
|
||||
meta_path = os.path.join(out_dir, "__default__", "meta.json")
|
||||
assert os.path.exists(meta_path)
|
||||
with open(meta_path, "r") as f:
|
||||
meta = json.load(f)
|
||||
assert "sequence" in meta
|
||||
assert "loss_mask" in meta
|
||||
|
||||
def test_full_text_pipeline(self, temp_dir, test_tokenizer):
|
||||
import tempfile as tmp
|
||||
|
||||
tokenizer_dir = os.path.join(temp_dir, "tok")
|
||||
os.makedirs(tokenizer_dir, exist_ok=True)
|
||||
|
||||
test_tokenizer._tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json"))
|
||||
with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w") as f:
|
||||
json.dump(
|
||||
{"special_tokens": {"pad_token": "<pad>", "unk_token": "<unk>"}}, f
|
||||
)
|
||||
|
||||
jsonl_path = os.path.join(temp_dir, "text.jsonl")
|
||||
with open(jsonl_path, "w", encoding="utf-8") as f:
|
||||
f.write(
|
||||
json.dumps(
|
||||
{
|
||||
"text": "Hello world this is a test document with enough characters to pass the minimum length filter."
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
f.write(
|
||||
json.dumps(
|
||||
{
|
||||
"text": "Another document for testing purposes with sufficient length to be processed."
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(type="text", text_key="text"),
|
||||
preprocessing=ProcessingConfig(
|
||||
max_seq_len=2048, min_chars=10, deduplicate=True
|
||||
),
|
||||
output=OutputConfig(storage_format="bin"),
|
||||
)
|
||||
|
||||
out_dir = os.path.join(temp_dir, "output")
|
||||
Pipeline(
|
||||
config=config,
|
||||
input_paths=[jsonl_path],
|
||||
output_dir=out_dir,
|
||||
tokenizer_path=tokenizer_dir,
|
||||
).run()
|
||||
|
||||
meta_path = os.path.join(out_dir, "__default__", "meta.json")
|
||||
assert os.path.exists(meta_path)
|
||||
with open(meta_path, "r") as f:
|
||||
meta = json.load(f)
|
||||
assert "sequence" in meta
|
||||
assert "loss_mask" not in meta
|
||||
|
||||
def test_full_instruction_pipeline(self, temp_dir, test_tokenizer):
|
||||
tokenizer_dir = os.path.join(temp_dir, "tok")
|
||||
os.makedirs(tokenizer_dir, exist_ok=True)
|
||||
test_tokenizer._tokenizer.save(os.path.join(tokenizer_dir, "tokenizer.json"))
|
||||
with open(os.path.join(tokenizer_dir, "tokenizer_config.json"), "w") as f:
|
||||
json.dump(
|
||||
{"special_tokens": {"pad_token": "<pad>", "unk_token": "<unk>"}}, f
|
||||
)
|
||||
|
||||
jsonl_path = os.path.join(temp_dir, "instruct.jsonl")
|
||||
with open(jsonl_path, "w", encoding="utf-8") as f:
|
||||
f.write(
|
||||
json.dumps(
|
||||
{
|
||||
"prompt": "Tell me a joke",
|
||||
"response": "Why did the chicken cross the road?",
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
f.write(
|
||||
json.dumps(
|
||||
{
|
||||
"prompt": "What is AI?",
|
||||
"response": "Artificial Intelligence is a field of computer science.",
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
|
||||
config = PipelineConfig(
|
||||
input=InputConfig(
|
||||
type="instruction", prompt_key="prompt", response_key="response"
|
||||
),
|
||||
mask={"prompt": "mask", "response": "train"},
|
||||
mask_default="mask",
|
||||
preprocessing=ProcessingConfig(max_seq_len=2048),
|
||||
output=OutputConfig(storage_format="bin"),
|
||||
)
|
||||
|
||||
out_dir = os.path.join(temp_dir, "output")
|
||||
Pipeline(
|
||||
config=config,
|
||||
input_paths=[jsonl_path],
|
||||
output_dir=out_dir,
|
||||
tokenizer_path=tokenizer_dir,
|
||||
).run()
|
||||
|
||||
meta_path = os.path.join(out_dir, "__default__", "meta.json")
|
||||
assert os.path.exists(meta_path)
|
||||
with open(meta_path, "r") as f:
|
||||
meta = json.load(f)
|
||||
assert "sequence" in meta
|
||||
assert "loss_mask" in meta
|
||||
|
||||
|
||||
class TestUtility:
|
||||
def test_filter_by_length(self):
|
||||
assert filter_by_length("hello world", min_len=5)
|
||||
assert not filter_by_length("hi", min_len=5)
|
||||
assert not filter_by_length("x" * 100, max_len=50)
|
||||
assert filter_by_length("just right", min_len=5, max_len=20)
|
||||
|
||||
def test_dedup_signature(self):
|
||||
a = {"key": "value", "number": 1}
|
||||
b = {"number": 1, "key": "value"}
|
||||
assert dedup_signature(a) == dedup_signature(b)
|
||||
c = {"key": "different"}
|
||||
assert dedup_signature(a) != dedup_signature(c)
|
||||
|
||||
|
||||
class TestFactoryRegistration:
|
||||
def test_registered_builders(self):
|
||||
names = MaskBuilderFactory._registry.list_names()
|
||||
assert "chat" in names
|
||||
assert "instruction" in names
|
||||
assert "text" in names
|
||||
|
||||
def test_create_chat_builder(self):
|
||||
builder = MaskBuilderFactory.create("chat")
|
||||
assert isinstance(builder, ChatMaskBuilder)
|
||||
|
||||
def test_create_instruction_builder(self):
|
||||
builder = MaskBuilderFactory.create("instruction")
|
||||
assert isinstance(builder, InstructionMaskBuilder)
|
||||
|
||||
def test_create_text_builder(self):
|
||||
builder = MaskBuilderFactory.create("text")
|
||||
assert isinstance(builder, TextMaskBuilder)
|
||||
Reference in New Issue
Block a user