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.

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 disponibles). Cree un collator de datos personalizado con:

DataCollatorWithPadding cuando necesitas el relleno estándar basado en el tokenizador más campos adicionales.
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 y extienda su método __call__. El ejemplo de abajo añade un campo "score".

Elimina primero el campo personalizado porque pad() no lo reconoce.
Llama a la clase padre para manejar input_ids y attention_mask.
Añada el campo "score" de nuevo al 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

Pasa el collator de datos personalizado a Trainer como cualquier otro collator de datos.

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 a continuación utiliza DataCollatorMixin porque cada muestra de entrenamiento tiene una respuesta elegida y una rechazada, y el modelo necesita ver ambas.

Separe chosen_ids y rejected_ids porque pad espera listas planas.
Concatena el par de entrada en una única lista.
Genera attention_mask con torch.ones_like en lugar del tokenizador, ya que el collator trabaja con listas de IDs de tokens sin procesar.
Rellena input_ids y 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óximos pasos

Vea todos los data collators disponibles para tareas comunes como la clasificación de tokens.

Update on GitHub