Teknosys learning library

12. Evaluate PEGASUS on CNN/DailyMail

CNN/DailyMail pairs news article inputs with highlights reference summaries. This smoke test evaluates ten held-out examples using google/pegasus-cnn_dailymail. Ten examples are not a benchmark; expand evaluation and inspect generated summaries manually. Truncation may discard important facts.

python -m pip install "transformers>=4.46,<5" torch datasets evaluate rouge_score sentencepiece accelerate

import torch
import evaluate
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

checkpoint = "google/pegasus-cnn_dailymail"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()
test = load_dataset("abisee/cnn_dailymail", "3.0.0", split="test[:10]")
predictions = []
with torch.inference_mode():
    for row in test:
        inputs = tokenizer(
            row["article"], return_tensors="pt",
            truncation=True, max_length=1024,
        ).to(device)
        output = model.generate(
            **inputs, max_new_tokens=128, num_beams=4, do_sample=False,
        )
        predictions.append(tokenizer.decode(output[0], skip_special_tokens=True))

rouge = evaluate.load("rouge")
scores = rouge.compute(
    predictions=predictions,
    references=test["highlights"],
    use_stemmer=True,
)
print(scores)  # Computed scores, not claimed benchmark results.

Requires substantial model memory and downloads. Pin exact package versions and model/dataset revisions for reproducible experiments. Do not use test results to repeatedly tune your model.

Educational material. Verify examples against current documentation and applicable licenses. Report a correction.