ViperEkura/DataPipeline
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: 修复特殊token 问题

Browse Source
This commit is contained in:
ViperEkura
2026-04-02 16:16:02 +08:00
parent e01ec081b3
commit f44ad6912e
28 changed files with 334 additions and 163 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/__init__.py
+1 -1
View File
@@ -1 +1 @@
# Test suite for DataPipeline
# Test suite for DataPipeline
tests/test_cache.py
+25 -11
View File
@@ -30,7 +30,6 @@ class DummyProcessor(BaseProcessor):
class TestCacheJsonl:
def test_basic_cache_functionality(self):
with tempfile.TemporaryDirectory() as tmpdir:
jsonl_path = os.path.join(tmpdir, "test.jsonl")
@@ -41,8 +40,11 @@ class TestCacheJsonl:
processor = DummyProcessor()
output_files = cache_jsonl(
files=[jsonl_path], output_dir=tmpdir,
processor=processor, pack_size=-1, pad_value=0,
files=[jsonl_path],
output_dir=tmpdir,
processor=processor,
pack_size=-1,
pad_value=0,
)
assert len(output_files) == 1
assert os.path.exists(output_files[0])
@@ -57,8 +59,11 @@ class TestCacheJsonl:
processor = DummyProcessor()
output_files = cache_jsonl(
files=[jsonl_path], output_dir=tmpdir,
processor=processor, pack_size=10, pad_value=0,
files=[jsonl_path],
output_dir=tmpdir,
processor=processor,
pack_size=10,
pad_value=0,
)
assert len(output_files) == 1
assert os.path.exists(output_files[0])
@@ -73,8 +78,11 @@ class TestCacheJsonl:
processor = DummyProcessor()
output_files = cache_jsonl(
files=[jsonl_path], output_dir=tmpdir,
processor=processor, pack_size=0, pad_value=-1,
files=[jsonl_path],
output_dir=tmpdir,
processor=processor,
pack_size=0,
pad_value=-1,
)
assert len(output_files) == 1
assert os.path.exists(output_files[0])
@@ -91,8 +99,11 @@ class TestCacheJsonl:
processor = DummyProcessor()
output_files = cache_jsonl(
files=files, output_dir=tmpdir,
processor=processor, pack_size=-1, pad_value=0,
files=files,
output_dir=tmpdir,
processor=processor,
pack_size=-1,
pad_value=0,
)
assert len(output_files) == 2
@@ -103,7 +114,10 @@ class TestCacheJsonl:
processor = DummyProcessor()
output_files = cache_jsonl(
files=[jsonl_path], output_dir=tmpdir,
processor=processor, pack_size=-1, pad_value=0,
files=[jsonl_path],
output_dir=tmpdir,
processor=processor,
pack_size=-1,
pad_value=0,
)
assert len(output_files) == 1
tests/test_io.py
+18 -9
View File
@@ -11,7 +11,6 @@ from pipeline.io import IOHandler
class TestIOHandler:
def test_fetch_files_in_directory(self):
with tempfile.TemporaryDirectory() as tmpdir:
Path(tmpdir, "file1.txt").touch()
@@ -41,7 +40,9 @@ class TestIOHandler:
os.makedirs(os.path.join(tmpdir, "folder1"))
os.makedirs(os.path.join(tmpdir, "folder2"))
folders = IOHandler.fetch_folders(tmpdir, filter_func=lambda x: "folder1" in x)
folders = IOHandler.fetch_folders(
tmpdir, filter_func=lambda x: "folder1" in x
)
assert len(folders) == 1
def test_save_and_load_h5(self):
@@ -56,8 +57,12 @@ class TestIOHandler:
loaded = IOHandler.load_h5(tmpdir, share_memory=False)
assert "sequence" in loaded
assert "labels" in loaded
assert torch.equal(loaded["sequence"][0], torch.tensor([1, 2, 3], dtype=torch.int32))
assert torch.equal(loaded["labels"][0], torch.tensor([4, 5], dtype=torch.int32))
assert torch.equal(
loaded["sequence"][0], torch.tensor([1, 2, 3], dtype=torch.int32)
)
assert torch.equal(
loaded["labels"][0], torch.tensor([4, 5], dtype=torch.int32)
)
def test_save_h5_creates_directory(self):
with tempfile.TemporaryDirectory() as tmpdir:
@@ -70,9 +75,9 @@ class TestIOHandler:
with tempfile.TemporaryDirectory() as tmpdir:
for i, data in enumerate([[1, 2, 3], [4, 5, 6]]):
h5_path = os.path.join(tmpdir, f"file{i}.h5")
with h5py.File(h5_path, 'w') as f:
with h5py.File(h5_path, "w") as f:
grp = f.create_group("data")
grp.create_dataset('data_0', data=data)
grp.create_dataset("data_0", data=data)
loaded = IOHandler.load_h5(tmpdir, share_memory=False)
assert len(loaded["data"]) == 2
@@ -82,9 +87,9 @@ class TestIOHandler:
subdir = os.path.join(tmpdir, "subdir")
os.makedirs(subdir)
h5_path = os.path.join(subdir, "nested.h5")
with h5py.File(h5_path, 'w') as f:
with h5py.File(h5_path, "w") as f:
grp = f.create_group("test")
grp.create_dataset('data_0', data=[1, 2])
grp.create_dataset("data_0", data=[1, 2])
loaded = IOHandler.load_h5(tmpdir, share_memory=False)
assert "test" in loaded
@@ -93,7 +98,11 @@ class TestIOHandler:
def test_save_h5_multiple_tensors_per_key(self):
with tempfile.TemporaryDirectory() as tmpdir:
tensor_group = {
"batch": [torch.tensor([1, 2]), torch.tensor([3, 4, 5]), torch.tensor([6])],
"batch": [
torch.tensor([1, 2]),
torch.tensor([3, 4, 5]),
torch.tensor([6]),
],
}
IOHandler.save_h5(tmpdir, "multi", tensor_group)
loaded = IOHandler.load_h5(tmpdir, share_memory=False)
tests/test_packing.py
+27 -18
View File
@@ -6,7 +6,6 @@ from pipeline.packing import SequencePacker
class TestSequencePacker:
def test_normal_packing(self):
packer = SequencePacker(pack_size=10, pad_value=0)
sequences = [
@@ -37,17 +36,21 @@ class TestSequencePacker:
def test_long_sequence_split_across_chunks(self):
"""Sequences longer than pack_size are split across multiple chunks."""
packer = SequencePacker(pack_size=5, pad_value=0)
packages = packer.pack([torch.tensor([1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.int32)])
packages = packer.pack(
[torch.tensor([1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.int32)]
)
assert len(packages) == 2
assert packages[0].tolist() == [1, 2, 3, 4, 5]
assert packages[1].tolist() == [6, 7, 8, 0, 0]
def test_padding_value(self):
packer = SequencePacker(pack_size=8, pad_value=99)
packages = packer.pack([
torch.tensor([1, 2], dtype=torch.int32),
torch.tensor([3], dtype=torch.int32),
])
packages = packer.pack(
[
torch.tensor([1, 2], dtype=torch.int32),
torch.tensor([3], dtype=torch.int32),
]
)
assert packages[0][:3].tolist() == [1, 2, 3]
assert packages[0][3:].tolist() == [99] * 5
@@ -85,10 +88,12 @@ class TestSequencePacker:
def test_exact_pack_size_fit(self):
packer = SequencePacker(pack_size=5, pad_value=0)
packages = packer.pack([
torch.tensor([1, 2, 3, 4, 5], dtype=torch.int32),
torch.tensor([6, 7, 8, 9, 10], dtype=torch.int32),
])
packages = packer.pack(
[
torch.tensor([1, 2, 3, 4, 5], dtype=torch.int32),
torch.tensor([6, 7, 8, 9, 10], dtype=torch.int32),
]
)
assert len(packages) == 2
assert packages[0].tolist() == [1, 2, 3, 4, 5]
assert packages[1].tolist() == [6, 7, 8, 9, 10]
@@ -127,10 +132,12 @@ class TestSequencePacker:
def test_stream_split_across_chunks(self):
"""Sequences are split across chunks in streaming mode."""
packer = SequencePacker(pack_size=5, pad_value=0)
packages = packer.pack([
torch.tensor([1, 2, 3], dtype=torch.int32),
torch.tensor([4, 5, 6, 7, 8], dtype=torch.int32),
])
packages = packer.pack(
[
torch.tensor([1, 2, 3], dtype=torch.int32),
torch.tensor([4, 5, 6, 7, 8], dtype=torch.int32),
]
)
assert len(packages) == 2
# First chunk: [1, 2, 3, 4, 5] — first seq + part of second
assert packages[0].tolist() == [1, 2, 3, 4, 5]
@@ -151,9 +158,11 @@ class TestSequencePacker:
"""Streaming concat preserves input order, no sorting."""
packer = SequencePacker(pack_size=4, pad_value=-1)
# short then long (fits in 2 chunks)
packages = packer.pack([
torch.tensor([1], dtype=torch.int32),
torch.tensor([2, 3, 4, 5, 6, 7], dtype=torch.int32),
])
packages = packer.pack(
[
torch.tensor([1], dtype=torch.int32),
torch.tensor([2, 3, 4, 5, 6, 7], dtype=torch.int32),
]
)
assert packages[0].tolist() == [1, 2, 3, 4]
assert packages[1].tolist() == [5, 6, 7, -1]
tests/test_processors.py
+21 -7
View File
@@ -44,18 +44,24 @@ class TestSFTProcessor:
assert SFTProcessor(DummyTokenizer()).output_keys == ["sequence", "loss_mask"]
def test_process_returns_both_keys(self):
result = SFTProcessor(DummyTokenizer()).process({"query": "hello", "response": "world"})
result = SFTProcessor(DummyTokenizer()).process(
{"query": "hello", "response": "world"}
)
assert "sequence" in result
assert "loss_mask" in result
assert isinstance(result["sequence"], torch.Tensor)
assert isinstance(result["loss_mask"], torch.Tensor)
def test_loss_mask_correct_length(self):
result = SFTProcessor(DummyTokenizer()).process({"query": "hi", "response": "bye"})
result = SFTProcessor(DummyTokenizer()).process(
{"query": "hi", "response": "bye"}
)
assert len(result["sequence"]) == len(result["loss_mask"])
def test_loss_mask_is_bool(self):
result = SFTProcessor(DummyTokenizer()).process({"query": "ab", "response": "cd"})
result = SFTProcessor(DummyTokenizer()).process(
{"query": "ab", "response": "cd"}
)
assert result["loss_mask"].dtype == torch.bool
@@ -89,13 +95,19 @@ class TestDPOProcessor:
class TestProcessorFactory:
def test_create_pre_train_processor(self):
assert isinstance(ProcessorFactory.create("pt", DummyTokenizer()), PreTrainProcessor)
assert isinstance(
ProcessorFactory.create("pt", DummyTokenizer()), PreTrainProcessor
)
def test_create_sft_processor(self):
assert isinstance(ProcessorFactory.create("sft", DummyTokenizer()), SFTProcessor)
assert isinstance(
ProcessorFactory.create("sft", DummyTokenizer()), SFTProcessor
)
def test_create_dpo_processor(self):
assert isinstance(ProcessorFactory.create("dpo", DummyTokenizer()), DPOProcessor)
assert isinstance(
ProcessorFactory.create("dpo", DummyTokenizer()), DPOProcessor
)
def test_create_invalid_processor_raises_error(self):
with pytest.raises(ValueError, match="Unknown processor type"):
@@ -114,4 +126,6 @@ class TestProcessorFactory:
return {"custom": torch.tensor([1, 2, 3])}
ProcessorFactory.register("custom")(CustomProcessor)
assert isinstance(ProcessorFactory.create("custom", DummyTokenizer()), CustomProcessor)
assert isinstance(
ProcessorFactory.create("custom", DummyTokenizer()), CustomProcessor
)
tests/test_strategies.py
+10 -7
View File
@@ -28,7 +28,7 @@ class DummyStrategy(PromptStrategy):
return prefix + query_tokens
def assemble_response(self, response_tokens):
suffix = self._encode_format("<eos>")
suffix = self._encode_format("<｜end▁of▁sentence｜>")
return response_tokens + suffix
@@ -46,9 +46,9 @@ class TestChatMLStrategy:
query_tokens = tk.encode("hello")
prompt = strategy.assemble_prompt(query_tokens)
text = _decode(prompt)
assert "<|im_start|>user" in text
assert "<｜im▁start｜>user" in text
assert "hello" in text
assert "<|im_start|>assistant" in text
assert "<｜im▁start｜>assistant" in text
def test_assemble_response(self):
tk = DummyTokenizer()
@@ -57,15 +57,18 @@ class TestChatMLStrategy:
response = strategy.assemble_response(response_tokens)
text = _decode(response)
assert "world" in text
assert "<|im_end|>" in text
assert "<eos>" in text
assert "<｜im▁end｜>" in text
assert "<｜end▁of▁sentence｜>" in text
def test_prompt_ends_with_assistant_start(self):
tk = DummyTokenizer()
strategy = ChatMLStrategy(tk)
prompt = strategy.assemble_prompt(tk.encode("hi"))
# prompt 末尾应该是 assistant_start 的 token ids
assert prompt[-len(strategy._assistant_start_ids):] == strategy._assistant_start_ids
assert (
prompt[-len(strategy._assistant_start_ids) :]
== strategy._assistant_start_ids
)
class TestAlpacaStrategy:
@@ -89,7 +92,7 @@ class TestAlpacaStrategy:
response = strategy.assemble_response(response_tokens)
text = _decode(response)
assert "world" in text
assert "<eos>" in text
assert "<｜end▁of▁sentence｜>" in text
class TestStrategyFactory: