Pre-training with Extracted Gap-sentences for Abstractive SUmmarization Sequence-to-sequence models, or PEGASUS, uses self-supervised objective Gap Sentences Generation (GSG) to train a transformer encoder-decoder model. The paper can be found on arXiv. ICML 2020 accepted.
If you use this code or these models, please cite the following paper:
@misc{zhang2019pegasus,
title={PEGASUS: Pre-training with Extracted Gap-sentences for Abstractive Summarization},
author={Jingqing Zhang and Yao Zhao and Mohammad Saleh and Peter J. Liu},
year={2019},
eprint={1912.08777},
archivePrefix={arXiv},
primaryClass={cs.CL}
}
Update (2022/08): Go to pegasus/flax for PEGASUS-X models
We train a pegasus model with sampled gap sentence ratios on both C4 and HugeNews, and stochastically sample important sentences. The updated the results are reported in this table.
dataset | C4 | HugeNews | Mixed & Stochastic |
---|---|---|---|
xsum | 45.20/22.06/36.99 | 47.21/24.56/39.25 | 47.60/24.83/39.64 |
cnn_dailymail | 43.90/21.20/40.76 | 44.17/21.47/41.11 | 44.16/21.56/41.30 |
newsroom | 45.07/33.39/41.28 | 45.15/33.51/41.33 | 45.98/34.20/42.18 |
multi_news | 46.74/17.95/24.26 | 47.52/18.72/24.91 | 47.65/18.75/24.95 |
gigaword | 38.75/19.96/36.14 | 39.12/19.86/36.24 | 39.65/20.47/36.76 |
wikihow | 43.07/19.70/34.79 | 41.35/18.51/33.42 | 46.39/22.12/38.41 * |
reddit_tifu | 26.54/8.94/21.64 | 26.63/9.01/21.60 | 27.99/9.81/22.94 |
big_patent | 53.63/33.16/42.25 | 53.41/32.89/42.07 | 52.29/33.08/41.66 * |
arxiv | 44.70/17.27/25.80 | 44.67/17.18/25.73 | 44.21/16.95/25.67 |
pubmed | 45.49/19.90/27.69 | 45.09/19.56/27.42 | 45.97/20.15/28.25 |
aeslc | 37.69/21.85/36.84 | 37.40/21.22/36.45 | 37.68/21.25/36.51 |
billsum | 57.20/39.56/45.80 | 57.31/40.19/45.82 | 59.67/41.58/47.59 |
The "Mixed & Stochastic" model has the following changes:
(*) the numbers of wikihow and big_patent datasets are not comparable because of change in tokenization and data:
Please create a project first and create an instance
gcloud compute instances create \
${VM_NAME} \
--zone=${ZONE} \
--machine-type=n1-highmem-8 \
--accelerator type=nvidia-tesla-v100,count=1 \
--boot-disk-size=500GB \
--image-project=ml-images \
--image-family=tf-1-15 \
--maintenance-policy TERMINATE --restart-on-failure
Clone library on github and install requirements.
git clone https://github.com/google-research/pegasus
cd pegasus
export PYTHONPATH=.
pip3 install -r requirements.txt
Download vocab, pretrained and fine-tuned checkpoints of all experiments from Google Cloud.
Alternatively in terminal, follow the instruction and install gsutil. Then
mkdir ckpt
gsutil cp -r gs://pegasus_ckpt/ ckpt/
Finetune on an existing dataset aeslc
.
python3 pegasus/bin/train.py --params=aeslc_transformer \
--param_overrides=vocab_filename=ckpt/pegasus_ckpt/c4.unigram.newline.10pct.96000.model \
--train_init_checkpoint=ckpt/pegasus_ckpt/model.ckpt-1500000 \
--model_dir=ckpt/pegasus_ckpt/aeslc
If you would like to finetune on a subset of dataset, please refer to the example of input pattern.
Evaluate on the finetuned dataset.
python3 pegasus/bin/evaluate.py --params=aeslc_transformer \
--param_overrides=vocab_filename=ckpt/pegasus_ckpt/c4.unigram.newline.10pct.96000.model,batch_size=1,beam_size=5,beam_alpha=0.6 \
--model_dir=ckpt/pegasus_ckpt/aeslc
Note that the above example is using a single GPU so the batch_size is much smaller than the results reported in the paper.
Two types of dataset format are supported: TensorFlow Datasets (TFDS) or TFRecords.
This tutorial shows how to add a new dataset in TFDS.
(The fine-tuning dataset is expected to be supervised, please provide
supervised_keys
in dataset info).
Tfrecords format requires each record to be a tf example of {"inputs":tf.string, "targets":tf.string}
.
For example, if you registered a TFDS dataset called new_tfds_dataset
for training and evaluation, and have some files in tfrecord format called new_dataset_files.tfrecord*
for test, they can be registered in /pegasus/params/public_params.py
.
@registry.register("new_params")
def my_param(param_overrides):
return public_params.transformer_params(
{
"train_pattern": "tfds:new_tfds_dataset,train",
"dev_pattern": "tfds:new_tfds_dataset,validation",
"test_pattern": "tfrecord:new_dataset_files.tfrecord*",
"max_input_len": 512,
"max_output_len": 128,
"train_steps": 10000,
"learning_rate": 0.0001,
"batch_size": 8,
}, param_overrides)
Evaluation results can be found in mode_dir
. Summarization metrics are automatically
calculated for each evaluation point.
ROUGE is the main metric for summarization quality.
BLEU is an alternative quality metric for language generation.
Extractive Fragments Coverage & Density are metrics that measures the abstractiveness of the summary.
Repetition Rates measures generation repetition failure modes.
Length statistics measures the length distribution of decodes comparing to gold summary.
Several types of output files can be found in model_dir
Pretraining (on C4 or any other corpus) requires a customly built tensorflow that includes ops for on-the-fly parsing that processes raw text document into model inputs and targets ids. Please refer to pegasus/ops/pretrain_parsing_ops.cc and pegasus/data/parsers.py for details.
Contains parts of code and design for training and evaluation of summarization models originally by Ben Goodrich [email protected].