wehub-resource-sync
254 lines
9.3 KiB
Python
254 lines
9.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import contextlib
|
|
import copy
|
|
import queue
|
|
from pathlib import Path
|
|
from typing import TypeAlias, TypeVar
|
|
|
|
from transformers import AutoTokenizer, PythonBackend, TokenizersBackend
|
|
|
|
from vllm.transformers_utils.config import get_sentence_transformer_tokenizer_config
|
|
|
|
from .protocol import TokenizerLike
|
|
|
|
HfTokenizer: TypeAlias = PythonBackend | TokenizersBackend
|
|
_T = TypeVar("_T", bound=TokenizerLike)
|
|
|
|
|
|
class ThreadSafeHFTokenizerMixin:
|
|
"""Mixin class for thread-safe HF fast tokenizers."""
|
|
|
|
pass
|
|
|
|
|
|
def maybe_make_thread_pool(tokenizer: _T, copies: int = 1):
|
|
"""
|
|
If `tokenizer` is a `TokenizersBackend`, modify the tokenizer
|
|
in-place to make the public interface thread-safe by routing calls
|
|
through a deep-copied tokenizer pool.
|
|
|
|
Note that:
|
|
- Only ``TokenizerLike``'s public interface is thread-safe.
|
|
This doesn't include ``_tokenizer`` property nor any mutation
|
|
methods like ``add_special_tokens`` or ``add_tokens``.
|
|
- Adjacent method calls could happen on different deep copies.
|
|
"""
|
|
if not isinstance(tokenizer, TokenizersBackend) or isinstance(
|
|
tokenizer, ThreadSafeHFTokenizerMixin
|
|
):
|
|
return tokenizer
|
|
|
|
og_tokenizer = copy.copy(tokenizer)
|
|
|
|
tokenizer_pool: queue.Queue[TokenizersBackend] = queue.Queue()
|
|
for _ in range(copies):
|
|
tokenizer_pool.put(copy.deepcopy(og_tokenizer))
|
|
|
|
@contextlib.contextmanager
|
|
def _borrow_from_pool():
|
|
try:
|
|
tok = tokenizer_pool.get_nowait()
|
|
yield tok
|
|
except queue.Empty:
|
|
tok = copy.deepcopy(og_tokenizer)
|
|
yield tok
|
|
finally:
|
|
tokenizer_pool.put(tok)
|
|
|
|
class TokenizerPool(tokenizer.__class__, ThreadSafeHFTokenizerMixin): # type: ignore
|
|
def apply_chat_template(self, *args, **kwargs):
|
|
with _borrow_from_pool() as tok:
|
|
return tok.apply_chat_template(*args, **kwargs)
|
|
|
|
def batch_decode(self, *args, **kwargs):
|
|
with _borrow_from_pool() as tok:
|
|
return tok.batch_decode(*args, **kwargs)
|
|
|
|
def batch_encode(self, *args, **kwargs):
|
|
with _borrow_from_pool() as tok:
|
|
return tok.batch_encode(*args, **kwargs)
|
|
|
|
def convert_tokens_to_ids(self, *args, **kwargs):
|
|
with _borrow_from_pool() as tok:
|
|
return tok.convert_tokens_to_ids(*args, **kwargs)
|
|
|
|
def convert_ids_to_tokens(self, *args, **kwargs):
|
|
with _borrow_from_pool() as tok:
|
|
return tok.convert_ids_to_tokens(*args, **kwargs)
|
|
|
|
def convert_tokens_to_string(self, *args, **kwargs):
|
|
with _borrow_from_pool() as tok:
|
|
return tok.convert_tokens_to_string(*args, **kwargs)
|
|
|
|
def decode(self, *args, **kwargs):
|
|
with _borrow_from_pool() as tok:
|
|
return tok.decode(*args, **kwargs)
|
|
|
|
def encode(self, *args, **kwargs):
|
|
with _borrow_from_pool() as tok:
|
|
return tok.encode(*args, **kwargs)
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
with _borrow_from_pool() as tok:
|
|
return tok(*args, **kwargs)
|
|
|
|
def __reduce__(self):
|
|
return maybe_make_thread_pool, (og_tokenizer, copies)
|
|
|
|
TokenizerPool.__name__ = f"TokenizerPool{og_tokenizer.__class__.__name__}"
|
|
|
|
tokenizer.__class__ = TokenizerPool
|
|
# Return the tokenizer: TokenizerPool.__reduce__ reconstructs through this
|
|
# function, so falling off the end would unpickle to None (issue #45433).
|
|
return tokenizer
|
|
|
|
|
|
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.
|
|
This proxy caches these properties for faster access.
|
|
"""
|
|
cached_tokenizer = copy.copy(tokenizer)
|
|
|
|
tokenizer_all_special_ids = tokenizer.all_special_ids
|
|
tokenizer_all_special_tokens = tokenizer.all_special_tokens
|
|
tokenizer_vocab = tokenizer.get_vocab()
|
|
tokenizer_len = len(tokenizer)
|
|
# The underlying tokenizer class could be MistralCommonBackend,
|
|
# which does not implement is_fast in Transformers
|
|
tokenizer_is_fast = getattr(tokenizer, "is_fast", True)
|
|
|
|
# MistralCommonBackend is tekken-backed and needs byte-fallback-aware tokenization.
|
|
mistral_tekkenizer = None
|
|
if getattr(getattr(tokenizer, "tokenizer", None), "instruct_tokenizer", None):
|
|
from vllm.tokenizers.mistral import mistral_common_tekkenizer
|
|
|
|
mistral_tekkenizer = mistral_common_tekkenizer(tokenizer)
|
|
|
|
max_token_id = max(tokenizer_vocab.values())
|
|
max_chars_per_token = max(len(tok) for tok in tokenizer_vocab)
|
|
|
|
# Some tokenizers (e.g., QwenTokenizer) have special tokens that
|
|
# are added and included in the implementation of the vocab_size
|
|
# property, but not in get_vocab(); if there is an implementation
|
|
# of vocab size, we should take the greater value.
|
|
if hasattr(tokenizer, "vocab_size"):
|
|
with contextlib.suppress(NotImplementedError):
|
|
max_token_id = max(max_token_id, tokenizer.vocab_size)
|
|
|
|
class CachedTokenizer(tokenizer.__class__): # type: ignore
|
|
@property
|
|
def all_special_ids(self) -> list[int]:
|
|
return tokenizer_all_special_ids
|
|
|
|
@property
|
|
def all_special_tokens(self) -> list[str]:
|
|
return tokenizer_all_special_tokens
|
|
|
|
@property
|
|
def max_token_id(self) -> int:
|
|
return max_token_id
|
|
|
|
@property
|
|
def max_chars_per_token(self) -> int:
|
|
return max_chars_per_token
|
|
|
|
@property
|
|
def is_fast(self) -> bool:
|
|
return tokenizer_is_fast
|
|
|
|
def convert_ids_to_tokens(self, ids, skip_special_tokens: bool = False):
|
|
if mistral_tekkenizer is not None:
|
|
from vllm.tokenizers.mistral import tekken_convert_ids_to_tokens
|
|
|
|
return tekken_convert_ids_to_tokens(mistral_tekkenizer, ids)
|
|
return super().convert_ids_to_tokens(
|
|
ids, skip_special_tokens=skip_special_tokens
|
|
)
|
|
|
|
def convert_tokens_to_string(self, tokens: list[str]) -> str:
|
|
if mistral_tekkenizer is not None:
|
|
from vllm.tokenizers.mistral import tekken_convert_tokens_to_string
|
|
|
|
return tekken_convert_tokens_to_string(mistral_tekkenizer, tokens)
|
|
try:
|
|
return super().convert_tokens_to_string(tokens)
|
|
except NotImplementedError:
|
|
# The underlying tokenizer class could be MistralCommonBackend,
|
|
# which does not implement convert_tokens_to_string in Transformers
|
|
return "".join(tokens)
|
|
|
|
def get_vocab(self) -> dict[str, int]:
|
|
return tokenizer_vocab
|
|
|
|
def __len__(self) -> int:
|
|
return tokenizer_len
|
|
|
|
def __reduce__(self):
|
|
return get_cached_tokenizer, (tokenizer,)
|
|
|
|
CachedTokenizer.__name__ = f"Cached{tokenizer.__class__.__name__}"
|
|
|
|
cached_tokenizer.__class__ = CachedTokenizer
|
|
return cached_tokenizer
|
|
|
|
|
|
class CachedHfTokenizer(TokenizerLike):
|
|
@classmethod
|
|
def from_pretrained(
|
|
cls,
|
|
path_or_repo_id: str | Path,
|
|
*args,
|
|
trust_remote_code: bool = False,
|
|
revision: str | None = None,
|
|
download_dir: str | None = None,
|
|
**kwargs,
|
|
) -> HfTokenizer:
|
|
try:
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
path_or_repo_id,
|
|
*args,
|
|
trust_remote_code=trust_remote_code,
|
|
revision=revision,
|
|
cache_dir=download_dir,
|
|
**kwargs,
|
|
)
|
|
except ValueError as e:
|
|
# If the error pertains to the tokenizer class not existing or not
|
|
# currently being imported,
|
|
# suggest using the --trust-remote-code flag.
|
|
if not trust_remote_code and (
|
|
"does not exist or is not currently imported." in str(e)
|
|
or "requires you to execute the tokenizer file" in str(e)
|
|
):
|
|
err_msg = (
|
|
"Failed to load the tokenizer. If the tokenizer "
|
|
"is a custom tokenizer not yet available in the "
|
|
"HuggingFace transformers library, consider "
|
|
"setting `trust_remote_code=True` in LLM or using "
|
|
"the `--trust-remote-code` flag in the CLI. If the "
|
|
"model was created with a newer version of "
|
|
"transformers, consider upgrading: "
|
|
"`uv pip install --upgrade transformers`"
|
|
)
|
|
raise RuntimeError(err_msg) from e
|
|
else:
|
|
raise e
|
|
|
|
# The special_tokens in tokenizer should also be
|
|
# controlled by do_lower_case in encoder_config
|
|
encoder_config = get_sentence_transformer_tokenizer_config(
|
|
path_or_repo_id, revision
|
|
)
|
|
if isinstance(encoder_config, dict) and encoder_config.get(
|
|
"do_lower_case", False
|
|
):
|
|
special_tokens_map = {
|
|
k: v.lower() for k, v in tokenizer.special_tokens_map.items()
|
|
}
|
|
tokenizer.add_special_tokens(special_tokens_map)
|
|
|
|
return get_cached_tokenizer(tokenizer)
|