229 lines
7.6 KiB
Python
229 lines
7.6 KiB
Python
"""单元测试：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}") |