Teknosys learning library

13. Fine-tune PEGASUS with paired summaries

Continue in the Python environment above, reusing its model and tokenizer. This small training run starts from an already CNN/DailyMail-adapted checkpoint; it demonstrates the workflow, not training from scratch or a promised quality improvement. Article tokens are inputs and summary tokens are labels. Padding labels use -100 so the loss ignores them.

from transformers import (
    DataCollatorForSeq2Seq, Seq2SeqTrainer, Seq2SeqTrainingArguments,
)

train = load_dataset("abisee/cnn_dailymail", "3.0.0", split="train[:100]")
validation = load_dataset("abisee/cnn_dailymail", "3.0.0", split="validation[:20]")

def preprocess(batch):
    inputs = tokenizer(batch["article"], max_length=512, truncation=True)
    labels = tokenizer(
        text_target=batch["highlights"], max_length=128, truncation=True,
    )
    inputs["labels"] = labels["input_ids"]
    return inputs

train_tokens = train.map(preprocess, batched=True, remove_columns=train.column_names)
validation_tokens = validation.map(
    preprocess, batched=True, remove_columns=validation.column_names,
)
args = Seq2SeqTrainingArguments(
    output_dir="./pegasus-demo",
    learning_rate=2e-5,
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,
    gradient_accumulation_steps=4,
    num_train_epochs=1,
    eval_strategy="epoch",
    save_strategy="epoch",
    report_to="none",
    seed=42,
)
trainer = Seq2SeqTrainer(
    model=model,
    args=args,
    train_dataset=train_tokens,
    eval_dataset=validation_tokens,
    data_collator=DataCollatorForSeq2Seq(tokenizer, model=model),
)
trainer.train()
print(trainer.evaluate())  # Validation loss; this does not compute ROUGE.
trainer.save_model("./pegasus-demo/final")
tokenizer.save_pretrained("./pegasus-demo/final")

Training requires much more memory than inference. Reduce sequence lengths or use a smaller model if needed. This example targets the Transformers 4.x APIs installed above. Evaluate generated summaries separately against your unchanged baseline and reserve the test set for the final comparison. Review data licensing before redistribution.

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