text_plain_output_1.png
from rich import box

from rich.console import Console

from rich.table import Column, Table

from torch import cuda

from torch.utils.data import Dataset, DataLoader, RandomSampler, SequentialSampler

from transformers import T5Tokenizer, T5ForConditionalGeneration

import numpy as np

import os

import os

import pandas as pd

import pandas as pd

import torch

import pandas as pd
path = '/kaggle/input/news-headlines/news_summary.csv'
df = pd.read_csv(path)

console = Console(record=True)
def display_df(df):
    """display dataframe in ASCII format"""
    console = Console()
    table = Table(Column('source_text', justify='center'), Column('target_text', justify='center'), title='Sample Data', pad_edge=False, box=box.ASCII)
    for i, row in enumerate(df.values.tolist()):
        table.add_row(row[0], row[1])
training_logger = Table(Column('Epoch', justify='center'), Column('Steps', justify='center'), Column('Loss', justify='center'), title='Training Status', pad_edge=False, box=box.ASCII)

from torch import cuda
device = 'cuda' if cuda.is_available() else 'cpu'
device

class CustomDataSetClass(Dataset):
    """
    Creating a custom dataset for reading the dataset and
    loading it into the dataloader to pass it to the
    transformer for finetuning the model
    """
    def __init__(self, dataframe, tokenizer, source_len, target_len, source_text, target_text):
        """
        Initializes a Dataset class
        Args:
            dataframe (pandas.DataFrame): Input dataframe
            tokenizer (transformers.tokenizer): Transformers tokenizer
            source_len (int): Max length of source text
            target_len (int): Max length of target text
            source_text (str): column name of source text
            target_text (str): column name of target text
        """
        self.tokenizer = tokenizer
        self.data = dataframe
        self.source_len = source_len
        self.summ_len = target_len
        self.target_text = self.data[target_text]
        self.source_text = self.data[source_text]
    def __len__(self):
        """returns the length of dataframe"""
        return len(self.target_text)
    def __getitem__(self, index):
        """return the input ids, attention masks and target ids"""
        source_text = str(self.source_text[index])
        target_text = str(self.target_text[index])
        source_text = ' '.join(source_text.split())
        target_text = ' '.join(target_text.split())
        source = self.tokenizer.batch_encode_plus([source_text], max_length=self.source_len, pad_to_max_length=True, truncation=True, padding='max_length', return_tensors='pt')
        target = self.tokenizer.batch_encode_plus([target_text], max_length=self.summ_len, pad_to_max_length=True, truncation=True, padding='max_length', return_tensors='pt')
        source_ids = source['input_ids'].squeeze()
        source_mask = source['attention_mask'].squeeze()
        target_ids = target['input_ids'].squeeze()
        target_mask = target['attention_mask'].squeeze()
        return {'source_ids': source_ids.to(dtype=torch.long), 'source_mask': source_mask.to(dtype=torch.long), 'target_ids': target_ids.to(dtype=torch.long), 'target_ids_y': target_ids.to(dtype=torch.long)}

def train(epoch, tokenizer, model, device, loader, optimizer):
    """
    Function to be called for training with the parameters passed from main function
    """
    model.train()
    for _, data in enumerate(loader, 0):
        y = data['target_ids'].to(device, dtype=torch.long)
        y_ids = y[:, :-1].contiguous()
        lm_labels = y[:, 1:].clone().detach()
        lm_labels[y[:, 1:] == tokenizer.pad_token_id] = -100
        ids = data['source_ids'].to(device, dtype=torch.long)
        mask = data['source_mask'].to(device, dtype=torch.long)
        outputs = model(input_ids=ids, attention_mask=mask, decoder_input_ids=y_ids, labels=lm_labels)
        loss = outputs[0]
        if _ % 100 == 0:
            training_logger.add_row(str(epoch), str(_), str(loss))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

