tokenizer.py 10.8 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import contextlib
5
import copy
6
import importlib.util
7
import os
8
from functools import lru_cache
9
from pathlib import Path
10
from typing import TYPE_CHECKING, Any, TypeAlias
11

12
import huggingface_hub
13
from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
14
from typing_extensions import assert_never
15

16
from vllm import envs
Woosuk Kwon's avatar
Woosuk Kwon committed
17
from vllm.logger import init_logger
18
from vllm.transformers_utils.config import get_sentence_transformer_tokenizer_config
19
from vllm.transformers_utils.gguf_utils import get_gguf_file_path_from_hf
20
from vllm.transformers_utils.repo_utils import list_filtered_repo_files
21
from vllm.transformers_utils.tokenizers import MistralTokenizer
22
23
24
25
26
27
from vllm.transformers_utils.utils import (
    check_gguf_file,
    is_gguf,
    is_remote_gguf,
    split_remote_gguf,
)
28

29
30
if TYPE_CHECKING:
    from vllm.config import ModelConfig
31
32
33
34
    from vllm.transformers_utils.tokenizer_base import TokenizerBase
else:
    ModelConfig = Any
    TokenizerBase = Any
35

36
37
logger = init_logger(__name__)

38
AnyTokenizer: TypeAlias = PreTrainedTokenizer | PreTrainedTokenizerFast | TokenizerBase
39

40

41
42
43
44
def decode_tokens(
    tokenizer: AnyTokenizer,
    token_ids: list[int],
    *,
45
    skip_special_tokens: bool | None = None,
46
47
48
) -> str:
    """
    Backend-agnostic equivalent of HF's
49
    `tokenizer.decode(token_ids, ...)`.
50

51
    `skip_special_tokens=None` means to use the backend's default
52
    settings.
53
    """
54
    if skip_special_tokens is not None:
55
        return tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)
56

57
    return tokenizer.decode(token_ids)
58
59


60
61
62
63
def encode_tokens(
    tokenizer: AnyTokenizer,
    text: str,
    *,
64
65
66
    truncation: bool | None = None,
    max_length: int | None = None,
    add_special_tokens: bool | None = None,
67
68
69
) -> list[int]:
    """
    Backend-agnostic equivalent of HF's
70
    `tokenizer.encode(text, ...)`.
71

72
    `add_special_tokens=None` means to use the backend's default
73
    settings.
74
    """
75
76
77
78
79
80
81
82

    kw_args: dict[str, Any] = {}
    if max_length is not None:
        kw_args["max_length"] = max_length

    if truncation is not None:
        kw_args["truncation"] = truncation

83
    if add_special_tokens is not None:
84
        kw_args["add_special_tokens"] = add_special_tokens
85

86
    return tokenizer.encode(text, **kw_args)
87
88


89
def get_cached_tokenizer(tokenizer: AnyTokenizer) -> AnyTokenizer:
90
    """
91
    By default, transformers will recompute multiple tokenizer properties
92
93
94
95
    each time they are called, leading to a significant slowdown.
    This proxy caches these properties for faster access.
    """
    cached_tokenizer = copy.copy(tokenizer)
96

97
98
    tokenizer_all_special_ids = tokenizer.all_special_ids
    tokenizer_all_special_tokens = tokenizer.all_special_tokens
99
    tokenizer_all_special_tokens_extended = tokenizer.all_special_tokens_extended
100
    tokenizer_vocab = tokenizer.get_vocab()
101
    tokenizer_len = len(tokenizer)
102

103
    max_token_id = max(tokenizer_vocab.values())
104
105
106
107
108
109
110
    # 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)
111

112
    class CachedTokenizer(tokenizer.__class__):  # type: ignore
113
        @property
114
        def all_special_ids(self) -> list[int]:
115
116
117
            return tokenizer_all_special_ids

        @property
118
        def all_special_tokens(self) -> list[str]:
119
120
121
            return tokenizer_all_special_tokens

        @property
122
        def all_special_tokens_extended(self) -> list[str]:
123
124
            return tokenizer_all_special_tokens_extended

125
        @property
126
        def max_token_id(self) -> int:
127
128
            return max_token_id

129
        def get_vocab(self) -> dict[str, int]:
130
131
            return tokenizer_vocab

132
        def __len__(self) -> int:
133
134
            return tokenizer_len

135
        def __reduce__(self):
136
            return get_cached_tokenizer, (tokenizer,)
137

138
139
    CachedTokenizer.__name__ = f"Cached{tokenizer.__class__.__name__}"

140
141
    cached_tokenizer.__class__ = CachedTokenizer
    return cached_tokenizer
142
143


144
def get_tokenizer(
145
    tokenizer_name: str | Path,
146
    *args,
147
    tokenizer_mode: str = "auto",
148
    trust_remote_code: bool = False,
149
150
    revision: str | None = None,
    download_dir: str | None = None,
151
    **kwargs,
152
) -> AnyTokenizer:
153
    """Gets a tokenizer for the given model name via HuggingFace or ModelScope."""
154
    if envs.VLLM_USE_MODELSCOPE:
155
156
157
158
159
        # download model from ModelScope hub,
        # lazy import so that modelscope is not required for normal use.
        # pylint: disable=C.
        from modelscope.hub.snapshot_download import snapshot_download

160
161
162
        # avoid circuit import
        from vllm.model_executor.model_loader.weight_utils import get_lock

163
164
        # Only set the tokenizer here, model will be downloaded on the workers.
        if not os.path.exists(tokenizer_name):
