# Collators de dados

Um data collator monta amostras individuais do conjunto de dados em um lote para o modelo. Ele também pode fazer o padding dinamicamente das amostras até a sequência mais longa em *cada lote*, o que é mais eficiente do que fazer o padding até um comprimento máximo 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])
}
```

O Transformers fornece data collators para diversas tarefas (veja todos os [data collators](./main_classes/data_collator) disponíveis). Crie um data collator personalizado com:

- [DataCollatorWithPadding](#datacollatorwithpadding) quando você precisa de padding padrão baseado em tokenizer, além de campos extras.
- [DataCollatorMixin](#datacollatormixin) quando você precisa de lógica de padding personalizada, múltiplas entradas em pares por amostra ou uma estrutura de lote que o tokenizer não consegue produzir sozinho.

## DataCollatorWithPadding

Para casos de uso simples, como adicionar um campo extra, crie uma subclasse de [DataCollatorWithPadding](/docs/transformers/main/pt/main_classes/data_collator#transformers.DataCollatorWithPadding) e estenda seu método `__call__`. O exemplo abaixo adiciona um campo `"score"`.

1. Remova primeiro o campo personalizado, pois [pad()](/docs/transformers/main/pt/internal/tokenization_utils#transformers.PreTrainedTokenizerBase.pad) não o reconhece.
2. Chame a classe pai para tratar `input_ids` e `attention_mask`.
3. Adicione o campo `"score"` novamente ao 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
```

Passe o data collator personalizado para o [Trainer](/docs/transformers/main/pt/main_classes/trainer#transformers.Trainer) como qualquer outro data collator.

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

## DataCollatorMixin

Subclasse `DataCollatorMixin` para ter controle total sobre a montagem do lote e implemente seu próprio método `__call__`. Crie uma lógica de padding personalizada, trate múltiplos tipos de entrada ou crie estruturas de lote inteiramente novas. O exemplo [DataCollatorForPreference](https://github.com/huggingface/trl/blob/cfbdd3bea4448cde878c0da0de49551f553c61fe/trl/trainer/reward_trainer.py#L126) abaixo usa `DataCollatorMixin` porque cada amostra de treinamento possui uma resposta escolhida e uma rejeitada, e o modelo precisa ver ambas.

1. Separe `chosen_ids` e `rejected_ids` porque `pad` espera listas planas.
2. Concatene o par de entrada em uma única lista.
3. Gere `attention_mask` com [torch.ones_like](https://docs.pytorch.org/docs/stable/generated/torch.ones_like.html) em vez do tokenizer, pois o collator trabalha com listas brutas de IDs de tokens.
4. Preencha `input_ids` e `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óximas etapas

- Veja todos os [data collators](./main_classes/data_collator) disponíveis para tarefas comuns, como classificação de tokens.