|
|
|
|
|
import torch |
|
import torch.nn as nn |
|
import torch.optim as optim |
|
from torch.utils.data import DataLoader, random_split |
|
import pandas as pd |
|
from tqdm import tqdm |
|
import time |
|
from .Utilities import LanguageDataset |
|
|
|
class Seq2Seq(): |
|
""" |
|
Base class for Seq2Seq (text-generation models). This class will be inherited by wrappers of transformers like GPT2 |
|
and T5. |
|
|
|
Attributes: |
|
|
|
Methods: |
|
|
|
""" |
|
|
|
def __init__(self, gpu=0, max_length=0, model_path=None): |
|
|
|
|
|
if torch.cuda.is_available(): |
|
self.device = torch.device('cuda') |
|
else: |
|
try: |
|
self.device = torch.device('mps') |
|
except Exception: |
|
self.device = torch.device('cpu') |
|
|
|
|
|
self.gpu = gpu |
|
|
|
|
|
if model_path: self.model = torch.load(model_path).to(self.device) |
|
else: self.model = None |
|
self.model_name = "" |
|
self.tokenizer = None |
|
self.max_length = max_length |
|
|
|
|
|
self.train_loader = None |
|
self.valid_loader = None |
|
self.results = pd.DataFrame(columns=['epoch', 'model_arch', 'batch_size', 'gpu', 'training_loss', 'validation_loss', 'epoch_duration_sec']) |
|
|
|
def load_data(self, df, batch_size, train_ratio=0.8): |
|
self.batch_size = batch_size |
|
dataset = LanguageDataset(df, self.tokenizer) |
|
train_size = int(0.8*len(dataset)) |
|
valid_size = len(dataset) - train_size |
|
train_data, valid_data = random_split(dataset, [train_size, valid_size]) |
|
self.max_length = dataset.max_length |
|
self.train_loader = DataLoader(train_data, batch_size=self.batch_size, shuffle=True) |
|
self.valid_loader = DataLoader(valid_data, batch_size=self.batch_size) |
|
|
|
""" Return training results """ |
|
def summary(self): |
|
return self.results |
|
|
|
""" Save model to path """ |
|
def to_pt(self, path): |
|
torch.save(self.model, path) |
|
|
|
|
|
class GPT2(Seq2Seq): |
|
""" |
|
This is the GPT2 implementation of Seq2Seq. |
|
""" |
|
|
|
def __init__(self, gpu, model_name, batch_size=16): |
|
super().__init__(gpu, max_length=0) |
|
from transformers import GPT2Tokenizer, GPT2LMHeadModel |
|
self.model_name = model_name |
|
self.model = GPT2LMHeadModel.from_pretrained(self.model_name).to(self.device) |
|
self.tokenizer = GPT2Tokenizer.from_pretrained(self.model_name) |
|
self.tokenizer.pad_token = self.tokenizer.eos_token |
|
|
|
def train(self, num_epochs=3, train_ratio=0.8): |
|
criterion = nn.CrossEntropyLoss(ignore_index=self.tokenizer.pad_token_id) |
|
optimizer = optim.Adam(self.model.parameters(), lr=5e-4) |
|
|
|
|
|
|
|
results = pd.DataFrame(columns=['epoch', 'transformer', 'batch_size', 'gpu', |
|
'training_loss', 'validation_loss', 'epoch_duration_sec']) |
|
|
|
for epoch in range(num_epochs): |
|
start_time = time.time() |
|
|
|
|
|
|
|
self.model.train() |
|
epoch_training_loss = 0 |
|
train_iterator = tqdm(self.train_loader, |
|
desc=f"Training Epoch {epoch + 1}/{num_epochs} Batch Size: {self.batch_size}, Transformer: {self.model_name}") |
|
for batch in train_iterator: |
|
optimizer.zero_grad() |
|
inputs = batch['input_ids'].squeeze(1).to(self.device) |
|
targets = inputs.clone() |
|
outputs = self.model(input_ids=inputs, labels=targets) |
|
loss = outputs.loss |
|
loss.backward() |
|
optimizer.step() |
|
train_iterator.set_postfix({'Training Loss': loss.item()}) |
|
epoch_training_loss += loss.item() |
|
avg_epoch_training_loss = epoch_training_loss / len(train_iterator) |
|
|
|
|
|
|
|
self.model.eval() |
|
epoch_validation_loss = 0 |
|
total_loss = 0 |
|
valid_iterator = tqdm(self.valid_loader, desc=f"Validation Epoch {epoch + 1}/{num_epochs}") |
|
with torch.no_grad(): |
|
for batch in valid_iterator: |
|
inputs = batch['input_ids'].squeeze(1).to(self.device) |
|
targets = inputs.clone() |
|
outputs = self.model(input_ids=inputs, labels=targets) |
|
loss = outputs.loss |
|
total_loss += loss |
|
valid_iterator.set_postfix({'Validation Loss': loss.item()}) |
|
epoch_validation_loss += loss.item() |
|
|
|
avg_epoch_validation_loss = epoch_validation_loss / len(self.valid_loader) |
|
|
|
end_time = time.time() |
|
epoch_duration_sec = end_time - start_time |
|
|
|
new_row = {'transformer': self.model_name, |
|
'batch_size': self.batch_size, |
|
'gpu': self.gpu, |
|
'epoch': epoch + 1, |
|
'training_loss': avg_epoch_training_loss, |
|
'validation_loss': avg_epoch_validation_loss, |
|
'epoch_duration_sec': epoch_duration_sec} |
|
|
|
self.results.loc[len(self.results)] = new_row |
|
print(f"Epoch: {epoch + 1}, Validation Loss: {total_loss / len(self.valid_loader)}") |
|
|
|
def generate_text(self, input_str, top_k=16, top_p=0.95, temperature=1.0, repetition_penalty=1.2): |
|
|
|
input_ids= self.tokenizer.encode(input_str, return_tensors='pt').to(self.device) |
|
|
|
|
|
output = self.model.generate( |
|
input_ids, |
|
max_length=self.max_length, |
|
num_return_sequences=1, |
|
do_sample=True, |
|
top_k=top_k, |
|
top_p=top_p, |
|
temperature=temperature, |
|
repetition_penalty=repetition_penalty |
|
) |
|
|
|
|
|
decoded_output = self.tokenizer.decode(output[0], skip_special_tokens=True) |
|
return decoded_output |
|
|
|
class FlanT5(Seq2Seq): |
|
""" |
|
This is the T5 implementation of Seq2Seq - it is designed to support T5 models of various sizes. |
|
""" |
|
def __init__(self, gpu, model_name, batch_size=16): |
|
super().__init__(gpu, max_length=0) |
|
from transformers import T5ForConditionalGeneration, T5Tokenizer |
|
self.model_name = model_name |
|
self.model = T5ForConditionalGeneration.from_pretrained(self.model_name).to(self.device) |
|
self.tokenizer = T5Tokenizer.from_pretrained(self.model_name) |
|
self.tokenizer.pad_token = self.tokenizer.eos_token |
|
|
|
def train(self, num_epochs=3, train_ratio=0.8): |
|
criterion = nn.CrossEntropyLoss(ignore_index=self.tokenizer.pad_token_id) |
|
optimizer = optim.Adam(self.model.parameters(), lr=5e-4) |
|
|
|
|
|
|
|
self.results = pd.DataFrame(columns=['epoch', 'transformer', 'batch_size', 'gpu', |
|
'training_loss', 'validation_loss', 'epoch_duration_sec']) |
|
|
|
for epoch in range(num_epochs): |
|
start_time = time.time() |
|
|
|
|
|
|
|
self.model.train() |
|
epoch_training_loss = 0 |
|
train_iterator = tqdm(self.train_loader, |
|
desc=f"Training Epoch {epoch + 1}/{num_epochs} Batch Size: {self.batch_size}, Transformer: {self.model_name}") |
|
for batch in train_iterator: |
|
optimizer.zero_grad() |
|
inputs = batch['input_ids'].squeeze(1).to(self.device) |
|
targets = batch['labels'].squeeze(1).to(self.device) |
|
outputs = self.model(input_ids=inputs, labels=targets) |
|
loss = outputs.loss |
|
loss.backward() |
|
optimizer.step() |
|
train_iterator.set_postfix({'Training Loss': loss.item()}) |
|
epoch_training_loss += loss.item() |
|
avg_epoch_training_loss = epoch_training_loss / len(train_iterator) |
|
|
|
|
|
|
|
self.model.eval() |
|
epoch_validation_loss = 0 |
|
total_loss = 0 |
|
valid_iterator = tqdm(self.valid_loader, desc=f"Validation Epoch {epoch + 1}/{num_epochs}") |
|
with torch.no_grad(): |
|
for batch in valid_iterator: |
|
inputs = batch['input_ids'].squeeze(1).to(self.device) |
|
targets = batch['labels'].squeeze(1).to(self.device) |
|
outputs = self.model(input_ids=inputs, labels=targets) |
|
loss = outputs.loss |
|
total_loss += loss |
|
valid_iterator.set_postfix({'Validation Loss': loss.item()}) |
|
epoch_validation_loss += loss.item() |
|
|
|
avg_epoch_validation_loss = epoch_validation_loss / len(self.valid_loader) |
|
|
|
end_time = time.time() |
|
epoch_duration_sec = end_time - start_time |
|
|
|
new_row = {'transformer': self.model_name, |
|
'batch_size': self.batch_size, |
|
'gpu': self.gpu, |
|
'epoch': epoch + 1, |
|
'training_loss': avg_epoch_training_loss, |
|
'validation_loss': avg_epoch_validation_loss, |
|
'epoch_duration_sec': epoch_duration_sec} |
|
|
|
self.results.loc[len(self.results)] = new_row |
|
print(f"Epoch: {epoch + 1}, Validation Loss: {total_loss / len(self.valid_loader)}") |
|
|
|
def generate_text(self, input_str, top_k=16, top_p=0.95, temperature=1.0, repetition_penalty=1.2): |
|
|
|
input_ids = self.tokenizer.encode(input_str, return_tensors='pt', max_length=self.max_length, truncation=True).to(self.device) |
|
|
|
output_ids = self.model.generate(input_ids, |
|
max_length=self.max_length, |
|
do_sample=True, |
|
top_k=top_k, |
|
top_p=top_p, |
|
temperature=temperature, |
|
repetition_penalty=repetition_penalty) |
|
|
|
output_str = self.tokenizer.decode(output_ids[0], skip_special_tokens=True) |
|
return output_str |