Home

Awesome

PRIMERA

The official code for PRIMERA: Pyramid-based Masked Sentence Pre-training for Multi-document Summarization.

PRIMERA is a pre-trained model for multi-document representation with focus on summarization that reduces the need for dataset-specific architectures and large amounts of fine-tuning labeled data. With extensive experiments on 6 multi-document summarization datasets from 3 different domains on the zero-shot, few-shot and full-supervised settings, PRIMER outperforms current state-of-the-art models on most of these settings with large margins.

Updates (2022-MAR-09)

For better usage of the model, we convert our trained models to the Huggingface version, which will be loaded to the Huggingface Model Hub soon. (The code for model conversion can be found Convert_to_hf_LED.ipynb, where the input is the state_dict() of our model)

We update the scripts and (example) bash files to run the Huggingface version of PRIMERA in the ./script/primer_hf_main.py and ./run_bash/, respectively. We also create a notebook as an example usage for evaluating our fine-tuned model on the multi-news dataset (Evaluation_Example.ipynb).

ModelRouge-1Rouge-2Rouge-L
PRIMERA42.013.620.8
PRIMERA-hf41.713.620.5
PRIMERA(finetuned)49.921.125.9
PRIMERA-hf(finetuned)49.920.925.8

Set up

  1. Create new virtual environment by
conda create --name primer python=3.7
conda activate primer
conda install cudatoolkit=10.0
  1. Install Longformer by
pip install git+https://github.com/allenai/longformer.git
  1. Install requirements to run the summarization scripts and data generation scripts by
pip install -r requirements.txt

Usage of PRIMERA

  1. Download the pre-trained PRIMERA model here to ./PRIMERA_model
  2. Load the tokenizer and model by
from transformers import AutoTokenizer
from longformer import LongformerEncoderDecoderForConditionalGeneration
from longformer import LongformerEncoderDecoderConfig

tokenizer = AutoTokenizer.from_pretrained('./PRIMERA_model/')
config = LongformerEncoderDecoderConfig.from_pretrained('./PRIMERA_model/')
model = LongformerEncoderDecoderForConditionalGeneration.from_pretrained(
            './PRIMERA_model/', config=config)

Make sure the documents separated with <doc-sep> in the input.

Summarization Scripts

You can use script/primer_main.py for pre-train/train/test PRIMERA, and script/compared_model_main.py for train/test BART/PEGASUS/LED.

Sample usages of both scripts can be found in run_bash/.

Datasets

Fully Supervised Models

We provide all the fully supervised models below.

Pre-training Data Generation

Newshead: we crawled the newshead dataset using the original code, and cleaned up the crawled data, the final newshead dataset can be found here.

You can use utils/pretrain_preprocess.py to generate pre-training data.

  1. Generate data with scores and entities with --mode compute_all_scores (The processed data with scores and entities can be found here)
  2. Generate pre-training data with --mode pretraining_data_with_score:
    • Pegasus: --strategy greedy --metric pegasus_score
    • Entity_Pyramid: --strategy greedy_entity_pyramid --metric pyramid_rouge (The processed data that could directly be used for pre-training can be found here)