Implement data pipeline module (#2)

- Implement data pipeline module (a2d6c00f06d350e44f367ae2a14caee908f11454)


Co-authored-by: xinyuanxie <xiexinyuan341@users.noreply.huggingface.co>

.gitignore

@@ -21,7 +21,7 @@ env/
 *.swo
 
 # Data
-data/
+/data/
 *.tsv
 *.csv
 

requirements.txt

@@ -20,7 +20,7 @@ rouge-score>=0.1.2
 nltk>=3.8.0
 
 # Utilities
-numpy>=1.24.0
+numpy>=1.24.0,<2.0
 pandas>=2.0.0
 tqdm>=4.66.0
 pyyaml>=6.0.0

setup.py

@@ -11,9 +11,11 @@ setup(
         "torch>=2.1.0",
         "transformers>=4.36.0",
         "datasets>=2.16.0",
+        "tokenizers>=0.15.0",
         "sentencepiece>=0.1.99",
         "accelerate>=0.25.0",
         "sacrebleu>=2.4.0",
+        "numpy>=1.24.0,<2.0",
         "omegaconf>=2.3.0",
         "rich>=13.0.0",
         "tqdm>=4.66.0",

src/easytranslate/data/README.md

@@ -0,0 +1,126 @@
+# EasyTranslate Data 模块使用说明
+
+本目录负责英中翻译任务的数据加载、清洗、分词、样本构造和动态批处理。推荐主数据集使用 WMT19 zh-en，调试或小样本实验可使用 OPUS-100 en-zh。
+
+## 文件职责
+
+- `dataset.py`: 加载 WMT/OPUS/custom 数据，并提供 `TranslationDataset`。
+- `tokenizer.py`: 训练 BPE tokenizer，并用 `TokenizerWrapper` 统一 tokenizer 接口。
+- `preprocessing.py`: 文本标准化、长度过滤、去重。
+- `collator.py`: batch padding、attention mask、动态 token batch。
+
+## 推荐数据集
+
+正式实验建议使用：
+
+```python
+from easytranslate.data import load_wmt_dataset
+
+raw = load_wmt_dataset(year="19", language_pair="zh-en")
+```
+
+快速调试建议使用：
+
+```python
+from easytranslate.data import load_opus_dataset
+
+raw = load_opus_dataset(subset="en-zh")
+```
+
+两个加载函数都会把样本统一成：
+
+```python
+{"src": "English sentence", "tgt": "中文句子"}
+```
+
+## 从文本到 DataLoader
+
+```python
+from torch.utils.data import DataLoader
+
+from easytranslate.data import (
+    DynamicBatchSampler,
+    TranslationCollator,
+    TranslationDataset,
+    preprocess_pipeline,
+    train_bpe_tokenizer,
+)
+
+src_texts = raw["train"]["src"]
+tgt_texts = raw["train"]["tgt"]
+
+src_texts, tgt_texts = preprocess_pipeline(
+    src_texts,
+    tgt_texts,
+    lowercase_src=False,
+    remove_punctuation=False,
+    max_src_len=256,
+    max_tgt_len=256,
+    length_ratio_threshold=3.0,
+)
+
+tokenizer = train_bpe_tokenizer(
+    list(src_texts) + list(tgt_texts),
+    vocab_size=32000,
+    min_frequency=2,
+    save_path="outputs/tokenizer/bpe.json",
+)
+
+train_dataset = TranslationDataset(
+    src_texts,
+    tgt_texts,
+    tokenizer=tokenizer,
+    max_src_len=256,
+    max_tgt_len=256,
+)
+
+lengths = [
+    (len(tokenizer.encode(src, add_special_tokens=True)), len(tokenizer.encode(tgt, add_special_tokens=True)))
+    for src, tgt in zip(src_texts, tgt_texts)
+]
+
+batch_sampler = DynamicBatchSampler(lengths, max_tokens_per_batch=8192)
+collator = TranslationCollator(pad_token_id=tokenizer.pad_token_id)
+
+loader = DataLoader(
+    train_dataset,
+    batch_sampler=batch_sampler,
+    collate_fn=collator,
+    num_workers=4,
+    pin_memory=True,
+)
+```
+
+## Batch 字段
+
+`TranslationCollator` 输出：
+
+- `src_ids`: `[B, S]`
+- `tgt_input_ids`: `[B, T]`，以 `<s>` 开头，用于 teacher forcing
+- `labels`: `[B, T]`，以 `</s>` 结尾，padding 为 `-100`
+- `src_padding_mask`: `[B, S]`，padding 位置为 `True`
+- `tgt_padding_mask`: `[B, T]`，padding 位置为 `True`
+- `src_attention_mask` / `tgt_attention_mask`: 有效 token 为 `1`
+- `src_lens` / `tgt_lens`: 原始长度
+
+## 自定义平行语料
+
+```python
+from easytranslate.data import load_custom_dataset
+
+data = load_custom_dataset(
+    train_src="data/train.en",
+    train_tgt="data/train.zh",
+    val_src="data/val.en",
+    val_tgt="data/val.zh",
+    test_src="data/test.en",
+    test_tgt="data/test.zh",
+)
+```
+
+## 处理原则
+
+- 保留英文大小写和中英文标点，默认不做 lowercase、不去标点。
+- 清洗只做 Unicode NFKC、控制字符删除、空白合并。
+- tokenizer 只用训练集训练，不使用验证集或测试集。
+- 从零训练 Transformer 时建议使用共享 bilingual BPE；微调 NLLB 时直接使用 NLLB tokenizer。

src/easytranslate/data/__init__.py

@@ -0,0 +1,32 @@
+"""Data utilities for EasyTranslate."""
+
+from easytranslate.data.collator import DynamicBatchSampler, TranslationCollator
+from easytranslate.data.dataset import (
+    TranslationDataset,
+    load_custom_dataset,
+    load_opus_dataset,
+    load_wmt_dataset,
+)
+from easytranslate.data.preprocessing import (
+    clean_text,
+    deduplicate_pairs,
+    filter_by_length,
+    preprocess_pipeline,
+)
+from easytranslate.data.tokenizer import TokenizerWrapper, build_tokenizer, train_bpe_tokenizer
+
+__all__ = [
+    "DynamicBatchSampler",
+    "TokenizerWrapper",
+    "TranslationCollator",
+    "TranslationDataset",
+    "build_tokenizer",
+    "clean_text",
+    "deduplicate_pairs",
+    "filter_by_length",
+    "load_custom_dataset",
+    "load_opus_dataset",
+    "load_wmt_dataset",
+    "preprocess_pipeline",
+    "train_bpe_tokenizer",
+]

src/easytranslate/data/collator.py

@@ -0,0 +1,114 @@
+"""Batch collation and dynamic batching for translation training."""
+
+from __future__ import annotations
+
+import random
+from typing import Iterator, Sequence
+
+import torch
+from torch.nn.utils.rnn import pad_sequence
+from torch.utils.data import Sampler
+
+
+class TranslationCollator:
+    """Pad variable-length translation examples into one batch."""
+
+    def __init__(self, pad_token_id: int = 0, label_pad_token_id: int = -100):
+        self.pad_token_id = pad_token_id
+        self.label_pad_token_id = label_pad_token_id
+
+    def __call__(self, batch: Sequence[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
+        src_ids = pad_sequence(
+            [item["src_ids"] for item in batch],
+            batch_first=True,
+            padding_value=self.pad_token_id,
+        )
+        tgt_input_ids = pad_sequence(
+            [item["tgt_input_ids"] for item in batch],
+            batch_first=True,
+            padding_value=self.pad_token_id,
+        )
+        labels = pad_sequence(
+            [item["labels"] for item in batch],
+            batch_first=True,
+            padding_value=self.label_pad_token_id,
+        )
+
+        src_padding_mask = src_ids.eq(self.pad_token_id)
+        tgt_padding_mask = tgt_input_ids.eq(self.pad_token_id)
+
+        return {
+            "src_ids": src_ids,
+            "tgt_input_ids": tgt_input_ids,
+            "labels": labels,
+            "src_padding_mask": src_padding_mask,
+            "tgt_padding_mask": tgt_padding_mask,
+            "src_attention_mask": (~src_padding_mask).long(),
+            "tgt_attention_mask": (~tgt_padding_mask).long(),
+            "src_lens": torch.stack([item["src_len"] for item in batch]),
+            "tgt_lens": torch.stack([item["tgt_len"] for item in batch]),
+        }
+
+
+class DynamicBatchSampler(Sampler[list[int]]):
+    """Create batches constrained by an approximate max token budget."""
+
+    def __init__(
+        self,
+        lengths: Sequence[int | tuple[int, int]],
+        max_tokens_per_batch: int = 8192,
+        shuffle: bool = True,
+        drop_last: bool = False,
+    ):
+        self.lengths = [max(length) if isinstance(length, tuple) else int(length) for length in lengths]
+        self.max_tokens_per_batch = max_tokens_per_batch
+        self.shuffle = shuffle
+        self.drop_last = drop_last
+
+    def __iter__(self) -> Iterator[list[int]]:
+        indices = list(range(len(self.lengths)))
+        if self.shuffle:
+            random.shuffle(indices)
+
+        indices.sort(key=lambda idx: self.lengths[idx])
+        batches: list[list[int]] = []
+        batch: list[int] = []
+        max_len = 0
+
+        for idx in indices:
+            candidate_max_len = max(max_len, self.lengths[idx])
+            candidate_tokens = candidate_max_len * (len(batch) + 1)
+
+            if batch and candidate_tokens > self.max_tokens_per_batch:
+                batches.append(batch)
+                batch = []
+                max_len = 0
+
+            batch.append(idx)
+            max_len = max(max_len, self.lengths[idx])
+
+        if batch and not self.drop_last:
+            batches.append(batch)
+
+        if self.shuffle:
+            random.shuffle(batches)
+
+        yield from batches
+
+    def __len__(self) -> int:
+        count = 0
+        batch_size = 0
+        max_len = 0
+
+        for length in sorted(self.lengths):
+            candidate_max_len = max(max_len, length)
+            if batch_size and candidate_max_len * (batch_size + 1) > self.max_tokens_per_batch:
+                count += 1
+                batch_size = 0
+                max_len = 0
+            batch_size += 1
+            max_len = max(max_len, length)
+
+        if batch_size and not self.drop_last:
+            count += 1
+        return count

src/easytranslate/data/dataset.py

@@ -0,0 +1,157 @@
+"""Dataset loading and PyTorch dataset classes for English-Chinese translation."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Mapping, Sequence
+
+import torch
+from torch.utils.data import Dataset
+
+from easytranslate.data.preprocessing import preprocess_pipeline
+from easytranslate.data.tokenizer import TokenizerWrapper
+
+
+class TranslationDataset(Dataset):
+    """PyTorch dataset that builds seq2seq inputs for teacher forcing."""
+
+    def __init__(
+        self,
+        src_texts: Sequence[str],
+        tgt_texts: Sequence[str],
+        tokenizer: TokenizerWrapper | None = None,
+        src_tokenizer: TokenizerWrapper | None = None,
+        tgt_tokenizer: TokenizerWrapper | None = None,
+        max_src_len: int = 256,
+        max_tgt_len: int = 256,
+    ):
+        if len(src_texts) != len(tgt_texts):
+            raise ValueError("src_texts and tgt_texts must have the same length")
+
+        if tokenizer is not None:
+            src_tokenizer = src_tokenizer or tokenizer
+            tgt_tokenizer = tgt_tokenizer or tokenizer
+        if src_tokenizer is None or tgt_tokenizer is None:
+            raise ValueError("Provide tokenizer or both src_tokenizer and tgt_tokenizer")
+
+        self.src_texts = list(src_texts)
+        self.tgt_texts = list(tgt_texts)
+        self.src_tokenizer = src_tokenizer
+        self.tgt_tokenizer = tgt_tokenizer
+        self.max_src_len = max_src_len
+        self.max_tgt_len = max_tgt_len
+
+    def __len__(self) -> int:
+        return len(self.src_texts)
+
+    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
+        src_ids = self.src_tokenizer.encode(
+            self.src_texts[index],
+            add_special_tokens=True,
+            max_length=self.max_src_len,
+        )
+
+        target_core = self.tgt_tokenizer.encode(
+            self.tgt_texts[index],
+            add_special_tokens=False,
+            max_length=max(1, self.max_tgt_len - 1),
+        )
+        tgt_input_ids = [self.tgt_tokenizer.bos_token_id] + target_core
+        labels = target_core + [self.tgt_tokenizer.eos_token_id]
+
+        return {
+            "src_ids": torch.tensor(src_ids, dtype=torch.long),
+            "tgt_input_ids": torch.tensor(tgt_input_ids, dtype=torch.long),
+            "labels": torch.tensor(labels, dtype=torch.long),
+            "src_len": torch.tensor(len(src_ids), dtype=torch.long),
+            "tgt_len": torch.tensor(len(labels), dtype=torch.long),
+        }
+
+
+def _translation_to_columns(dataset, src_lang: str, tgt_lang: str):
+    def convert(example):
+        translation = example["translation"]
+        return {"src": translation[src_lang], "tgt": translation[tgt_lang]}
+
+    return dataset.map(convert, remove_columns=dataset.column_names)
+
+
+def load_wmt_dataset(
+    year: str = "19",
+    language_pair: str = "zh-en",
+    src_lang: str = "en",
+    tgt_lang: str = "zh",
+    split: str | None = None,
+    cache_dir: str | None = None,
+):
+    """Load WMT zh-en and normalize rows to {'src', 'tgt'}."""
+    from datasets import load_dataset
+
+    dataset_name = f"wmt/wmt{year}"
+    try:
+        dataset = load_dataset(dataset_name, language_pair, split=split, cache_dir=cache_dir)
+    except Exception:
+        dataset = load_dataset(f"wmt{year}", language_pair, split=split, cache_dir=cache_dir)
+
+    if split is not None:
+        return _translation_to_columns(dataset, src_lang, tgt_lang)
+
+    return dataset.map(
+        lambda example: {"src": example["translation"][src_lang], "tgt": example["translation"][tgt_lang]},
+        remove_columns=next(iter(dataset.values())).column_names,
+    )
+
+
+def load_opus_dataset(
+    subset: str = "en-zh",
+    src_lang: str = "en",
+    tgt_lang: str = "zh",
+    split: str | None = None,
+    cache_dir: str | None = None,
+):
+    """Load OPUS-100 and normalize rows to {'src', 'tgt'}."""
+    from datasets import load_dataset
+
+    dataset = load_dataset("Helsinki-NLP/opus-100", subset, split=split, cache_dir=cache_dir)
+    if split is not None:
+        return _translation_to_columns(dataset, src_lang, tgt_lang)
+
+    return dataset.map(
+        lambda example: {"src": example["translation"][src_lang], "tgt": example["translation"][tgt_lang]},
+        remove_columns=next(iter(dataset.values())).column_names,
+    )
+
+
+def _read_lines(path: str | Path) -> list[str]:
+    with Path(path).open("r", encoding="utf-8") as f:
+        return [line.rstrip("\n") for line in f]
+
+
+def load_custom_dataset(
+    train_src: str | Path,
+    train_tgt: str | Path,
+    val_src: str | Path | None = None,
+    val_tgt: str | Path | None = None,
+    test_src: str | Path | None = None,
+    test_tgt: str | Path | None = None,
+    preprocessing_config: Mapping | None = None,
+) -> dict[str, dict[str, list[str]]]:
+    """Load parallel text files and return split dictionaries."""
+    preprocessing_config = dict(preprocessing_config or {})
+
+    def load_split(src_path: str | Path, tgt_path: str | Path) -> dict[str, list[str]]:
+        src_texts, tgt_texts = preprocess_pipeline(
+            _read_lines(src_path),
+            _read_lines(tgt_path),
+            **preprocessing_config,
+        )
+        return {"src": src_texts, "tgt": tgt_texts}
+
+    result = {"train": load_split(train_src, train_tgt)}
+
+    if val_src and val_tgt:
+        result["validation"] = load_split(val_src, val_tgt)
+    if test_src and test_tgt:
+        result["test"] = load_split(test_src, test_tgt)
+
+    return result

src/easytranslate/data/preprocessing.py

@@ -0,0 +1,125 @@
+"""Text cleaning and filtering utilities for translation corpora."""
+
+from __future__ import annotations
+
+import re
+import unicodedata
+from typing import Callable, Iterable, Sequence
+
+CONTROL_OR_ZERO_WIDTH_RE = re.compile(r"[\u0000-\u001f\u007f-\u009f\u200b\u200c\u200d\ufeff]")
+SPACE_RE = re.compile(r"\s+")
+PUNCT_RE = re.compile(r"[^\w\s\u4e00-\u9fff]", flags=re.UNICODE)
+
+
+def clean_text(text: str, lowercase: bool = False, remove_punctuation: bool = False) -> str:
+    """Normalize a single sentence without changing its meaning aggressively."""
+    if text is None:
+        return ""
+
+    text = unicodedata.normalize("NFKC", str(text))
+    text = CONTROL_OR_ZERO_WIDTH_RE.sub("", text)
+    text = SPACE_RE.sub(" ", text).strip()
+
+    if lowercase:
+        text = text.lower()
+    if remove_punctuation:
+        text = PUNCT_RE.sub("", text)
+        text = SPACE_RE.sub(" ", text).strip()
+
+    return text
+
+
+def _default_length(text: str) -> int:
+    """Use whitespace tokens for Latin text and character count for CJK-heavy text."""
+    cjk_chars = sum(1 for ch in text if "\u4e00" <= ch <= "\u9fff")
+    if cjk_chars >= max(1, len(text) // 3):
+        return len(text.replace(" ", ""))
+    return len(text.split())
+
+
+def filter_by_length(
+    src: str,
+    tgt: str,
+    min_src_len: int = 1,
+    min_tgt_len: int = 1,
+    max_src_len: int = 256,
+    max_tgt_len: int = 256,
+    length_ratio_threshold: float = 3.0,
+    length_fn: Callable[[str], int] | None = None,
+) -> bool:
+    """Return True when a sentence pair passes basic length and ratio checks."""
+    length_fn = length_fn or _default_length
+    src_len = length_fn(src)
+    tgt_len = length_fn(tgt)
+
+    if src_len < min_src_len or tgt_len < min_tgt_len:
+        return False
+    if src_len > max_src_len or tgt_len > max_tgt_len:
+        return False
+
+    shorter = max(1, min(src_len, tgt_len))
+    longer = max(src_len, tgt_len)
+    return longer / shorter <= length_ratio_threshold
+
+
+def deduplicate_pairs(pairs: Iterable[tuple[str, str]]) -> list[tuple[str, str]]:
+    """Deduplicate by exact cleaned source-target pair while preserving order."""
+    seen: set[tuple[str, str]] = set()
+    result: list[tuple[str, str]] = []
+
+    for src, tgt in pairs:
+        key = (src, tgt)
+        if key in seen:
+            continue
+        seen.add(key)
+        result.append(key)
+
+    return result
+
+
+def preprocess_pipeline(
+    src_texts: Sequence[str],
+    tgt_texts: Sequence[str],
+    lowercase_src: bool = False,
+    lowercase_tgt: bool = False,
+    remove_punctuation: bool = False,
+    max_src_len: int = 256,
+    max_tgt_len: int = 256,
+    min_src_len: int = 1,
+    min_tgt_len: int = 1,
+    filter_by_length_enabled: bool = True,
+    length_ratio_threshold: float = 3.0,
+    deduplicate: bool = True,
+) -> tuple[list[str], list[str]]:
+    """Clean, filter, and optionally deduplicate parallel source-target texts."""
+    if len(src_texts) != len(tgt_texts):
+        raise ValueError("src_texts and tgt_texts must have the same length")
+
+    pairs: list[tuple[str, str]] = []
+    for raw_src, raw_tgt in zip(src_texts, tgt_texts):
+        src = clean_text(raw_src, lowercase=lowercase_src, remove_punctuation=remove_punctuation)
+        tgt = clean_text(raw_tgt, lowercase=lowercase_tgt, remove_punctuation=remove_punctuation)
+
+        if not src or not tgt:
+            continue
+        if filter_by_length_enabled and not filter_by_length(
+            src,
+            tgt,
+            min_src_len=min_src_len,
+            min_tgt_len=min_tgt_len,
+            max_src_len=max_src_len,
+            max_tgt_len=max_tgt_len,
+            length_ratio_threshold=length_ratio_threshold,
+        ):
+            continue
+
+        pairs.append((src, tgt))
+
+    if deduplicate:
+        pairs = deduplicate_pairs(pairs)
+
+    if not pairs:
+        return [], []
+
+    src_clean, tgt_clean = zip(*pairs)
+    return list(src_clean), list(tgt_clean)

src/easytranslate/data/tokenizer.py

@@ -0,0 +1,197 @@
+"""Tokenizer wrappers and BPE training helpers."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Iterable, Mapping, Sequence
+
+
+class TokenizerWrapper:
+    """Small adapter that gives HF tokenizers and tokenizers.Tokenizer one API."""
+
+    def __init__(
+        self,
+        tokenizer,
+        pad_token: str = "<pad>",
+        unk_token: str = "<unk>",
+        bos_token: str = "<s>",
+        eos_token: str = "</s>",
+    ):
+        self.tokenizer = tokenizer
+        self.pad_token = pad_token
+        self.unk_token = unk_token
+        self.bos_token = bos_token
+        self.eos_token = eos_token
+
+    @property
+    def pad_token_id(self) -> int:
+        return self.token_to_id(self.pad_token)
+
+    @property
+    def unk_token_id(self) -> int:
+        return self.token_to_id(self.unk_token)
+
+    @property
+    def bos_token_id(self) -> int:
+        return self.token_to_id(self.bos_token)
+
+    @property
+    def eos_token_id(self) -> int:
+        return self.token_to_id(self.eos_token)
+
+    @property
+    def vocab_size(self) -> int:
+        if hasattr(self.tokenizer, "get_vocab_size"):
+            return int(self.tokenizer.get_vocab_size())
+        return int(len(self.tokenizer))
+
+    def token_to_id(self, token: str) -> int:
+        if hasattr(self.tokenizer, "token_to_id"):
+            idx = self.tokenizer.token_to_id(token)
+        elif hasattr(self.tokenizer, "convert_tokens_to_ids"):
+            idx = self.tokenizer.convert_tokens_to_ids(token)
+        else:
+            raise TypeError("Unsupported tokenizer type")
+
+        if idx is None:
+            raise ValueError(f"Token {token!r} is not in the tokenizer vocabulary")
+        return int(idx)
+
+    def encode(self, text: str, add_special_tokens: bool = False, max_length: int | None = None) -> list[int]:
+        if hasattr(self.tokenizer, "encode") and self.tokenizer.__class__.__module__.startswith("tokenizers"):
+            ids = self.tokenizer.encode(text).ids
+        else:
+            ids = self.tokenizer.encode(text, add_special_tokens=add_special_tokens)
+            add_special_tokens = False
+
+        if add_special_tokens:
+            ids = [self.bos_token_id] + list(ids) + [self.eos_token_id]
+
+        if max_length is not None:
+            ids = list(ids)[:max_length]
+
+        return list(ids)
+
+    def decode(self, ids: Sequence[int], skip_special_tokens: bool = True) -> str:
+        if hasattr(self.tokenizer, "decode"):
+            try:
+                return self.tokenizer.decode(list(ids), skip_special_tokens=skip_special_tokens)
+            except TypeError:
+                return self.tokenizer.decode(list(ids))
+        raise TypeError("Unsupported tokenizer type")
+
+    def save(self, path: str | Path) -> None:
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        if hasattr(self.tokenizer, "save"):
+            self.tokenizer.save(str(path))
+            return
+        if hasattr(self.tokenizer, "save_pretrained"):
+            self.tokenizer.save_pretrained(str(path))
+            return
+        raise TypeError("Unsupported tokenizer type")
+
+
+def _special_tokens(config: Mapping | None = None) -> dict[str, str]:
+    tokens = {
+        "pad": "<pad>",
+        "unk": "<unk>",
+        "bos": "<s>",
+        "eos": "</s>",
+    }
+    if config:
+        tokens.update(dict(config))
+    return tokens
+
+
+def train_bpe_tokenizer(
+    texts: Iterable[str],
+    vocab_size: int = 32000,
+    min_frequency: int = 2,
+    special_tokens: Mapping[str, str] | None = None,
+    save_path: str | Path | None = None,
+) -> TokenizerWrapper:
+    """Train a byte-level BPE tokenizer on source and target training text."""
+    from tokenizers import Tokenizer
+    from tokenizers.decoders import ByteLevel as ByteLevelDecoder
+    from tokenizers.models import BPE
+    from tokenizers.normalizers import NFKC, Sequence as NormalizerSequence
+    from tokenizers.pre_tokenizers import ByteLevel
+    from tokenizers.trainers import BpeTrainer
+
+    tokens = _special_tokens(special_tokens)
+    ordered_specials = [tokens["pad"], tokens["unk"], tokens["bos"], tokens["eos"]]
+
+    tokenizer = Tokenizer(BPE(unk_token=tokens["unk"]))
+    tokenizer.normalizer = NormalizerSequence([NFKC()])
+    tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)
+    tokenizer.decoder = ByteLevelDecoder()
+
+    trainer = BpeTrainer(
+        vocab_size=vocab_size,
+        min_frequency=min_frequency,
+        special_tokens=ordered_specials,
+        show_progress=True,
+    )
+    tokenizer.train_from_iterator((text for text in texts if text), trainer=trainer)
+
+    wrapper = TokenizerWrapper(
+        tokenizer,
+        pad_token=tokens["pad"],
+        unk_token=tokens["unk"],
+        bos_token=tokens["bos"],
+        eos_token=tokens["eos"],
+    )
+
+    if save_path is not None:
+        wrapper.save(save_path)
+
+    return wrapper
+
+
+def build_tokenizer(config: Mapping, train_texts: Iterable[str] | None = None) -> TokenizerWrapper:
+    """Build a tokenizer from project config."""
+    tokenizer_type = config.get("type", "bpe")
+    tokens = _special_tokens(config.get("special_tokens"))
+
+    if tokenizer_type == "pretrained":
+        from transformers import AutoTokenizer
+
+        model_name = config.get("model_name") or config.get("pretrained_model_name")
+        if not model_name:
+            raise ValueError("pretrained tokenizer requires config['model_name']")
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        return TokenizerWrapper(
+            tokenizer,
+            pad_token=tokenizer.pad_token or tokens["pad"],
+            unk_token=tokenizer.unk_token or tokens["unk"],
+            bos_token=tokenizer.bos_token or tokens["bos"],
+            eos_token=tokenizer.eos_token or tokens["eos"],
+        )
+
+    if tokenizer_type in {"bpe", "sentencepiece"}:
+        tokenizer_path = config.get("path") or config.get("tokenizer_path")
+        if tokenizer_path and Path(tokenizer_path).exists():
+            from tokenizers import Tokenizer
+
+            return TokenizerWrapper(
+                Tokenizer.from_file(str(tokenizer_path)),
+                pad_token=tokens["pad"],
+                unk_token=tokens["unk"],
+                bos_token=tokens["bos"],
+                eos_token=tokens["eos"],
+            )
+
+        if train_texts is None:
+            raise ValueError("BPE tokenizer requires train_texts when no tokenizer path is provided")
+
+        return train_bpe_tokenizer(
+            train_texts,
+            vocab_size=int(config.get("vocab_size", 32000)),
+            min_frequency=int(config.get("min_frequency", 2)),
+            special_tokens=tokens,
+            save_path=tokenizer_path,
+        )
+
+    raise ValueError(f"Unsupported tokenizer type: {tokenizer_type}")

tests/test_data.py

@@ -5,43 +5,83 @@
 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):
+    def test_dataset_length(self, tiny_tokenizer):
         """测试数据集长度。"""
-        # TODO: 创建 mock tokenizer，构建小型数据集，验证 len() 正确
-        pass
+        dataset = TranslationDataset(["Hello world"], ["你好 世界"], tokenizer=tiny_tokenizer)
+        assert len(dataset) == 1
 
-    def test_getitem_returns_correct_keys(self):
+    def test_getitem_returns_correct_keys(self, tiny_tokenizer):
         """测试 __getitem__ 返回正确的字段。"""
-        # TODO: 验证返回 dict 包含 src_ids, tgt_input_ids, labels, src_len, tgt_len
-        pass
+        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):
+    def test_getitem_tensor_types(self, tiny_tokenizer):
         """测试返回的 tensor 类型正确。"""
-        pass
+        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):
         """测试源目标数量不匹配时抛出异常。"""
