| 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 | |
| 333 | class SortishSampler(Sampler): |