bleurt-tiny-128-pt

PyTorch and safetensors conversion of Google's BLEURT checkpoint bleurt-tiny-128, 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-tiny-128 from google-research/bleurt.
• Weights: taken from the PyTorch weights of lucadiliello/bleurt-tiny-128 (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 bert. 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 128 (read from the bleurt_config.json of the TensorFlow original), num_labels is 1, do_lower_case is true, and unused keys of the port (for example _name_or_path and id2label) are removed.
• Tokenizer: WordPiece (vocab.txt, lowercased), copied from the port without changes.

Why this conversion exists

The published PyTorch port differs from the TensorFlow original in three ways:

• Activation: its config uses hidden_act: gelu, which is the exact (erf) GELU. TensorFlow BLEURT uses the tanh approximation.
• Sequence length: its config only has max_position_embeddings: 512. The original checkpoint was exported with max_seq_length: 128. Without an explicit limit of 128, nothing truncates at 128 tokens, so longer pairs are scored on different input than in TensorFlow.
• Truncation: for truncated pairs whose remaining sides have equal length, the fast WordPiece tokenizer of transformers removes a token from the other side than TensorFlow does (TensorFlow BLEURT removes it from the candidate).

Measured against TensorFlow on the test set below, the port as documented (exact GELU, truncation by transformers, no explicit length limit) gives for bleurt-tiny-128: Pearson 0.986785, maximum absolute difference 6.68e-01, mean absolute difference 1.87e-02, and 179 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: 1.43e-06
• Mean absolute difference: 3.56e-07
• Pairs with a difference above 1e-3: 0 of 1,096

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: bert and architectures: [BertForSequenceClassification]. The weights use the prefix bert. and classifier., so the checkpoint loads into transformers.BertForSequenceClassification with one output label. This was checked with transformers 5.18.0 on bleurt-tiny-128, bleurt-base-512 and bleurt-large-128: the logits differ from the MATEA loader by at most 1.1e-06 on 48 test pairs. hidden_act is gelu_new in config.json. The stock tokenizer does not apply the truncation rule described below.

Config keys that affect scoring:

• hidden_act: gelu_new (tanh approximation).
• max_seq_length: 128. 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"
}