A dataset that calls prepare_seq2seq_batch.
| 248 | |
| 249 | |
| 250 | class Seq2SeqDataset(AbstractSeq2SeqDataset): |
| 251 | """A dataset that calls prepare_seq2seq_batch.""" |
| 252 | |
| 253 | def __getitem__(self, index) -> Dict[str, str]: |
| 254 | index = index + 1 # linecache starts at 1 |
| 255 | source_line = self.prefix + linecache.getline(str(self.src_file), index).rstrip("\n") |
| 256 | tgt_line = linecache.getline(str(self.tgt_file), index).rstrip("\n") |
| 257 | assert source_line, f"empty source line for index {index}" |
| 258 | assert tgt_line, f"empty tgt line for index {index}" |
| 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