Repository: huggingface/transformers
Pull request: Patch MistralCommonTokenizer

# What does this PR do?

This PR patches the `MistralCommonTokenizer`:
1. [BUG FIX] spm now is correctly supported as previous usage would result in an error for `_piece_to_id` and `_is_control_token`.
2. [BUG FIX] previous implementation of `_piece_to_id` was incorrect for Tekkenizer.

a) the special tokens were not supported
b) the normal tokens were not shifted by adding the number of special tokens

4. [MAYBE BUG FIX] Changed `get_vocab` to a function that should better mimic what happens in `Transformers`: now the mapping is based on the real ids but some ids are missing due to conversion loss of some tokens in Tekken.
5. [FEATURE] `add_generation_prompt` has been added. This is to match signature of Transformers. In practice this value is ignored except if:

a) `continue_final_message` and `add_generation_prompt` are `True` an error is raised.
b) if `add_generation_prompt` is `True` and the last message is `assistant` then an error is raised as the user should have passed `continue_final_message`.

6. ~[FEATURE] Now the tokenizer is set to fast because:~

a) ~it is true~(Edit: not so sure about that after discussion)
b) it removes annoying message when initializing the tokenizer.

Edit 2: this has been reverted as it doesn't bring value and is misleading.

7. [OPTIMIZATION] now image tensors are converted without slowness warnings from torch.
8. [DOCS] updated some docs to remove unused args.

Also added minimal tests to ensure SPM also works in the future.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

@patrickvonplaten @ArthurZucker @itazap 

Changed files:
filename: src/transformers/tokenization_mistral_common.py
patch: @@ -115,11 +115,6 @@
                 of returning overflowing tokens.
             return_special_tokens_mask (`bool`, *optional*, defaults to `False`):
                 Whether or not to return special tokens mask information.
-            return_offsets_mapping (`bool`, *optional*, defaults to `False`):
-                Whether or not to return `(char_start, char_end)` for each token.
-
-                This is only available on fast tokenizers inheriting from [`PreTrainedTokenizerFast`], if using
-                Python's tokenizer, this method will raise `NotImplementedError`.
             return_length  (`bool`, *optional*, defaults to `False`):
                 Whether or not to return the lengths of the encoded inputs.
             verbose (`bool`, *optional*, defaults to `True`):
@@ -176,6 +171,7 @@ class MistralCommonTokenizer(PushToHubMixin):
     Supports the following methods from the `PreTrainedTokenizerBase` class:
 
     - [`~MistralCommonTokenizer.get_vocab`]: Returns the vocabulary as a dictionary of token to index.
+        This is a lossy conversion for Tekkenizer as some decoding errors are collapsed into the same token.
     - [`~MistralCommonTokenizer.encode`]: Encode a string to a list of integers.
     - [`~MistralCommonTokenizer.decode`]: Decode a list of integers to a string.
     - [`~MistralCommonTokenizer.batch_decode`]: Decode a batch of list of integers to a list of strings.
