| |
| """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) |
|
|
| |
| |
|
|
|
|
| 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, |
|
|
| ) |
|
|
| |
| |
|
|
|
|
| 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) |
