| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import contextlib |
| 4 | import copy |
| 5 | import queue |
| 6 | from pathlib import Path |
| 7 | from typing import TypeAlias, TypeVar |
| 8 | |
| 9 | from transformers import AutoTokenizer, PythonBackend, TokenizersBackend |
| 10 | |
| 11 | from vllm.transformers_utils.config import get_sentence_transformer_tokenizer_config |
| 12 | |
| 13 | from .protocol import TokenizerLike |
| 14 | |
| 15 | HfTokenizer: TypeAlias = PythonBackend | TokenizersBackend |
| 16 | _T = TypeVar("_T", bound=TokenizerLike) |
| 17 | |
| 18 | |
| 19 | class ThreadSafeHFTokenizerMixin: |
| 20 | """Mixin class for thread-safe HF fast tokenizers.""" |
| 21 | |
| 22 | pass |
| 23 | |
| 24 | |
| 25 | def maybe_make_thread_pool(tokenizer: _T, copies: int = 1): |
| 26 | """If `tokenizer` is a `TokenizersBackend`, modify the tokenizer |
| 27 | in-place to make the public interface thread-safe by routing calls |
| 28 | through a deep-copied tokenizer pool. |
| 29 | |
| 30 | Note that: |
| 31 | - Only ``TokenizerLike``'s public interface is thread-safe. |
| 32 | This doesn't include ``_tokenizer`` property nor any mutation |
| 33 | methods like ``add_special_tokens`` or ``add_tokens``. |
| 34 | - Adjacent method calls could happen on different deep copies. |
| 35 | """ |
| 36 | if not isinstance(tokenizer, TokenizersBackend) or isinstance( |
| 37 | tokenizer, ThreadSafeHFTokenizerMixin |
| 38 | ): |
| 39 | return tokenizer |
| 40 | |
| 41 | og_tokenizer = copy.copy(tokenizer) |
| 42 | |
| 43 | tokenizer_pool: queue.Queue[TokenizersBackend] = queue.Queue() |
| 44 | for _ in range(copies): |
| 45 | tokenizer_pool.put(copy.deepcopy(og_tokenizer)) |
| 46 | |
| 47 | @contextlib.contextmanager |
| 48 | def _borrow_from_pool(): |
| 49 | try: |
| 50 | tok = tokenizer_pool.get_nowait() |
| 51 | yield tok |
| 52 | except queue.Empty: |
| 53 | tok = copy.deepcopy(og_tokenizer) |
| 54 | yield tok |
| 55 | finally: |
| 56 | tokenizer_pool.put(tok) |
| 57 | |
| 58 | class TokenizerPool(tokenizer.__class__, ThreadSafeHFTokenizerMixin): # type: ignore |
| 59 | def apply_chat_template(self, *args, **kwargs): |
| 60 | with _borrow_from_pool() as tok: |
| 61 | return tok.apply_chat_template(*args, **kwargs) |
| 62 | |
| 63 | def batch_decode(self, *args, **kwargs): |
| 64 | with _borrow_from_pool() as tok: |
| 65 | return tok.batch_decode(*args, **kwargs) |
| 66 | |
| 67 | def batch_encode(self, *args, **kwargs): |
| 68 | with _borrow_from_pool() as tok: |
| 69 | return tok.batch_encode(*args, **kwargs) |
| 70 | |
| 71 | def convert_tokens_to_ids(self, *args, **kwargs): |
| 72 | with _borrow_from_pool() as tok: |
| 73 | return tok.convert_tokens_to_ids(*args, **kwargs) |
| 74 | |
| 75 | def convert_ids_to_tokens(self, *args, **kwargs): |
| 76 | with _borrow_from_pool() as tok: |
| 77 | return tok.convert_ids_to_tokens(*args, **kwargs) |
| 78 | |
| 79 | def convert_tokens_to_string(self, *args, **kwargs): |
| 80 | with _borrow_from_pool() as tok: |
| 81 | return tok.convert_tokens_to_string(*args, **kwargs) |
| 82 | |
| 83 | def decode(self, *args, **kwargs): |
| 84 | with _borrow_from_pool() as tok: |
| 85 | return tok.decode(*args, **kwargs) |
| 86 | |
| 87 | def encode(self, *args, **kwargs): |
| 88 | with _borrow_from_pool() as tok: |
| 89 | return tok.encode(*args, **kwargs) |
| 90 | |
| 91 | def __call__(self, *args, **kwargs): |
| 92 | with _borrow_from_pool() as tok: |
| 93 | return tok(*args, **kwargs) |
| 94 | |
| 95 | def __reduce__(self): |
| 96 | return maybe_make_thread_pool, (og_tokenizer, copies) |
| 97 | |
| 98 | TokenizerPool.__name__ = f"TokenizerPool{og_tokenizer.__class__.__name__}" |
| 99 | |
| 100 | tokenizer.__class__ = TokenizerPool |
| 101 | # Return the tokenizer: TokenizerPool.__reduce__ reconstructs through this |
| 102 | # function, so falling off the end would unpickle to None (issue #45433). |
| 103 | return tokenizer |
| 104 | |
| 105 | |
| 106 | def get_cached_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer: |
| 107 | """By default, transformers will recompute multiple tokenizer properties |
| 108 | each time they are called, leading to a significant slowdown. |
| 109 | This proxy caches these properties for faster access. |
| 110 | """ |
| 111 | cached_tokenizer = copy.copy(tokenizer) |
| 112 | |
| 113 | tokenizer_all_special_ids = tokenizer.all_special_ids |
| 114 | tokenizer_all_special_tokens = tokenizer.all_special_tokens |
| 115 | tokenizer_vocab = tokenizer.get_vocab() |
| 116 | tokenizer_len = len(tokenizer) |
| 117 | # The underlying tokenizer class could be a specific backend, |
| 118 | # which does not always implement is_fast in Transformers |
| 119 | tokenizer_is_fast = getattr(tokenizer, "is_fast", True) |
| 120 | |
| 121 | max_token_id = max(tokenizer_vocab.values()) |
| 122 | max_chars_per_token = max(len(tok) for tok in tokenizer_vocab) |
| 123 | |
| 124 | # Some tokenizers (e.g., QwenTokenizer) have special tokens that |
| 125 | # are added and included in the implementation of the vocab_size |
| 126 | # property, but not in get_vocab(); if there is an implementation |
| 127 | # of vocab size, we should take the greater value. vocab_size is a |
| 128 | # count, so the largest id it implies is vocab_size - 1. |
| 129 | if hasattr(tokenizer, "vocab_size"): |
| 130 | with contextlib.suppress(NotImplementedError): |
| 131 | max_token_id = max(max_token_id, tokenizer.vocab_size - 1) |
| 132 | |
| 133 | class CachedTokenizer(tokenizer.__class__): # type: ignore |
| 134 | @property |
| 135 | def all_special_ids(self) -> list[int]: |
| 136 | return tokenizer_all_special_ids |
| 137 | |
| 138 | @property |
| 139 | def all_special_tokens(self) -> list[str]: |
| 140 | return tokenizer_all_special_tokens |
| 141 | |
| 142 | @property |
| 143 | def max_token_id(self) -> int: |
| 144 | return max_token_id |
| 145 | |
| 146 | @property |
| 147 | def max_chars_per_token(self) -> int: |
| 148 | return max_chars_per_token |
| 149 | |
| 150 | @property |
| 151 | def is_fast(self) -> bool: |
| 152 | return tokenizer_is_fast |
| 153 | |
| 154 | def get_vocab(self) -> dict[str, int]: |
| 155 | return tokenizer_vocab |
| 156 | |
| 157 | def __len__(self) -> int: |
| 158 | return tokenizer_len |
| 159 | |
| 160 | def __reduce__(self): |
| 161 | return get_cached_tokenizer, (tokenizer,) |
| 162 | |
| 163 | CachedTokenizer.__name__ = f"Cached{tokenizer.__class__.__name__}" |
| 164 | |
| 165 | cached_tokenizer.__class__ = CachedTokenizer |
| 166 | return cached_tokenizer |
| 167 | |
| 168 | |
| 169 | class CachedHfTokenizer(TokenizerLike): |
| 170 | @classmethod |
| 171 | def from_pretrained( |
| 172 | cls, |
| 173 | path_or_repo_id: str | Path, |
| 174 | *args, |
| 175 | trust_remote_code: bool = False, |
| 176 | revision: str | None = None, |
| 177 | download_dir: str | None = None, |
| 178 | **kwargs, |
| 179 | ) -> HfTokenizer: |
| 180 | try: |
| 181 | tokenizer = AutoTokenizer.from_pretrained( |
| 182 | path_or_repo_id, |
| 183 | *args, |
| 184 | trust_remote_code=trust_remote_code, |
| 185 | revision=revision, |
| 186 | cache_dir=download_dir, |
| 187 | **kwargs, |
| 188 | ) |
| 189 | except ValueError as e: |
| 190 | # If the error pertains to the tokenizer class not existing or not |
| 191 | # currently being imported, |
| 192 | # suggest using the --trust-remote-code flag. |
| 193 | if not trust_remote_code and ( |
| 194 | "does not exist or is not currently imported." in str(e) |
| 195 | or "requires you to execute the tokenizer file" in str(e) |
| 196 | ): |
| 197 | err_msg = ( |
| 198 | "Failed to load the tokenizer. If the tokenizer " |
| 199 | "is a custom tokenizer not yet available in the " |
| 200 | "HuggingFace transformers library, consider " |
| 201 | "setting `trust_remote_code=True` in LLM or using " |
| 202 | "the `--trust-remote-code` flag in the CLI. If the " |
| 203 | "model was created with a newer version of " |
| 204 | "transformers, consider upgrading: " |
| 205 | "`uv pip install --upgrade transformers`" |
| 206 | ) |
| 207 | raise RuntimeError(err_msg) from e |
| 208 | else: |
| 209 | raise e |
| 210 | |
| 211 | # The special_tokens in tokenizer should also be |
| 212 | # controlled by do_lower_case in encoder_config |
| 213 | encoder_config = get_sentence_transformer_tokenizer_config( |
| 214 | path_or_repo_id, revision |
| 215 | ) |
| 216 | if isinstance(encoder_config, dict) and encoder_config.get( |
| 217 | "do_lower_case", False |
| 218 | ): |
| 219 | special_tokens_map = { |
| 220 | k: v.lower() for k, v in tokenizer.special_tokens_map.items() |
| 221 | } |
| 222 | tokenizer.add_special_tokens(special_tokens_map) |
| 223 | |
| 224 | return get_cached_tokenizer(tokenizer) |