# Coladores de datos

Un data collator ensambla muestras individuales del conjunto de datos en un lote para el modelo. También puede rellenar dinámicamente las muestras hasta la secuencia más larga de *cada lote*, lo cual es más eficiente que rellenar hasta una longitud máxima global.

```md
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 proporciona collators de datos para diversas tareas (consulte todos los [collators de datos](./main_classes/data_collator) disponibles). Cree un collator de datos personalizado con:

- [DataCollatorWithPadding](#datacollatorwithpadding) cuando necesitas el relleno estándar basado en el tokenizador más campos adicionales.
- [DataCollatorMixin](#datacollatormixin) cuando necesitas una lógica de relleno personalizada, múltiples entradas emparejadas por muestra o una estructura de lote que el tokenizador no puede generar por sí solo.

## DataCollatorWithPadding

Para casos de uso simples, como añadir un campo adicional, subclase [DataCollatorWithPadding](/docs/transformers/main/es/main_classes/data_collator#transformers.DataCollatorWithPadding) y extienda su método `__call__`. El ejemplo de abajo añade un campo `"score"`.

1. Elimina primero el campo personalizado, ya que [pad()](/docs/transformers/main/es/internal/tokenization_utils#transformers.PreTrainedTokenizerBase.pad) no lo reconoce.
2. Llama a la clase padre para manejar `input_ids` y `attention_mask`.
3. Añada el campo `"score"` de nuevo al lote.

```py
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
```

Pasa el collator de datos personalizado a [Trainer](/docs/transformers/main/es/main_classes/trainer#transformers.Trainer) como cualquier otro collator de datos.

```py
trainer = Trainer(
    ...,
    data_collator=DataCollatorWithScore(tokenizer=tokenizer),
)
```

## DataCollatorMixin

Subclase `DataCollatorMixin` para tener control total sobre la ensamblaje de lotes e implemente su propio método `__call__`. Cree lógica de relleno personalizada, maneje múltiples tipos de entrada o cree estructuras de lote completamente nuevas. El ejemplo [DataCollatorForPreference](https://github.com/huggingface/trl/blob/cfbdd3bea4448cde878c0da0de49551f553c61fe/trl/trainer/reward_trainer.py#L126) a continuación utiliza `DataCollatorMixin` porque cada muestra de entrenamiento tiene una respuesta elegida y una rechazada, y el modelo necesita ver ambas.

1. Separe `chosen_ids` y `rejected_ids` porque `pad` espera listas planas.
2. Concatena el par de entrada en una única lista.
3. Genera `attention_mask` con [torch.ones_like](https://docs.pytorch.org/docs/stable/generated/torch.ones_like.html) en lugar del tokenizador, ya que el collator trabaja con listas de IDs de tokens sin procesar.
4. Rellena `input_ids` y `attention_mask`.

```py
import torch
from dataclasses import dataclass
from transformers.data.data_collator import DataCollatorMixin
from trl.trainer.utils import pad

@dataclass
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
```

## Próximos pasos

- Vea todos los [data collators](./main_classes/data_collator) disponibles para tareas comunes como la clasificación de tokens.