データコラテーター

データコラレーターは、モデル用のバッチに個別のデータセットサンプルをまとめる役割を担います。また、各バッチ内の最長シーケンスに合わせてサンプルを動的にパディングすることもでき、グローバルな最大長でパディングするよりも効率的です。

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を使用します。

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

カスタムデータコラテーターを、他のデータコラテーターと同様にTrainerに渡します。

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

DataCollatorMixin

DataCollatorMixin をサブクラス化してバッチの組み立てを完全に制御し、独自の __call__ メソッドを実装します。カスタムのパディングロジックを構築したり、複数の入力タイプを処理したり、まったく新しいバッチ構造を作成したりできます。以下の DataCollatorForPreference の例では、各トレーニングサンプルに選択された応答と拒否された応答の両方が含まれ、モデルが両方を参照する必要があるため、DataCollatorMixin が使用されています。

pad はフラットなリストを期待するため、chosen_ids と rejected_ids を分離します。
入力ペアを単一のリストに連結します。
コラテーターは生のトークン ID リストで動作するため、トークナイザーではなく torch.ones_like を使用して attention_mask を生成します。
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

次のステップ

トークン分類などの一般的なタスクで利用可能なすべてのデータコラテーターを確認してください。

Update on GitHub