File size: 2,636 Bytes
662f2cf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 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) |