Spaces:
Runtime error
Runtime error
import torch | |
import pandas as pd | |
from transformers import ( | |
AdamW, | |
T5ForConditionalGeneration, | |
T5TokenizerFast as T5Tokenizer, MT5Tokenizer, MT5ForConditionalGeneration, ByT5Tokenizer, | |
) | |
from torch.utils.data import Dataset, DataLoader | |
import pytorch_lightning as pl | |
from pytorch_lightning.loggers import MLFlowLogger, WandbLogger | |
from pytorch_lightning import Trainer | |
from pytorch_lightning.callbacks.early_stopping import EarlyStopping | |
from pytorch_lightning import LightningDataModule | |
from pytorch_lightning import LightningModule | |
from datasets import load_metric | |
from tqdm.auto import tqdm | |
# from dagshub.pytorch_lightning import DAGsHubLogger | |
torch.cuda.empty_cache() | |
pl.seed_everything(42) | |
class DataModule(Dataset): | |
""" | |
Data Module for pytorch | |
""" | |
def __init__( | |
self, | |
data: pd.DataFrame, | |
tokenizer: T5Tokenizer, | |
source_max_token_len: int = 512, | |
target_max_token_len: int = 512, | |
): | |
""" | |
:param data: | |
:param tokenizer: | |
:param source_max_token_len: | |
:param target_max_token_len: | |
""" | |
self.data = data | |
self.target_max_token_len = target_max_token_len | |
self.source_max_token_len = source_max_token_len | |
self.tokenizer = tokenizer | |
def __len__(self): | |
return len(self.data) | |
def __getitem__(self, index: int): | |
data_row = self.data.iloc[index] | |
input_encoding = self.tokenizer( | |
data_row["input_text"], | |
max_length=self.source_max_token_len, | |
padding="max_length", | |
truncation=True, | |
return_attention_mask=True, | |
add_special_tokens=True, | |
return_tensors="pt", | |
) | |
output_encoding = self.tokenizer( | |
data_row["output_text"], | |
max_length=self.target_max_token_len, | |
padding="max_length", | |
truncation=True, | |
return_attention_mask=True, | |
add_special_tokens=True, | |
return_tensors="pt", | |
) | |
labels = output_encoding["input_ids"] | |
labels[ | |
labels == 0 | |
] = -100 | |
return dict( | |
keywords=data_row["input_text"], | |
text=data_row["output_text"], | |
keywords_input_ids=input_encoding["input_ids"].flatten(), | |
keywords_attention_mask=input_encoding["attention_mask"].flatten(), | |
labels=labels.flatten(), | |
labels_attention_mask=output_encoding["attention_mask"].flatten(), | |
) | |
class PLDataModule(LightningDataModule): | |
def __init__( | |
self, | |
train_df: pd.DataFrame, | |
test_df: pd.DataFrame, | |
tokenizer: T5Tokenizer, | |
source_max_token_len: int = 512, | |
target_max_token_len: int = 512, | |
batch_size: int = 4, | |
split: float = 0.1, | |
num_workers: int = 2 | |
): | |
""" | |
:param data_df: | |
:param tokenizer: | |
:param source_max_token_len: | |
:param target_max_token_len: | |
:param batch_size: | |
:param split: | |
""" | |
super().__init__() | |
self.train_df = train_df | |
self.test_df = test_df | |
self.split = split | |
self.batch_size = batch_size | |
self.target_max_token_len = target_max_token_len | |
self.source_max_token_len = source_max_token_len | |
self.tokenizer = tokenizer | |
self.num_workers = num_workers | |
def setup(self, stage=None): | |
self.train_dataset = DataModule( | |
self.train_df, | |
self.tokenizer, | |
self.source_max_token_len, | |
self.target_max_token_len, | |
) | |
self.test_dataset = DataModule( | |
self.test_df, | |
self.tokenizer, | |
self.source_max_token_len, | |
self.target_max_token_len, | |
) | |
def train_dataloader(self): | |
""" training dataloader """ | |
return DataLoader( | |
self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers | |
) | |
def test_dataloader(self): | |
""" test dataloader """ | |
return DataLoader( | |
self.test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers | |
) | |
def val_dataloader(self): | |
""" validation dataloader """ | |
return DataLoader( | |
self.test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers | |
) | |
class LightningModel(LightningModule): | |
""" PyTorch Lightning Model class""" | |
def __init__(self, tokenizer, model, learning_rate, adam_epsilon, weight_decay, output: str = "outputs"): | |
""" | |
initiates a PyTorch Lightning Model | |
Args: | |
tokenizer : T5 tokenizer | |
model : T5 model | |
output (str, optional): output directory to save model checkpoints. Defaults to "outputs". | |
""" | |
super().__init__() | |
self.model = model | |
self.tokenizer = tokenizer | |
self.output = output | |
self.learning_rate = learning_rate | |
self.adam_epsilon = adam_epsilon | |
self.weight_decay = weight_decay | |
def forward(self, input_ids, attention_mask, decoder_attention_mask, labels=None): | |
""" forward step """ | |
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_size): | |
""" training step """ | |
input_ids = batch["keywords_input_ids"] | |
attention_mask = batch["keywords_attention_mask"] | |
labels = batch["labels"] | |
labels_attention_mask = batch["labels_attention_mask"] | |
loss, outputs = self( | |
input_ids=input_ids, | |
attention_mask=attention_mask, | |
decoder_attention_mask=labels_attention_mask, | |
labels=labels, | |
) | |
self.log("train_loss", loss, prog_bar=True, logger=True) | |
return loss | |
def validation_step(self, batch, batch_size): | |
""" validation step """ | |
input_ids = batch["keywords_input_ids"] | |
attention_mask = batch["keywords_attention_mask"] | |
labels = batch["labels"] | |
labels_attention_mask = batch["labels_attention_mask"] | |
loss, outputs = self( | |
input_ids=input_ids, | |
attention_mask=attention_mask, | |
decoder_attention_mask=labels_attention_mask, | |
labels=labels, | |
) | |
self.log("val_loss", loss, prog_bar=True, logger=True) | |
return loss | |
def test_step(self, batch, batch_size): | |
""" test step """ | |
input_ids = batch["keywords_input_ids"] | |
attention_mask = batch["keywords_attention_mask"] | |
labels = batch["labels"] | |
labels_attention_mask = batch["labels_attention_mask"] | |
loss, outputs = self( | |
input_ids=input_ids, | |
attention_mask=attention_mask, | |
decoder_attention_mask=labels_attention_mask, | |
labels=labels, | |
) | |
self.log("test_loss", loss, prog_bar=True, logger=True) | |
return loss | |
def configure_optimizers(self): | |
""" configure optimizers """ | |
model = self.model | |
no_decay = ["bias", "LayerNorm.weight"] | |
optimizer_grouped_parameters = [ | |
{ | |
"params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], | |
"weight_decay": self.weight_decay, | |
}, | |
{ | |
"params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], | |
"weight_decay": 0.0, | |
}, | |
] | |
optimizer = AdamW(optimizer_grouped_parameters, lr=self.learning_rate, eps=self.adam_epsilon) | |
self.opt = optimizer | |
return [optimizer] | |
class Summarization: | |
""" Custom Summarization class """ | |
def __init__(self) -> None: | |
""" initiates Summarization class """ | |
pass | |
def from_pretrained(self, model_type="t5", model_name="t5-base") -> None: | |
""" | |
loads T5/MT5 Model model for training/finetuning | |
Args: | |
model_name (str, optional): exact model architecture name, "t5-base" or "t5-large". Defaults to "t5-base". | |
:param model_type: | |
""" | |
if model_type == "t5": | |
self.tokenizer = T5Tokenizer.from_pretrained(f"{model_name}") | |
self.model = T5ForConditionalGeneration.from_pretrained( | |
f"{model_name}", return_dict=True | |
) | |
elif model_type == "mt5": | |
self.tokenizer = MT5Tokenizer.from_pretrained(f"{model_name}") | |
self.model = MT5ForConditionalGeneration.from_pretrained( | |
f"{model_name}", return_dict=True | |
) | |
elif model_type == "byt5": | |
self.tokenizer = ByT5Tokenizer.from_pretrained(f"{model_name}") | |
self.model = T5ForConditionalGeneration.from_pretrained( | |
f"{model_name}", return_dict=True | |
) | |
def train( | |
self, | |
train_df: pd.DataFrame, | |
eval_df: pd.DataFrame, | |
source_max_token_len: int = 512, | |
target_max_token_len: int = 512, | |
batch_size: int = 8, | |
max_epochs: int = 5, | |
use_gpu: bool = True, | |
outputdir: str = "models", | |
early_stopping_patience_epochs: int = 0, # 0 to disable early stopping feature | |
learning_rate: float = 0.0001, | |
adam_epsilon: float = 0.01, | |
num_workers: int = 2, | |
weight_decay: float = 0.0001 | |
): | |
""" | |
trains T5/MT5 model on custom dataset | |
Args: | |
train_df (pd.DataFrame): training datarame. Dataframe must have 2 column --> "input_text" and "output_text" | |
eval_df ([type], optional): validation datarame. Dataframe must have 2 column --> "input_text" and | |
"output_text" | |
source_max_token_len (int, optional): max token length of source text. Defaults to 512. | |
target_max_token_len (int, optional): max token length of target text. Defaults to 512. | |
batch_size (int, optional): batch size. Defaults to 8. | |
max_epochs (int, optional): max number of epochs. Defaults to 5. | |
use_gpu (bool, optional): if True, model uses gpu for training. Defaults to True. | |
outputdir (str, optional): output directory to save model checkpoints. Defaults to "outputs". | |
early_stopping_patience_epochs (int, optional): monitors val_loss on epoch end and stops training, | |
if val_loss does not improve after the specied number of epochs. set 0 to disable early stopping. | |
Defaults to 0 (disabled) | |
:param learning_rate: | |
:param adam_epsilon: | |
""" | |
self.target_max_token_len = target_max_token_len | |
self.data_module = PLDataModule( | |
train_df, | |
eval_df, | |
self.tokenizer, | |
batch_size=batch_size, | |
source_max_token_len=source_max_token_len, | |
target_max_token_len=target_max_token_len, | |
num_workers=num_workers, | |
) | |
self.T5Model = LightningModel( | |
tokenizer=self.tokenizer, model=self.model, output=outputdir, | |
learning_rate=learning_rate, adam_epsilon=adam_epsilon, weight_decay=weight_decay | |
) | |
# MLlogger = MLFlowLogger(experiment_name="Summarization", | |
# tracking_uri="https://dagshub.com/gagan3012/summarization.mlflow")\ | |
WandLogger = WandbLogger(project="summarization-dagshub") | |
# logger = DAGsHubLogger(metrics_path='reports/training_metrics.txt') | |
early_stop_callback = ( | |
[ | |
EarlyStopping( | |
monitor="val_loss", | |
min_delta=0.00, | |
patience=early_stopping_patience_epochs, | |
verbose=True, | |
mode="min", | |
) | |
] | |
if early_stopping_patience_epochs > 0 | |
else None | |
) | |
gpus = -1 if use_gpu and torch.cuda.is_available() else 0 | |
trainer = Trainer( | |
logger=[WandLogger], | |
callbacks=early_stop_callback, | |
max_epochs=max_epochs, | |
gpus=gpus, | |
progress_bar_refresh_rate=5, | |
) | |
trainer.fit(self.T5Model, self.data_module) | |
def load_model( | |
self, model_type: str = 't5', model_dir: str = "models", use_gpu: bool = False | |
): | |
""" | |
loads a checkpoint for inferencing/prediction | |
Args: | |
model_type (str, optional): "t5" or "mt5". Defaults to "t5". | |
model_dir (str, optional): path to model directory. Defaults to "outputs". | |
use_gpu (bool, optional): if True, model uses gpu for inferencing/prediction. Defaults to True. | |
""" | |
if model_type == "t5": | |
self.tokenizer = T5Tokenizer.from_pretrained(f"{model_dir}") | |
self.model = T5ForConditionalGeneration.from_pretrained( | |
f"{model_dir}", return_dict=True | |
) | |
elif model_type == "mt5": | |
self.tokenizer = MT5Tokenizer.from_pretrained(f"{model_dir}") | |
self.model = MT5ForConditionalGeneration.from_pretrained( | |
f"{model_dir}", return_dict=True | |
) | |
elif model_type == "byt5": | |
self.tokenizer = ByT5Tokenizer.from_pretrained(f"{model_dir}") | |
self.model = T5ForConditionalGeneration.from_pretrained( | |
f"{model_dir}", return_dict=True | |
) | |
if use_gpu: | |
if torch.cuda.is_available(): | |
self.device = torch.device("cuda") | |
else: | |
raise Exception("exception ---> no gpu found. set use_gpu=False, to use CPU") | |
else: | |
self.device = torch.device("cpu") | |
self.model = self.model.to(self.device) | |
def save_model( | |
self, | |
model_dir="../../models" | |
): | |
""" | |
Save model to dir | |
:param model_dir: | |
:return: model is saved | |
""" | |
path = f"{model_dir}" | |
self.tokenizer.save_pretrained(path) | |
self.model.save_pretrained(path) | |
def predict( | |
self, | |
source_text: str, | |
max_length: int = 512, | |
num_return_sequences: int = 1, | |
num_beams: int = 2, | |
top_k: int = 50, | |
top_p: float = 0.95, | |
do_sample: bool = True, | |
repetition_penalty: float = 2.5, | |
length_penalty: float = 1.0, | |
early_stopping: bool = True, | |
skip_special_tokens: bool = True, | |
clean_up_tokenization_spaces: bool = True, | |
): | |
""" | |
generates prediction for T5/MT5 model | |
Args: | |
source_text (str): any text for generating predictions | |
max_length (int, optional): max token length of prediction. Defaults to 512. | |
num_return_sequences (int, optional): number of predictions to be returned. Defaults to 1. | |
num_beams (int, optional): number of beams. Defaults to 2. | |
top_k (int, optional): Defaults to 50. | |
top_p (float, optional): Defaults to 0.95. | |
do_sample (bool, optional): Defaults to True. | |
repetition_penalty (float, optional): Defaults to 2.5. | |
length_penalty (float, optional): Defaults to 1.0. | |
early_stopping (bool, optional): Defaults to True. | |
skip_special_tokens (bool, optional): Defaults to True. | |
clean_up_tokenization_spaces (bool, optional): Defaults to True. | |
Returns: | |
list[str]: returns predictions | |
""" | |
input_ids = self.tokenizer.encode( | |
source_text, return_tensors="pt", add_special_tokens=True | |
) | |
input_ids = input_ids.to(self.device) | |
generated_ids = self.model.generate( | |
input_ids=input_ids, | |
num_beams=num_beams, | |
max_length=max_length, | |
repetition_penalty=repetition_penalty, | |
length_penalty=length_penalty, | |
early_stopping=early_stopping, | |
top_p=top_p, | |
top_k=top_k, | |
num_return_sequences=num_return_sequences, | |
) | |
preds = self.tokenizer.decode( | |
generated_ids[0], | |
skip_special_tokens=skip_special_tokens, | |
clean_up_tokenization_spaces=clean_up_tokenization_spaces, | |
) | |
return preds | |
def evaluate( | |
self, | |
test_df: pd.DataFrame, | |
metrics: str = "rouge" | |
): | |
metric = load_metric(metrics) | |
input_text = test_df['input_text'] | |
references = test_df['output_text'] | |
references = references.to_list() | |
predictions = [self.predict(x) for x in tqdm(input_text)] | |
results = metric.compute(predictions=predictions, references=references) | |
output = { | |
'Rouge 1': { | |
'Rouge_1 Low Precision': results["rouge1"].low.precision, | |
'Rouge_1 Low recall': results["rouge1"].low.recall, | |
'Rouge_1 Low F1': results["rouge1"].low.fmeasure, | |
'Rouge_1 Mid Precision': results["rouge1"].mid.precision, | |
'Rouge_1 Mid recall': results["rouge1"].mid.recall, | |
'Rouge_1 Mid F1': results["rouge1"].mid.fmeasure, | |
'Rouge_1 High Precision': results["rouge1"].high.precision, | |
'Rouge_1 High recall': results["rouge1"].high.recall, | |
'Rouge_1 High F1': results["rouge1"].high.fmeasure, | |
}, | |
'Rouge 2': { | |
'Rouge_2 Low Precision': results["rouge2"].low.precision, | |
'Rouge_2 Low recall': results["rouge2"].low.recall, | |
'Rouge_2 Low F1': results["rouge2"].low.fmeasure, | |
'Rouge_2 Mid Precision': results["rouge2"].mid.precision, | |
'Rouge_2 Mid recall': results["rouge2"].mid.recall, | |
'Rouge_2 Mid F1': results["rouge2"].mid.fmeasure, | |
'Rouge_2 High Precision': results["rouge2"].high.precision, | |
'Rouge_2 High recall': results["rouge2"].high.recall, | |
'Rouge_2 High F1': results["rouge2"].high.fmeasure, | |
}, | |
'Rouge L': { | |
'Rouge_L Low Precision': results["rougeL"].low.precision, | |
'Rouge_L Low recall': results["rougeL"].low.recall, | |
'Rouge_L Low F1': results["rougeL"].low.fmeasure, | |
'Rouge_L Mid Precision': results["rougeL"].mid.precision, | |
'Rouge_L Mid recall': results["rougeL"].mid.recall, | |
'Rouge_L Mid F1': results["rougeL"].mid.fmeasure, | |
'Rouge_L High Precision': results["rougeL"].high.precision, | |
'Rouge_L High recall': results["rougeL"].high.recall, | |
'Rouge_L High F1': results["rougeL"].high.fmeasure, | |
}, | |
'rougeLsum': { | |
'rougeLsum Low Precision': results["rougeLsum"].low.precision, | |
'rougeLsum Low recall': results["rougeLsum"].low.recall, | |
'rougeLsum Low F1': results["rougeLsum"].low.fmeasure, | |
'rougeLsum Mid Precision': results["rougeLsum"].mid.precision, | |
'rougeLsum Mid recall': results["rougeLsum"].mid.recall, | |
'rougeLsum Mid F1': results["rougeLsum"].mid.fmeasure, | |
'rougeLsum High Precision': results["rougeLsum"].high.precision, | |
'rougeLsum High recall': results["rougeLsum"].high.recall, | |
'rougeLsum High F1': results["rougeLsum"].high.fmeasure, | |
} | |
} | |
return output | |