Коллаторы данных

Коллатор данных собирает отдельные примеры из набора данных в батч для модели. Он также может динамически дополнять примеры до длины самой длинной последовательности в каждом батче, что эффективнее, чем дополнять до глобальной максимальной длины.

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 предоставляет коллабораторы данных для различных задач (см. все доступные коллабораторы данных). Создайте собственный коллаборатор данных с помощью:

DataCollatorWithPadding, если вам нужен стандартный padding на основе токенизатора плюс дополнительные поля.
DataCollatorMixin — если вам нужна собственная логика дополнения (padding), несколько парных входных данных на один пример или структура батча, которую токенизатор не может сформировать самостоятельно.

DataCollatorWithPadding

Для простых сценариев, например, добавления дополнительного поля, создайте подкласс DataCollatorWithPadding и расширьте его метод __call__. В примере ниже добавляется поле "score".

Сначала удалите пользовательское поле, так как pad() его не распознаёт.
Вызовите родительский класс для обработки input_ids и attention_mask.
Добавьте поле "score" обратно в батч.

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

Передайте собственный data collator в Trainer так же, как любой другой data collator.

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

DataCollatorMixin

Наследуйте DataCollatorMixin, чтобы полностью контролировать формирование батчей, и реализуйте собственный метод __call__. Создавайте собственную логику отступов, обрабатывайте несколько типов входных данных или формируйте совершенно новые структуры батчей. Пример DataCollatorForPreference ниже использует DataCollatorMixin, так как каждая обучающая выборка содержит выбранный и отклонённый ответы, и модели необходимо видеть оба.

Разделите chosen_ids и rejected_ids, так как pad ожидает плоские списки.
Соедините входную пару в единый список.
Создавайте attention_mask с помощью torch.ones_like, а не токенизатора, так как коллабор работает со списками сырых ID токенов.
Выровняйте input_ids и 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

Дальнейшие шаги

См. все доступные data collators для типичных задач, таких как классификация токенов.

Update on GitHub