(self, batch)
| 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 |
nothing calls this directly
no test coverage detected