(self, batch: dict)
| 118 | return lmap(str.strip, gen_text) |
| 119 | |
| 120 | def _step(self, batch: dict) -> Tuple: |
| 121 | pad_token_id = self.tokenizer.pad_token_id |
| 122 | source_ids, source_mask, y = batch["input_ids"], batch["attention_mask"], batch["decoder_input_ids"] |
| 123 | y_ids = y[:, :-1].contiguous() |
| 124 | lm_labels = y[:, 1:].clone() |
| 125 | lm_labels[y[:, 1:] == pad_token_id] = -100 |
| 126 | outputs = self(source_ids, attention_mask=source_mask, decoder_input_ids=y_ids, labels=lm_labels,) |
| 127 | loss = outputs[0] |
| 128 | return (loss,) |
| 129 | |
| 130 | def training_step(self, batch, batch_idx) -> Dict: |
| 131 | loss_tensors = self._step(batch) |
no outgoing calls
no test coverage detected