-        pass
+        with pytest.raises(ValueError):
+            TranslationDataset(["a", "b"], ["甲"], tokenizer=object())
 
 
 class TestTokenizer:
     """测试分词器。"""
 
-    def test_bpe_train_and_encode(self):
+    def test_bpe_train_and_encode(self, tiny_tokenizer):
         """测试 BPE 训练和编码。"""
-        pass
+        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):
+    def test_encode_decode_roundtrip(self, tiny_tokenizer):
         """测试编码-解码往返一致性。"""
-        pass
+        ids = tiny_tokenizer.encode("Hello world")
+        decoded = tiny_tokenizer.decode(ids)
+        assert "Hello" in decoded
+        assert "world" in decoded
 
-    def test_special_tokens(self):
+    def test_special_tokens(self, tiny_tokenizer):
         """测试特殊 token 正确。"""
-        pass
+        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:
@@ -49,24 +89,55 @@ class TestPreprocessing:
 
     def test_clean_text_unicode(self):
         """测试 Unicode 标准化。"""
-        pass
+        assert clean_text("ＡＢＣ\u200b  １２３") == "ABC 123"
 
     def test_filter_by_length(self):
         """测试按长度过滤。"""
-        pass
+        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):
         """测试去重。"""
-        pass
+        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):
+    def test_padding(self, tiny_tokenizer):
         """测试 padding 正确。"""
-        pass
-
-    def test_attention_mask(self):
+        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 正确。"""
-        pass
+        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