Upload summarizer.py

summarizer.py

@@ -0,0 +1,74 @@
+import os
+import torch
+import pytorch_lightning as pl
+from torch import nn
+from transformers import AdamW
+from transformers import T5ForConditionalGeneration
+from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning.loggers import TensorBoardLogger
+
+class SummarizerModel(pl.LightningModule):
+    def __init__(self, model_name = None):
+        super().__init__()
+        self.model = T5ForConditionalGeneration.from_pretrained(model_name, return_dict = True)
+        
+    def forward(self, 
+                input_ids, 
+                attention_mask, 
+                decoder_attention_mask, 
+                labels = None):
+        output = self.model(
+            input_ids,
+            attention_mask = attention_mask,
+            labels = labels,
+            decoder_attention_mask = decoder_attention_mask
+        )
+        return output.loss, output.logits
+    
+    def training_step(self, batch, batch_idx):
+        input_ids = batch['text_input_ids']
+        attention_mask = batch['text_attention_mask']
+        labels = batch['labels']
+        decoder_attention_mask = batch['labels_attention_mask']
+
+        loss, outputs = self.forward(
+            input_ids = input_ids,
+            attention_mask = attention_mask,
+            decoder_attention_mask = decoder_attention_mask,
+            labels = labels
+        )
+        self.log("train_loss", loss, prog_bar = True, logger = True)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        input_ids = batch['text_input_ids']
+        attention_mask = batch['text_attention_mask']
+        labels = batch['labels']
+        decoder_attention_mask = batch['labels_attention_mask']
+
+        loss, outputs = self.forward(
+            input_ids = input_ids,
+            attention_mask = attention_mask,
+            decoder_attention_mask = decoder_attention_mask,
+            labels = labels
+        )
+        self.log("val_loss", loss, prog_bar = True, logger = True)
+        return loss
+    
+    def test_step(self, batch, batch_idx):
+        input_ids = batch['text_input_ids']
+        attention_mask = batch['text_attention_mask']
+        labels = batch['labels']
+        decoder_attention_mask = batch['labels_attention_mask']
+
+        loss, outputs = self.forward(
+            input_ids = input_ids,
+            attention_mask = attention_mask,
+            decoder_attention_mask = decoder_attention_mask,
+            labels = labels
+        )
+        self.log("test_loss", loss, prog_bar = True, logger = True)
+        return loss
+    
+    def configure_optimizers(self):
+        return AdamW(self.model.parameters(), lr = 0.0001)