# -*- coding: utf-8 -*-
"""hf_workshop_project.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/16rr3KcHT3lyfI2QjUDm720EZjpP8Jw28
"""

from datasets import load_dataset
from transformers import T5Tokenizer,T5ForConditionalGeneration
from transformers import Seq2SeqTrainingArguments,Seq2SeqTrainer
import evaluate
import numpy as np

df = load_dataset("knkarthick/samsum")

tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-small")
model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-small",device_map = "auto")

def tokenize(data):
  input = ["summarize: "+ text for text in data['dialogue']]
  model_inputs =  tokenizer(input,max_length=128,padding='max_length',truncation=True)
  label = tokenizer(data['summary'],max_length=128,padding='max_length',truncation=True)
  model_inputs['labels'] = label['input_ids']
  return model_inputs

tokenized_train_data = df['train'].map(tokenize,batched= True)
tokenized_validation_data = df['validation'].map(tokenize,batched= True)

# tokenized_train_data = df['train'].select(range(2000)).map(tokenize,batched= True)
# tokenized_validation_data = df['validation'].select(range(600)).map(tokenize,batched= True)


training_args = Seq2SeqTrainingArguments(
    output_dir = './results',
    eval_strategy = 'epoch',
    learning_rate = 3e-5,
    per_device_train_batch_size = 8,
    per_device_eval_batch_size = 8,
    num_train_epochs = 10,
    weight_decay = 0.01,
    report_to = "none",
    logging_dir = './logs',
    fp16 = False,
    predict_with_generate= True,
    generation_max_length= 128,

)

# !pip install evaluate
# !pip install rouge_score


metric = evaluate.load('rouge')

def compute_metrics(eval_pred) :
  preds,labels = eval_pred

  if isinstance(preds,tuple):
    preds = preds[0]

  if preds.ndim == 3:
        preds = np.argmax(preds, axis=-1)

  preds = np.where(preds < 0, tokenizer.pad_token_id, preds)
  decoded_preds = tokenizer.batch_decode(preds,skip_special_tokens= True)

  labels = np.where(labels !=-100,labels,tokenizer.pad_token_id)
  decoded_labels = tokenizer.batch_decode(labels,skip_special_tokens= True)

  return metric.compute(predictions=decoded_preds,
                 references = decoded_labels,
                 use_stemmer = True)

trainer = Seq2SeqTrainer(
    model = model,
    train_dataset= tokenized_train_data,
    eval_dataset= tokenized_validation_data,
    args = training_args,
    compute_metrics= compute_metrics
)

trainer.train()

save_dir = './summary_model'
trainer.save_model(save_dir)
tokenizer.save_pretrained(save_dir)