@@ -354,9 +350,13 @@ def get_vocab(self) -> dict[str, int]:
             `Dict[str, int]`: The vocabulary.
         """
         if self._cache_get_vocab is None:
-            self._cache_get_vocab = {
-                token: idx for idx, token in enumerate(self.tokenizer.instruct_tokenizer.tokenizer.vocab())
-            }
+            # We reverse the order to make sure that the first token is the one to be returned when there are multiple tokens with the same string representation.
+            vocab = self.tokenizer.instruct_tokenizer.tokenizer.vocab()
+            self._cache_get_vocab = {token: self._piece_to_id(token, False) for token in vocab}
+            # Order the dict.
+            self._cache_get_vocab = dict(
+                sorted(((k, v) for k, v in self._cache_get_vocab.items()), key=lambda x: x[1])
+            )
         return self._cache_get_vocab
 
     def __len__(self):
@@ -517,7 +517,7 @@ def batch_decode(
 
     def _is_control_token(self, token_id: int) -> bool:
         if self._tokenizer_type == MistralTokenizerType.spm:
-            return token_id in self.tokenizer.instruct_tokenizer.tokenizer._control_tokens()
+            return token_id in self.tokenizer.instruct_tokenizer.tokenizer._control_tokens
         elif self._tokenizer_type == MistralTokenizerType.tekken:
             return token_id < self.tokenizer.instruct_tokenizer.tokenizer.num_special_tokens
         else:
@@ -563,15 +563,27 @@ def convert_ids_to_tokens(
             return tokens[0]
         return tokens
 
-    def _piece_to_id(self, piece: str) -> int:
+    def _tekken_piece_to_id(self, piece: str, warn: bool) -> int:
+        tekken_tokenizer = self.tokenizer.instruct_tokenizer.tokenizer
+        assert isinstance(tekken_tokenizer, Tekkenizer), type(tekken_tokenizer)
+
+        piece_bytes = piece.encode("utf-8")
+        shift = tekken_tokenizer.num_special_tokens
+        try:
+            return shift + tekken_tokenizer._tekken_token2id_nospecial[piece_bytes]
+        except KeyError:
+            piece_str = piece_bytes.decode("utf-8")
+            if piece_str in tekken_tokenizer._special_tokens_reverse_vocab:
+                return tekken_tokenizer._special_tokens_reverse_vocab[piece_str]
+            if warn:
+                logger.warning("Failed to convert token %s to id, replacing with <unk>", piece_bytes)
+            return tekken_tokenizer.unk_id
+
+    def _piece_to_id(self, piece: str, warn: bool) -> int:
         if self._tokenizer_type == MistralTokenizerType.spm:
             return self.tokenizer.instruct_tokenizer.tokenizer._model.piece_to_id(piece)
         elif self._tokenizer_type == MistralTokenizerType.tekken:
-            pieces = self.tokenizer.instruct_tokenizer.tokenizer._model.encode(
-                piece, allowed_special="all", disallowed_special=set()
-            )
-            assert len(pieces) == 1, f"Expected to decode 1 token, got {len(pieces)}"
-            return pieces[0]
+            return self._tekken_piece_to_id(piece, warn)
         else:
             raise ValueError(f"Unknown tokenizer type: {self._tokenizer_type}")
 
@@ -595,7 +607,7 @@ def convert_tokens_to_ids(self, tokens: Union[str, list[str]]) -> Union[int, lis
 
         ids: list[int] = []
         for token in tokens:
-            ids.append(self._piece_to_id(token))
+            ids.append(self._piece_to_id(token, True))
 
         if one_token:
             return ids[0]
@@ -647,13 +659,7 @@ def _encode_plus(
         return_special_tokens_mask: bool = False,
         return_length: bool = False,
         verbose: bool = True,
-        **kwargs,
     ) -> BatchEncoding:
-        if kwargs:
-            raise ValueError(
-                f"Kwargs {list(kwargs.keys())} are not supported by `MistralCommonTokenizer._encode_plus`."
-            )
-
         def get_input_ids(text):
             if isinstance(text, str):
                 return self._text_to_ids(text, add_special_tokens)
@@ -699,10 +705,8 @@ def _batch_encode_plus(
         return_attention_mask: Optional[bool] = None,
         return_overflowing_tokens: bool = False,
         return_special_tokens_mask: bool = False,
-        return_offsets_mapping: bool = False,
         return_length: bool = False,
         verbose: bool = True,
-        **kwargs,
     ) -> BatchEncoding:
         def get_input_ids(text):
             if isinstance(text, str):
@@ -712,13 +716,6 @@ def get_input_ids(text):
             else:
                 raise ValueError("Input is not valid. Should be a string or a list/tuple of integers.")
 
-        if return_offsets_mapping:
-            raise NotImplementedError(
-                "return_offset_mapping is not available when using Python tokenizers. "
-                "To use this feature, change your tokenizer to one deriving from "
-                "transformers.PreTrainedTokenizerFast."
-            )
-
         input_ids = []
         for ids in batch_text:
             input_ids.append(get_input_ids(ids))
@@ -746,7 +743,7 @@ def _all_special_ids(self) -> set[int]:
         if self._tokenizer_type == MistralTokenizerType.tekken:
             return {t["rank"] for t in self.tokenizer.instruct_tokenizer.tokenizer._all_special_tokens}
         elif self._tokenizer_type == MistralTokenizerType.spm:
-            return self.tokenizer.instruct_tokenizer.tokenizer._control_tokens()
+            return self.tokenizer.instruct_tokenizer.tokenizer._control_tokens
         else:
             raise ValueError(f"Unknown tokenizer type: {self._tokenizer_type}")
 
@@ -966,9 +963,7 @@ def _get_padding_truncation_strategies(
                     logger.warning(
                         "Truncation was not explicitly activated but `max_length` is provided a specific value, please"
                         " use `truncation=True` to explicitly truncate examples to max length. Defaulting to"
-                        " 'longest_first' truncation strategy. If you encode pairs of sequences (GLUE-style) with the"
-                        " tokenizer you can select this strategy more precisely by providing a specific strategy to"
-                        " `truncation`."
+                        " 'longest_first' truncation strategy."
                     )
                 self.deprecation_warnings["Truncation-not-explicitly-activated"] = True
             truncation = "longest_first"
@@ -1376,6 +1371,7 @@ def apply_chat_template(
         self,
         conversation: Union[list[dict[str, str]], list[list[dict[str, str]]]],
         tools: Optional[list[Union[dict, Callable]]] = None,
+        add_generation_prompt: bool = False,
         continue_final_message: bool = False,
         tokenize: bool = True,
         padding: Union[bool, str, PaddingStrategy] = False,
@@ -1398,6 +1394,9 @@ def apply_chat_template(
                 giving the name, description and argument types for the tool. See our
                 [chat templating guide](https://huggingface.co/docs/transformers/main/en/chat_templating#automated-function-conversion-for-tool-use)
                 for more information.
+            add_generation_prompt (`bool`, *optional*):
+                This argument is a no-op for `MistralCommonTokenizer`. However it cannot be used at the same time as `continue_final_message` to keep the API consistent and
+                if any conversation ends with an assistant message, it will raise an error. In such case, use `continue_final_message` instead.
             continue_final_message (bool, *optional*):
                 If this is set, the chat will be formatted so that the final
                 message in the chat is open-ended, without any EOS tokens. The model will continue this message
@@ -1442,6 +1441,9 @@ def apply_chat_template(
         if not isinstance(truncation, bool):
             raise TypeError("`truncation` must be a boolean for `apply_chat_template` method.")
 
+        if add_generation_prompt and continue_final_message:
+            raise ValueError("Cannot use both `add_generation_prompt` and `continue_final_message`.")
+
         if isinstance(conversation, (list, tuple)) and (
             isinstance(conversation[0], (list, tuple)) or hasattr(conversation[0], "messages")
         ):
@@ -1451,6 +1453,14 @@ def apply_chat_template(
             conversations = [conversation]
             is_batched = False
 
+        if add_generation_prompt:
+            for conversation in conversations:
+                last_message = conversation[-1]
+                if last_message.get("role") == "assistant":
+                    raise ValueError(
+                        "The last message in the conversation is already an assistant message. Consider using `continue_final_message` instead."
+                    )
+
         def _maybe_adapt_message(message: dict[str, Any]) -> None:
             """Adapt message to `mistral-common` format and leave validation to `mistral-common`."""
             if not isinstance(message, dict):
@@ -1543,7 +1553,7 @@ def _maybe_adapt_message(message: dict[str, Any]) -> None:
                                 "Unable to convert output to PyTorch tensors format, PyTorch is not installed."
                             )
 
-                        pixel_values = torch.tensor(images)
+                        pixel_values = torch.from_numpy(np.stack(images))
                     elif return_tensors == "np":
                         pixel_values = np.array(images)
                     elif return_tensors is None:
@@ -1667,7 +1677,6 @@ def _is_valid_text_input(t):
                 return_special_tokens_mask=return_special_tokens_mask,
                 return_length=return_length,
                 verbose=verbose,
-                **kwargs,
             )
         else:
             return self._encode_plus(
@@ -1685,7 +1694,6 @@ def _is_valid_text_input(t):
                 return_special_tokens_mask=return_special_tokens_mask,
                 return_length=return_length,
                 verbose=verbose,
-                **kwargs,
             )
 
     @classmethod
@@ -1760,9 +1768,10 @@ def from_pretrained(
             raise ValueError("`init_inputs` are not supported by `MistralCommonTokenizer.from_pretrained`.")
 
         # Handle kwargs and AutoTokenizer case
-        if kwargs and not set(kwargs.keys()).issubset({"_from_auto", "trust_remote_code"}):
+        ignore_subset = {"_from_auto", "trust_remote_code"}
+        if kwargs and not (set_kwargs := set(kwargs.keys())).issubset(ignore_subset):
             raise ValueError(
-                f"Kwargs {list(kwargs.keys())} are not supported by `MistralCommonTokenizer.from_pretrained`."
+                f"Kwargs {list(set_kwargs - ignore_subset)} are not supported by `MistralCommonTokenizer.from_pretrained`."
             )
 
         if not os.path.isdir(pretrained_model_name_or_path):

filename: tests/test_tokenization_mistral_common.py
patch: @@ -34,6 +34,7 @@
     from mistral_common.exceptions import InvalidMessageStructureException
     from mistral_common.protocol.instruct.request import ChatCompletionRequest
     from mistral_common.protocol.transcription.request import TranscriptionRequest
+    from mistral_common.tokens.tokenizers.base import SpecialTokenPolicy
     from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
     from mistral_common.tokens.tokenizers.utils import list_local_hf_repo_files
 
@@ -82,6 +83,10 @@ def setUpClass(cls):
         cls.ref_tokenizer: MistralTokenizer = MistralTokenizer.from_hf_hub(
             cls.repo_id, local_files_only=cls.local_files_only
         )
+
+        # Define SPM tokenizer to test the private methods that handle SPM and Tekken differencies.
+        cls.spm_repo_id = "mistralai/Mistral-7B-v0.3"
+
         # cls.tokenizer_audio: MistralCommonTokenizer = AutoTokenizer.from_pretrained(
         #     "hf-internal-testing/namesspace-mistralai-repo_name-Voxtral-Mini-3B-2507"
         # )
@@ -127,12 +132,53 @@ def tearDownClass(cls):
         del cls.ref_special_ids
         gc.collect()
 
+    # Copy paste of `MistralCommonTokenizer._tekken_piece_to_id`
     def _ref_piece_to_id(self, piece: str) -> int:
-        pieces = self.ref_tokenizer.instruct_tokenizer.tokenizer._model.encode(
-            piece, allowed_special="all", disallowed_special=set()
+        tekken_tokenizer = self.ref_tokenizer.instruct_tokenizer.tokenizer
+
+        piece_bytes = piece.encode("utf-8")
+        shift = tekken_tokenizer.num_special_tokens
+        try:
+            return shift + tekken_tokenizer._tekken_token2id_nospecial[piece_bytes]
+        except KeyError:
+            piece_str = piece_bytes.decode("utf-8")
+            if piece_str in tekken_tokenizer._special_tokens_reverse_vocab:
+                return tekken_tokenizer._special_tokens_reverse_vocab[piece_str]
+            return tekken_tokenizer.unk_id
+
+    def _get_spm_tokenizer(self) -> MistralCommonTokenizer:
+        local_files_only = len(list_local_hf_repo_files(self.spm_repo_id, revision=None)) > 0
+        return AutoTokenizer.from_pretrained(
+            self.spm_repo_id,
+            local_files_only=local_files_only,
+            revision=None,
+            tokenizer_type="mistral",
         )
-        assert len(pieces) == 1, f"Expected to decode 1 token, got {len(pieces)}"
-        return pieces[0]
+
+    def test_spm_vs_tekken_is_control_token(self):
+        spm_tokenizer = self._get_spm_tokenizer()
+        self.assertTrue(spm_tokenizer._is_control_token(1))
+        self.assertTrue(spm_tokenizer._is_control_token(768))
+        self.assertFalse(spm_tokenizer._is_control_token(2000))
+
+        self.assertTrue(self.tokenizer._is_control_token(0))
+        self.assertTrue(self.tokenizer._is_control_token(768))
+        self.assertTrue(self.tokenizer._is_control_token(999))
+        self.assertFalse(self.tokenizer._is_control_token(1000))
+
+    def test_spm_vs_tekken_piece_to_id(self):
+        spm_tokenizer = self._get_spm_tokenizer()
+        self.assertEqual(spm_tokenizer._piece_to_id("
", False), 1)
+        self.assertEqual(spm_tokenizer._piece_to_id("h", False), 29484)
+
+        self.assertEqual(self.tokenizer._piece_to_id("
", False), 1)
+        self.assertEqual(self._ref_piece_to_id("
"), 1)
+        self.assertEqual(self.tokenizer._piece_to_id("\u0000", False), 1000)
+        self.assertEqual(self._ref_piece_to_id("\u0000"), 1000)
+        self.assertEqual(self.tokenizer._piece_to_id(" String", False), 3000)
+        self.assertEqual(self._ref_piece_to_id(" String"), 3000)
+        self.assertEqual(self.tokenizer._piece_to_id("后汉书", False), 131071)
+        self.assertEqual(self._ref_piece_to_id("后汉书"), 131071)
 
     def test_vocab_size(self):
         self.assertEqual(self.tokenizer.vocab_size, self.ref_tokenizer.instruct_tokenizer.tokenizer.n_words)
@@ -785,6 +831,40 @@ def test_apply_chat_template_continue_final_message(self):
         with self.assertRaises(InvalidMessageStructureException):
             self.tokenizer.apply_chat_template(conversation, tokenize=False, continue_final_message=False)
 
+    def test_apply_chat_template_with_add_generation_prompt(self):
+        conversation = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "Hi!"},
+        ]
+
+        # Test 1:
+        # with add_generation_prompt
+        for add_generation_prompt in [False, True]:
+            expected_tokenized = self.ref_tokenizer.encode_chat_completion(
+                ChatCompletionRequest.from_openai(conversation)
+            )
+            token_outputs = self.tokenizer.apply_chat_template(
+                conversation, tokenize=True, add_generation_prompt=add_generation_prompt
+            )
+            self.assertEqual(token_outputs, expected_tokenized.tokens)
+
+        # Test 2:
+        # with continue_final_message
+        with self.assertRaises(ValueError):
+            self.tokenizer.apply_chat_template(
+                conversation, tokenize=True, add_generation_prompt=True, continue_final_message=True
+            )
+
+        # Test 3:
+        # with last message with assistant role
+        conversation = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "Hi!"},
+            {"role": "asistant", "content": "Hey!"},
+        ]
+        with self.assertRaises(ValueError):
+            self.tokenizer.apply_chat_template(conversation, tokenize=True, add_generation_prompt=True)
+
     def test_apply_chat_template_with_tools(self):
         conversation = [
             {"role": "system", "content": "You are a helpful assistant."},
@@ -891,7 +971,8 @@ def test_apply_chat_template_with_image(self):
             conversation, tokenize=True, return_dict=True, return_tensors="pt"
         )
         self.assertEqual(output_dict["input_ids"].tolist()[0], expected_tokenized.tokens)
-        self.assertTrue(torch.allclose(output_dict["pixel_values"], torch.tensor(expected_tokenized.images)))
+        expected_images_pt_tensor = torch.from_numpy(np.stack(expected_tokenized.images))
+        self.assertTrue(torch.allclose(output_dict["pixel_values"], expected_images_pt_tensor))
 
     def test_apply_chat_template_with_audio(self):
         ref_conversation = conversation = [
@@ -1063,7 +1144,7 @@ def test_batch_apply_chat_template(self):
         ):
             self.tokenizer.apply_chat_template(conversations, tools=tools, tokenize=True, unk_args="")
 
-    def test_batch_apply_images(self):
+    def test_batch_apply_chat_template_images(self):
         conversations = [
             [
                 {"role": "system", "content": "You are a helpful assistant."},
@@ -1135,7 +1216,8 @@ def test_batch_apply_images(self):
         )
         self.assertEqual(output["input_ids"].tolist(), [expected_tokenized.tokens] * 3)
         self.assertEqual(output["input_ids"].shape[0], len(expected_tokenized.images) * 3)
-        self.assertTrue(torch.allclose(output["pixel_values"], torch.tensor([expected_tokenized.images] * 3)))
+        expected_images_pt_tensor = torch.from_numpy(np.stack([expected_tokenized.images] * 3))
+        self.assertTrue(torch.allclose(output["pixel_values"], expected_images_pt_tensor))
 
         output = self.tokenizer.apply_chat_template(
             conversations, tokenize=True, return_dict=True, return_tensors="np"
@@ -1194,6 +1276,54 @@ def test_batch_apply_chat_template_with_continue_final_message(self):
                 continue_final_message=True,
             )
 
+    def test_batch_apply_chat_template_with_add_generation_prompt(self):
+        conversations = [
+            [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "Hi!"},
+            ],
+            [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "Hi!"},
+            ],
+        ]
+
+        # Test 1:
+        # with add_generation_prompt
+        for add_generation_prompt in [False, True]:
+            expected_tokenized = [
+                self.ref_tokenizer.encode_chat_completion(ChatCompletionRequest.from_openai(conversation))
+                for conversation in conversations
+            ]
+            token_outputs = self.tokenizer.apply_chat_template(
+                conversations, tokenize=True, add_generation_prompt=add_generation_prompt
+            )
+            for output, expected in zip(token_outputs, expected_tokenized):
+                self.assertEqual(output, expected.tokens)
+
+        # Test 2:
+        # with continue_final_message
+        with self.assertRaises(ValueError):
+            self.tokenizer.apply_chat_template(
+                conversations, tokenize=True, add_generation_prompt=True, continue_final_message=True
+            )
+
+        # Test 3:
+        # with last message with assistant role
+        conversations = [
+            [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "Hi!"},
+                {"role": "asistant", "content": "Hey!"},
+            ],
+            [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "Hi!"},
+            ],
+        ]
+        with self.assertRaises(ValueError):
+            self.tokenizer.apply_chat_template(conversations, tokenize=True, add_generation_prompt=True)
+
     def test_batch_apply_chat_template_with_truncation(
         self,
     ):
@@ -1799,3 +1929,16 @@ def test_batch_call_with_padding_and_truncation(self):
                         for i, ids in enumerate(expected_tokens)
                     ],
                 )
+
+    def test_get_vocab(self):
+        vocab = self.tokenizer.get_vocab()
+        # loss of some tokens due to conversion
+        self.assertNotEqual(len(vocab), len(self.tokenizer))
+        for token, id_token in vocab.items():
+            # Issue during conversion
+            if id_token == 0 and token != "<unk>":
+                continue
+            self.assertEqual(self.tokenizer.convert_tokens_to_ids(token), id_token)
+            self.assertEqual(
+                self.ref_tokenizer.decode([id_token], special_token_policy=SpecialTokenPolicy.KEEP), token
+            )