Trình tổng hợp dữ liệu

Data collator tập hợp các mẫu dữ liệu riêng lẻ thành một batch cho mô hình. Nó cũng có thể đệm động các mẫu đến độ dài chuỗi dài nhất trong mỗi batch, hiệu quả hơn so với việc đệm đến độ dài tối đa toàn cục.

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 cung cấp các data collator cho nhiều tác vụ khác nhau (xem tất cả các data collator có sẵn). Tạo một data collator tùy chỉnh với:

DataCollatorWithPadding khi bạn cần padding dựa trên tokenizer chuẩn cùng với các trường bổ sung.
DataCollatorMixin khi bạn cần logic padding tùy chỉnh, nhiều đầu vào ghép cặp cho mỗi mẫu, hoặc cấu trúc batch mà tokenizer không thể tự tạo ra.

DataCollatorWithPadding

Đối với các trường hợp sử dụng đơn giản như thêm một trường bổ sung, hãy kế thừa DataCollatorWithPadding và mở rộng phương thức __call__ của nó. Ví dụ bên dưới thêm một trường "score".

Đầu tiên hãy xóa trường tùy chỉnh vì pad() không nhận diện được nó.
Gọi lớp cha để xử lý input_ids và attention_mask.
Thêm trường "score" trở lại vào batch.

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

Truyền data collator tùy chỉnh vào Trainer như bất kỳ data collator nào khác.

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

DataCollatorMixin

Kế thừa DataCollatorMixin để kiểm soát hoàn toàn quá trình lắp ráp batch và triển khai phương thức __call__ của riêng bạn. Xây dựng logic padding tùy chỉnh, xử lý nhiều loại đầu vào, hoặc tạo ra các cấu trúc batch hoàn toàn mới. Ví dụ DataCollatorForPreference bên dưới sử dụng DataCollatorMixin vì mỗi mẫu huấn luyện có một phản hồi được chọn và một phản hồi bị từ chối, và mô hình cần nhìn thấy cả hai.

Tách chosen_ids và rejected_ids vì pad yêu cầu các danh sách phẳng.
Nối cặp đầu vào thành một danh sách duy nhất.
Tạo attention_mask bằng torch.ones_like thay vì dùng tokenizer vì collator hoạt động với các danh sách token ID thô.
Đệm input_ids và 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

Các bước tiếp theo

Xem tất cả các data collator có sẵn cho các tác vụ phổ biến như phân loại token.

Update on GitHub