Call prepare_seq2seq_batch.
(self, batch)
| 259 | return {"tgt_texts": tgt_line, "src_texts": source_line, "id": index - 1} |
| 260 | |
| 261 | def collate_fn(self, batch) -> Dict[str, torch.Tensor]: |
| 262 | """Call prepare_seq2seq_batch.""" |
| 263 | batch_encoding: Dict[str, torch.Tensor] = self.tokenizer.prepare_seq2seq_batch( |
| 264 | [x["src_texts"] for x in batch], |
| 265 | tgt_texts=[x["tgt_texts"] for x in batch], |
| 266 | max_length=self.max_source_length, |
| 267 | max_target_length=self.max_target_length, |
| 268 | return_tensors="pt", |
| 269 | **self.dataset_kwargs, |
| 270 | ).data |
| 271 | batch_encoding["ids"] = torch.tensor([x["id"] for x in batch]) |
| 272 | return batch_encoding |
| 273 | |
| 274 | |
| 275 | class Seq2SeqDataCollator: |
nothing calls this directly
no outgoing calls
no test coverage detected