165
166
167
168
169
170
171
172
173
            # Use file lock to prevent multiple processes from
            # downloading the same file at the same time.
            with get_lock(tokenizer_name, download_dir):
                tokenizer_path = snapshot_download(
                    model_id=tokenizer_name,
                    cache_dir=download_dir,
                    revision=revision,
                    local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
                    # Ignore weights - we only need the tokenizer.
174
175
                    ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"],
                )
176
                tokenizer_name = tokenizer_path
177

178
179
    if tokenizer_mode == "slow":
        if kwargs.get("use_fast", False):
180
            raise ValueError("Cannot use the fast tokenizer in slow tokenizer mode.")
181
182
        kwargs["use_fast"] = False

183
184
185
    if "truncation_side" not in kwargs:
        kwargs["truncation_side"] = "left"

186
    # Separate model folder from file path for GGUF models
187
188
189
190
191
    if is_gguf(tokenizer_name):
        if check_gguf_file(tokenizer_name):
            kwargs["gguf_file"] = Path(tokenizer_name).name
            tokenizer_name = Path(tokenizer_name).parent
        elif is_remote_gguf(tokenizer_name):
192
193
194
195
196
197
198
199
            tokenizer_name, quant_type = split_remote_gguf(tokenizer_name)
            # Get the HuggingFace Hub path for the GGUF file
            gguf_file = get_gguf_file_path_from_hf(
                tokenizer_name,
                quant_type,
                revision=revision,
            )
            kwargs["gguf_file"] = gguf_file
200

201
202
203
204
205
206
207
208
209
    # if `tokenizer_mode` == "auto", check if tokenizer can be loaded via Mistral format
    # first to use official Mistral tokenizer if possible.
    mistral_common_installed = importlib.util.find_spec("mistral_common") is not None
    if tokenizer_mode == "auto" and mistral_common_installed:
        allow_patterns = ["tekken.json", "tokenizer.model.v*"]
        files_list = list_filtered_repo_files(
            model_name_or_path=str(tokenizer_name),
            allow_patterns=allow_patterns,
            revision=revision,
210
        )
211
212
        if len(files_list) > 0:
            tokenizer_mode = "mistral"
213
214

    tokenizer: AnyTokenizer
215
    if tokenizer_mode == "mistral":
216
        logger.debug_once(f"Loading MistralTokenizer from {tokenizer_name}")
217
218
219
        tokenizer = MistralTokenizer.from_pretrained(
            str(tokenizer_name), revision=revision
        )
220
    elif tokenizer_mode == "custom":
221
        from vllm.transformers_utils.tokenizer_base import TokenizerRegistry
222

223
        logger.debug_once(f"Loading CustomTokenizer from {tokenizer_name}")
224
225
226
227
228
229
230
        tokenizer = TokenizerRegistry.get_tokenizer(
            str(tokenizer_name),
            *args,
            revision=revision,
            download_dir=download_dir,
            **kwargs,
        )
231
232
    else:
        try:
233
            logger.debug_once(f"Loading AutoTokenizer from {tokenizer_name}")
234
            tokenizer = AutoTokenizer.from_pretrained(
235
236
237
                tokenizer_name,
                *args,
                trust_remote_code=trust_remote_code,
238
                revision=revision,
239
240
241
242
243
244
245
                **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 (
246
247
248
249
250
251
252
253
254
255
                "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."
                )
256
257
258
259
                raise RuntimeError(err_msg) from e
            else:
                raise e

260
261
262
        # The special_tokens in tokenizer should also be
        # controlled by do_lower_case in encoder_config
        encoder_config = get_sentence_transformer_tokenizer_config(
263
264
            tokenizer_name, revision
        )
265
        if isinstance(encoder_config, dict) and encoder_config.get(
266
267
            "do_lower_case", False
        ):
268
            special_tokens_map = {
269
                k: v.lower() for k, v in tokenizer.special_tokens_map.items()
270
271
272
            }
            tokenizer.add_special_tokens(special_tokens_map)

273
274
275
        if not isinstance(tokenizer, PreTrainedTokenizerFast):
            logger.warning(
                "Using a slow tokenizer. This might cause a significant "
276
277
                "slowdown. Consider using a fast tokenizer instead."
            )
278
        tokenizer = get_cached_tokenizer(tokenizer)
279

280
    return tokenizer
281
282


283
284
285
286
cached_get_tokenizer = lru_cache(get_tokenizer)


def cached_tokenizer_from_config(
287
    model_config: ModelConfig,
288
289
290
291
292
    **kwargs: Any,
):
    return cached_get_tokenizer(
        model_config.tokenizer,
        tokenizer_mode=model_config.tokenizer_mode,
293
        revision=model_config.tokenizer_revision,
294
295
296
297
298
        trust_remote_code=model_config.trust_remote_code,
        **kwargs,
    )


299
300
301
302
303
304
305
306
def init_tokenizer_from_configs(model_config: ModelConfig):
    runner_type = model_config.runner_type
    if runner_type == "generate" or runner_type == "draft":
        truncation_side = "left"
    elif runner_type == "pooling":
        truncation_side = "right"
    else:
        assert_never(runner_type)
307

308
309
310
311
312
313
314
    return get_tokenizer(
        model_config.tokenizer,
        tokenizer_mode=model_config.tokenizer_mode,
        trust_remote_code=model_config.trust_remote_code,
        revision=model_config.tokenizer_revision,
        truncation_side=truncation_side,
    )