text_html_output_82.png text_html_output_255.png text_html_output_149.png text_html_output_277.png text_html_output_338.png text_html_output_282.png text_html_output_219.png text_html_output_130.png text_html_output_320.png text_html_output_155.png text_html_output_328.png text_html_output_327.png text_html_output_307.png text_html_output_314.png text_html_output_103.png text_html_output_321.png text_html_output_197.png text_html_output_88.png text_html_output_80.png text_html_output_29.png text_html_output_89.png text_html_output_146.png text_html_output_261.png text_html_output_60.png text_html_output_227.png text_html_output_159.png text_html_output_237.png text_html_output_68.png text_html_output_341.png text_html_output_178.png text_html_output_117.png text_html_output_27.png text_html_output_115.png text_html_output_239.png text_html_output_235.png text_html_output_28.png text_html_output_209.png text_html_output_356.png text_html_output_263.png text_html_output_310.png text_html_output_141.png text_html_output_323.png text_html_output_45.png text_html_output_90.png text_html_output_274.png text_html_output_329.png text_html_output_353.png text_html_output_304.png text_html_output_10.png text_html_output_183.png text_html_output_165.png text_html_output_59.png text_html_output_319.png text_html_output_180.png text_html_output_116.png text_html_output_161.png text_html_output_347.png text_html_output_190.png text_html_output_258.png text_html_output_142.png text_html_output_218.png text_html_output_224.png text_html_output_181.png text_html_output_186.png text_html_output_22.png text_html_output_230.png text_html_output_350.png text_html_output_214.png text_html_output_125.png text_html_output_150.png text_html_output_234.png text_html_output_16.png text_html_output_337.png text_html_output_40.png text_html_output_223.png text_html_output_91.png text_html_output_100.png text_html_output_293.png text_html_output_317.png text_html_output_176.png text_html_output_199.png text_html_output_144.png text_html_output_173.png text_html_output_305.png text_html_output_153.png text_html_output_52.png text_html_output_122.png text_html_output_166.png text_html_output_53.png text_html_output_4.png text_html_output_194.png text_html_output_41.png text_html_output_81.png text_html_output_205.png text_html_output_300.png text_html_output_132.png text_html_output_57.png text_html_output_73.png text_html_output_147.png text_html_output_203.png text_html_output_158.png text_html_output_254.png text_html_output_35.png text_html_output_306.png text_html_output_244.png text_html_output_6.png text_html_output_47.png text_html_output_139.png text_html_output_113.png text_html_output_241.png text_html_output_269.png text_html_output_339.png text_html_output_168.png text_html_output_61.png text_html_output_253.png text_html_output_311.png text_html_output_288.png text_html_output_63.png text_html_output_123.png text_html_output_94.png text_html_output_136.png text_html_output_222.png text_html_output_140.png text_html_output_248.png text_html_output_333.png text_html_output_266.png text_html_output_352.png text_html_output_211.png text_html_output_26.png text_html_output_316.png text_html_output_346.png text_html_output_37.png text_html_output_270.png text_html_output_2.png text_html_output_38.png text_html_output_246.png text_html_output_15.png text_html_output_5.png text_html_output_271.png text_html_output_281.png text_html_output_358.png text_html_output_124.png text_html_output_295.png text_html_output_252.png text_html_output_163.png text_html_output_200.png text_html_output_75.png text_html_output_208.png text_html_output_64.png text_html_output_204.png text_html_output_110.png text_html_output_206.png text_html_output_14.png text_html_output_191.png text_html_output_210.png text_html_output_332.png text_html_output_71.png text_html_output_98.png text_html_output_120.png text_html_output_157.png text_html_output_65.png text_html_output_43.png text_html_output_104.png text_html_output_101.png text_html_output_95.png text_html_output_268.png text_html_output_108.png text_html_output_175.png text_html_output_336.png text_html_output_131.png text_html_output_127.png text_html_output_121.png text_html_output_170.png text_html_output_23.png text_html_output_278.png text_html_output_118.png text_html_output_309.png text_html_output_249.png text_html_output_290.png text_html_output_291.png text_html_output_19.png text_html_output_283.png text_html_output_196.png text_html_output_137.png text_html_output_238.png text_html_output_48.png text_html_output_324.png text_html_output_298.png text_html_output_72.png text_html_output_354.png text_html_output_251.png text_html_output_169.png text_html_output_9.png text_html_output_109.png text_html_output_46.png text_html_output_201.png text_html_output_49.png text_html_output_13.png text_html_output_20.png text_html_output_250.png text_html_output_114.png text_html_output_70.png text_html_output_51.png text_html_output_151.png text_html_output_192.png text_html_output_259.png text_html_output_357.png text_html_output_129.png text_html_output_105.png text_html_output_164.png text_html_output_326.png text_html_output_134.png text_html_output_111.png text_html_output_188.png text_html_output_126.png text_html_output_256.png text_html_output_233.png text_html_output_148.png text_html_output_135.png text_html_output_212.png text_html_output_184.png text_html_output_225.png text_html_output_21.png text_html_output_69.png text_html_output_285.png text_html_output_128.png text_html_output_349.png text_html_output_83.png text_html_output_215.png text_html_output_79.png text_html_output_308.png text_html_output_289.png text_html_output_276.png text_html_output_335.png text_html_output_302.png text_html_output_187.png text_html_output_172.png text_html_output_202.png text_html_output_348.png text_html_output_257.png text_html_output_322.png text_html_output_343.png text_html_output_267.png text_html_output_92.png text_html_output_1.png text_html_output_93.png text_html_output_66.png text_html_output_36.png text_html_output_78.png text_html_output_17.png text_html_output_33.png text_html_output_42.png text_html_output_189.png text_html_output_44.png text_html_output_301.png text_html_output_160.png text_html_output_96.png text_html_output_87.png text_html_output_247.png text_html_output_299.png text_html_output_242.png text_html_output_50.png text_html_output_243.png text_html_output_30.png text_html_output_18.png text_html_output_236.png text_html_output_85.png text_html_output_334.png text_html_output_280.png text_html_output_330.png text_html_output_77.png text_html_output_275.png text_html_output_351.png text_html_output_54.png text_html_output_286.png text_html_output_162.png text_html_output_229.png text_html_output_179.png text_html_output_156.png text_html_output_76.png text_html_output_325.png text_html_output_345.png text_html_output_12.png text_html_output_213.png text_html_output_303.png text_html_output_292.png text_html_output_97.png text_html_output_34.png text_html_output_221.png text_html_output_193.png text_html_output_74.png text_html_output_272.png text_html_output_62.png text_html_output_273.png text_html_output_198.png text_html_output_56.png text_html_output_297.png text_html_output_55.png text_html_output_216.png text_html_output_240.png text_html_output_220.png text_html_output_106.png text_html_output_86.png text_html_output_217.png text_html_output_145.png text_html_output_102.png text_html_output_264.png application_vnd.jupyter.stderr_output_1.png text_html_output_11.png text_html_output_342.png text_html_output_195.png text_html_output_84.png text_html_output_24.png text_html_output_355.png text_html_output_284.png text_html_output_182.png text_html_output_167.png text_html_output_287.png text_html_output_31.png text_html_output_260.png text_html_output_313.png text_html_output_228.png text_html_output_344.png text_html_output_99.png text_html_output_143.png text_html_output_154.png text_html_output_177.png text_html_output_8.png text_html_output_67.png text_html_output_25.png text_html_output_231.png text_html_output_331.png text_html_output_58.png text_html_output_294.png text_html_output_171.png text_html_output_265.png text_html_output_232.png text_html_output_152.png text_html_output_39.png text_html_output_279.png text_html_output_207.png text_html_output_312.png text_html_output_174.png text_html_output_315.png text_html_output_112.png text_html_output_245.png text_html_output_32.png text_html_output_119.png text_html_output_296.png text_html_output_3.png text_html_output_185.png text_html_output_133.png text_html_output_107.png text_html_output_138.png text_html_output_226.png text_html_output_262.png text_html_output_7.png text_html_output_318.png text_html_output_340.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')

predictions.sample(20)