File size: 3,548 Bytes
a535e2c 206fa83 37ac195 a535e2c 37ac195 |
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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
import pytorch_lightning as pl
from transformers import AutoModelForSeq2SeqLM
import torch
class T5(pl.LightningModule):
def __init__(self, lr=5e-5, num_train_epochs=15, warmup_steps=1000):
super().__init__()
self.model = AutoModelForSeq2SeqLM.from_pretrained("t5-small")
self.train_losses=[]
self.validation_losses=[]
self.train_losses_epoch=[]
self.validation_losses_epoch=[]
self.save_hyperparameters()
def forward(self, input_ids, attention_mask, labels=None):
outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
return outputs
def common_step(self, batch, batch_idx):
outputs = self(**batch)
loss = outputs.loss
return loss
def training_step(self, batch, batch_idx):
loss = self.common_step(batch, batch_idx)
# logs metrics for each training_step,
# and the average across the epoch
self.log("training_loss", loss)
self.train_losses.append(loss)
return loss
def validation_step(self, batch, batch_idx):
loss = self.common_step(batch, batch_idx)
self.log("validation_loss", loss, on_epoch=True)
self.validation_losses.append(loss)
return loss
return loss
def on_train_epoch_end(self):
# Calculate average loss for the epoch and append to the list
avg_train_loss = sum(self.train_losses)/ len(self.train_losses)
self.train_losses_epoch.append(avg_train_loss.item())
# Reset epoch loss accumulator
self.train_losses = []
def on_validation_epoch_end(self):
# Calculate average loss for the epoch and append to the list
avg_val_loss = sum(self.validation_losses) / len(self.validation_losses)
self.validation_losses_epoch.append(avg_val_loss.item())
# Reset epoch loss accumulator
self.validation_losses = []
# Reset epoch loss accumulator
self.test_losses = []
def configure_optimizers(self):
# create optimizer
optimizer = AdamW(self.model.parameters(), lr=self.hparams.lr)
# create learning rate scheduler
num_train_optimization_steps = self.hparams.num_train_epochs * len(train_dataloader)
lr_scheduler = {'scheduler': get_linear_schedule_with_warmup(optimizer,
num_warmup_steps=self.hparams.warmup_steps,
num_training_steps=num_train_optimization_steps),
'name': 'learning_rate',
'interval':'step',
'frequency': 1}
return {"optimizer": optimizer, "lr_scheduler": lr_scheduler}
def generate(self, input_ids, max_new_tokens=30, device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')):
input_ids=input_ids.clone().detach().reshape((1,-1)).to(device)
return self.model.generate(input_ids)
def push_to_hub(self, model_name, organization):
# Save the model
self.model.push_to_hub(model_name, organization)
def from_pretrained(self, model_path):
AutoModelForSeq2SeqLM.from_pretrained(model_path)
def train_dataloader(self):
return train_dataloader
def val_dataloader(self):
return valid_dataloader
def test_dataloader(self):
return test_dataloader
|