Collator data

Data collator merangkai sampel dataset individual menjadi batch untuk model. Ia juga dapat melakukan padding secara dinamis pada sampel hingga panjang urutan terpanjang di setiap batch, yang lebih efisien daripada melakukan padding hingga panjang maksimum 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 menyediakan data collator untuk berbagai tugas (lihat semua data collator yang tersedia). Buat data collator kustom dengan:

DataCollatorWithPadding ketika Anda memerlukan padding berbasis tokenizer standar ditambah bidang tambahan.
DataCollatorMixin ketika Anda memerlukan logika padding kustom, beberapa input berpasangan per sampel, atau struktur batch yang tidak dapat dihasilkan oleh tokenizer secara mandiri.

DataCollatorWithPadding

Untuk kasus penggunaan sederhana seperti menambahkan bidang tambahan, turunkan DataCollatorWithPadding dan perluas metode __call__-nya. Contoh di bawah ini menambahkan bidang "score".

Hapus field kustom terlebih dahulu karena pad() tidak mengenalinya.
Panggil kelas induk untuk menangani input_ids dan attention_mask.
Tambahkan kembali field "score" ke 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

Lewatkan data collator kustom ke Trainer seperti data collator lainnya.

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

DataCollatorMixin

Subclass DataCollatorMixin untuk kontrol penuh atas perakitan batch dan implementasikan metode __call__ Anda sendiri. Bangun logika padding kustom, tangani berbagai jenis input, atau buat struktur batch yang sepenuhnya baru. Contoh DataCollatorForPreference di bawah menggunakan DataCollatorMixin karena setiap sampel pelatihan memiliki respons yang dipilih dan ditolak, dan model perlu melihat keduanya.

Pisahkan chosen_ids dan rejected_ids karena pad mengharapkan daftar datar.
Gabungkan pasangan input menjadi satu daftar.
Buat attention_mask dengan torch.ones_like alih-alih tokenizer karena collator bekerja dengan daftar ID token mentah.
Pad input_ids dan 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

Langkah selanjutnya

Lihat semua data collator yang tersedia untuk tugas umum seperti klasifikasi token.

Update on GitHub