def validate(epoch, tokenizer, model, device, loader):
    """
  Function to evaluate model for predictions
  """
    model.eval()
    predictions = []
    actuals = []
    with torch.no_grad():
        for _, data in enumerate(loader, 0):
            y = data['target_ids'].to(device, dtype=torch.long)
            ids = data['source_ids'].to(device, dtype=torch.long)
            mask = data['source_mask'].to(device, dtype=torch.long)
            generated_ids = model.generate(input_ids=ids, attention_mask=mask, max_length=150, num_beams=2, repetition_penalty=2.5, length_penalty=1.0, early_stopping=True)
            preds = [tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True) for g in generated_ids]
            target = [tokenizer.decode(t, skip_special_tokens=True, clean_up_tokenization_spaces=True) for t in y]
            predictions.extend(preds)
            actuals.extend(target)
    return (predictions, actuals)

def T5Trainer(dataframe, source_text, target_text, model_params, output_dir='/kaggle/working/'):
    """
    T5 trainer
    """
    torch.manual_seed(model_params['SEED'])
    np.random.seed(model_params['SEED'])
    torch.backends.cudnn.deterministic = True
    console.log(f"[Model]: Loading {model_params['MODEL']}...\n")
    tokenizer = T5Tokenizer.from_pretrained(model_params['MODEL'])
    model = T5ForConditionalGeneration.from_pretrained(model_params['MODEL'])
    model = model.to(device)
    console.log(f'[Data]: Reading data...\n')
    dataframe = dataframe[[source_text, target_text]]
    train_size = 0.8
    train_dataset = dataframe.sample(frac=train_size, random_state=model_params['SEED'])
    val_dataset = dataframe.drop(train_dataset.index).reset_index(drop=True)
    train_dataset = train_dataset.reset_index(drop=True)
    training_set = CustomDataSetClass(train_dataset, tokenizer, model_params['MAX_SOURCE_TEXT_LENGTH'], model_params['MAX_TARGET_TEXT_LENGTH'], source_text, target_text)
    val_set = CustomDataSetClass(val_dataset, tokenizer, model_params['MAX_SOURCE_TEXT_LENGTH'], model_params['MAX_TARGET_TEXT_LENGTH'], source_text, target_text)
    train_params = {'batch_size': model_params['TRAIN_BATCH_SIZE'], 'shuffle': True, 'num_workers': 0}
    val_params = {'batch_size': model_params['VALID_BATCH_SIZE'], 'shuffle': False, 'num_workers': 0}
    training_loader = DataLoader(training_set, **train_params)
    val_loader = DataLoader(val_set, **val_params)
    optimizer = torch.optim.Adam(params=model.parameters(), lr=model_params['LEARNING_RATE'])
    console.log(f'[Initiating Fine Tuning]...\n')
    for epoch in range(model_params['TRAIN_EPOCHS']):
        train(epoch, tokenizer, model, device, training_loader, optimizer)
    console.log(f'[Saving Model]...\n')
    path = os.path.join(output_dir, 'model_files')
    model.save_pretrained(path)
    tokenizer.save_pretrained(path)
    console.log(f'[Initiating Validation]...\n')
    for epoch in range(model_params['VAL_EPOCHS']):
        predictions, actuals = validate(epoch, tokenizer, model, device, val_loader)
        final_df = pd.DataFrame({'Generated Text': predictions, 'Actual Text': actuals})
        final_df.to_csv(os.path.join(output_dir, 'predictions.csv'))
    console.save_text(os.path.join(output_dir, 'logs.txt'))
    console.log(f'[Validation Completed.]\n')
    return final_df

model_params = {'MODEL': 't5-base', 'TRAIN_BATCH_SIZE': 8, 'VALID_BATCH_SIZE': 8, 'TRAIN_EPOCHS': 1, 'VAL_EPOCHS': 1, 'LEARNING_RATE': 0.0001, 'MAX_SOURCE_TEXT_LENGTH': 512, 'MAX_TARGET_TEXT_LENGTH': 50, 'SEED': 42}

df['text'] = 'summarize: ' + df['text']
predictions = T5Trainer(dataframe=df, source_text='text', target_text='headlines', model_params=model_params, output_dir='outputs')