| 12 | from torchmetrics.classification import F1Score |
| 13 | |
| 14 | class NaiveSeqClassModule(pl.LightningModule): |
| 15 | # Instantiate the model |
| 16 | def __init__(self, tokenizer, model, use_question_stance_approach=True, learning_rate=1e-3): |
| 17 | super().__init__() |
| 18 | self.tokenizer = tokenizer |
| 19 | self.model = model |
| 20 | self.learning_rate = learning_rate |
| 21 | |
| 22 | self.train_acc = torchmetrics.Accuracy() |
| 23 | self.val_acc = torchmetrics.Accuracy() |
| 24 | self.test_acc = torchmetrics.Accuracy() |
| 25 | |
| 26 | self.train_f1 = F1Score(num_classes=4, average="macro") |
| 27 | self.val_f1 = F1Score(num_classes=4, average=None) |
| 28 | self.test_f1 = F1Score(num_classes=4, average=None) |
| 29 | |
| 30 | self.use_question_stance_approach = use_question_stance_approach |
| 31 | |
| 32 | |
| 33 | # Do a forward pass through the model |
| 34 | def forward(self, input_ids, **kwargs): |
| 35 | return self.model(input_ids, **kwargs) |
| 36 | |
| 37 | def configure_optimizers(self): |
| 38 | optimizer = AdamW(self.parameters(), lr = self.learning_rate) |
| 39 | return optimizer |
| 40 | |
| 41 | def training_step(self, batch, batch_idx): |
| 42 | x, x_mask, y = batch |
| 43 | |
| 44 | outputs = self(x, attention_mask=x_mask, labels=y) |
| 45 | logits = outputs.logits |
| 46 | loss = outputs.loss |
| 47 | |
| 48 | #cross_entropy = torch.nn.CrossEntropyLoss() |
| 49 | #loss = cross_entropy(logits, y) |
| 50 | |
| 51 | preds = torch.argmax(logits, axis=1) |
| 52 | |
| 53 | self.train_acc(preds.cpu(), y.cpu()) |
| 54 | self.train_f1(preds.cpu(), y.cpu()) |
| 55 | |
| 56 | self.log("train_loss", loss) |
| 57 | |
| 58 | return {'loss': loss} |
| 59 | |
| 60 | def training_epoch_end(self, outs): |
| 61 | self.log('train_acc_epoch', self.train_acc) |
| 62 | self.log('train_f1_epoch', self.train_f1) |
| 63 | |
| 64 | def validation_step(self, batch, batch_idx): |
| 65 | x, x_mask, y = batch |
| 66 | |
| 67 | outputs = self(x, attention_mask=x_mask, labels=y) |
| 68 | logits = outputs.logits |
| 69 | loss = outputs.loss |
| 70 | |
| 71 | preds = torch.argmax(logits, axis=1) |