fix train_new_from_iterator in the case of byte-level tokenizers

src/transformers/tokenization_utils_fast.py

@@ -21,6 +21,7 @@
 from collections import defaultdict
 from typing import Any, Dict, List, Optional, Tuple, Union
 
+import tokenizers.pre_tokenizers as pre_tokenizers_fast
 from tokenizers import Encoding as EncodingFast
 from tokenizers import Tokenizer as TokenizerFast
 from tokenizers.decoders import Decoder as DecoderFast
@@ -699,6 +700,8 @@ def train_new_from_iterator(
             kwargs["end_of_word_suffix"] = tokenizer_json["model"]["end_of_word_suffix"]
         if tokenizer_json["model"]["type"] == "Unigram" and unk_token is not None:
             kwargs["unk_token"] = unk_token
+        if tokenizer_json["pre_tokenizer"]["type"] == "ByteLevel":
+            kwargs["initial_alphabet"] = pre_tokenizers_fast.ByteLevel.alphabet()
 
         trainer_class = MODEL_TO_TRAINER_MAPPING[tokenizer_json["model"]["type"]]
         trainer = trainer_class(vocab_size=vocab_size, special_tokens=special_tokens, **kwargs)

tests/pipelines/test_pipelines_common.py

@@ -143,6 +143,18 @@ def gen_test(ModelClass, checkpoint, tiny_config, tokenizer_class, feature_extra
             @skipIf(tiny_config is None, "TinyConfig does not exist")
             @skipIf(checkpoint is None, "checkpoint does not exist")
             def test(self):
+                if tokenizer_class is not None:
+                    try:
+                        tokenizer = get_tiny_tokenizer_from_checkpoint(checkpoint)
+                        tiny_config.vocab_size = len(tokenizer)
+                    # Rust Panic exception are NOT Exception subclass
+                    # Some test tokenizer contain broken vocabs or custom PreTokenizer, so we
+                    # provide some default tokenizer and hope for the best.
+                    except:  # noqa: E722
+                        self.skipTest(f"Ignoring {ModelClass}, cannot create a simple tokenizer")
+                else:
+                    tokenizer = None
+
                 if ModelClass.__name__.endswith("ForCausalLM"):
                     tiny_config.is_encoder_decoder = False
                     if hasattr(tiny_config, "encoder_no_repeat_ngram_size"):
@@ -160,24 +172,14 @@ def test(self):
                     )
                 if hasattr(model, "eval"):
                     model = model.eval()
-                if tokenizer_class is not None:
-                    try:
-                        tokenizer = get_tiny_tokenizer_from_checkpoint(checkpoint)
-                        # XLNet actually defines it as -1.
-                        if isinstance(model.config, (RobertaConfig, IBertConfig)):
-                            tokenizer.model_max_length = model.config.max_position_embeddings - 2
-                        elif (
-                            hasattr(model.config, "max_position_embeddings")
-                            and model.config.max_position_embeddings > 0
-                        ):
-                            tokenizer.model_max_length = model.config.max_position_embeddings
-                    # Rust Panic exception are NOT Exception subclass
-                    # Some test tokenizer contain broken vocabs or custom PreTokenizer, so we
-                    # provide some default tokenizer and hope for the best.
-                    except:  # noqa: E722
-                        self.skipTest(f"Ignoring {ModelClass}, cannot create a simple tokenizer")
-                else:
-                    tokenizer = None
+
+                if tokenizer is not None:
+                    # XLNet actually defines it as -1.
+                    if isinstance(model.config, (RobertaConfig, IBertConfig)):
+                        tokenizer.model_max_length = model.config.max_position_embeddings - 2
+                    elif hasattr(model.config, "max_position_embeddings") and model.config.max_position_embeddings > 0:
+                        tokenizer.model_max_length = model.config.max_position_embeddings
+
                 feature_extractor = get_tiny_feature_extractor_from_checkpoint(
                     checkpoint, tiny_config, feature_extractor_class
                 )

tests/tokenization/test_tokenization_fast.py

@@ -39,6 +39,7 @@ def setUp(self):
         self.test_rust_tokenizer = True
 
         model_paths = ["robot-test/dummy-tokenizer-fast", "robot-test/dummy-tokenizer-wordlevel"]
+        self.bytelevel_bpe_model_name = "SaulLu/dummy-tokenizer-bytelevel-bpe"
 
         # Inclusion of 2 tokenizers to test different types of models (Unigram and WordLevel for the moment)
         self.tokenizers_list = [(PreTrainedTokenizerFast, model_path, {}) for model_path in model_paths]
@@ -99,6 +100,15 @@ def test_training_new_tokenizer_with_special_tokens_change(self):
                     shutil.rmtree(self.tmpdirname)
                     self.tmpdirname = tmpdirname_orig
 
+    def test_training_new_tokenizer_with_bytelevel(self):
+        tokenizer = self.rust_tokenizer_class.from_pretrained(self.bytelevel_bpe_model_name)
+
+        toy_text_iterator = ("a" for _ in range(1000))
+        new_tokenizer = tokenizer.train_new_from_iterator(text_iterator=toy_text_iterator, length=1000, vocab_size=50)
+
+        encoding_ids = new_tokenizer.encode("a🤗")
+        self.assertEqual(encoding_ids, [64, 172, 253, 97, 245])
+
 
 @require_tokenizers
 class TokenizerVersioningTest(unittest.TestCase):