collator.py (1431B)
1 """Module for batching socio-political datasets.""" 2 3 from __future__ import annotations 4 5 from typing import TYPE_CHECKING, Any 6 7 import transformers 8 9 10 if TYPE_CHECKING: 11 import t5.dataset 12 import t5.model 13 14 15 class DataCollator: 16 """Group socio-political event samples into a batch.""" 17 18 # We need a custom collator in order to keep all features in the batch, including some that cannot be batched automatically. 19 def __init__(self: DataCollator, dataset: t5.dataset.Dataset, model: t5.model.Model) -> None: 20 """Initialize socio-political event data collator.""" 21 self.dataset = dataset 22 self.model = model 23 self.text_collator = transformers.data.DataCollatorForSeq2Seq(tokenizer=self.dataset.tokenizer, model=self.model) 24 self.text_keys: list[str] = ["input_ids", "attention_mask", "labels"] 25 26 def __call__(self: DataCollator, samples: list[dict[str, Any]]) -> dict[str, Any]: 27 """Collate a list of samples into a batch.""" 28 text_inputs: list[dict[str, Any]] = [{key: sample[key] for key in self.text_keys} for sample in samples] 29 batch: dict[str, Any] = {key: [] for key in samples[0] if key not in self.text_keys} 30 batch.update(self.text_collator(text_inputs)) 31 for sample in samples: 32 for key, value in sample.items(): 33 if key not in self.text_keys: 34 batch[key].append(value) 35 return batch