BLEURT Documentation

A BLEURT checkpoint is a self-contained folder containing a TensorFlow regression model and necessary resources. It is technically a Tensorflow SavedModel accompanied by a bleurt_config.json file.

Checkpoints also include files required by BERT, such as a bert_config.json and a WordPiece dictionary or SentencePiece model for tokenization.

Important: Each checkpoint represents a different model; therefore, results produced by different checkpoints are not directly comparable.

If you want to train a metric from a 'fresh' BERT checkpoint (one not yet fine-tuned on ratings), use the bleurt.finetune module with the following parameters. If you are using RemBERT, you may also specify a sentence_piece_model.

BERT_DIR=bleurt/test_checkpoint
BERT_CKPT=variables/variables
python -m bleurt.finetune \
  -init_checkpoint=${BERT_DIR}/${BERT_CKPT} \
  -bert_config_file=${BERT_DIR}/bert_config.json \
  -vocab_file=${BERT_DIR}/vocab.txt \
  -model_dir=my_new_bleurt_checkpoint \
  -train_set=bleurt/test_data/ratings_train.jsonl \
  -dev_set=bleurt/test_data/ratings_dev.jsonl \
  -num_train_steps=500

The recommended checkpoint for accurate, multilingual results is BLEURT-20. It is a self-contained folder containing the regression model.

To use it, download the zip, unzip it, and point your command or API to the resulting folder.

# Download and unzip BLEURT-20
wget https://storage.googleapis.com/bleurt-oss-21/BLEURT-20.zip .
unzip BLEURT-20.zip

# Run scoring with the new checkpoint
python -m bleurt.score_files \
  -candidate_file=bleurt/test_data/candidates \
  -reference_file=bleurt/test_data/references \
  -bleurt_checkpoint=BLEURT-20

Use the bleurt.wmt.db_builder module to download WMT Metrics Shared Task archives and aggregate them into a single JSONL file. This is useful for working with ratings that are otherwise spread across multiple archives.

Supported years range from 2015 to 2019.

python -m bleurt.wmt.db_builder \
  -target_language="en" \
  -rating_years="2015 2016" \
  -target_file=wmt.jsonl

BLEURT requires Python 3, Tensorflow (>=1.15), and tf-slim (>=1.1). To install, clone the repository and install the local package using pip.

You can verify your installation by running the included unit tests.

pip install --upgrade pip
git clone https://github.com/google-research/bleurt.git
cd bleurt
pip install .

# Verify installation
python -m unittest bleurt.score_test
python -m unittest bleurt.score_not_eager_test
python -m unittest bleurt.finetune_test
python -m unittest bleurt.score_files_test

You can fine-tune previous generation checkpoints (those based on English BERT) on your own custom ratings data using the bleurt.finetune module.

The script tokenizes the input, serializes it into TFRecord files, and runs a training/evaluation loop. It saves the best model and exports it as a BLEURT checkpoint.

Note: It is highly recommended to use a GPU for fine-tuning. You can visualize training progress using Tensorboard by pointing it to your model_dir.

python -m bleurt.finetune \
  -init_bleurt_checkpoint=bleurt/test_checkpoint \
  -model_dir=my_new_bleurt_checkpoint \
  -train_set=bleurt/test_data/ratings_train.jsonl \
  -dev_set=bleurt/test_data/ratings_dev.jsonl \
  -num_train_steps=500

For embedding BLEURT into a TensorFlow computation graph (e.g., for Tensorboard visualization during training), use score.create_bleurt_ops() to generate the necessary operations.

import tensorflow as tf
# Set tf.enable_eager_execution() if using TF 1.x.

from bleurt import score

references = tf.constant(["This is a test."])
candidates = tf.constant(["This is the test."])

bleurt_ops = score.create_bleurt_ops()
bleurt_out = bleurt_ops(references=references, candidates=candidates)

assert bleurt_out["predictions"].shape == (1,)
print(bleurt_out["predictions"])

To use BLEURT within a Python script, import the score module and use the BleurtScorer class. This is the recommended way for programmatic access.

Note: BLEURT works in both eager_mode (default in TF 2.0) and tf.Session (TF 1.0).

from bleurt import score

checkpoint = "bleurt/test_checkpoint"
references = ["This is a test."]
candidates = ["This is the test."]

scorer = score.BleurtScorer(checkpoint)
scores = scorer.score(references=references, candidates=candidates)
assert isinstance(scores, list) and len(scores) == 1
print(scores)

The bleurt.finetune script includes many parameters for controlling the training process. You can view the full list of available flags by running the help command. Many parameters are inherited from the BERT codebase.

python finetune.py -helpfull