| 273 | |
| 274 | |
| 275 | class Seq2SeqDataCollator: |
| 276 | def __init__(self, tokenizer, data_args, tpu_num_cores=None): |
| 277 | self.tokenizer = tokenizer |
| 278 | self.pad_token_id = tokenizer.pad_token_id |
| 279 | assert ( |
| 280 | self.pad_token_id is not None |
| 281 | ), f"pad_token_id is not defined for ({self.tokenizer.__class__.__name__}), it must be defined." |
| 282 | self.data_args = data_args |
| 283 | self.tpu_num_cores = tpu_num_cores |
| 284 | self.dataset_kwargs = {"add_prefix_space": True} if isinstance(tokenizer, BartTokenizer) else {} |
| 285 | if data_args.src_lang is not None: |
| 286 | self.dataset_kwargs["src_lang"] = data_args.src_lang |
| 287 | if data_args.tgt_lang is not None: |
| 288 | self.dataset_kwargs["tgt_lang"] = data_args.tgt_lang |
| 289 | |
| 290 | def __call__(self, batch) -> Dict[str, torch.Tensor]: |
| 291 | if hasattr(self.tokenizer, "prepare_seq2seq_batch"): |
| 292 | batch = self._encode(batch) |
| 293 | input_ids, attention_mask, labels = ( |
| 294 | batch["input_ids"], |
| 295 | batch["attention_mask"], |
| 296 | batch["labels"], |
| 297 | ) |
| 298 | else: |
| 299 | input_ids = torch.stack([x["input_ids"] for x in batch]) |
| 300 | attention_mask = torch.stack([x["attention_mask"] for x in batch]) |
| 301 | labels = torch.stack([x["labels"] for x in batch]) |
| 302 | |
| 303 | labels = trim_batch(labels, self.pad_token_id) |
| 304 | input_ids, attention_mask = trim_batch(input_ids, self.pad_token_id, attention_mask=attention_mask) |
| 305 | |
| 306 | batch = { |
| 307 | "input_ids": input_ids, |
| 308 | "attention_mask": attention_mask, |
| 309 | "labels": labels, |
| 310 | } |
| 311 | return batch |
| 312 | |
| 313 | def _shift_right_t5(self, input_ids): |
| 314 | # shift inputs to the right |
| 315 | shifted_input_ids = input_ids.new_zeros(input_ids.shape) |
| 316 | shifted_input_ids[..., 1:] = input_ids[..., :-1].clone() |
| 317 | shifted_input_ids[..., 0] = self.pad_token_id |
| 318 | return shifted_input_ids |
| 319 | |
| 320 | def _encode(self, batch) -> Dict[str, torch.Tensor]: |
| 321 | batch_encoding = self.tokenizer.prepare_seq2seq_batch( |
| 322 | [x["src_texts"] for x in batch], |
| 323 | tgt_texts=[x["tgt_texts"] for x in batch], |
| 324 | max_length=self.data_args.max_source_length, |
| 325 | max_target_length=self.data_args.max_target_length, |
| 326 | padding="max_length" if self.tpu_num_cores is not None else "longest", # TPU hack |
| 327 | return_tensors="pt", |
| 328 | **self.dataset_kwargs, |
| 329 | ) |
| 330 | return batch_encoding.data |
| 331 | |
| 332 |
nothing calls this directly
no outgoing calls
no test coverage detected