[Refactor] TokenizerRegistry only uses lazy imports (vllm-project#30609)

Signed-off-by: DarkLight1337 <[email protected]>

Author: DarkLight1337

tests/test_inputs.py

@@ -7,7 +7,7 @@
 from vllm.inputs import zip_enc_dec_prompts
 from vllm.inputs.parse import parse_raw_prompts
 from vllm.inputs.preprocess import InputPreprocessor
-from vllm.tokenizers import init_tokenizer_from_config
+from vllm.tokenizers import cached_tokenizer_from_config
 
 pytestmark = pytest.mark.cpu_test
 
@@ -108,7 +108,7 @@ def test_zip_enc_dec_prompts(mm_processor_kwargs, expected_mm_kwargs):
 )
 def test_preprocessor_always_mm_code_path(model_id, prompt):
     model_config = ModelConfig(model=model_id)
-    tokenizer = init_tokenizer_from_config(model_config)
+    tokenizer = cached_tokenizer_from_config(model_config)
     input_preprocessor = InputPreprocessor(model_config, tokenizer)
 
     # HF processor adds sep token

tests/tokenizers_/test_basic.py

@@ -3,38 +3,39 @@
 from typing import _get_protocol_attrs  # type: ignore
 
 import pytest
-from transformers import PreTrainedTokenizerBase
+from transformers import (
+    PreTrainedTokenizer,
+    PreTrainedTokenizerBase,
+    PreTrainedTokenizerFast,
+)
 
 from vllm.tokenizers import TokenizerLike, get_tokenizer
+from vllm.tokenizers.mistral import MistralTokenizer
 
 
 def _get_missing_attrs(obj: object, target: type):
     return [k for k in _get_protocol_attrs(target) if not hasattr(obj, k)]
 
 
+def _assert_tokenizer_like(tokenizer: object):
+    missing_attrs = _get_missing_attrs(tokenizer, TokenizerLike)
+    assert not missing_attrs, f"Missing attrs: {missing_attrs}"
+
+
 def test_tokenizer_like_protocol():
-    assert not (
-        missing_attrs := _get_missing_attrs(
-            get_tokenizer("gpt2", use_fast=False),
-            TokenizerLike,
-        )
-    ), f"Missing attrs: {missing_attrs}"
-
-    assert not (
-        missing_attrs := _get_missing_attrs(
-            get_tokenizer("gpt2", use_fast=True),
-            TokenizerLike,
-        )
-    ), f"Missing attrs: {missing_attrs}"
-
-    assert not (
-        missing_attrs := _get_missing_attrs(
-            get_tokenizer(
-                "mistralai/Mistral-7B-Instruct-v0.3", tokenizer_mode="mistral"
-            ),
-            TokenizerLike,
-        )
-    ), f"Missing attrs: {missing_attrs}"
+    tokenizer = get_tokenizer("gpt2", use_fast=False)
+    assert isinstance(tokenizer, PreTrainedTokenizer)
+    _assert_tokenizer_like(tokenizer)
+
+    tokenizer = get_tokenizer("gpt2", use_fast=True)
+    assert isinstance(tokenizer, PreTrainedTokenizerFast)
+    _assert_tokenizer_like(tokenizer)
+
+    tokenizer = get_tokenizer(
+        "mistralai/Mistral-7B-Instruct-v0.3", tokenizer_mode="mistral"
+    )
+    assert isinstance(tokenizer, MistralTokenizer)
+    _assert_tokenizer_like(tokenizer)
 
 
 @pytest.mark.parametrize("tokenizer_name", ["facebook/opt-125m", "gpt2"])

tests/tokenizers_/test_registry.py

@@ -2,7 +2,14 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from pathlib import Path
 
-from vllm.tokenizers import TokenizerLike, TokenizerRegistry, get_tokenizer
+import pytest
+
+from vllm.tokenizers import TokenizerLike
+from vllm.tokenizers.registry import (
+    TokenizerRegistry,
+    get_tokenizer,
+    resolve_tokenizer_args,
+)
 
 
 class TestTokenizer(TokenizerLike):
@@ -40,10 +47,22 @@ def is_fast(self) -> bool:
         return True
 
 
