데이터 콜레이터

데이터 콜레이터 (Data collators)

데이터 콜레이터(data collator)는 데이터셋의 개별 샘플들을 모델에 넣을 수 있는 배치(batch) 하나로 묶어주는 역할을 해요. 각 배치에서 가장 긴 시퀀스에 맞춰 샘플을 동적으로 패딩할 수도 있는데, 이는 전역(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는 다양한 작업을 위한 데이터 콜레이터를 제공해요 (사용 가능한 모든 data collators를 확인해 보세요). 커스텀 데이터 콜레이터는 다음 두 가지 방법으로 만들 수 있어요:

  • 표준적인 tokenizer 기반 패딩에 추가 필드까지 필요한 경우에는 DataCollatorWithPadding을 사용해요.
  • 커스텀 패딩 로직, 샘플당 여러 개의 쌍을 이루는 입력, 또는 tokenizer가 스스로는 만들어낼 수 없는 배치 구조가 필요하다면 DataCollatorMixin을 사용해요.

DataCollatorWithPadding

추가 필드를 하나 더하는 것 같은 간단한 용도라면 DataCollatorWithPadding을 서브클래싱하고 __call__ 메서드를 확장하면 돼요. 아래 예시는 "score" 필드를 추가하는 코드랍니다.

  1. 먼저 커스텀 필드를 제거해요. pad()는 그 필드를 인식하지 못하기 때문이에요.
  2. 부모 클래스를 호출해서 input_ids와 attention_mask를 처리해요.
  3. 배치에 "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 예시는 각 학습 샘플에 selected 응답과 rejected 응답이 있고 모델이 둘 다 봐야 하기 때문에 DataCollatorMixin을 사용하고 있어요.

  1. chosen_ids와 rejected_ids를 분리해요. pad는 평평한 리스트(flat list)를 기대하기 때문이에요.
  2. 입력 쌍을 하나의 리스트로 이어붙여요.
  3. tokenizer 대신 torch.ones_like로 attention_mask를 생성해요. 콜레이터가 원시 토큰 ID 리스트를 다루기 때문이에요.
  4. 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

더 알아보기 (Learn more)

  • 토큰 분류 같은 일반적인 작업에 사용할 수 있는 모든 data collators를 확인해 보세요.