Repository: huggingface/transformers (Python) — Bug Fix

Issue:
mBART: incorrect <mask> token id
# 🐛 Bug

## Information

Model I am using: mBART

## To reproduce

```
from transformers import MBartTokenizer
tokenizer = MBartTokenizer.from_pretrained('facebook/mbart-large-cc25')
print(tokenizer.convert_tokens_to_ids(['<mask>', 'ar_AR']))
```
The output for the above code is `[250001, 250001]`  - two different special tokens are mapped to the same id.

## Expected behavior

As far as I can tell, `<mask>` token should be mapped to id 250026.

I've checked [fairseq implementation](https://github.com/pytorch/fairseq/blob/master/fairseq/tasks/multilingual_denoising.py) and it seems that `<mask>` token is added after all the language codes, so it should be the last token in the vocab.

Currently, when I try to use mBART to denoise text with `<mask>` tokens, it mostly just ignores them, but if I replace mask ids with 250026, the model actually generates new text in place of `<mask>` tokens:

```
from transformers import MBartTokenizer, BartForConditionalGeneration

tokenizer = MBartTokenizer.from_pretrained('facebook/mbart-large-cc25')
model = BartForConditionalGeneration.from_pretrained('facebook/mbart-large-cc25')

text = 'I highly recommend <mask> - it is one of the best <mask> ever read!'
inputs = tokenizer.prepare_translation_batch([text], src_lang='en_XX')

outputs = model.generate(inputs['input_ids'], decoder_start_token_id=tokenizer.lang_code_to_id['en_XX'], 
                         num_beams=5)
print(tokenizer.batch_decode(outputs)[0])
```

The output is:
```
en_XX
 highly recommend - it is one of the best ever read!
```
Replacing mask ids:

```
where = (inputs['input_ids'] == 250001)
inputs['input_ids'][where] = 250026

outputs = model.generate(inputs['input_ids'], decoder_start_token_id=tokenizer.lang_code_to_id['en_XX'], 
                         num_beams=5)
print(tokenizer.batch_decode(outputs)[0])
```

The output is:
```
en_XX
 highly recommend this book - it is one of the best books I have ever read!
```

(In both cases, the model also skips the first input token when generating output, as discussed in #5755.)

I've also noticed that fairseq is using [language code tokens](https://github.com/pytorch/fairseq/blob/108bb2560b1ec01524ba723bc7c69186875afa0a/fairseq/tasks/multilingual_denoising.py#L62) of the form `[en_XX]` rather than just `en_XX`, which can lead to different tokenization if words like  `en_XX` appear in the text, but that's a rather contrived case.


## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: 3.0.2


@sshleifer



Discussion:

Resolving patch:
diff --git a/src/transformers/tokenization_bart.py b/src/transformers/tokenization_bart.py
--- a/src/transformers/tokenization_bart.py
+++ b/src/transformers/tokenization_bart.py
@@ -122,6 +122,7 @@ def __init__(self, *args, **kwargs):
         }
         self.id_to_lang_code = {v: k for k, v in self.lang_code_to_id.items()}
         self.cur_lang_code = self.lang_code_to_id["en_XX"]
+        self.fairseq_tokens_to_ids["<mask>"] = len(self.sp_model) + len(self.lang_code_to_id) + self.fairseq_offset
 
         self.fairseq_tokens_to_ids.update(self.lang_code_to_id)
         self.fairseq_ids_to_tokens = {v: k for k, v in self.fairseq_tokens_to_ids.items()}


Test patch:
diff --git a/tests/test_modeling_mbart.py b/tests/test_modeling_mbart.py
--- a/tests/test_modeling_mbart.py
+++ b/tests/test_modeling_mbart.py
@@ -123,6 +123,7 @@ def test_mbart_fast_forward(self):
         self.assertEqual(logits.shape, expected_shape)
 
 
+@require_torch
 class MBartCC25IntegrationTest(AbstractMBartIntegrationTest):
     checkpoint_name = "facebook/mbart-large-cc25"
     src_text = [
@@ -140,3 +141,14 @@ def test_cc25_generate(self):
         )
         decoded = self.tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)
         self.assertEqual(self.tgt_text[0], decoded[0])
+
+    @slow
+    def test_fill_mask(self):
+        inputs = self.tokenizer.prepare_translation_batch(["One of the best <mask> I ever read!"]).to(torch_device)
+        outputs = self.model.generate(
+            inputs["input_ids"], decoder_start_token_id=self.tokenizer.lang_code_to_id["en_XX"], num_beams=1
+        )
+        prediction: str = self.tokenizer.batch_decode(
+            outputs, clean_up_tokenization_spaces=True, skip_special_tokens=True
+        )[0]
+        self.assertEqual(prediction, "of the best books I ever read!")
diff --git a/tests/test_tokenization_mbart.py b/tests/test_tokenization_mbart.py
--- a/tests/test_tokenization_mbart.py
+++ b/tests/test_tokenization_mbart.py
@@ -1,3 +1,4 @@
+import tempfile
 import unittest
 
 from transformers import AutoTokenizer, BatchEncoding, MBartTokenizer
@@ -171,3 +172,13 @@ def test_enro_tokenizer_truncation(self):
         self.assertEqual(ids[-2], 2)
         self.assertEqual(ids[-1], EN_CODE)
         self.assertEqual(len(ids), desired_max_length)
+
+    def test_mask_token(self):
+        self.assertListEqual(self.tokenizer.convert_tokens_to_ids(["<mask>", "ar_AR"]), [250026, 250001])
+
+    def test_special_tokens_unaffacted_by_save_load(self):
+        tmpdirname = tempfile.mkdtemp()
+        original_special_tokens = self.tokenizer.fairseq_tokens_to_ids
+        self.tokenizer.save_pretrained(tmpdirname)
+        new_tok = MBartTokenizer.from_pretrained(tmpdirname)
+        self.assertDictEqual(new_tok.fairseq_tokens_to_ids, original_special_tokens)