| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import contextlib |
| 4 | from dataclasses import dataclass, field |
| 5 | from functools import lru_cache |
| 6 | from pathlib import Path |
| 7 | from typing import TYPE_CHECKING |
| 8 | |
| 9 | import huggingface_hub |
| 10 | from typing_extensions import TypeVar, assert_never |
| 11 | |
| 12 | from vllm import envs |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.transformers_utils.config import _maybe_register_hf_config, get_config |
| 15 | from vllm.transformers_utils.repo_utils import ( |
| 16 | any_pattern_in_repo_files, |
| 17 | is_mistral_model_repo, |
| 18 | ) |
| 19 | from vllm.utils.import_utils import resolve_obj_by_qualname |
| 20 | |
| 21 | from .hf import CachedHfTokenizer |
| 22 | from .protocol import TokenizerLike |
| 23 | |
| 24 | if TYPE_CHECKING: |
| 25 | from vllm.config.model import ModelConfig, RunnerType |
| 26 | |
| 27 | logger = init_logger(__name__) |
| 28 | |
| 29 | |
| 30 | # Model types whose hub tokenizer_class is incorrect and should be overridden with |
| 31 | # TokenizersBackend (the generic fast tokenizer). Adding a model type here is always a |
| 32 | # temporary workaround and better long term solutions are: |
| 33 | # - Add model type to MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS in transformers (better) |
| 34 | # - Fix tokenizer_class on the hub for the affected models (best) |
| 35 | _MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = { |
| 36 | "internlm2", |
| 37 | "step3_vl", |
| 38 | "step3p7", |
| 39 | "unlimited-ocr", |
| 40 | } |
| 41 | |
| 42 | _VLLM_TOKENIZERS = { |
| 43 | # ``cohere`` mode uses the standard cached HF tokenizer; only the |
| 44 | # renderer (template stage) is replaced with a melody-based one. |
| 45 | "cohere": ("hf", "CachedHfTokenizer"), |
| 46 | "deepseek_v32": ("deepseek_v32", "DeepseekV32Tokenizer"), |
| 47 | "deepseek_v4": ("deepseek_v4", "DeepseekV4Tokenizer"), |
| 48 | "deepseek_v41": ("deepseek_v41", "DeepseekV41Tokenizer"), |
| 49 | "hf": ("hf", "CachedHfTokenizer"), |
| 50 | "kimi_audio": ("kimi_audio", "KimiAudioTokenizer"), |
| 51 | "kimi_k3": ("hf", "CachedHfTokenizer"), |
| 52 | "mistral": ("mistral", "MistralTokenizer"), |
| 53 | # Inkling uses the plain HF tokenizer for token operations; the "inkling" |
| 54 | # mode exists to select the InklingRenderer, which renders chat to |
| 55 | # token ids natively (Inkling has no Jinja chat template). |
| 56 | "inkling": ("hf", "CachedHfTokenizer"), |
| 57 | } |
| 58 | |
| 59 | |
| 60 | @dataclass |
| 61 | class _TokenizerRegistry: |
| 62 | # Tokenizer mode -> (tokenizer module, tokenizer class) |
| 63 | tokenizers: dict[str, tuple[str, str]] = field(default_factory=dict) |
| 64 | |
| 65 | def register(self, tokenizer_mode: str, module: str, class_name: str) -> None: |
| 66 | if tokenizer_mode in self.tokenizers: |
| 67 | logger.warning( |
| 68 | "%s.%s is already registered for tokenizer_mode=%r. " |
| 69 | "It is overwritten by the new one.", |
| 70 | module, |
| 71 | class_name, |
| 72 | tokenizer_mode, |
| 73 | ) |
| 74 | |
| 75 | self.tokenizers[tokenizer_mode] = (module, class_name) |
| 76 | |
| 77 | def load_tokenizer_cls(self, tokenizer_mode: str) -> type[TokenizerLike]: |
| 78 | if tokenizer_mode not in self.tokenizers: |
| 79 | raise ValueError(f"No tokenizer registered for {tokenizer_mode=!r}.") |
| 80 | |
| 81 | module, class_name = self.tokenizers[tokenizer_mode] |
| 82 | logger.debug_once(f"Loading {class_name} for {tokenizer_mode=!r}") |
| 83 | |
| 84 | return resolve_obj_by_qualname(f"{module}.{class_name}") |
| 85 | |
| 86 | def load_tokenizer(self, tokenizer_mode: str, *args, **kwargs) -> TokenizerLike: |
| 87 | tokenizer_cls = self.load_tokenizer_cls(tokenizer_mode) |
| 88 | return tokenizer_cls.from_pretrained(*args, **kwargs) |
| 89 | |
| 90 | |
| 91 | TokenizerRegistry = _TokenizerRegistry( |
| 92 | { |
| 93 | mode: (f"vllm.tokenizers.{mod_relname}", cls_name) |
| 94 | for mode, (mod_relname, cls_name) in _VLLM_TOKENIZERS.items() |
| 95 | } |
| 96 | ) |
| 97 | |
| 98 | |
| 99 | def resolve_tokenizer_args( |
| 100 | tokenizer_name: str | Path, |
| 101 | *args, |
| 102 | runner_type: "RunnerType" = "generate", |
| 103 | tokenizer_mode: str = "auto", |
| 104 | **kwargs, |
| 105 | ): |
| 106 | revision: str | None = kwargs.get("revision") |
| 107 | download_dir: str | None = kwargs.get("download_dir") |
| 108 | |
| 109 | if envs.VLLM_USE_MODELSCOPE: |
| 110 | # download model from ModelScope hub, |
| 111 | # lazy import so that modelscope is not required for normal use. |
| 112 | from modelscope.hub.snapshot_download import snapshot_download |
| 113 | |
| 114 | # avoid circular import |
| 115 | from vllm.model_executor.model_loader.weight_utils import get_lock |
| 116 | |
| 117 | # Only set the tokenizer here, model will be downloaded on the workers. |
| 118 | if not Path(tokenizer_name).exists(): |
| 119 | # Use file lock to prevent multiple processes from |
| 120 | # downloading the same file at the same time. |
| 121 | with get_lock(tokenizer_name, download_dir): |
| 122 | tokenizer_path = snapshot_download( |
| 123 | model_id=str(tokenizer_name), |
| 124 | cache_dir=download_dir, |
| 125 | revision=revision, |
| 126 | local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, |
| 127 | # Ignore weights - we only need the tokenizer. |
| 128 | ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"], |
| 129 | ) |
| 130 | tokenizer_name = tokenizer_path |
| 131 | |
| 132 | if "truncation_side" not in kwargs: |
| 133 | if runner_type == "generate" or runner_type == "draft": |
| 134 | kwargs["truncation_side"] = "left" |
| 135 | elif runner_type == "pooling": |
| 136 | kwargs["truncation_side"] = "right" |
| 137 | else: |
| 138 | assert_never(runner_type) |
| 139 | |
| 140 | # Try to use official Mistral tokenizer if possible |
| 141 | if ( |
| 142 | tokenizer_mode == "auto" |
| 143 | and is_mistral_model_repo( |
| 144 | model_name_or_path=str(tokenizer_name), |
| 145 | revision=revision, |
| 146 | token=kwargs.get("token"), |
| 147 | ) |
| 148 | and any_pattern_in_repo_files( |
| 149 | model_name_or_path=str(tokenizer_name), |
| 150 | allow_patterns=["tekken.json", "tokenizer.model.v*"], |
| 151 | revision=revision, |
| 152 | token=kwargs.get("token"), |
| 153 | ) |
| 154 | ): |
| 155 | tokenizer_mode = "mistral" |
| 156 | |
| 157 | # Fallback to HF tokenizer |
| 158 | if tokenizer_mode == "auto": |
| 159 | tokenizer_mode = "hf" |
| 160 | |
| 161 | if tokenizer_mode == "hf": |
| 162 | if kwargs.pop("mistral_format", False): |
| 163 | raise ValueError( |
| 164 | "mistral_format=True is not supported with tokenizer_mode='hf'" |
| 165 | ) |
| 166 | kwargs["mistral_format"] = False |
| 167 | |
| 168 | return tokenizer_mode, tokenizer_name, args, kwargs |
| 169 | |
| 170 | |
| 171 | cached_resolve_tokenizer_args = lru_cache(resolve_tokenizer_args) |
| 172 | |
| 173 | |
| 174 | def tokenizer_args_from_config(config: "ModelConfig", **kwargs): |
| 175 | return cached_resolve_tokenizer_args( |
| 176 | config.tokenizer, |
| 177 | runner_type=config.runner_type, |
| 178 | tokenizer_mode=config.tokenizer_mode, |
| 179 | revision=config.tokenizer_revision, |
| 180 | trust_remote_code=config.trust_remote_code, |
| 181 | **kwargs, |
| 182 | ) |
| 183 | |
| 184 | |
| 185 | _T = TypeVar("_T", bound=TokenizerLike, default=TokenizerLike) |
| 186 | |
| 187 | |
| 188 | def get_tokenizer( |
| 189 | tokenizer_name: str | Path, |
| 190 | *args, |
| 191 | tokenizer_cls: type[_T] = TokenizerLike, # type: ignore[assignment] |
| 192 | trust_remote_code: bool = False, |
| 193 | revision: str | None = None, |
| 194 | download_dir: str | None = None, |
| 195 | **kwargs, |
| 196 | ) -> _T: |
| 197 | """Gets a tokenizer for the given model name via HuggingFace or ModelScope.""" |
| 198 | if envs.VLLM_USE_FASTOKENS: |
| 199 | # Process-global, idempotent patch that swaps the Rust BPE backend |
| 200 | # of any HF fast tokenizer loaded afterwards. No-op for non-HF modes. |
| 201 | from .fastokens import apply_fastokens_patch |
| 202 | |
| 203 | apply_fastokens_patch() |
| 204 | |
| 205 | tokenizer_mode, tokenizer_name, args, kwargs = cached_resolve_tokenizer_args( |
| 206 | tokenizer_name, |
| 207 | *args, |
| 208 | trust_remote_code=trust_remote_code, |
| 209 | revision=revision, |
| 210 | download_dir=download_dir, |
| 211 | **kwargs, |
| 212 | ) |
| 213 | |
| 214 | if tokenizer_cls == TokenizerLike: |
| 215 | tokenizer_cls_ = TokenizerRegistry.load_tokenizer_cls(tokenizer_mode) |
| 216 | else: |
| 217 | tokenizer_cls_ = tokenizer_cls |
| 218 | |
| 219 | # Ensure that, if the config were to come from vllm.transformers_utils.config, it is |
| 220 | # registered with AutoConfig before the tokenizer is loaded. This is necessary since |
| 221 | # tokenizer_cls_.from_pretrained will call AutoConfig.from_pretrained internally. |
| 222 | # This may fail for paths that don't have a model config (e.g. LoRA adapters), |
| 223 | # which is fine — those don't need custom config registration. |
| 224 | # HF-backed tokenizers must receive the HF config. In a dual-format Mistral |
| 225 | # repository, auto detection intentionally prefers params.json, but passing |
| 226 | # that generic config to AutoTokenizer can select the wrong tokenizer class. |
| 227 | config_format = "hf" if tokenizer_cls_ is CachedHfTokenizer else "auto" |
| 228 | config = None |
| 229 | with contextlib.suppress(ValueError, OSError): |
| 230 | config = get_config( |
| 231 | tokenizer_name, |
| 232 | trust_remote_code=trust_remote_code, |
| 233 | revision=revision, |
| 234 | config_format=config_format, |
| 235 | token=kwargs.get("token"), |
| 236 | ) |
| 237 | |
| 238 | # Some models have an incorrect tokenizer_class on the hub. |
| 239 | # For these model types, bypass AutoTokenizer and use TokenizersBackend directly. |
| 240 | model_type = getattr(config, "model_type", None) if config else None |
| 241 | if model_type in _MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: |
| 242 | from transformers.tokenization_utils_tokenizers import TokenizersBackend |
| 243 | |
| 244 | logger.debug( |
| 245 | "Overriding tokenizer_class to TokenizersBackend for model_type=%r", |
| 246 | model_type, |
| 247 | ) |
| 248 | tokenizer_cls_ = TokenizersBackend |
| 249 | |
| 250 | if config is not None and tokenizer_cls_ is CachedHfTokenizer: |
| 251 | # AutoTokenizer otherwise reloads config.json internally. Reuse the |
| 252 | # config that get_config just loaded successfully so a concurrent Hub |
| 253 | # cache refresh cannot invalidate the file between the two reads. |
| 254 | kwargs.setdefault("config", config) |
| 255 | |
| 256 | tokenizer = tokenizer_cls_.from_pretrained(tokenizer_name, *args, **kwargs) |
| 257 | if model_type in _MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: |
| 258 | from vllm.tokenizers.hf import get_cached_tokenizer |
| 259 | |
| 260 | tokenizer = get_cached_tokenizer(tokenizer) |
| 261 | if not tokenizer.is_fast: |
| 262 | logger.warning( |
| 263 | "Using a slow tokenizer. This might cause a significant " |
| 264 | "slowdown. Consider using a fast tokenizer instead." |
| 265 | ) |
| 266 | |
| 267 | return tokenizer # type: ignore |
| 268 | |
| 269 | |
| 270 | cached_get_tokenizer = lru_cache(get_tokenizer) |
| 271 | |
| 272 | |
| 273 | def cached_tokenizer_from_config(model_config: "ModelConfig", **kwargs): |
| 274 | if model_config.skip_tokenizer_init: |
| 275 | return None |
| 276 | |
| 277 | _maybe_register_hf_config(getattr(model_config, "hf_config", None)) |
| 278 | |
| 279 | return cached_get_tokenizer( |
| 280 | model_config.tokenizer, |
| 281 | runner_type=model_config.runner_type, |
| 282 | tokenizer_mode=model_config.tokenizer_mode, |
| 283 | revision=model_config.tokenizer_revision, |
| 284 | trust_remote_code=model_config.trust_remote_code, |
| 285 | **kwargs, |
| 286 | ) |