File size: 3,717 Bytes
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 | """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
|