LaserTagger

LaserTagger is a text-editing model which predicts a sequence of token-level edit operations to transform a source text into a target text. The model currently supports four different edit operations:

Keep the token.
Delete the token.
Add a phrase before the token.
Swap the order of input sentences (if there are two of them).

Operation 3 can be combined with 1 and 2. Compared to sequence-to-sequence models, LaserTagger is (1) less prone to hallucination, (2) more data efficient, and (3) faster at inference time.

A detailed method description and evaluation can be found in our EMNLP'19 paper: https://arxiv.org/abs/1909.01187

LaserTagger is built on Python 3, Tensorflow and BERT. It works with CPU, GPU, and Cloud TPU.

Usage Instructions

Running an experiment with LaserTagger consists of the following steps:

Optimize the vocabulary of phrases that can be added by LaserTagger.
Convert target texts into target tag sequences.
Finetune a pretrained BERT model to predict the tags.
Compute predictions.
Evaluate the predictions.

Next we go through these steps, using the Split-and-Rephrase (WikiSplit) task as a running example.

You can run all of the steps with

sh run_wikisplit_experiment.sh

after setting the paths in the beginning of the script.

Note: Text should be tokenized with spaces separating the tokens before applying LaserTagger.

1. Phrase Vocabulary Optimization

Download the WikiSplit dataset and run the following command to find a set of phrases that the model is allowed to add.

export WIKISPLIT_DIR=/path/to/wikisplit
export OUTPUT_DIR=/path/to/output

python phrase_vocabulary_optimization.py \
  --input_file=${WIKISPLIT_DIR}/train.tsv \
  --input_format=wikisplit \
  --vocabulary_size=500 \
  --max_input_examples=1000000 \
  --output_file=${OUTPUT_DIR}/label_map.txt

Note that you can also set max_input_examples to a smaller value to get a reasonable vocabulary, but then you should sort the dataset rows in the case of WikiSplit. The rows are in an alphabetical order so taking first k of them might not give you a representative sample of the data.

2. Converting Target Texts to Tags

Download a pretrained BERT model from the official repository. We've used the 12-layer ''BERT-Base, Cased'' model for all of our experiments. Then convert the original TSV datasets into TFRecord format.

export BERT_BASE_DIR=/path/to/cased_L-12_H-768_A-12

python preprocess_main.py \
  --input_file=${WIKISPLIT_DIR}/tune.tsv \
  --input_format=wikisplit \
  --output_tfrecord=${OUTPUT_DIR}/tune.tf_record \
  --label_map_file=${OUTPUT_DIR}/label_map.txt \
  --vocab_file=${BERT_BASE_DIR}/vocab.txt \
  --output_arbitrary_targets_for_infeasible_examples=true

python preprocess_main.py \
    --input_file=${WIKISPLIT_DIR}/train.tsv \
    --input_format=wikisplit \
    --output_tfrecord=${OUTPUT_DIR}/train.tf_record \
    --label_map_file=${OUTPUT_DIR}/label_map.txt \
    --vocab_file=${BERT_BASE_DIR}/vocab.txt \
    --output_arbitrary_targets_for_infeasible_examples=false

3. Model Training

Model hyperparameters are specified in lasertagger_config.json. This configuration file extends bert_config.json which comes with the zipped pretrained BERT model.

Note that if you want to switch from using LaserTagger_FF to LaserTagger_AR, you should set "use_t2t_decoder": true in the LaserTagger config. The latter is usually more accurate, whereas the former runs inference faster.

Train the model on CPU/GPU.

# Check these numbers from the "*.num_examples" files created in step 2.
export NUM_TRAIN_EXAMPLES=310922
export NUM_EVAL_EXAMPLES=5000
export CONFIG_FILE=configs/lasertagger_config.json
export EXPERIMENT=wikisplit_experiment_name

python run_lasertagger.py \
  --training_file=${OUTPUT_DIR}/train.tf_record \
  --eval_file=${OUTPUT_DIR}/tune.tf_record \
  --label_map_file=${OUTPUT_DIR}/label_map.txt \
  --model_config_file=${CONFIG_FILE} \
  --output_dir=${OUTPUT_DIR}/models/${EXPERIMENT} \
  --init_checkpoint=${BERT_BASE_DIR}/bert_model.ckpt \
  --do_train=true \
  --do_eval=true \
  --train_batch_size=256 \
  --save_checkpoints_steps=500 \
  --num_train_examples=${NUM_TRAIN_EXAMPLES} \
  --num_eval_examples=${NUM_EVAL_EXAMPLES}

To train on Cloud TPU, you should additionally set:

  --use_tpu=true \
  --tpu_name=${TPU_NAME}

Please see BERT TPU instructions and the Google Cloud TPU tutorial for how to use Cloud TPUs.

4. Prediction

First you need to export your model.

python run_lasertagger.py \
  --label_map_file=${OUTPUT_DIR}/label_map.txt \
  --model_config_file=${CONFIG_FILE} \
  --output_dir=${OUTPUT_DIR}/models/${EXPERIMENT} \
  --do_export=true \
  --export_path=${OUTPUT_DIR}/models/${EXPERIMENT}/export

You can additionally use init_checkpoint to specify which checkpoint to export (the default is to export the latest).

Compute the predicted tags and realize the output text with:

export SAVED_MODEL_DIR=/path/to/exported/model
export PREDICTION_FILE=${OUTPUT_DIR}/models/${EXPERIMENT}/pred.tsv

python predict_main.py \
  --input_file=${WIKISPLIT_DIR}/validation.tsv \
  --input_format=wikisplit \
  --output_file=${PREDICTION_FILE} \
  --label_map_file=${OUTPUT_DIR}/label_map.txt \
  --vocab_file=${BERT_BASE_DIR}/vocab.txt \
  --saved_model=${SAVED_MODEL_DIR}

Note that the above will run inference with batch size of 1 so it's not optimal in terms of inference time.

5. Evaluation

Compute the evaluation scores.

python score_main.py --prediction_file=${PREDICTION_FILE}

Example output:

Exact score:     15.220
SARI score:      61.668
 KEEP score:     93.059
 ADDITION score: 32.168
 DELETION score: 59.778

How to Cite LaserTagger

@inproceedings{malmi2019lasertagger,
  title={Encode, Tag, Realize: High-Precision Text Editing},
  author={Eric Malmi and Sebastian Krause and Sascha Rothe and Daniil Mirylenka and Aliaksei Severyn},
  booktitle={EMNLP-IJCNLP},
  year={2019}
}

License

Apache 2.0; see LICENSE for details.

Disclaimer

This repository contains a Python reimplementation of our original C++ code used for the paper and thus some discrepancies compared to the paper results are possible. However, we've verified that we get the similar results on the WikiSplit dataset.

This is not an official Google product.

kasnerz/lasertagger