+@pytest.mark.parametrize("runner_type", ["generate", "pooling"])
+def test_resolve_tokenizer_args_idempotent(runner_type):
+    tokenizer_mode, tokenizer_name, args, kwargs = resolve_tokenizer_args(
+        "facebook/opt-125m",
+        runner_type=runner_type,
+    )
+
+    assert (tokenizer_mode, tokenizer_name, args, kwargs) == resolve_tokenizer_args(
+        tokenizer_name, *args, **kwargs
+    )
+
+
 def test_customized_tokenizer():
     TokenizerRegistry.register("test_tokenizer", __name__, TestTokenizer.__name__)
 
-    tokenizer = TokenizerRegistry.get_tokenizer("test_tokenizer", "abc")
+    tokenizer = TokenizerRegistry.load_tokenizer("test_tokenizer", "abc")
     assert isinstance(tokenizer, TestTokenizer)
     assert tokenizer.path_or_repo_id == "abc"
     assert tokenizer.bos_token_id == 0

vllm/entrypoints/chat_utils.py

@@ -50,7 +50,6 @@
 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalDataDict, MultiModalUUIDDict
 from vllm.multimodal.utils import MEDIA_CONNECTOR_REGISTRY, MediaConnector
 from vllm.tokenizers import TokenizerLike
-from vllm.tokenizers.mistral import MistralTokenizer
 from vllm.transformers_utils.chat_templates import get_chat_template_fallback_path
 from vllm.transformers_utils.processor import cached_get_processor
 from vllm.utils import random_uuid
@@ -60,6 +59,8 @@
 
 if TYPE_CHECKING:
     import torch
+
+    from vllm.tokenizers.mistral import MistralTokenizer
 else:
     torch = LazyLoader("torch", globals(), "torch")
 
