Data collators · Hugging Face Transformers documentation Data collators Transformers 🏡 View all docsAWS Trainium & InferentiaAccelerateArgillaAutoTrainBitsandbytesCLIChat UIDataset viewerDatasetsDeploying on AWSDiffusersDistilabelEvaluateGoogle CloudGoogle TPUsGradioHubHub Python LibraryHuggingface.jsInference Endpoints (dedicated)Inference ProvidersKernelsLeRobotLeaderboardsLightevalMicrosoft AzureOpenEnvOptimumPEFTReachy MiniSafetensorsSentence TransformersTRLTasksText Embeddings InferenceText Generation InferenceTokenizersTrackioTransformersTransformers.jsXetsmolagentstimm Search documentation mainv5.14.0v5.13.1v5.12.0v5.11.0v5.10.4v5.9.0v5.8.1v5.7.0v5.6.2v5.5.4v5.4.0v5.3.0v5.2.0v5.1.0v5.0.0v4.57.6v4.56.2v4.55.4v4.53.3v4.52.3v4.51.3v4.50.0v4.49.0v4.48.2v4.47.1v4.46.3v4.45.2v4.44.2v4.43.4v4.42.4v4.41.2v4.40.2v4.39.3v4.38.2v4.37.2v4.36.1v4.35.2v4.34.1v4.33.3v4.32.1v4.31.0v4.30.0v4.29.1v4.28.1v4.27.2v4.26.1v4.25.1v4.24.0v4.23.1v4.22.2v4.21.3v4.20.1v4.19.4v4.18.0v4.17.0v4.16.2v4.15.0v4.14.1v4.13.0v4.12.5v4.11.3v4.10.1v4.9.2v4.8.2v4.7.0v4.6.0v4.5.1v4.4.2v4.3.3v4.2.2v4.1.1v4.0.1v3.5.1v3.4.0v3.3.1v3.2.0v3.1.0v3.0.2v2.11.0v2.10.0v2.9.1v2.8.0v2.7.0v2.6.0v2.5.1v2.4.1v2.3.0v2.2.2v2.1.1v2.0.0v1.2.0v1.1.0v1.0.0doc-builder-html ARDEENESFRHIITJAKOPTROTRZH Join the Hugging Face community and get access to the augmented documentation experience Collaborate on models, datasets and Spaces Faster examples with accelerated inference Switch between documentation themes Sign Up to get started Copy page Data collators A data collator assembles individual dataset samples into a batch for the model. It can also dynamically pad samples to the longest sequence in each batch, which is more efficient than padding to a global maximum length. Copied Dataset[0] → {"input_ids": [101, 2003], "labels": 1} Dataset[1] → {"input_ids": [101, 2003, 1996], "labels": 0} Dataset[2] → {"input_ids": [101, 7592], "labels": 1} ↓ collator { "input_ids": tensor([[101, 2003, 0], # padded to longest [101, 2003, 1996], [101, 7592, 0]]), "labels": tensor([1, 0, 1]) } Transformers provides data collators for various tasks (see all available data collators). Create a custom data collator with: DataCollatorWithPadding when you need standard tokenizer-based padding plus extra fields. DataCollatorMixin when you need custom padding logic, multiple paired inputs per sample, or a batch structure the tokenizer can’t produce on its own. DataCollatorWithPadding For simple use cases like adding an extra field, subclass DataCollatorWithPadding and extend its __call__ method. The example below adds a "score" field. Remove the custom field first because pad() doesn’t recognize it. Call the parent class to handle input_ids and attention_mask. Add the "score" field back to the batch. Copied import torch from dataclasses import dataclass from transformers import DataCollatorWithPadding, PreTrainedTokenizerBase @dataclass class DataCollatorWithScore(DataCollatorWithPadding): tokenizer: PreTrainedTokenizerBase def __call__(self, features): scores = [f.pop("score") for f in features] batch = super().__call__(features) batch["score"] = torch.tensor(scores, dtype=torch.float) return batch Pass the custom data collator to Trainer like any other data collator. Copied trainer = Trainer( ..., data_collator=DataCollatorWithScore(tokenizer=tokenizer), ) DataCollatorMixin Subclass DataCollatorMixin for full control over batch assembly and implement your own __call__ method. Build custom padding logic, handle multiple input types, or create entirely new batch structures. The DataCollatorForPreference example below uses DataCollatorMixin because each training sample has a chosen and rejected response, and the model needs to see both. Separate chosen_ids and rejected_ids because pad expects flat lists. Concatenate the input pair into a single list. Generate attention_mask with torch.ones_like instead of the tokenizer because the collator works with raw token ID lists. Pad input_ids and attention_mask. Copied import torch from transformers import DataCollatorMixin from trl.trainer.utils import pad class DataCollatorForPreference(DataCollatorMixin): pad_token_id: int pad_to_multiple_of: int | None = None def __call__(self, examples: list[dict]) -> dict: chosen_input_ids = [torch.tensor(ex["chosen_ids"]) for ex in examples] rejected_input_ids = [torch.tensor(ex["rejected_ids"]) for ex in examples] input_ids = chosen_input_ids + rejected_input_ids attention_mask = [torch.ones_like(ids) for ids in input_ids] output = { "input_ids": pad( input_ids, padding_value=self.pad_token_id, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of, ), "attention_mask": pad( attention_mask, padding_value=0, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of, ), } ... return output Next steps See all available data collators for common tasks like token classification. Update on GitHub ←Callbacks Optimizers and schedulers→