ViperEkura
598e1ce4ae
- 将 pipeline/packing.py 拆分为 packing/ 子包 (base/stream/binpack) - 新增 BfdPacker(默认)/FfDPacker/GreedyPacker，移除 StreamingPacker - 超长序列直接截断至 pack_size - group_size 语义改为"每 N 个 chunk 合并为一块"，默认 1000 - 新增 AutoTokenizer.token_to_id()，修复 ChatML 中 hacky 的 nl_id 获取 - pad_value 默认改为 2(pad_token_id)，position_ids pad=0, loss_mask pad=False - 新增 position_ids 打包后归零一致性测试 - scripts/cache_h5.py 新增 --pack-algo 参数
124 lines
4.1 KiB
Python
124 lines
4.1 KiB
Python
"""Tests for pipeline.cache module."""
|
|
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import torch
|
|
from pathlib import Path
|
|
|
|
from pipeline.io import cache_jsonl
|
|
from pipeline.processors import BaseProcessor
|
|
|
|
|
|
class DummyProcessor(BaseProcessor):
|
|
"""Dummy processor for testing."""
|
|
|
|
def __init__(self):
|
|
self._output_keys = ["sequence", "loss_mask"]
|
|
|
|
@property
|
|
def output_keys(self):
|
|
return self._output_keys
|
|
|
|
def process(self, item):
|
|
text = item.get("text", "")
|
|
tokens = [ord(c) for c in text[:10]]
|
|
return {
|
|
"sequence": torch.tensor(tokens, dtype=torch.int32),
|
|
"loss_mask": torch.ones(len(tokens), dtype=torch.int32),
|
|
}
|
|
|
|
|
|
class TestCacheJsonl:
|
|
def test_basic_cache_functionality(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
jsonl_path = os.path.join(tmpdir, "test.jsonl")
|
|
test_data = [{"text": "hello"}, {"text": "world"}, {"text": "test"}]
|
|
with open(jsonl_path, "w", encoding="utf-8") as f:
|
|
for item in test_data:
|
|
f.write(json.dumps(item) + "\n")
|
|
|
|
processor = DummyProcessor()
|
|
output_files = cache_jsonl(
|
|
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])
|
|
|
|
def test_packer_state_independence(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
jsonl_path = os.path.join(tmpdir, "test.jsonl")
|
|
test_data = [{"text": "ab"}, {"text": "abcde"}, {"text": "abc"}]
|
|
with open(jsonl_path, "w", encoding="utf-8") as f:
|
|
for item in test_data:
|
|
f.write(json.dumps(item) + "\n")
|
|
|
|
processor = DummyProcessor()
|
|
output_files = cache_jsonl(
|
|
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])
|
|
|
|
def test_no_packing_mode(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
jsonl_path = os.path.join(tmpdir, "test.jsonl")
|
|
test_data = [{"text": "hello"}, {"text": "world"}]
|
|
with open(jsonl_path, "w", encoding="utf-8") as f:
|
|
for item in test_data:
|
|
f.write(json.dumps(item) + "\n")
|
|
|
|
processor = DummyProcessor()
|
|
output_files = cache_jsonl(
|
|
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])
|
|
|
|
def test_multiple_files(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
files = []
|
|
for i in range(2):
|
|
jsonl_path = os.path.join(tmpdir, f"test{i}.jsonl")
|
|
test_data = [{"text": f"data{i}"}]
|
|
with open(jsonl_path, "w", encoding="utf-8") as f:
|
|
f.write(json.dumps(test_data[0]) + "\n")
|
|
files.append(jsonl_path)
|
|
|
|
processor = DummyProcessor()
|
|
output_files = cache_jsonl(
|
|
files=files,
|
|
output_dir=tmpdir,
|
|
processor=processor,
|
|
pack_size=-1,
|
|
pad_value=0,
|
|
)
|
|
assert len(output_files) == 2
|
|
|
|
def test_empty_file_handling(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
jsonl_path = os.path.join(tmpdir, "empty.jsonl")
|
|
Path(jsonl_path).touch()
|
|
|
|
processor = DummyProcessor()
|
|
output_files = cache_jsonl(
|
|
files=[jsonl_path],
|
|
output_dir=tmpdir,
|
|
processor=processor,
|
|
pack_size=-1,
|
|
pad_value=0,
|
|
)
|
|
assert len(output_files) == 0
|