데이터 콜레이터

Data collator는 개별 데이터셋 샘플을 모델용 배치로 조립합니다. 또한 각 배치 내 가장 긴 시퀀스 길이에 맞춰 샘플을 동적으로 패딩할 수 있으며, 이는 전역 최대 길이에 패딩하는 것보다 효율적입니다.

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를 사용합니다.
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를 분리합니다.
입력 쌍을 단일 리스트로 연결합니다.
collator가 원시 토큰 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

다음 단계

토큰 분류와 같은 일반적인 작업을 위한 모든 사용 가능한 data collators를 확인하세요.

Update on GitHub