| import pytorch_lightning as pl |
| import torch |
| from transformers.optimization import AdamW |
| import torchmetrics |
|
|
|
|
| class DualEncoderModule(pl.LightningModule): |
|
|
| def __init__(self, tokenizer, model, learning_rate=1e-3): |
| super().__init__() |
| self.tokenizer = tokenizer |
| self.model = model |
| self.learning_rate = learning_rate |
|
|
| self.train_acc = torchmetrics.Accuracy( |
| task="multiclass", num_classes=model.num_labels |
| ) |
| self.val_acc = torchmetrics.Accuracy( |
| task="multiclass", num_classes=model.num_labels |
| ) |
| self.test_acc = torchmetrics.Accuracy( |
| task="multiclass", num_classes=model.num_labels |
| ) |
|
|
| def forward(self, input_ids, **kwargs): |
| return self.model(input_ids, **kwargs) |
|
|
| def configure_optimizers(self): |
| optimizer = AdamW(self.parameters(), lr=self.learning_rate) |
| return optimizer |
|
|
| def training_step(self, batch, batch_idx): |
| pos_ids, pos_mask, neg_ids, neg_mask = batch |
|
|
| neg_ids = neg_ids.view(-1, neg_ids.shape[-1]) |
| neg_mask = neg_mask.view(-1, neg_mask.shape[-1]) |
|
|
| pos_outputs = self( |
| pos_ids, |
| attention_mask=pos_mask, |
| labels=torch.ones(pos_ids.shape[0], dtype=torch.uint8).to( |
| pos_ids.get_device() |
| ), |
| ) |
| neg_outputs = self( |
| neg_ids, |
| attention_mask=neg_mask, |
| labels=torch.zeros(neg_ids.shape[0], dtype=torch.uint8).to( |
| neg_ids.get_device() |
| ), |
| ) |
|
|
| loss_scale = 1.0 |
| loss = pos_outputs.loss + loss_scale * neg_outputs.loss |
|
|
| pos_logits = pos_outputs.logits |
| pos_preds = torch.argmax(pos_logits, axis=1) |
| self.train_acc( |
| pos_preds.cpu(), torch.ones(pos_ids.shape[0], dtype=torch.uint8).cpu() |
| ) |
|
|
| neg_logits = neg_outputs.logits |
| neg_preds = torch.argmax(neg_logits, axis=1) |
| self.train_acc( |
| neg_preds.cpu(), torch.zeros(neg_ids.shape[0], dtype=torch.uint8).cpu() |
| ) |
|
|
| return {"loss": loss} |
|
|
| def validation_step(self, batch, batch_idx): |
| pos_ids, pos_mask, neg_ids, neg_mask = batch |
|
|
| neg_ids = neg_ids.view(-1, neg_ids.shape[-1]) |
| neg_mask = neg_mask.view(-1, neg_mask.shape[-1]) |
|
|
| pos_outputs = self( |
| pos_ids, |
| attention_mask=pos_mask, |
| labels=torch.ones(pos_ids.shape[0], dtype=torch.uint8).to( |
| pos_ids.get_device() |
| ), |
| ) |
| neg_outputs = self( |
| neg_ids, |
| attention_mask=neg_mask, |
| labels=torch.zeros(neg_ids.shape[0], dtype=torch.uint8).to( |
| neg_ids.get_device() |
| ), |
| ) |
|
|
| loss_scale = 1.0 |
| loss = pos_outputs.loss + loss_scale * neg_outputs.loss |
|
|
| pos_logits = pos_outputs.logits |
| pos_preds = torch.argmax(pos_logits, axis=1) |
| self.val_acc( |
| pos_preds.cpu(), torch.ones(pos_ids.shape[0], dtype=torch.uint8).cpu() |
| ) |
|
|
| neg_logits = neg_outputs.logits |
| neg_preds = torch.argmax(neg_logits, axis=1) |
| self.val_acc( |
| neg_preds.cpu(), torch.zeros(neg_ids.shape[0], dtype=torch.uint8).cpu() |
| ) |
|
|
| self.log("val_acc", self.val_acc) |
|
|
| return {"loss": loss} |
|
|
| def test_step(self, batch, batch_idx): |
| pos_ids, pos_mask, neg_ids, neg_mask = batch |
|
|
| neg_ids = neg_ids.view(-1, neg_ids.shape[-1]) |
| neg_mask = neg_mask.view(-1, neg_mask.shape[-1]) |
|
|
| pos_outputs = self( |
| pos_ids, |
| attention_mask=pos_mask, |
| labels=torch.ones(pos_ids.shape[0], dtype=torch.uint8).to( |
| pos_ids.get_device() |
| ), |
| ) |
| neg_outputs = self( |
| neg_ids, |
| attention_mask=neg_mask, |
| labels=torch.zeros(neg_ids.shape[0], dtype=torch.uint8).to( |
| neg_ids.get_device() |
| ), |
| ) |
|
|
| pos_logits = pos_outputs.logits |
| pos_preds = torch.argmax(pos_logits, axis=1) |
| self.test_acc( |
| pos_preds.cpu(), torch.ones(pos_ids.shape[0], dtype=torch.uint8).cpu() |
| ) |
|
|
| neg_logits = neg_outputs.logits |
| neg_preds = torch.argmax(neg_logits, axis=1) |
| self.test_acc( |
| neg_preds.cpu(), torch.zeros(neg_ids.shape[0], dtype=torch.uint8).cpu() |
| ) |
|
|
| self.log("test_acc", self.test_acc) |
|
|
