(self, batch, batch_idx)
| 32 | return optimizer |
| 33 | |
| 34 | def training_step(self, batch, batch_idx): |
| 35 | pos_ids, pos_mask, neg_ids, neg_mask = batch |
| 36 | |
| 37 | neg_ids = neg_ids.view(-1, neg_ids.shape[-1]) |
| 38 | neg_mask = neg_mask.view(-1, neg_mask.shape[-1]) |
| 39 | |
| 40 | pos_outputs = self(pos_ids, attention_mask=pos_mask, labels=torch.ones(pos_ids.shape[0], dtype=torch.uint8).to(pos_ids.get_device())) |
| 41 | neg_outputs = self(neg_ids, attention_mask=neg_mask, labels=torch.zeros(neg_ids.shape[0], dtype=torch.uint8).to(neg_ids.get_device())) |
| 42 | |
| 43 | loss_scale = 1.0 |
| 44 | loss = pos_outputs.loss + loss_scale * neg_outputs.loss |
| 45 | |
| 46 | pos_logits = pos_outputs.logits |
| 47 | pos_preds = torch.argmax(pos_logits, axis=1) |
| 48 | self.train_acc(pos_preds.cpu(), torch.ones(pos_ids.shape[0], dtype=torch.uint8).cpu()) |
| 49 | |
| 50 | neg_logits = neg_outputs.logits |
| 51 | neg_preds = torch.argmax(neg_logits, axis=1) |
| 52 | self.train_acc(neg_preds.cpu(), torch.zeros(neg_ids.shape[0], dtype=torch.uint8).cpu()) |
| 53 | |
| 54 | |
| 55 | return {'loss': loss} |
| 56 | |
| 57 | def validation_step(self, batch, batch_idx): |
| 58 | pos_ids, pos_mask, neg_ids, neg_mask = batch |
nothing calls this directly
no outgoing calls
no test coverage detected