🚨🚨🚨 [NLLB Tokenizer] Fix the prefix tokens 🚨🚨🚨

src/transformers/models/nllb/tokenization_nllb.py

@@ -390,11 +390,11 @@ def _switch_to_target_mode(self):
     def set_src_lang_special_tokens(self, src_lang) -> None:
         """Reset the special tokens to the source lang setting. No prefix and suffix=[eos, src_lang_code]."""
         self.cur_lang_code = self.lang_code_to_id[src_lang]
-        self.prefix_tokens = []
-        self.suffix_tokens = [self.eos_token_id, self.cur_lang_code]
+        self.prefix_tokens = [self.cur_lang_code]
+        self.suffix_tokens = [self.eos_token_id]
 
     def set_tgt_lang_special_tokens(self, lang: str) -> None:
         """Reset the special tokens to the target language setting. No prefix and suffix=[eos, tgt_lang_code]."""
         self.cur_lang_code = self.lang_code_to_id[lang]
-        self.prefix_tokens = []
-        self.suffix_tokens = [self.eos_token_id, self.cur_lang_code]
+        self.prefix_tokens = [self.cur_lang_code]
+        self.suffix_tokens = [self.eos_token_id]

src/transformers/models/nllb/tokenization_nllb_fast.py

@@ -289,8 +289,8 @@ def _switch_to_target_mode(self):
     def set_src_lang_special_tokens(self, src_lang) -> None:
         """Reset the special tokens to the source lang setting. No prefix and suffix=[eos, src_lang_code]."""
         self.cur_lang_code = self.convert_tokens_to_ids(src_lang)
-        self.prefix_tokens = []
-        self.suffix_tokens = [self.eos_token_id, self.cur_lang_code]
+        self.prefix_tokens = [self.cur_lang_code]
+        self.suffix_tokens = [self.eos_token_id]
 
         prefix_tokens_str = self.convert_ids_to_tokens(self.prefix_tokens)
         suffix_tokens_str = self.convert_ids_to_tokens(self.suffix_tokens)
@@ -304,8 +304,8 @@ def set_src_lang_special_tokens(self, src_lang) -> None:
     def set_tgt_lang_special_tokens(self, lang: str) -> None:
         """Reset the special tokens to the target language setting. No prefix and suffix=[eos, tgt_lang_code]."""
         self.cur_lang_code = self.convert_tokens_to_ids(lang)
-        self.prefix_tokens = []
-        self.suffix_tokens = [self.eos_token_id, self.cur_lang_code]
+        self.prefix_tokens = [self.cur_lang_code]
+        self.suffix_tokens = [self.eos_token_id]
 
         prefix_tokens_str = self.convert_ids_to_tokens(self.prefix_tokens)
         suffix_tokens_str = self.convert_ids_to_tokens(self.suffix_tokens)

tests/models/nllb/test_tokenization_nllb.py

@@ -305,6 +305,7 @@ class NllbDistilledIntegrationTest(unittest.TestCase):
         " face decât să înrăutăţească violenţele şi mizeria pentru milioane de oameni.",
     ]
     expected_src_tokens = [
+        256047,
         16297,
         134408,
         8165,
@@ -319,7 +320,6 @@ class NllbDistilledIntegrationTest(unittest.TestCase):
         108,
         49486,
         2,
-        256047,
     ]
 
     @classmethod
@@ -355,8 +355,8 @@ def test_enro_tokenizer_truncation(self):
         assert isinstance(src_text[0], str)
         desired_max_length = 10
         ids = self.tokenizer(src_text, max_length=desired_max_length, truncation=True).input_ids[0]
-        self.assertEqual(ids[-2], 2)
-        self.assertEqual(ids[-1], EN_CODE)
+        self.assertEqual(ids[-1], 2)
+        self.assertEqual(ids[0], EN_CODE)
         self.assertEqual(len(ids), desired_max_length)
 
     def test_mask_token(self):
@@ -389,10 +389,10 @@ def test_enro_tokenizer_prepare_batch(self):
         self.assertEqual((2, 15), batch.attention_mask.shape)
         result = batch.input_ids.tolist()[0]
         self.assertListEqual(self.expected_src_tokens, result)
-        self.assertEqual(2, batch.decoder_input_ids[0, -1])  # EOS
+        self.assertEqual(RO_CODE, batch.decoder_input_ids[0, 0])  # EOS
         # Test that special tokens are reset
-        self.assertEqual(self.tokenizer.prefix_tokens, [])
-        self.assertEqual(self.tokenizer.suffix_tokens, [self.tokenizer.eos_token_id, EN_CODE])
+        self.assertEqual(self.tokenizer.prefix_tokens, [EN_CODE])
+        self.assertEqual(self.tokenizer.suffix_tokens, [self.tokenizer.eos_token_id])
 
     def test_seq2seq_max_length(self):
         batch = self.tokenizer(self.src_text, padding=True, truncation=True, max_length=3, return_tensors="pt")
@@ -419,7 +419,7 @@ def test_tokenizer_translation(self):
             nested_simplify(inputs),
             {
                 # A, test, EOS, en_XX
-                "input_ids": [[70, 7356, 2, 256047]],
+                "input_ids": [[256047, 70, 7356, 2]],
                 "attention_mask": [[1, 1, 1, 1]],
                 # ar_AR
                 "forced_bos_token_id": 256057,