(self, batch)
| 233 | ) |
| 234 | |
| 235 | def collate_fn(self, batch) -> Dict[str, torch.Tensor]: |
| 236 | input_ids = torch.stack([x["input_ids"] for x in batch]) |
| 237 | masks = torch.stack([x["attention_mask"] for x in batch]) |
| 238 | target_ids = torch.stack([x["labels"] for x in batch]) |
| 239 | pad_token_id = self.pad_token_id |
| 240 | y = trim_batch(target_ids, pad_token_id) |
| 241 | source_ids, source_mask = trim_batch(input_ids, pad_token_id, attention_mask=masks) |
| 242 | batch = { |
| 243 | "input_ids": source_ids, |
| 244 | "attention_mask": source_mask, |
| 245 | "labels": y, |
| 246 | } |
| 247 | return batch |
| 248 | |
| 249 | |
| 250 | class Seq2SeqDataset(AbstractSeq2SeqDataset): |
nothing calls this directly
no test coverage detected