File size: 5,434 Bytes
6b73a07 | 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 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | """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
|