(self, batch)
| 125 | return source_ids, source_mask, y |
| 126 | |
| 127 | def collate_fn(self, batch) -> dict: |
| 128 | input_ids = torch.stack([x["input_ids"] for x in batch]) |
| 129 | masks = torch.stack([x["attention_mask"] for x in batch]) |
| 130 | target_ids = torch.stack([x["decoder_input_ids"] for x in batch]) |
| 131 | pad_token_id = self.pad_token_id |
| 132 | y = trim_batch(target_ids, pad_token_id) |
| 133 | source_ids, source_mask = trim_batch(input_ids, pad_token_id, attention_mask=masks) |
| 134 | batch = {"input_ids": source_ids, "attention_mask": source_mask, "decoder_input_ids": y} |
| 135 | return batch |
| 136 | |
| 137 | @property |
| 138 | def src_lens(self): # Can delete? |
nothing calls this directly
no test coverage detected