AmitPandit175/Text_Summarization

1from transformers import BartTokenizer, BartForConditionalGeneration2from datasets import load_dataset, load_metric3 4def evaluate_model(model_dir="./models/bart", dataset_name="cnn_dailymail"):5    tokenizer = BartTokenizer.from_pretrained(model_dir)6    model = BartForConditionalGeneration.from_pretrained(model_dir)7    dataset = load_dataset(dataset_name, '3.0.0', split='test[:100]')8    rouge = load_metric("rouge")9 10    for sample in dataset:11        inputs = tokenizer(sample['article'], return_tensors="pt", truncation=True, max_length=512)12        summary_ids = model.generate(inputs["input_ids"], max_length=128, num_beams=4)13        output = tokenizer.decode(summary_ids[0], skip_special_tokens=True)14        rouge.add(prediction=output, reference=sample["highlights"])15 16    final_score = rouge.compute()17    print(final_score)18 19if __name__ == "__main__":20    evaluate_model()21