Collateurs de données

Un collateur de données assemble des échantillons individuels du jeu de données en un lot pour le modèle. Il peut également effectuer un bourrage dynamique des échantillons jusqu'à la séquence la plus longue de chaque lot, ce qui est plus efficace que le bourrage jusqu'à une longueur maximale globale.

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 fournit des collateurs de données pour diverses tâches (voir tous les collateurs de données disponibles). Créez un collateur de données personnalisé avec :

• DataCollatorWithPadding lorsque vous avez besoin d’un padding standard basé sur le tokenizer ainsi que de champs supplémentaires.
• DataCollatorMixin lorsque vous avez besoin d’une logique de padding personnalisée, de plusieurs entrées appariées par échantillon ou d’une structure de lot que le tokenizer ne peut pas produire seul.

DataCollatorWithPadding

Pour des cas d’utilisation simples comme l’ajout d’un champ supplémentaire, créez une sous-classe de DataCollatorWithPadding et étendez sa méthode __call__. L’exemple ci-dessous ajoute un champ "score".

1. Supprimez d’abord le champ personnalisé, car pad() ne le reconnaît pas.
2. Appeler la classe parente pour gérer input_ids et attention_mask.
3. Ajoutez le champ "score" au lot.

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

Passez le collateur de données personnalisé à Trainer comme n'importe quel autre collateur de données.

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

DataCollatorMixin

Sous-classez DataCollatorMixin pour un contrôle total sur l’assemblage des lots et implémentez votre propre méthode __call__. Construisez une logique de bourrage personnalisée, gérez plusieurs types d’entrées ou créez de nouvelles structures de lots. L’exemple DataCollatorForPreference ci-dessous utilise DataCollatorMixin car chaque échantillon d’entraînement comporte une réponse choisie et une réponse rejetée, et le modèle doit voir les deux.

1. Séparez chosen_ids et rejected_ids car pad attend des listes plates.
2. Concaténer la paire d'entrées en une seule liste.
3. Générez attention_mask avec torch.ones_like plutôt qu’avec le tokenizer, car le collator fonctionne avec des listes brutes d’identifiants de jetons.
4. Effectuer le remplissage de input_ids et 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

Étapes suivantes

• Consultez tous les data collators disponibles pour des tâches courantes telles que la classification de jetons.