From 2f2401a40664973ff8187bb9fef44c244f9092c2 Mon Sep 17 00:00:00 2001 From: Patrick Ribbsaeter Date: Sun, 18 Oct 2026 22:20:03 +0200 Subject: [PATCH] fix(tokenizers): align Mistral and Gemma end-token control with deprecation warning (#2977) - Standardize tokenize_messages kwargs on add_start_tokens and add_end_tokens for Mistral and Gemma tokenizers. - Preserve backward-compatibility for legacy add_eos parameter with DeprecationWarning (stacklevel=2). - Fix __call__ end-token forwarding during evaluation and inference. - Add comprehensive test coverage asserting special token control across all permutations (including off-diagonal cases), warning-free canonical path, and legacy equivalence. Fixes #2479. --- .../models/gemma/test_gemma_tokenizer.py | 43 +++++++++++++--- .../models/mistral/test_mistral_tokenizer.py | 51 ++++++++++++++++--- torchtune/models/gemma/_tokenizer.py | 27 ++++++++-- torchtune/models/mistral/_tokenizer.py | 27 ++++++++-- 4 files changed, 126 insertions(+), 22 deletions(-) diff --git a/tests/torchtune/models/gemma/test_gemma_tokenizer.py b/tests/torchtune/models/gemma/test_gemma_tokenizer.py index 920f5f6..f691b31 100644 --- a/tests/torchtune/models/gemma/test_gemma_tokenizer.py +++ b/tests/torchtune/models/gemma/test_gemma_tokenizer.py @@ -4,6 +4,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import warnings + import pytest from tests.common import ASSETS from torchtune.data import Message @@ -55,7 +57,18 @@ class TestGemmaTokenizer: assert expected_tokens == tokens assert expected_mask == mask - def test_tokenize_messages_drop_eos(self, tokenizer, expected_tokens): + @pytest.mark.parametrize( + "add_start_tokens, add_end_tokens", + [ + (True, True), + (True, False), + (False, True), + (False, False), + ], + ) + def test_tokenize_messages_special_token_control( + self, tokenizer, expected_tokens, add_start_tokens, add_end_tokens + ): messages = [ Message( role="user", @@ -74,10 +87,28 @@ class TestGemmaTokenizer: "good conversation over coffee.", ), ] - tokens, mask = tokenizer.tokenize_messages(messages, add_eos=False) - # Drop eos token. - expected_tokens = expected_tokens[:-1] - # On 1 less then with eos - expected_mask = [True] * 75 + [False] * 124 + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + tokens, mask = tokenizer.tokenize_messages( + messages, + add_start_tokens=add_start_tokens, + add_end_tokens=add_end_tokens, + ) + if not add_start_tokens: + expected_tokens = expected_tokens[1:] + if not add_end_tokens: + expected_tokens = expected_tokens[:-1] + expected_mask = [True] * (75 if add_start_tokens else 74) + [False] * ( + 125 if add_end_tokens else 124 + ) assert expected_tokens == tokens assert expected_mask == mask + + with pytest.deprecated_call(): + legacy_tokens, legacy_mask = tokenizer.tokenize_messages( + messages, + add_start_tokens=add_start_tokens, + add_eos=add_end_tokens, + ) + assert legacy_tokens == tokens + assert legacy_mask == mask diff --git a/tests/torchtune/models/mistral/test_mistral_tokenizer.py b/tests/torchtune/models/mistral/test_mistral_tokenizer.py index 9958fda..126f7a5 100644 --- a/tests/torchtune/models/mistral/test_mistral_tokenizer.py +++ b/tests/torchtune/models/mistral/test_mistral_tokenizer.py @@ -4,6 +4,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import warnings + import pytest from tests.common import ASSETS from torchtune.data import Message @@ -86,13 +88,50 @@ class TestMistralTokenizer: assert expected_tokens == tokens assert expected_mask == mask - def test_tokenize_message_drop_eos(self, messages, expected_tokens): + @pytest.mark.parametrize( + "add_start_tokens, add_end_tokens", + [ + (True, True), + (True, False), + (False, True), + (False, False), + ], + ) + def test_tokenize_messages_special_token_control( + self, messages, expected_tokens, add_start_tokens, add_end_tokens + ): tokenizer = self.tokenizer(template=False) - tokens, mask = tokenizer.tokenize_messages(messages, add_eos=False) + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + tokens, mask = tokenizer.tokenize_messages( + messages, + add_start_tokens=add_start_tokens, + add_end_tokens=add_end_tokens, + ) + + if not add_start_tokens: + expected_tokens = expected_tokens[1:] + if not add_end_tokens: + expected_tokens = expected_tokens[:-1] - # Drop eos token. - expected_tokens = expected_tokens[:-1] - # On 1 less then with eos - expected_mask = [True] * 75 + [False] * 124 + expected_mask = [True] * (75 if add_start_tokens else 74) + [False] * ( + 125 if add_end_tokens else 124 + ) assert expected_tokens == tokens assert expected_mask == mask + + with pytest.deprecated_call(): + legacy_tokens, legacy_mask = tokenizer.tokenize_messages( + messages, + add_start_tokens=add_start_tokens, + add_eos=add_end_tokens, + ) + assert legacy_tokens == tokens + assert legacy_mask == mask + + def test_inference_drops_end_token(self, messages, expected_tokens): + tokenizer = self.tokenizer(template=False) + sample = tokenizer({"messages": messages}, inference=True) + + assert sample["tokens"] == expected_tokens[:-1] + assert sample["mask"] == [True] * 75 + [False] * 124 diff --git a/torchtune/models/gemma/_tokenizer.py b/torchtune/models/gemma/_tokenizer.py index ed8c814..329b768 100644 --- a/torchtune/models/gemma/_tokenizer.py +++ b/torchtune/models/gemma/_tokenizer.py @@ -5,6 +5,7 @@ # LICENSE file in the root directory of this source tree. from typing import Any, Mapping, Optional +from warnings import warn from torchtune.data import Message, PromptTemplate from torchtune.modules.transforms import Transform @@ -104,7 +105,9 @@ class GemmaTokenizer(ModelTokenizer, Transform): self, messages: list[Message], *, - add_eos: bool = True, + add_start_tokens: bool = True, + add_end_tokens: bool = True, + add_eos: bool | None = None, ) -> tuple[list[int], list[bool]]: r"""Tokenize a list of messages one at a time then concatenate them, returning a list of tokens and a list of masks. @@ -131,11 +134,25 @@ class GemmaTokenizer(ModelTokenizer, Transform): Args: messages (list[Message]): A list of messages, each containing role, content, and masked attributes. - add_eos (bool): Whether to append EOS after assistant message, default to True + add_start_tokens (bool): Whether to add BOS token to the beginning of the + first message. Default True. + add_end_tokens (bool): Whether to add EOS token to the end of the last + message. Default True. + add_eos (bool | None): Deprecated compatibility alias for add_end_tokens. + When provided, it takes precedence. Returns: tuple[list[int], list[bool]]: The tokenized messages """ + if add_eos is not None: + warn( + "add_eos is deprecated and will be removed in a future release. " + "Please use add_end_tokens instead.", + DeprecationWarning, + stacklevel=2, + ) + add_end_tokens = add_eos + templated_messages = ( self.prompt_template(messages) if self.prompt_template is not None @@ -144,8 +161,8 @@ class GemmaTokenizer(ModelTokenizer, Transform): return tokenize_messages_no_special_tokens( tokenizer=self, messages=templated_messages, - bos_id=self.bos_id, - eos_id=self.eos_id if add_eos else None, + bos_id=self.bos_id if add_start_tokens else None, + eos_id=self.eos_id if add_end_tokens else None, truncation_type=self.truncation_type, ) @@ -165,7 +182,7 @@ class GemmaTokenizer(ModelTokenizer, Transform): and the "messages" field removed. """ messages = sample.pop("messages") - tokens, mask = self.tokenize_messages(messages) + tokens, mask = self.tokenize_messages(messages, add_end_tokens=not inference) sample["tokens"] = tokens sample["mask"] = mask return sample diff --git a/torchtune/models/mistral/_tokenizer.py b/torchtune/models/mistral/_tokenizer.py index 5a4d5fc..61edf47 100644 --- a/torchtune/models/mistral/_tokenizer.py +++ b/torchtune/models/mistral/_tokenizer.py @@ -5,6 +5,7 @@ # LICENSE file in the root directory of this source tree. from typing import Any, Mapping, Optional +from warnings import warn from torchtune.data import Message, PromptTemplate from torchtune.models.mistral._prompt_template import MistralChatTemplate @@ -129,7 +130,9 @@ class MistralTokenizer(ModelTokenizer, Transform): self, messages: list[Message], *, - add_eos: bool = True, + add_start_tokens: bool = True, + add_end_tokens: bool = True, + add_eos: bool | None = None, ) -> tuple[list[int], list[bool]]: r"""Tokenize a list of messages one at a time then concatenate them, returning a list of tokens and a list of masks. @@ -161,11 +164,25 @@ class MistralTokenizer(ModelTokenizer, Transform): Args: messages (list[Message]): A list of messages, each containing role, content, and masked attributes. - add_eos (bool): Whether to append EOS after assistant message, default to True + add_start_tokens (bool): Whether to add BOS token to the beginning of the + first message. Default True. + add_end_tokens (bool): Whether to add EOS token to the end of the last + message. Default True. + add_eos (bool | None): Deprecated compatibility alias for add_end_tokens. + When provided, it takes precedence. Returns: tuple[list[int], list[bool]]: The tokenized messages """ + if add_eos is not None: + warn( + "add_eos is deprecated and will be removed in a future release. " + "Please use add_end_tokens instead.", + DeprecationWarning, + stacklevel=2, + ) + add_end_tokens = add_eos + templated_messages = ( self.prompt_template(messages) if self.prompt_template is not None @@ -174,8 +191,8 @@ class MistralTokenizer(ModelTokenizer, Transform): return tokenize_messages_no_special_tokens( tokenizer=self, messages=templated_messages, - bos_id=self.bos_id, - eos_id=self.eos_id if add_eos else None, + bos_id=self.bos_id if add_start_tokens else None, + eos_id=self.eos_id if add_end_tokens else None, truncation_type=self.truncation_type, ) @@ -196,7 +213,7 @@ class MistralTokenizer(ModelTokenizer, Transform): inference (bool): Whether the template is being used for inference or not. """ messages = sample.pop("messages") - tokens, mask = self.tokenize_messages(messages) + tokens, mask = self.tokenize_messages(messages, add_end_tokens=not inference) sample["tokens"] = tokens sample["mask"] = mask return sample -- 2.55.0.windows.3