File size: 3,439 Bytes
0861a59 | 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 | """
Custom data collators for transformer training.
"""
import torch
import random
from typing import Dict, List, Any, Union
from dataclasses import dataclass
@dataclass
class DataCollatorForLanguageModeling:
"""
Data collator for language modeling.
This collator will tokenize inputs and dynamically mask tokens
for masked language modeling tasks.
"""
tokenizer: Any
mlm: bool = True # Whether to use masked language modeling
mlm_probability: float = 0.15 # Probability of masking a token
def __call__(self, examples: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
"""
Collate a batch of examples.
Args:
examples: List of examples from dataset
Returns:
Batch dictionary for model
"""
# Extract input_ids
input_ids = [example["input_ids"] for example in examples]
# Concatenate inputs
batch = self.tokenizer.pad(
{"input_ids": input_ids},
return_tensors="pt"
)
# If masked language modeling is enabled
if self.mlm:
inputs, labels = self.mask_tokens(batch["input_ids"])
return {"input_ids": inputs, "labels": labels}
else:
labels = batch["input_ids"].clone()
return {
"input_ids": batch["input_ids"],
"labels": labels,
"attention_mask": batch.get("attention_mask", None)
}
def mask_tokens(
self, inputs: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
"""
Prepare masked tokens inputs/labels for masked language modeling.
Args:
inputs: Input tensor
Returns:
Tuple of (masked inputs, labels)
"""
labels = inputs.clone()
# Get probability mask
probability_matrix = torch.full(labels.shape, self.mlm_probability)
# Create special tokens mask
if hasattr(self.tokenizer, "get_special_tokens_mask"):
special_tokens_mask = [
self.tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True)
for val in labels.tolist()
]
special_tokens_mask = torch.tensor(special_tokens_mask, dtype=torch.bool)
else:
special_tokens_mask = torch.tensor(
[
[self._is_special_token(x) for x in val]
for val in labels.tolist()
],
dtype=torch.bool,
)
# Don't mask special tokens
probability_matrix.masked_fill_(special_tokens_mask, value=0.0)
# Get mask indices
masked_indices = torch.bernoulli(probability_matrix).bool()
# Set labels for non-masked tokens to -100 (ignored in loss)
labels[~masked_indices] = -100
# Set 80% of masked tokens to [MASK]
indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices
if hasattr(self.tokenizer, "mask_token_id") and self.tokenizer.mask_token_id is not None:
inputs[indices_replaced] = self.tokenizer.mask_token_id
# Set |