BLEURT-20-D12-pt

PyTorch and safetensors conversion of Google's BLEURT checkpoint BLEURT-20-D12, made for MATEA, a library for machine translation evaluation. BLEURT scores a candidate text against a reference text with a fine tuned regression model. A higher score means closer agreement with the reference.

Source and conversion

• Original: the TensorFlow checkpoint BLEURT-20-D12 from google-research/bleurt.
• Weights: taken from the PyTorch weights of lucadiliello/BLEURT-20-D12 (the bleurt-pytorch port by Luca Di Liello). They were not re-exported from the TensorFlow files. Parity with the TensorFlow original was measured afterwards (see below).
• Changes to the weights: position_ids buffers are dropped, parameter names carry the bleurt. prefix, and the file is stored as model.safetensors. Tensor values are unchanged.
• Changes to config.json: hidden_act is gelu_new (the tanh approximation of GELU that TensorFlow BLEURT uses), max_seq_length is 512 (read from the bleurt_config.json of the TensorFlow original), num_labels is 1, do_lower_case is false, and unused keys of the port (for example _name_or_path and id2label) are removed.
• Tokenizer: SentencePiece (spm.model), copied from the port without changes.

Why this conversion exists

The published PyTorch port differs from the TensorFlow original in one way:

• Activation: its config uses hidden_act: gelu, which is the exact (erf) GELU. TensorFlow BLEURT uses the tanh approximation.

The sequence length is not a difference for this checkpoint: the original was exported with max_seq_length: 512, which equals the port's max_position_embeddings. The SentencePiece tokenizer of this checkpoint gave identical token ids to TensorFlow on all 1,096 test pairs. MATEA applies the TensorFlow truncation rule regardless of the tokenizer.

Measured against TensorFlow on the test set below, the port as documented (exact GELU, truncation by transformers, no explicit length limit) gives for BLEURT-20-D12: Pearson 0.999999, maximum absolute difference 1.56e-03, mean absolute difference 2.20e-04, and 7 of 1,096 pairs with a difference above 1e-3.

Parity with TensorFlow

Test set: 1,096 pairs. These are 1,040 WMT22 segment pairs (en-de and zh-en, 4 systems each), 32 long pairs (45 consecutive segments concatenated, above 512 tokens) and 24 edge cases (empty hypothesis or reference, identical texts, punctuation, emoji, non Latin scripts). The converted checkpoint was scored with the MATEA loader (tanh GELU, TensorFlow style truncation, max_seq_length from the config) and compared with the original TensorFlow checkpoint (TensorFlow 2.15.1, google-research/bleurt at HEAD).

CPU, fp32, batch size 16:

• Pearson: 1.000000
• Maximum absolute difference: 6.56e-07
• Mean absolute difference: 1.05e-07
• Pairs with a difference above 1e-3: 0 of 1,096

On GPU (fp32, batch size 16), the maximum absolute difference to TensorFlow over the same 1,096 pairs is 1.1e-06, and 0 pairs differ by more than 1e-3.

Usage

The checkpoint is intended for use with MATEA, which implements the scoring rules below. A different loader reproduces the parity numbers above only if it follows the same rules.

The config has model_type: bleurt, which is not a model type in transformers. The weights use the prefix bleurt. (plus classifier.) and an extra encoder.embedding_projection layer that maps the embedding size (128) to the hidden size (1024). bleurt-pytorch implements this architecture, and MATEA ships a self-contained loader.

Config keys that affect scoring:

• hidden_act: gelu_new (tanh approximation).
• max_seq_length: 512. The input is [CLS] reference [SEP] candidate [SEP] and its total length must not exceed this value.
• max_position_embeddings: 512, the size of the position embedding table.

TensorFlow style truncation: tokenize reference and candidate without special tokens, then remove the last token of the longer sequence (of the candidate if both have equal length) until max_seq_length - 3 tokens remain, and add [CLS] and [SEP] afterwards. Token type ids are 0 for [CLS], the reference and the first [SEP], and 1 for the candidate and the last [SEP]. The model returns one score per pair from the classifier layer.

License and attribution

This conversion is released under the Apache License 2.0, the license of the original BLEURT repository.

• Original BLEURT models and code: Google Research, https://github.com/google-research/bleurt
• PyTorch port and the weights used as conversion source: Luca Di Liello, https://github.com/lucadiliello/bleurt-pytorch

Citation

@inproceedings{sellam-etal-2020-bleurt,
    title = "{BLEURT}: Learning Robust Metrics for Text Generation",
    author = "Sellam, Thibault  and
      Das, Dipanjan  and
      Parikh, Ankur",
    editor = "Jurafsky, Dan  and
      Chai, Joyce  and
      Schluter, Natalie  and
      Tetreault, Joel",
    booktitle = "Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics",
    month = jul,
    year = "2020",
    address = "Online",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2020.acl-main.704/",
    doi = "10.18653/v1/2020.acl-main.704",
    pages = "7881--7892"
}

@inproceedings{pu-etal-2021-learning,
    title = "Learning Compact Metrics for {MT}",
    author = "Pu, Amy  and
      Chung, Hyung Won  and
      Parikh, Ankur  and
      Gehrmann, Sebastian  and
      Sellam, Thibault",
    editor = "Moens, Marie-Francine  and
      Huang, Xuanjing  and
      Specia, Lucia  and
      Yih, Scott Wen-tau",
    booktitle = "Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing",
    month = nov,
    year = "2021",
    address = "Online and Punta Cana, Dominican Republic",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2021.emnlp-main.58/",
    doi = "10.18653/v1/2021.emnlp-main.58",
    pages = "751--762"
}