reafactor: 重构项目
This commit is contained in:
+31
-87
@@ -1,4 +1,4 @@
|
||||
"""单元测试：pipeline.cache 模块中的 cache_jsonl 函数"""
|
||||
"""Tests for pipeline.cache module."""
|
||||
|
||||
import json
|
||||
import os
|
||||
@@ -6,24 +6,23 @@ import tempfile
|
||||
import torch
|
||||
from pathlib import Path
|
||||
|
||||
from pipeline.cache import cache_jsonl
|
||||
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]] # 简单模拟tokenize
|
||||
|
||||
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),
|
||||
@@ -31,99 +30,57 @@ class DummyProcessor(BaseProcessor):
|
||||
|
||||
|
||||
class TestCacheJsonl:
|
||||
"""cache_jsonl 函数的测试套件"""
|
||||
|
||||
|
||||
def test_basic_cache_functionality(self):
|
||||
"""测试基本缓存功能：处理简单JSONL文件并生成HDF5"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 创建测试JSONL文件
|
||||
jsonl_path = os.path.join(tmpdir, "test.jsonl")
|
||||
test_data = [
|
||||
{"text": "hello"},
|
||||
{"text": "world"},
|
||||
{"text": "test"},
|
||||
]
|
||||
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()
|
||||
|
||||
# 调用 cache_jsonl
|
||||
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])
|
||||
|
||||
|
||||
def test_packer_state_independence(self):
|
||||
"""测试打包器状态独立性：验证不同 output_key 的打包结果是否独立"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 创建测试JSONL文件，包含不同长度的文本
|
||||
jsonl_path = os.path.join(tmpdir, "test.jsonl")
|
||||
test_data = [
|
||||
{"text": "ab"}, # 2 chars
|
||||
{"text": "abcde"}, # 5 chars
|
||||
{"text": "abc"}, # 3 chars
|
||||
]
|
||||
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()
|
||||
|
||||
# 调用 cache_jsonl，使用打包模式
|
||||
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])
|
||||
|
||||
|
||||
def test_no_packing_mode(self):
|
||||
"""测试无打包模式（pack_size <= 0）"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 创建测试JSONL文件
|
||||
jsonl_path = os.path.join(tmpdir, "test.jsonl")
|
||||
test_data = [
|
||||
{"text": "hello"},
|
||||
{"text": "world"},
|
||||
]
|
||||
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()
|
||||
|
||||
# 打包大小设为0表示不打包
|
||||
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])
|
||||
|
||||
|
||||
def test_multiple_files(self):
|
||||
"""测试处理多个文件"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 创建两个测试JSONL文件
|
||||
files = []
|
||||
for i in range(2):
|
||||
jsonl_path = os.path.join(tmpdir, f"test{i}.jsonl")
|
||||
@@ -131,35 +88,22 @@ class TestCacheJsonl:
|
||||
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,
|
||||
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文件
|
||||
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,
|
||||
files=[jsonl_path], output_dir=tmpdir,
|
||||
processor=processor, pack_size=-1, pad_value=0,
|
||||
)
|
||||
|
||||
assert len(output_files) == 1
|
||||
assert len(output_files) == 1
|
||||
|
||||
+35
-122
@@ -1,4 +1,4 @@
|
||||
"""单元测试：pipeline.io 模块中的 IOHandler 类"""
|
||||
"""Tests for pipeline.io module."""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
@@ -11,177 +11,90 @@ from pipeline.io import IOHandler
|
||||
|
||||
|
||||
class TestIOHandler:
|
||||
"""IOHandler 类的测试套件"""
|
||||
|
||||
|
||||
def test_fetch_files_in_directory(self):
|
||||
"""测试 fetch_files 方法能正确获取目录中的文件"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 创建测试文件
|
||||
test_file1 = os.path.join(tmpdir, "file1.txt")
|
||||
test_file2 = os.path.join(tmpdir, "file2.txt")
|
||||
Path(test_file1).touch()
|
||||
Path(test_file2).touch()
|
||||
|
||||
# 创建子目录和文件
|
||||
Path(tmpdir, "file1.txt").touch()
|
||||
Path(tmpdir, "file2.txt").touch()
|
||||
subdir = os.path.join(tmpdir, "subdir")
|
||||
os.makedirs(subdir)
|
||||
test_file3 = os.path.join(subdir, "file3.txt")
|
||||
Path(test_file3).touch()
|
||||
|
||||
# 获取文件列表
|
||||
Path(subdir, "file3.txt").touch()
|
||||
|
||||
files = IOHandler.fetch_files(tmpdir)
|
||||
|
||||
# 验证
|
||||
assert len(files) == 3
|
||||
assert any("file1.txt" in f for f in files)
|
||||
assert any("file2.txt" in f for f in files)
|
||||
assert any("file3.txt" in f for f in files)
|
||||
|
||||
|
||||
def test_fetch_files_empty_directory(self):
|
||||
"""测试空目录返回空列表"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
files = IOHandler.fetch_files(tmpdir)
|
||||
assert files == []
|
||||
|
||||
assert IOHandler.fetch_files(tmpdir) == []
|
||||
|
||||
def test_fetch_folders_in_directory(self):
|
||||
"""测试 fetch_folders 方法能正确获取子目录"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 创建子目录
|
||||
subdir1 = os.path.join(tmpdir, "folder1")
|
||||
subdir2 = os.path.join(tmpdir, "folder2")
|
||||
os.makedirs(subdir1)
|
||||
os.makedirs(subdir2)
|
||||
|
||||
# 创建嵌套子目录
|
||||
nested = os.path.join(subdir1, "nested")
|
||||
os.makedirs(nested)
|
||||
|
||||
# 获取文件夹列表
|
||||
os.makedirs(os.path.join(tmpdir, "folder1"))
|
||||
os.makedirs(os.path.join(tmpdir, "folder2"))
|
||||
os.makedirs(os.path.join(tmpdir, "folder1", "nested"))
|
||||
|
||||
folders = IOHandler.fetch_folders(tmpdir)
|
||||
|
||||
# 验证
|
||||
assert len(folders) == 3
|
||||
assert any("folder1" in f for f in folders)
|
||||
assert any("folder2" in f for f in folders)
|
||||
assert any("nested" in f for f in folders)
|
||||
|
||||
|
||||
def test_fetch_folders_with_filter(self):
|
||||
"""测试 fetch_folders 方法的过滤功能"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 创建子目录
|
||||
subdir1 = os.path.join(tmpdir, "folder1")
|
||||
subdir2 = os.path.join(tmpdir, "folder2")
|
||||
os.makedirs(subdir1)
|
||||
os.makedirs(subdir2)
|
||||
|
||||
# 使用过滤器只获取 folder1
|
||||
folders = IOHandler.fetch_folders(
|
||||
tmpdir,
|
||||
filter_func=lambda x: "folder1" in x
|
||||
)
|
||||
|
||||
# 验证
|
||||
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)
|
||||
assert len(folders) == 1
|
||||
assert "folder1" in folders[0]
|
||||
|
||||
|
||||
def test_save_and_load_h5(self):
|
||||
"""测试 save_h5 和 load_h5 方法的读写功能"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 创建测试数据
|
||||
tensor_group = {
|
||||
"sequence": [torch.tensor([1, 2, 3], dtype=torch.int32)],
|
||||
"labels": [torch.tensor([4, 5], dtype=torch.int32)],
|
||||
}
|
||||
|
||||
# 保存
|
||||
IOHandler.save_h5(tmpdir, "test", tensor_group)
|
||||
|
||||
# 验证文件已创建
|
||||
h5_path = os.path.join(tmpdir, "test.h5")
|
||||
assert os.path.exists(h5_path)
|
||||
|
||||
# 加载 - 传入目录而不是单个文件
|
||||
|
||||
assert os.path.exists(os.path.join(tmpdir, "test.h5"))
|
||||
loaded = IOHandler.load_h5(tmpdir, share_memory=False)
|
||||
|
||||
# 验证数据
|
||||
assert "sequence" in loaded
|
||||
assert "labels" in loaded
|
||||
assert len(loaded["sequence"]) == 1
|
||||
assert len(loaded["labels"]) == 1
|
||||
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):
|
||||
"""测试 save_h5 自动创建输出目录"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
output_dir = os.path.join(tmpdir, "nested", "output")
|
||||
|
||||
tensor_group = {
|
||||
"data": [torch.tensor([1, 2, 3])],
|
||||
}
|
||||
|
||||
# 保存到不存在的目录
|
||||
IOHandler.save_h5(output_dir, "test", tensor_group)
|
||||
|
||||
# 验证目录已创建
|
||||
IOHandler.save_h5(output_dir, "test", {"data": [torch.tensor([1, 2, 3])]})
|
||||
assert os.path.exists(output_dir)
|
||||
assert os.path.exists(os.path.join(output_dir, "test.h5"))
|
||||
|
||||
|
||||
def test_load_h5_multiple_files(self):
|
||||
"""测试 load_h5 方法能处理多个 H5 文件"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 创建第一个 H5 文件
|
||||
h5_path1 = os.path.join(tmpdir, "file1.h5")
|
||||
with h5py.File(h5_path1, 'w') as f:
|
||||
grp = f.create_group("data")
|
||||
grp.create_dataset('data_0', data=[1, 2, 3])
|
||||
|
||||
# 创建第二个 H5 文件
|
||||
h5_path2 = os.path.join(tmpdir, "file2.h5")
|
||||
with h5py.File(h5_path2, 'w') as f:
|
||||
grp = f.create_group("data")
|
||||
grp.create_dataset('data_0', data=[4, 5, 6])
|
||||
|
||||
# 加载目录
|
||||
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:
|
||||
grp = f.create_group("data")
|
||||
grp.create_dataset('data_0', data=data)
|
||||
|
||||
loaded = IOHandler.load_h5(tmpdir, share_memory=False)
|
||||
|
||||
# 验证
|
||||
assert "data" in loaded
|
||||
assert len(loaded["data"]) == 2
|
||||
|
||||
|
||||
def test_load_h5_with_rglob(self):
|
||||
"""测试 load_h5 能递归查找 H5 文件"""
|
||||
with tempfile.TemporaryDirectory() as tmpdir:
|
||||
# 在子目录中创建 H5 文件
|
||||
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:
|
||||
grp = f.create_group("test")
|
||||
grp.create_dataset('data_0', data=[1, 2])
|
||||
|
||||
# 加载根目录
|
||||
|
||||
loaded = IOHandler.load_h5(tmpdir, share_memory=False)
|
||||
|
||||
# 验证能找到子目录中的文件
|
||||
assert "test" in loaded
|
||||
assert len(loaded["test"]) == 1
|
||||
|
||||
|
||||
def test_save_h5_multiple_tensors_per_key(self):
|
||||
"""测试 save_h5 能保存多个张量到同一键"""
|
||||
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)
|
||||
|
||||
assert len(loaded["batch"]) == 3
|
||||
assert len(loaded["batch"]) == 3
|
||||
|
||||
+61
-191
@@ -1,4 +1,4 @@
|
||||
"""单元测试：pipeline.packing 模块中的 SequencePacker 类"""
|
||||
"""Tests for pipeline.packing module."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -6,252 +6,122 @@ from pipeline.packing import SequencePacker
|
||||
|
||||
|
||||
class TestSequencePacker:
|
||||
"""SequencePacker 类的测试套件"""
|
||||
|
||||
def test_normal_packing(self):
|
||||
"""测试正常打包场景：多个序列正确打包成固定长度的包"""
|
||||
packer = SequencePacker(pack_size=10, pad_value=0)
|
||||
|
||||
sequences = [
|
||||
torch.tensor([1, 2, 3], dtype=torch.int32),
|
||||
torch.tensor([4, 5], dtype=torch.int32),
|
||||
torch.tensor([6, 7, 8, 9], dtype=torch.int32),
|
||||
]
|
||||
|
||||
packages = packer.pack(sequences)
|
||||
|
||||
# 验证至少有包输出
|
||||
assert len(packages) >= 1
|
||||
|
||||
# 验证每个包的长度是正确的
|
||||
for pkg in packages:
|
||||
assert pkg.shape == (10,)
|
||||
# 验证填充值
|
||||
# 检查所有非零元素都在前几个位置，或者包是满的
|
||||
non_zero_count = (pkg != 0).sum().item()
|
||||
# 非零元素的数量应该等于原始序列元素的总和
|
||||
total_elements = sum(s.numel() for s in sequences)
|
||||
# 由于打包，第一个包包含3+2=5个元素，第二个包包含4个元素
|
||||
# 第一个包应该包含前两个序列
|
||||
pkg1 = packages[0]
|
||||
# 序列[1,2,3]和[4,5]按长度降序排序后是[1,2,3]在前，然后[4,5]
|
||||
# 但排序是原地修改...等等，我们已经修复了使用sorted()
|
||||
# 所以排序后的顺序是[6,7,8,9], [1,2,3], [4,5]
|
||||
# 第一个包包含[6,7,8,9]和部分[1,2,3] = 4+3=7，剩余3个位置放[4,5]
|
||||
# 所以第一个包应该是[6,7,8,9,1,2,3,4,5,0]
|
||||
|
||||
# 简化测试：验证打包后的张量包含所有原始数据
|
||||
|
||||
# Verify all original values are present
|
||||
all_values = []
|
||||
for pkg in packages:
|
||||
non_zero = pkg[pkg != 0].tolist()
|
||||
all_values.extend(non_zero)
|
||||
|
||||
# 检查所有原始数据是否都被包含
|
||||
original_values = [1, 2, 3, 4, 5, 6, 7, 8, 9]
|
||||
for val in original_values:
|
||||
assert val in all_values, f"Value {val} not found in packages"
|
||||
all_values.extend(pkg[pkg != 0].tolist())
|
||||
for val in [1, 2, 3, 4, 5, 6, 7, 8, 9]:
|
||||
assert val in all_values
|
||||
|
||||
def test_empty_list_input(self):
|
||||
"""测试空列表输入"""
|
||||
packer = SequencePacker(pack_size=10)
|
||||
|
||||
packages = packer.pack([])
|
||||
|
||||
assert packages == []
|
||||
|
||||
# 验证内部状态已正确初始化
|
||||
assert packer._current_pack is not None
|
||||
assert packer._current_pos == 0
|
||||
assert packer.pack([]) == []
|
||||
|
||||
def test_single_sequence_input(self):
|
||||
"""测试单个序列输入"""
|
||||
packer = SequencePacker(pack_size=10, pad_value=-1)
|
||||
|
||||
sequences = [torch.tensor([1, 2, 3], dtype=torch.int32)]
|
||||
|
||||
packages = packer.pack(sequences)
|
||||
|
||||
packages = packer.pack([torch.tensor([1, 2, 3], dtype=torch.int32)])
|
||||
assert len(packages) == 1
|
||||
pkg = packages[0]
|
||||
assert pkg.shape == (10,)
|
||||
assert pkg[:3].tolist() == [1, 2, 3]
|
||||
assert pkg[3:].tolist() == [-1] * 7
|
||||
assert packages[0][:3].tolist() == [1, 2, 3]
|
||||
assert packages[0][3:].tolist() == [-1] * 7
|
||||
|
||||
def test_truncate_long_sequence(self, caplog):
|
||||
"""测试超长序列截断，验证警告日志是否触发"""
|
||||
packer = SequencePacker(pack_size=5, pad_value=0)
|
||||
|
||||
sequences = [
|
||||
torch.tensor([1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.int32), # 长度8，超过pack_size=5
|
||||
]
|
||||
|
||||
packages = packer.pack(sequences)
|
||||
|
||||
packages = packer.pack([torch.tensor([1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.int32)])
|
||||
assert len(packages) == 1
|
||||
pkg = packages[0]
|
||||
assert pkg.shape == (5,)
|
||||
assert pkg.tolist() == [1, 2, 3, 4, 5] # 只保留前5个元素
|
||||
|
||||
# 验证警告日志已触发
|
||||
assert packages[0].tolist() == [1, 2, 3, 4, 5]
|
||||
assert "truncating" in caplog.text.lower() or "exceeds" in caplog.text.lower()
|
||||
|
||||
def test_padding_value(self):
|
||||
"""测试填充值正确应用"""
|
||||
packer = SequencePacker(pack_size=8, pad_value=99)
|
||||
|
||||
sequences = [
|
||||
packages = packer.pack([
|
||||
torch.tensor([1, 2], dtype=torch.int32),
|
||||
torch.tensor([3], dtype=torch.int32),
|
||||
]
|
||||
|
||||
packages = packer.pack(sequences)
|
||||
|
||||
assert len(packages) == 1
|
||||
pkg = packages[0]
|
||||
|
||||
# 前3个元素是数据
|
||||
assert pkg[:3].tolist() == [1, 2, 3]
|
||||
# 后5个元素是填充值
|
||||
assert pkg[3:].tolist() == [99] * 5
|
||||
])
|
||||
assert packages[0][:3].tolist() == [1, 2, 3]
|
||||
assert packages[0][3:].tolist() == [99] * 5
|
||||
|
||||
def test_different_dtypes(self):
|
||||
"""测试支持不同 dtype (int32, int64, float32)"""
|
||||
# int32
|
||||
packer_int32 = SequencePacker(pack_size=10, pad_value=0, dtype=torch.int32)
|
||||
sequences_int32 = [torch.tensor([1, 2, 3], dtype=torch.int32)]
|
||||
packages_int32 = packer_int32.pack(sequences_int32)
|
||||
assert packages_int32[0].dtype == torch.int32
|
||||
for dtype in [torch.int32, torch.int64, torch.float32]:
|
||||
packer = SequencePacker(pack_size=10, dtype=dtype)
|
||||
val = 1.0 if dtype == torch.float32 else 1
|
||||
packages = packer.pack([torch.tensor([val, 2, 3], dtype=dtype)])
|
||||
assert packages[0].dtype == dtype
|
||||
|
||||
# int64
|
||||
packer_int64 = SequencePacker(pack_size=10, pad_value=0, dtype=torch.int64)
|
||||
sequences_int64 = [torch.tensor([1, 2, 3], dtype=torch.int64)]
|
||||
packages_int64 = packer_int64.pack(sequences_int64)
|
||||
assert packages_int64[0].dtype == torch.int64
|
||||
|
||||
# float32
|
||||
packer_float32 = SequencePacker(pack_size=10, pad_value=0.0, dtype=torch.float32)
|
||||
sequences_float32 = [torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)]
|
||||
packages_float32 = packer_float32.pack(sequences_float32)
|
||||
assert packages_float32[0].dtype == torch.float32
|
||||
def test_dtype_conversion_on_mismatch(self, caplog):
|
||||
"""Tensors with mismatched dtype are silently converted."""
|
||||
packer = SequencePacker(pack_size=10, dtype=torch.int32)
|
||||
packages = packer.pack([torch.tensor([1, 2, 3], dtype=torch.int64)])
|
||||
assert packages[0].dtype == torch.int32
|
||||
assert packages[0][:3].tolist() == [1, 2, 3]
|
||||
|
||||
def test_non_1d_tensor_raises_error(self):
|
||||
"""测试非1D张量是否抛出异常"""
|
||||
packer = SequencePacker(pack_size=10)
|
||||
|
||||
# 2D 张量应该抛出异常
|
||||
sequences_2d = [torch.tensor([[1, 2], [3, 4]])] # shape: (2, 2)
|
||||
with pytest.raises(ValueError, match="Expected 1D tensor"):
|
||||
packer.pack(sequences_2d)
|
||||
|
||||
# 0D 张量 (标量) 应该抛出异常
|
||||
sequences_0d = [torch.tensor(5)] # shape: ()
|
||||
packer.pack([torch.tensor([[1, 2], [3, 4]])])
|
||||
with pytest.raises(ValueError, match="Expected 1D tensor"):
|
||||
packer.pack(sequences_0d)
|
||||
|
||||
# 3D 张量应该抛出异常
|
||||
sequences_3d = [torch.tensor([[[1, 2]]])] # shape: (1, 1, 2)
|
||||
with pytest.raises(ValueError, match="Expected 1D tensor"):
|
||||
packer.pack(sequences_3d)
|
||||
|
||||
def test_reset_method(self):
|
||||
"""测试 reset() 方法是否正确重置内部状态"""
|
||||
packer = SequencePacker(pack_size=10, pad_value=0)
|
||||
|
||||
# 第一次打包
|
||||
sequences1 = [torch.tensor([1, 2, 3], dtype=torch.int32)]
|
||||
packer.pack(sequences1)
|
||||
|
||||
# 验证内部状态已更新
|
||||
assert packer._current_pos == 0
|
||||
assert packer._current_pack is None # 最后一个包已发送，设置为None
|
||||
|
||||
# 重置
|
||||
packer.reset()
|
||||
|
||||
# 验证重置后的状态
|
||||
assert packer._current_pos == 0
|
||||
assert packer._current_pack is not None
|
||||
assert packer._current_pack.shape == (10,)
|
||||
assert packer._current_pack.tolist() == [0] * 10
|
||||
|
||||
# 验证重置后可以继续正常使用
|
||||
sequences2 = [torch.tensor([4, 5, 6], dtype=torch.int32)]
|
||||
packages = packer.pack(sequences2)
|
||||
|
||||
assert len(packages) == 1
|
||||
assert packages[0][:3].tolist() == [4, 5, 6]
|
||||
packer.pack([torch.tensor(5)])
|
||||
|
||||
def test_input_list_not_modified(self):
|
||||
"""测试输入列表是否未被修改（使用 sorted 而非 sort）"""
|
||||
packer = SequencePacker(pack_size=10)
|
||||
|
||||
# 创建原始序列列表（故意不按长度排序）
|
||||
original_sequences = [
|
||||
torch.tensor([3], dtype=torch.int32), # 长度1
|
||||
torch.tensor([1, 2], dtype=torch.int32), # 长度2
|
||||
torch.tensor([4, 5, 6, 7], dtype=torch.int32), # 长度4
|
||||
original = [
|
||||
torch.tensor([3], dtype=torch.int32),
|
||||
torch.tensor([1, 2], dtype=torch.int32),
|
||||
torch.tensor([4, 5, 6, 7], dtype=torch.int32),
|
||||
]
|
||||
|
||||
# 保存原始顺序的字符串表示
|
||||
original_repr = [seq.tolist() for seq in original_sequences]
|
||||
|
||||
# 打包
|
||||
packer.pack(original_sequences)
|
||||
|
||||
# 验证输入列表未被修改
|
||||
current_repr = [seq.tolist() for seq in original_sequences]
|
||||
assert current_repr == original_repr, "输入列表被修改了，应该使用 sorted() 而非 sort()"
|
||||
original_repr = [seq.tolist() for seq in original]
|
||||
packer.pack(original)
|
||||
assert [seq.tolist() for seq in original] == original_repr
|
||||
|
||||
def test_exact_pack_size_fit(self):
|
||||
"""测试序列长度恰好等于 pack_size 的情况"""
|
||||
packer = SequencePacker(pack_size=5, pad_value=0)
|
||||
|
||||
sequences = [
|
||||
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(sequences)
|
||||
|
||||
# 每个序列恰好占满一个包
|
||||
])
|
||||
assert len(packages) == 2
|
||||
|
||||
assert packages[0].tolist() == [1, 2, 3, 4, 5]
|
||||
assert packages[1].tolist() == [6, 7, 8, 9, 10]
|
||||
|
||||
def test_multiple_packs_full_utilization(self):
|
||||
"""测试多个包的高效利用"""
|
||||
packer = SequencePacker(pack_size=10, pad_value=-1)
|
||||
|
||||
# 创建多个小序列，确保高效打包
|
||||
sequences = [
|
||||
torch.tensor([1], dtype=torch.int32),
|
||||
torch.tensor([2], dtype=torch.int32),
|
||||
torch.tensor([3], dtype=torch.int32),
|
||||
torch.tensor([4], dtype=torch.int32),
|
||||
torch.tensor([5], dtype=torch.int32),
|
||||
torch.tensor([6], dtype=torch.int32),
|
||||
torch.tensor([7], dtype=torch.int32),
|
||||
torch.tensor([8], dtype=torch.int32),
|
||||
torch.tensor([9], dtype=torch.int32),
|
||||
torch.tensor([10], dtype=torch.int32),
|
||||
torch.tensor([11], dtype=torch.int32),
|
||||
]
|
||||
|
||||
sequences = [torch.tensor([i], dtype=torch.int32) for i in range(1, 12)]
|
||||
packages = packer.pack(sequences)
|
||||
|
||||
# 前10个序列打包成一个包，最后一个序列单独一个包
|
||||
assert len(packages) == 2
|
||||
assert packages[0].tolist() == list(range(1, 11))
|
||||
assert packages[1].tolist() == [11] + [-1] * 9
|
||||
|
||||
def test_dtype_mismatch_warning(self, caplog):
|
||||
"""测试 dtype 不匹配时的警告"""
|
||||
packer = SequencePacker(pack_size=10, pad_value=0, dtype=torch.int32)
|
||||
|
||||
sequences = [torch.tensor([1, 2, 3], dtype=torch.int64)]
|
||||
|
||||
packages = packer.pack(sequences)
|
||||
|
||||
# 应该触发 dtype 不匹配警告
|
||||
assert "dtype" in caplog.text.lower() or "converted" in caplog.text.lower()
|
||||
def test_cross_group_ordering(self):
|
||||
"""Tensor groups with identical per-item lengths are sorted identically."""
|
||||
packer = SequencePacker(pack_size=10, pad_value=0)
|
||||
# sequences: lengths [3, 1, 4] -> after sort desc: [4, 3, 1]
|
||||
seqs = [
|
||||
torch.tensor([1, 2, 3], dtype=torch.int32),
|
||||
torch.tensor([10], dtype=torch.int32),
|
||||
torch.tensor([4, 5, 6, 7], dtype=torch.int32),
|
||||
]
|
||||
masks = [
|
||||
torch.tensor([True, True, True], dtype=torch.bool),
|
||||
torch.tensor([True], dtype=torch.bool),
|
||||
torch.tensor([True, True, True, True], dtype=torch.bool),
|
||||
]
|
||||
packed_seqs = packer.pack(seqs)
|
||||
packer.reset()
|
||||
packed_masks = packer.pack(masks)
|
||||
|
||||
# Verify mask packer uses bool dtype
|
||||
assert packed_masks[0].dtype == torch.bool
|
||||
# Both groups should produce the same number of packages
|
||||
assert len(packed_seqs) == len(packed_masks)
|
||||
|
||||
+47
-143
@@ -1,4 +1,4 @@
|
||||
"""单元测试：pipeline.processors 模块中的处理器类"""
|
||||
"""Tests for pipeline.processors module."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -13,201 +13,105 @@ from pipeline.processors import (
|
||||
|
||||
|
||||
class DummyTokenizer:
|
||||
"""用于测试的虚拟分词器"""
|
||||
|
||||
def encode(self, text: str):
|
||||
# 简单模拟：返回文本字符的ASCII码列表
|
||||
def encode(self, text: str, add_special_tokens: bool = False):
|
||||
return [ord(c) for c in text]
|
||||
|
||||
|
||||
class TestBaseProcessor:
|
||||
"""BaseProcessor 抽象基类的测试"""
|
||||
|
||||
def test_abstract_class_cannot_be_instantiated(self):
|
||||
"""测试 BaseProcessor 不能直接实例化"""
|
||||
with pytest.raises(TypeError):
|
||||
BaseProcessor()
|
||||
|
||||
|
||||
class TestPreTrainProcessor:
|
||||
"""PreTrainProcessor 类的测试套件"""
|
||||
|
||||
def test_output_keys(self):
|
||||
"""测试 output_keys 属性"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = PreTrainProcessor(tokenizer)
|
||||
assert processor.output_keys == ["sequence"]
|
||||
|
||||
assert PreTrainProcessor(DummyTokenizer()).output_keys == ["sequence"]
|
||||
|
||||
def test_process_returns_tensor(self):
|
||||
"""测试 process 方法返回正确的张量"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = PreTrainProcessor(tokenizer)
|
||||
|
||||
processor = PreTrainProcessor(DummyTokenizer())
|
||||
result = processor.process({"text": "hello world"})
|
||||
|
||||
assert "sequence" in result
|
||||
assert isinstance(result["sequence"], torch.Tensor)
|
||||
assert result["sequence"].dtype == torch.int32
|
||||
|
||||
|
||||
def test_process_adds_eos(self):
|
||||
"""测试 process 方法添加 EOS 标记"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = PreTrainProcessor(tokenizer)
|
||||
|
||||
# 文本 "a" 的 ASCII 码是 97
|
||||
result = processor.process({"text": "a"})
|
||||
|
||||
# 应该包含文本的ASCII码 + <eos> (假设是 4)
|
||||
seq = result["sequence"]
|
||||
# 基本验证：返回的张量长度应该大于0
|
||||
assert len(seq) > 0
|
||||
result = PreTrainProcessor(DummyTokenizer()).process({"text": "a"})
|
||||
assert len(result["sequence"]) > 0
|
||||
|
||||
|
||||
class TestSFTProcessor:
|
||||
"""SFTProcessor 类的测试套件"""
|
||||
|
||||
def test_output_keys(self):
|
||||
"""测试 output_keys 属性"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = SFTProcessor(tokenizer)
|
||||
assert processor.output_keys == ["sequence", "loss_mask"]
|
||||
|
||||
assert SFTProcessor(DummyTokenizer()).output_keys == ["sequence", "loss_mask"]
|
||||
|
||||
def test_process_returns_both_keys(self):
|
||||
"""测试 process 方法返回所有键"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = SFTProcessor(tokenizer)
|
||||
|
||||
result = processor.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):
|
||||
"""测试 loss_mask 长度与 sequence 一致"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = SFTProcessor(tokenizer)
|
||||
|
||||
result = processor.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_after_query_is_true(self):
|
||||
"""测试 loss_mask 在响应部分为 True"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = SFTProcessor(tokenizer)
|
||||
|
||||
result = processor.process({
|
||||
"query": "ab", # 2 chars
|
||||
"response": "cd", # 2 chars
|
||||
})
|
||||
|
||||
# 验证 loss_mask 是 bool 类型
|
||||
|
||||
def test_loss_mask_is_bool(self):
|
||||
result = SFTProcessor(DummyTokenizer()).process({"query": "ab", "response": "cd"})
|
||||
assert result["loss_mask"].dtype == torch.bool
|
||||
|
||||
|
||||
class TestDPOProcessor:
|
||||
"""DPOProcessor 类的测试套件"""
|
||||
|
||||
def test_output_keys(self):
|
||||
"""测试 output_keys 属性"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = DPOProcessor(tokenizer)
|
||||
assert processor.output_keys == ["chosen", "chosen_mask", "rejected", "rejected_mask"]
|
||||
|
||||
keys = DPOProcessor(DummyTokenizer()).output_keys
|
||||
assert keys == ["chosen", "chosen_mask", "rejected", "rejected_mask"]
|
||||
|
||||
def test_process_returns_all_keys(self):
|
||||
"""测试 process 方法返回所有键"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = DPOProcessor(tokenizer)
|
||||
|
||||
result = processor.process({
|
||||
"query": "hello",
|
||||
"chosen": "response1",
|
||||
"rejected": "response2"
|
||||
})
|
||||
|
||||
expected_keys = ["chosen", "chosen_mask", "rejected", "rejected_mask"]
|
||||
for key in expected_keys:
|
||||
result = DPOProcessor(DummyTokenizer()).process(
|
||||
{"query": "hello", "chosen": "r1", "rejected": "r2"}
|
||||
)
|
||||
for key in ["chosen", "chosen_mask", "rejected", "rejected_mask"]:
|
||||
assert key in result
|
||||
assert isinstance(result[key], torch.Tensor)
|
||||
|
||||
def test_chosen_and_rejected_same_length_as_mask(self):
|
||||
"""测试 chosen/rejected 长度与 mask 一致"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = DPOProcessor(tokenizer)
|
||||
|
||||
result = processor.process({
|
||||
"query": "test",
|
||||
"chosen": "yes",
|
||||
"rejected": "no"
|
||||
})
|
||||
|
||||
|
||||
def test_masks_match_lengths(self):
|
||||
result = DPOProcessor(DummyTokenizer()).process(
|
||||
{"query": "test", "chosen": "yes", "rejected": "no"}
|
||||
)
|
||||
assert len(result["chosen"]) == len(result["chosen_mask"])
|
||||
assert len(result["rejected"]) == len(result["rejected_mask"])
|
||||
|
||||
|
||||
def test_masks_are_bool(self):
|
||||
"""测试 mask 张量是 bool 类型"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = DPOProcessor(tokenizer)
|
||||
|
||||
result = processor.process({
|
||||
"query": "test",
|
||||
"chosen": "yes",
|
||||
"rejected": "no"
|
||||
})
|
||||
|
||||
result = DPOProcessor(DummyTokenizer()).process(
|
||||
{"query": "test", "chosen": "yes", "rejected": "no"}
|
||||
)
|
||||
assert result["chosen_mask"].dtype == torch.bool
|
||||
assert result["rejected_mask"].dtype == torch.bool
|
||||
|
||||
|
||||
class TestProcessorFactory:
|
||||
"""ProcessorFactory 类的测试套件"""
|
||||
|
||||
def test_create_pre_train_processor(self):
|
||||
"""测试创建预训练处理器"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = ProcessorFactory.create("pt", tokenizer)
|
||||
assert isinstance(processor, PreTrainProcessor)
|
||||
|
||||
assert isinstance(ProcessorFactory.create("pt", DummyTokenizer()), PreTrainProcessor)
|
||||
|
||||
def test_create_sft_processor(self):
|
||||
"""测试创建 SFT 处理器"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = ProcessorFactory.create("sft", tokenizer)
|
||||
assert isinstance(processor, SFTProcessor)
|
||||
|
||||
assert isinstance(ProcessorFactory.create("sft", DummyTokenizer()), SFTProcessor)
|
||||
|
||||
def test_create_dpo_processor(self):
|
||||
"""测试创建 DPO 处理器"""
|
||||
tokenizer = DummyTokenizer()
|
||||
processor = ProcessorFactory.create("dpo", tokenizer)
|
||||
assert isinstance(processor, DPOProcessor)
|
||||
|
||||
assert isinstance(ProcessorFactory.create("dpo", DummyTokenizer()), DPOProcessor)
|
||||
|
||||
def test_create_invalid_processor_raises_error(self):
|
||||
"""测试创建无效处理器类型抛出异常"""
|
||||
tokenizer = DummyTokenizer()
|
||||
with pytest.raises(ValueError, match="Invalid processor type"):
|
||||
ProcessorFactory.create("invalid", tokenizer)
|
||||
|
||||
with pytest.raises(ValueError, match="Unknown processor type"):
|
||||
ProcessorFactory.create("invalid", DummyTokenizer())
|
||||
|
||||
def test_register_and_create_custom_processor(self):
|
||||
"""测试注册和创建自定义处理器"""
|
||||
class CustomProcessor(BaseProcessor):
|
||||
def __init__(self, tokenizer=None): # 接受 tokenizer 参数
|
||||
def __init__(self, tokenizer=None):
|
||||
self._tokenizer = tokenizer
|
||||
|
||||
|
||||
@property
|
||||
def output_keys(self):
|
||||
return ["custom"]
|
||||
|
||||
|
||||
def process(self, input_dict):
|
||||
return {"custom": torch.tensor([1, 2, 3])}
|
||||
|
||||
tokenizer = DummyTokenizer()
|
||||
ProcessorFactory.register("custom", CustomProcessor)
|
||||
processor = ProcessorFactory.create("custom", tokenizer)
|
||||
assert isinstance(processor, CustomProcessor)
|
||||
|
||||
ProcessorFactory.register("custom")(CustomProcessor)
|
||||
assert isinstance(ProcessorFactory.create("custom", DummyTokenizer()), CustomProcessor)
|
||||
|
||||
@@ -0,0 +1,118 @@
|
||||
"""Tests for strategy module."""
|
||||
|
||||
import pytest
|
||||
from pipeline.strategies import (
|
||||
PromptStrategy,
|
||||
ChatMLStrategy,
|
||||
AlpacaStrategy,
|
||||
StrategyFactory,
|
||||
)
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
|
||||
|
||||
class DummyTokenizer:
|
||||
def encode(self, text: str, add_special_tokens: bool = False):
|
||||
return [ord(c) for c in text]
|
||||
|
||||
|
||||
class DummyStrategy(PromptStrategy):
|
||||
def __init__(self, tokenizer):
|
||||
super().__init__(tokenizer)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "dummy"
|
||||
|
||||
def assemble_prompt(self, query_tokens):
|
||||
prefix = self._encode_format("Q: ")
|
||||
return prefix + query_tokens
|
||||
|
||||
def assemble_response(self, response_tokens):
|
||||
suffix = self._encode_format("<eos>")
|
||||
return response_tokens + suffix
|
||||
|
||||
|
||||
def _decode(tokens):
|
||||
return "".join(chr(t) for t in tokens)
|
||||
|
||||
|
||||
class TestChatMLStrategy:
|
||||
def test_name(self):
|
||||
assert ChatMLStrategy(DummyTokenizer()).name == "chatml"
|
||||
|
||||
def test_assemble_prompt(self):
|
||||
tk = DummyTokenizer()
|
||||
strategy = ChatMLStrategy(tk)
|
||||
query_tokens = tk.encode("hello")
|
||||
prompt = strategy.assemble_prompt(query_tokens)
|
||||
text = _decode(prompt)
|
||||
assert "<|im_start|>user" in text
|
||||
assert "hello" in text
|
||||
assert "<|im_start|>assistant" in text
|
||||
|
||||
def test_assemble_response(self):
|
||||
tk = DummyTokenizer()
|
||||
strategy = ChatMLStrategy(tk)
|
||||
response_tokens = tk.encode("world")
|
||||
response = strategy.assemble_response(response_tokens)
|
||||
text = _decode(response)
|
||||
assert "world" in text
|
||||
assert "<|im_end|>" in text
|
||||
assert "<eos>" 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
|
||||
|
||||
|
||||
class TestAlpacaStrategy:
|
||||
def test_name(self):
|
||||
assert AlpacaStrategy(DummyTokenizer()).name == "alpaca"
|
||||
|
||||
def test_assemble_prompt(self):
|
||||
tk = DummyTokenizer()
|
||||
strategy = AlpacaStrategy(tk)
|
||||
query_tokens = tk.encode("hello")
|
||||
prompt = strategy.assemble_prompt(query_tokens)
|
||||
text = _decode(prompt)
|
||||
assert "### Instruction:" in text
|
||||
assert "hello" in text
|
||||
assert "### Response:" in text
|
||||
|
||||
def test_assemble_response(self):
|
||||
tk = DummyTokenizer()
|
||||
strategy = AlpacaStrategy(tk)
|
||||
response_tokens = tk.encode("world")
|
||||
response = strategy.assemble_response(response_tokens)
|
||||
text = _decode(response)
|
||||
assert "world" in text
|
||||
assert "<eos>" in text
|
||||
|
||||
|
||||
class TestStrategyFactory:
|
||||
def test_create_chatml(self):
|
||||
tk = DummyTokenizer()
|
||||
assert isinstance(StrategyFactory.create("chatml", tk), ChatMLStrategy)
|
||||
|
||||
def test_create_alpaca(self):
|
||||
tk = DummyTokenizer()
|
||||
assert isinstance(StrategyFactory.create("alpaca", tk), AlpacaStrategy)
|
||||
|
||||
def test_create_invalid_raises_error(self):
|
||||
with pytest.raises(ValueError, match="Unknown strategy"):
|
||||
StrategyFactory.create("invalid_strategy", DummyTokenizer())
|
||||
|
||||
def test_register_and_create(self):
|
||||
StrategyFactory.register("dummy")(DummyStrategy)
|
||||
tk = DummyTokenizer()
|
||||
strategy = StrategyFactory.create("dummy", tk)
|
||||
assert isinstance(strategy, DummyStrategy)
|
||||
assert strategy.name == "dummy"
|
||||
|
||||
def test_available_strategies(self):
|
||||
strategies = StrategyFactory.available_types()
|
||||
assert "chatml" in strategies
|
||||
assert "alpaca" in strategies
|
||||
@@ -1,229 +0,0 @@
|
||||
"""单元测试：pipeline.tokenizer 模块中的 BpeTokenizer 类"""
|
||||
|
||||
import pytest
|
||||
from pathlib import Path
|
||||
import tempfile
|
||||
|
||||
from pipeline.tokenizer import BpeTokenizer
|
||||
|
||||
|
||||
class TestBpeTokenizer:
|
||||
"""BpeTokenizer 类的测试套件"""
|
||||
|
||||
def test_initialization_without_path(self):
|
||||
"""测试不加载外部文件初始化"""
|
||||
tokenizer = BpeTokenizer()
|
||||
assert tokenizer is not None
|
||||
assert hasattr(tokenizer, '_tokenizer')
|
||||
|
||||
def test_initialization_with_path(self):
|
||||
"""测试加载外部文件初始化"""
|
||||
# 这个测试假设没有预训练的分词器文件，所以只测试不抛出异常
|
||||
# 实际使用中需要提供有效的分词器文件路径
|
||||
try:
|
||||
tokenizer = BpeTokenizer(path="nonexistent.json")
|
||||
except Exception:
|
||||
# 预期会抛出异常，因为文件不存在
|
||||
pass
|
||||
|
||||
def test_vocab_size(self):
|
||||
"""测试获取词汇表大小"""
|
||||
tokenizer = BpeTokenizer()
|
||||
vocab_size = len(tokenizer)
|
||||
assert isinstance(vocab_size, int)
|
||||
assert vocab_size >= 0
|
||||
|
||||
def test_special_tokens_exist(self):
|
||||
"""测试特殊token是否存在"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
# 检查控制token
|
||||
assert hasattr(tokenizer, '_control_tokens')
|
||||
assert '<bos>' in tokenizer._control_tokens
|
||||
assert '<eos>' in tokenizer._control_tokens
|
||||
assert '<pad>' in tokenizer._control_tokens
|
||||
|
||||
# 检查特殊token
|
||||
assert hasattr(tokenizer, '_special_tokens')
|
||||
assert '<|im_start|>' in tokenizer._special_tokens
|
||||
assert '<|im_end|>' in tokenizer._special_tokens
|
||||
|
||||
def test_encode_string(self):
|
||||
"""测试编码单个字符串"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
# 使用简单的ASCII字符测试
|
||||
result = tokenizer.encode("hello")
|
||||
|
||||
# 返回应该是 token IDs 列表
|
||||
assert isinstance(result, list)
|
||||
|
||||
def test_encode_list(self):
|
||||
"""测试编码字符串列表"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
texts = ["hello", "world", "test"]
|
||||
result = tokenizer.encode(texts)
|
||||
|
||||
# 返回应该是列表的列表
|
||||
assert isinstance(result, list)
|
||||
assert len(result) == len(texts)
|
||||
for item in result:
|
||||
assert isinstance(item, list)
|
||||
|
||||
def test_encode_with_output_tokens(self):
|
||||
"""测试编码返回tokens而非ids"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
result = tokenizer.encode("hello", out_ids=False)
|
||||
|
||||
# 应该返回 token 字符串列表
|
||||
assert isinstance(result, list)
|
||||
|
||||
def test_encode_with_special_tokens(self):
|
||||
"""测试编码添加特殊token"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
result = tokenizer.encode("hello", add_special_tokens=True)
|
||||
|
||||
assert isinstance(result, list)
|
||||
|
||||
def test_decode(self):
|
||||
"""测试解码token IDs"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
# 解码空列表
|
||||
result = tokenizer.decode([])
|
||||
assert isinstance(result, str)
|
||||
|
||||
# 解码包含一些ID的列表（假设有 vocab）
|
||||
# 如果分词器未训练，可能无法正确解码
|
||||
result = tokenizer.decode([104, 101, 108, 108, 111]) # "hello" 的 ASCII
|
||||
assert isinstance(result, str)
|
||||
|
||||
def test_decode_with_special_tokens(self):
|
||||
"""测试解码保留特殊token"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
# 解码空列表
|
||||
result = tokenizer.decode([], skip_special_tokens=False)
|
||||
assert isinstance(result, str)
|
||||
|
||||
def test_stop_ids_property(self):
|
||||
"""测试 stop_ids 属性"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
stop_ids = tokenizer.stop_ids
|
||||
assert isinstance(stop_ids, list)
|
||||
|
||||
def test_special_token_properties(self):
|
||||
"""测试特殊token ID属性"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
# 这些属性可能返回 None 如果分词器未训练
|
||||
bos_id = tokenizer.bos_id
|
||||
eos_id = tokenizer.eos_id
|
||||
pad_id = tokenizer.pad_id
|
||||
|
||||
# 只验证属性存在且为 int 或 None
|
||||
assert isinstance(bos_id, (int, type(None)))
|
||||
assert isinstance(eos_id, (int, type(None)))
|
||||
assert isinstance(pad_id, (int, type(None)))
|
||||
|
||||
def test_save_method_exists(self):
|
||||
"""测试 save 方法存在"""
|
||||
tokenizer = BpeTokenizer()
|
||||
assert hasattr(tokenizer, 'save')
|
||||
assert callable(tokenizer.save)
|
||||
|
||||
def test_load_method_exists(self):
|
||||
"""测试 load 方法存在"""
|
||||
tokenizer = BpeTokenizer()
|
||||
assert hasattr(tokenizer, 'load')
|
||||
assert callable(tokenizer.load)
|
||||
|
||||
def test_train_method_exists(self):
|
||||
"""测试 train 方法存在"""
|
||||
tokenizer = BpeTokenizer()
|
||||
assert hasattr(tokenizer, 'train')
|
||||
assert callable(tokenizer.train)
|
||||
|
||||
def test_train_from_iterator_method_exists(self):
|
||||
"""测试 train_from_iterator 方法存在"""
|
||||
tokenizer = BpeTokenizer()
|
||||
assert hasattr(tokenizer, 'train_from_iterator')
|
||||
assert callable(tokenizer.train_from_iterator)
|
||||
|
||||
|
||||
class TestBpeTokenizerIntegration:
|
||||
"""BpeTokenizer 集成测试"""
|
||||
|
||||
def test_encode_decode_roundtrip(self):
|
||||
"""测试编码解码往返"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
original = "hello world"
|
||||
encoded = tokenizer.encode(original)
|
||||
decoded = tokenizer.decode(encoded)
|
||||
|
||||
# 往返后应该得到类似的结果
|
||||
# 注意：由于分词器可能未训练，结果可能不完全一致
|
||||
assert isinstance(encoded, list)
|
||||
assert isinstance(decoded, str)
|
||||
|
||||
def test_train_from_iterator_small_corpus(self, tmp_path):
|
||||
"""测试使用小语料库训练"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
# 创建临时训练文件
|
||||
train_file = tmp_path / "train.txt"
|
||||
train_content = "hello world\nthis is a test\nmachine learning\n"
|
||||
train_file.write_text(train_content)
|
||||
|
||||
# 训练分词器（使用较小的 vocab size 加快测试）
|
||||
try:
|
||||
tokenizer.train(
|
||||
files=[str(train_file)],
|
||||
vocab_size=100,
|
||||
min_freq=1,
|
||||
reserved_token_size=10
|
||||
)
|
||||
|
||||
# 验证训练后分词器可用
|
||||
result = tokenizer.encode("hello")
|
||||
assert isinstance(result, list)
|
||||
assert len(result) > 0
|
||||
except Exception as e:
|
||||
pytest.skip(f"Training failed: {e}")
|
||||
|
||||
def test_save_and_load_tokenizer(self, tmp_path):
|
||||
"""测试保存和加载分词器"""
|
||||
tokenizer = BpeTokenizer()
|
||||
|
||||
# 创建临时训练文件并训练
|
||||
train_file = tmp_path / "train.txt"
|
||||
train_content = "hello world\ntest data\n"
|
||||
train_file.write_text(train_content)
|
||||
|
||||
try:
|
||||
tokenizer.train(
|
||||
files=[str(train_file)],
|
||||
vocab_size=50,
|
||||
min_freq=1,
|
||||
reserved_token_size=5
|
||||
)
|
||||
|
||||
# 保存
|
||||
save_path = tmp_path / "tokenizer.json"
|
||||
tokenizer.save(str(save_path))
|
||||
|
||||
# 加载到新实例
|
||||
new_tokenizer = BpeTokenizer()
|
||||
new_tokenizer.load(str(save_path))
|
||||
|
||||
# 验证加载后分词器可用
|
||||
result = new_tokenizer.encode("hello")
|
||||
assert isinstance(result, list)
|
||||
|
||||
except Exception as e:
|
||||
pytest.skip(f"Save/load test failed: {e}")
|
||||
Reference in New Issue
Block a user