Spaces:
Sleeping
Sleeping
File size: 3,709 Bytes
34b369f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
import os
import pandas as pd
import torch
from datasets import load_dataset, load_from_disk, load_metric
from tqdm import tqdm
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from src.TextSummarizer.entity import entities
from src.TextSummarizer.logger import backend_logger
class ModelEvaluation:
def __init__(self, config: entities.ModelEvaluationConfig):
self.config = config
def generate_batch_sized_chunks(self,list_of_elements, batch_size):
"""split the dataset into smaller batches that we can process simultaneously
Yield successive batch-sized chunks from list_of_elements."""
for i in range(0, len(list_of_elements), batch_size):
yield list_of_elements[i : i + batch_size]
def calculate_metric_on_test_ds(self,
dataset,
metric,
model,
tokenizer,
batch_size=16,
device="cuda" if torch.cuda.is_available() else "cpu",
column_text="article",
column_summary="highlights"):
"""
Calculate the metrics.
"""
article_batches = list(self.generate_batch_sized_chunks(dataset[column_text], batch_size))
target_batches = list(self.generate_batch_sized_chunks(dataset[column_summary], batch_size))
for article_batch, target_batch in tqdm(
zip(article_batches, target_batches), total=len(article_batches)):
inputs = tokenizer(article_batch, max_length=1024, truncation=True,
padding="max_length", return_tensors="pt")
summaries = model.generate(input_ids=inputs["input_ids"].to(device),
attention_mask=inputs["attention_mask"].to(device),
length_penalty=0.8, num_beams=8, max_length=128)
''' parameter for length penalty ensures that the model does not generate sequences that are too long. '''
# Finally, we decode the generated texts,
# replace the token, and add the decoded texts with the references to the metric.
decoded_summaries = [tokenizer.decode(s, skip_special_tokens=True,
clean_up_tokenization_spaces=True)
for s in summaries]
decoded_summaries = [d.replace("", " ") for d in decoded_summaries]
metric.add_batch(predictions=decoded_summaries, references=target_batch)
# Finally compute and return the ROUGE scores.
score = metric.compute()
return score
def run(self):
"""
Run the model evaluation step.
"""
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained(self.config.tokenizer_path)
model_pegasus = AutoModelForSeq2SeqLM.from_pretrained(self.config.model_path).to(device)
#loading data
dataset_samsum_pt = load_from_disk(self.config.data_path)
rouge_names = ["rouge1", "rouge2", "rougeL", "rougeLsum"]
rouge_metric = load_metric('rouge')
score = self.calculate_metric_on_test_ds(
dataset_samsum_pt['test'][0:10], rouge_metric, model_pegasus, tokenizer, batch_size = 2, column_text = 'dialogue', column_summary= 'summary'
)
rouge_dict = dict((rn, score[rn].mid.fmeasure ) for rn in rouge_names )
df = pd.DataFrame(rouge_dict, index = ['pegasus'] )
df.to_csv(self.config.metric_file_name, index=False)
|