File size: 5,345 Bytes
c1a46f7 a2d6c00 c1a46f7 a2d6c00 c1a46f7 a2d6c00 c1a46f7 a2d6c00 c1a46f7 a2d6c00 c1a46f7 a2d6c00 c1a46f7 a2d6c00 c1a46f7 a2d6c00 c1a46f7 a2d6c00 c1a46f7 a2d6c00 c1a46f7 a2d6c00 c1a46f7 a2d6c00 c1a46f7 a2d6c00 c1a46f7 a2d6c00 c1a46f7 a2d6c00 c1a46f7 a2d6c00 c1a46f7 a2d6c00 c1a46f7 a2d6c00 c1a46f7 a2d6c00 c1a46f7 a2d6c00 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | """
单元测试 — 数据模块 (Person A 实现后需通过)
"""
import pytest
import torch
from easytranslate.data import (
DynamicBatchSampler,
TokenizerWrapper,
TranslationCollator,
TranslationDataset,
clean_text,
deduplicate_pairs,
filter_by_length,
preprocess_pipeline,
train_bpe_tokenizer,
)
@pytest.fixture()
def tiny_tokenizer() -> TokenizerWrapper:
texts = [
"Hello world",
"Machine translation is useful",
"I love natural language processing",
"你好 世界",
"机器翻译 很 有用",
"我 喜欢 自然语言处理",
]
return train_bpe_tokenizer(texts, vocab_size=80, min_frequency=1)
class TestTranslationDataset:
"""测试 TranslationDataset 类。"""
def test_dataset_length(self, tiny_tokenizer):
"""测试数据集长度。"""
dataset = TranslationDataset(["Hello world"], ["你好 世界"], tokenizer=tiny_tokenizer)
assert len(dataset) == 1
def test_getitem_returns_correct_keys(self, tiny_tokenizer):
"""测试 __getitem__ 返回正确的字段。"""
dataset = TranslationDataset(["Hello world"], ["你好 世界"], tokenizer=tiny_tokenizer)
item = dataset[0]
assert set(item) == {"src_ids", "tgt_input_ids", "labels", "src_len", "tgt_len"}
def test_getitem_tensor_types(self, tiny_tokenizer):
"""测试返回的 tensor 类型正确。"""
dataset = TranslationDataset(["Hello world"], ["你好 世界"], tokenizer=tiny_tokenizer)
item = dataset[0]
assert item["src_ids"].dtype == torch.long
assert item["tgt_input_ids"].dtype == torch.long
assert item["labels"].dtype == torch.long
def test_src_tgt_mismatch_raises(self):
"""测试源目标数量不匹配时抛出异常。"""
with pytest.raises(ValueError):
TranslationDataset(["a", "b"], ["甲"], tokenizer=object())
class TestTokenizer:
"""测试分词器。"""
def test_bpe_train_and_encode(self, tiny_tokenizer):
"""测试 BPE 训练和编码。"""
ids = tiny_tokenizer.encode("Hello world", add_special_tokens=True)
assert len(ids) >= 3
assert ids[0] == tiny_tokenizer.bos_token_id
assert ids[-1] == tiny_tokenizer.eos_token_id
def test_encode_decode_roundtrip(self, tiny_tokenizer):
"""测试编码-解码往返一致性。"""
ids = tiny_tokenizer.encode("Hello world")
decoded = tiny_tokenizer.decode(ids)
assert "Hello" in decoded
assert "world" in decoded
def test_special_tokens(self, tiny_tokenizer):
"""测试特殊 token 正确。"""
assert tiny_tokenizer.pad_token_id == 0
assert tiny_tokenizer.unk_token_id == 1
assert tiny_tokenizer.bos_token_id == 2
assert tiny_tokenizer.eos_token_id == 3
class TestPreprocessing:
"""测试预处理。"""
def test_clean_text_unicode(self):
"""测试 Unicode 标准化。"""
assert clean_text("ＡＢＣ\u200b １２３") == "ABC 123"
def test_filter_by_length(self):
"""测试按长度过滤。"""
assert filter_by_length("hello world", "你好世界", max_src_len=10, max_tgt_len=10)
assert not filter_by_length("hello " * 300, "你好", max_src_len=256)
def test_deduplicate(self):
"""测试去重。"""
pairs = deduplicate_pairs([("a", "甲"), ("a", "甲"), ("b", "乙")])
assert pairs == [("a", "甲"), ("b", "乙")]
def test_preprocess_pipeline(self):
"""测试完整预处理流水线。"""
src, tgt = preprocess_pipeline([" Hello world ", "", "Hello world"], [" 你好 世界 ", "空", "你好 世界"])
assert src == ["Hello world"]
assert tgt == ["你好 世界"]
class TestCollator:
"""测试数据整理器。"""
def test_padding(self, tiny_tokenizer):
"""测试 padding 正确。"""
dataset = TranslationDataset(
["Hello world", "Machine translation is useful"],
["你好 世界", "机器翻译 很 有用"],
tokenizer=tiny_tokenizer,
)
batch = TranslationCollator(pad_token_id=tiny_tokenizer.pad_token_id)([dataset[0], dataset[1]])
assert batch["src_ids"].ndim == 2
assert batch["tgt_input_ids"].ndim == 2
assert batch["labels"].shape == batch["tgt_input_ids"].shape
def test_attention_mask(self, tiny_tokenizer):
"""测试 attention mask 正确。"""
dataset = TranslationDataset(
["Hello", "Machine translation is useful"],
["你好", "机器翻译 很 有用"],
tokenizer=tiny_tokenizer,
)
batch = TranslationCollator(pad_token_id=tiny_tokenizer.pad_token_id)([dataset[0], dataset[1]])
assert batch["src_padding_mask"].dtype == torch.bool
assert torch.equal(batch["src_attention_mask"], (~batch["src_padding_mask"]).long())
def test_dynamic_batch_sampler(self):
"""测试动态 batch 不超过 token 预算。"""
sampler = DynamicBatchSampler([5, 6, 20, 21], max_tokens_per_batch=24, shuffle=False)
batches = list(sampler)
assert batches
for batch in batches:
assert max([5, 6, 20, 21][idx] for idx in batch) * len(batch) <= 24
|