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.

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 disponíveis). Crie um data collator personalizado com:

DataCollatorWithPadding quando você precisa de padding padrão baseado em tokenizer, além de campos extras.
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 e estenda seu método __call__. O exemplo abaixo adiciona um campo "score".

Remova primeiro o campo personalizado, pois pad() não o reconhece.
Chame a classe pai para tratar input_ids e attention_mask.
Adicione o campo "score" novamente ao lote.

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 como qualquer outro data collator.

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 abaixo usa DataCollatorMixin porque cada amostra de treinamento possui uma resposta escolhida e uma rejeitada, e o modelo precisa ver ambas.

Separe chosen_ids e rejected_ids porque pad espera listas planas.
Concatene o par de entrada em uma única lista.
Gere attention_mask com torch.ones_like em vez do tokenizer, pois o collator trabalha com listas brutas de IDs de tokens.
Preencha input_ids e attention_mask.

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 disponíveis para tarefas comuns, como classificação de tokens.

Update on GitHub