This works in Colab.
I only ran it for 3 epochs, on a small portion of a dataset, but it’s training and the metrics improved.
The main changes are:

passing in the decoder_input_ids parameter (no right shifting in this example)
handling the -100 pad_token_id value for the labels
using Adafactor instead of AdamW
using a label_smoothing_factor
setting the learning_rate to one of the values from the paper (appendix C)

I hope it helps.
! pip install transformers
! pip install datasets
! pip install sentencepiece
! pip install rouge_score

import torch
import datasets

from transformers import (
    AutoModelForSeq2SeqLM, AutoTokenizer,
    Seq2SeqTrainingArguments, Seq2SeqTrainer
  )

model = AutoModelForSeq2SeqLM.from_pretrained("sshleifer/distill-pegasus-xsum-16-8")
tokenizer = AutoTokenizer.from_pretrained("sshleifer/distill-pegasus-xsum-16-8")

# Use wiki_lingua for demonstration
language = "english"
data = datasets.load_dataset("wiki_lingua", name=language, split="train[:1%]")
source = [x["document"] for x in data["article"] if len(x["document"]) > 0]
target = [x["summary"] for x in data["article"] if len(x["document"]) > 0]
source = [item for sublist in source for item in sublist]
target = [item for sublist in target for item in sublist]
N_DEMO = 100  # Try on a small portion of the data
source = source[:N_DEMO]
target = target[:N_DEMO]

train_percentage = .8
test_percentage = 1-train_percentage
train_test_split_pct = int(len(source)*train_percentage)

train_summary = target[:train_test_split_pct]
train_webtext = source[:train_test_split_pct]
test_summary = target[train_test_split_pct:]
test_webtext = source[train_test_split_pct:]

encoder_max_length=512
decoder_max_length=64

# Tokenize our data
train_summary = tokenizer(train_summary, padding="max_length", 
                          truncation=True, max_length=decoder_max_length)
train_webtext = tokenizer(train_webtext, padding="max_length", 
                          truncation=True, max_length=encoder_max_length)
test_summary = tokenizer(test_summary, padding="max_length", 
                         truncation=True, max_length=decoder_max_length)
test_webtext = tokenizer(test_webtext, padding="max_length", 
                         truncation=True, max_length=encoder_max_length)

# Setup dataset objects
class Summary(torch.utils.data.Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        labels = self.labels['input_ids'][idx]
        labels100 = [-100 if token == tokenizer.pad_token_id 
                     else token for token in labels]
        item['labels'] = torch.tensor(labels100)
        item['decoder_input_ids'] = torch.tensor(labels)            
        return item

    def __len__(self):
        return len(self.labels)


# Get datasets
train_dataset = Summary(train_webtext, train_summary)
test_dataset = Summary(test_webtext, test_summary)

rouge = datasets.load_metric("rouge")

def compute_metrics(pred):
    """
    Borrowed from https://github.com/patrickvonplaten/notebooks/blob/master/RoBERTaShared_for_BBC_XSum.ipynb
    """
    labels_ids = pred.label_ids
    pred_ids = pred.predictions

    # all unnecessary tokens are removed
    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
    labels_ids[labels_ids == -100] = tokenizer.pad_token_id
    label_str = tokenizer.batch_decode(labels_ids, skip_special_tokens=True)

    rouge_output = rouge.compute(predictions=pred_str, references=label_str, rouge_types=["rouge2"])["rouge2"].mid

    return {
        "rouge2_precision": round(rouge_output.precision, 4),
        "rouge2_recall": round(rouge_output.recall, 4),
        "rouge2_fmeasure": round(rouge_output.fmeasure, 4),
    }

batch_size = 4  # demo

training_args = Seq2SeqTrainingArguments( 
    output_dir='./results',
    adafactor=True,     
    num_train_epochs=3, # 50, # demo
    per_device_train_batch_size=batch_size, 
    per_device_eval_batch_size=batch_size,
    predict_with_generate=True,
    do_train=True,
    do_eval=True,
    learning_rate=5e-4,
    label_smoothing_factor=0.1,
    overwrite_output_dir=True, 
    logging_dir='./logs'
)

trainer = Seq2SeqTrainer(
    model=model,
    tokenizer=tokenizer,
    args=training_args,
    compute_metrics=compute_metrics,
    train_dataset=train_dataset,
    eval_dataset=test_dataset
)

trainer.evaluate()

trainer.train()

trainer.evaluate()


image1012×1315 79.4 KB