@@ -1832,7 +1833,7 @@ def apply_hf_chat_template(
 
 
 def apply_mistral_chat_template(
-    tokenizer: MistralTokenizer,
+    tokenizer: "MistralTokenizer",
     messages: list[ChatCompletionMessageParam],
     chat_template: str | None,
     tools: list[dict[str, Any]] | None,

vllm/tokenizers/__init__.py

@@ -1,9 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from .deepseekv32 import DeepseekV32Tokenizer
-from .hf import HfTokenizer
-from .mistral import MistralTokenizer
 from .protocol import TokenizerLike
 from .registry import (
     TokenizerRegistry,
@@ -15,12 +12,9 @@
 
 __all__ = [
     "TokenizerLike",
-    "HfTokenizer",
-    "MistralTokenizer",
     "TokenizerRegistry",
     "cached_get_tokenizer",
     "get_tokenizer",
     "cached_tokenizer_from_config",
     "init_tokenizer_from_config",
-    "DeepseekV32Tokenizer",
 ]

vllm/tokenizers/deepseekv32.py

@@ -2,24 +2,18 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from pathlib import Path
+from typing import Any
 
 from transformers import BatchEncoding
 
-from .deepseek_v32_encoding import encode_messages
-from .hf import HfTokenizer, TokenizerLike
-from .registry import TokenizerRegistry
+from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
 
+from .deepseek_v32_encoding import encode_messages
+from .hf import CachedHfTokenizer
+from .protocol import TokenizerLike
 
-@TokenizerRegistry.register("deepseek_v32")
-class DeepseekV32Tokenizer(HfTokenizer):
-    def __init__(self, tokenizer: TokenizerLike):
-        self.tokenizer = tokenizer
-        self.name_or_path = (
-            tokenizer.name_or_path if hasattr(tokenizer, "name_or_path") else ""
-        )
-        self._added_vocab = self.tokenizer.get_added_vocab()
-        self._added_vocab_size = len(self._added_vocab)
 
+class DeepseekV32Tokenizer(CachedHfTokenizer):
     @classmethod
     def from_pretrained(
         cls,
@@ -40,7 +34,21 @@ def from_pretrained(
         )
         return DeepseekV32Tokenizer(tokenizer)
 
-    def apply_chat_template(self, messages, tools=None, **kwargs):
+    def __init__(self, tokenizer: TokenizerLike) -> None:
+        super().__init__()
+
+        self.tokenizer = tokenizer
+        self.name_or_path = getattr(tokenizer, "name_or_path", "")
+
+        self._added_vocab = self.tokenizer.get_added_vocab()
+        self._added_vocab_size = len(self._added_vocab)
+
+    def apply_chat_template(
+        self,
+        messages: list["ChatCompletionMessageParam"],
+        tools: list[dict[str, Any]] | None = None,
+        **kwargs,
+    ) -> str | list[int]:
         thinking = kwargs.get("thinking", False)
         thinking_mode = "thinking"
         if not thinking:
@@ -49,13 +57,24 @@ def apply_chat_template(self, messages, tools=None, **kwargs):
         messages = conversation.copy()
         if tools is not None and len(tools) > 0:
             messages.insert(0, {"role": "system"})
-            messages[0]["tools"] = tools
+            messages[0]["tools"] = tools  # type: ignore[typeddict-unknown-key]
 
         # Historical reasoning content is dropped when a new user message is introduced
         drop_thinking = messages[-1]["role"] == "user"
 
         encode_config = dict(thinking_mode=thinking_mode, drop_thinking=drop_thinking)
         prompt_str = encode_messages(messages, **encode_config)  # type: ignore
+
+        if kwargs.get("tokenize", True):
+            tokenizer_kwargs = {
+                k: kwargs[k] for k in ("truncation", "max_length") if k in kwargs
+            }
+            return self.encode(
+                prompt_str,
+                add_special_tokens=False,
+                **tokenizer_kwargs,
+            )
+
         return prompt_str
 
     def num_special_tokens_to_add(self) -> int:

vllm/tokenizers/hf.py

@@ -3,22 +3,18 @@
 import contextlib
 import copy
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TypeAlias
 
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
 
 from vllm.transformers_utils.config import get_sentence_transformer_tokenizer_config
 
 from .protocol import TokenizerLike
-from .registry import TokenizerRegistry
 
-if TYPE_CHECKING:
-    from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
+HfTokenizer: TypeAlias = PreTrainedTokenizer | PreTrainedTokenizerFast
 
 
-def get_cached_tokenizer(
-    tokenizer: "PreTrainedTokenizer | PreTrainedTokenizerFast",
-) -> TokenizerLike:
+def get_cached_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer:
     """
     By default, transformers will recompute multiple tokenizer properties
     each time they are called, leading to a significant slowdown.
@@ -65,11 +61,10 @@ def __reduce__(self):
     CachedTokenizer.__name__ = f"Cached{tokenizer.__class__.__name__}"
 
     cached_tokenizer.__class__ = CachedTokenizer
-    return cached_tokenizer  # type: ignore
+    return cached_tokenizer
 
 
-@TokenizerRegistry.register("hf")
-class HfTokenizer(TokenizerLike):
+class CachedHfTokenizer(TokenizerLike):
     @classmethod
     def from_pretrained(
         cls,
@@ -79,7 +74,7 @@ def from_pretrained(
         revision: str | None = None,
         download_dir: str | None = None,
         **kwargs,
-    ) -> "TokenizerLike":
+    ) -> HfTokenizer:
         try:
             tokenizer = AutoTokenizer.from_pretrained(
                 path_or_repo_id,

vllm/tokenizers/mistral.py

@@ -3,10 +3,11 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, cast
 
+from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
+from vllm.entrypoints.openai.protocol import ChatCompletionRequest
 from vllm.logger import init_logger
 
 from .protocol import TokenizerLike
-from .registry import TokenizerRegistry
 
 if TYPE_CHECKING:
     from mistral_common.protocol.instruct.request import (
@@ -15,9 +16,6 @@
     from mistral_common.tokens.tokenizers.tekken import Tekkenizer
     from transformers import BatchEncoding
 
-    from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
-    from vllm.entrypoints.openai.protocol import ChatCompletionRequest
-
     try:
         # Transformers v5
         from transformers.tokenization_mistral_common import MistralCommonBackend
@@ -201,7 +199,6 @@ def _tekken_token_to_id(tokenizer: "Tekkenizer", t: str | bytes) -> int:
         return tokenizer.unk_id
 
 
-@TokenizerRegistry.register("mistral")
 class MistralTokenizer(TokenizerLike):
     @classmethod
     def from_pretrained(

vllm/tokenizers/protocol.py

@@ -97,7 +97,7 @@ def apply_chat_template(
         messages: list["ChatCompletionMessageParam"],
         tools: list[dict[str, Any]] | None = None,
         **kwargs,
-    ) -> list[int]:
+    ) -> str | list[int]:
         raise NotImplementedError
 
     def convert_tokens_to_string(self, tokens: list[str]) -> str: