main
py 167 lines 5.12 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import json
4 from pathlib import Path
5 from types import SimpleNamespace
6 from unittest.mock import patch
7
8 import pytest
9 from transformers import AutoConfig
10 from transformers.models.auto.configuration_auto import CONFIG_MAPPING
11
12 from vllm.tokenizers import TokenizerLike
13 from vllm.tokenizers.registry import (
14 TokenizerRegistry,
15 cached_get_tokenizer,
16 cached_resolve_tokenizer_args,
17 cached_tokenizer_from_config,
18 get_tokenizer,
19 resolve_tokenizer_args,
20 )
21 from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeConfig
22
23
24 class TestTokenizer(TokenizerLike):
25 @classmethod
26 def from_pretrained(
27 cls,
28 path_or_repo_id: str | Path,
29 *args,
30 trust_remote_code: bool = False,
31 revision: str | None = None,
32 download_dir: str | None = None,
33 **kwargs,
34 ) -> "TestTokenizer":
35 return TestTokenizer(path_or_repo_id) # type: ignore
36
37 def __init__(self, path_or_repo_id: str | Path) -> None:
38 super().__init__()
39
40 self.path_or_repo_id = path_or_repo_id
41
42 @property
43 def bos_token_id(self) -> int:
44 return 0
45
46 @property
47 def eos_token_id(self) -> int:
48 return 1
49
50 @property
51 def pad_token_id(self) -> int:
52 return 2
53
54 @property
55 def is_fast(self) -> bool:
56 return True
57
58
59 @pytest.mark.parametrize("runner_type", ["generate", "pooling"])
60 def test_resolve_tokenizer_args_idempotent(runner_type):
61 tokenizer_mode, tokenizer_name, args, kwargs = resolve_tokenizer_args(
62 "facebook/opt-125m",
63 runner_type=runner_type,
64 )
65
66 assert (tokenizer_mode, tokenizer_name, args, kwargs) == resolve_tokenizer_args(
67 tokenizer_name, *args, **kwargs
68 )
69
70
71 @pytest.mark.parametrize("input_kwargs", [{}, {"mistral_format": False}])
72 def test_resolve_tokenizer_args_forces_hf_mistral_format_false(input_kwargs):
73 resolved_mode, _, _, kwargs = resolve_tokenizer_args(
74 "mistralai/Mistral-Nemo-Instruct-2407",
75 tokenizer_mode="hf",
76 **input_kwargs,
77 )
78
79 assert resolved_mode == "hf"
80 assert kwargs["mistral_format"] is False
81
82
83 def test_resolve_tokenizer_args_rejects_hf_mistral_format_true():
84 with pytest.raises(
85 ValueError,
86 match="mistral_format=True is not supported with tokenizer_mode='hf'",
87 ):
88 resolve_tokenizer_args(
89 "mistralai/Mistral-Nemo-Instruct-2407",
90 tokenizer_mode="hf",
91 mistral_format=True,
92 )
93
94
95 def test_customized_tokenizer():
96 TokenizerRegistry.register("test_tokenizer", __name__, TestTokenizer.__name__)
97
98 tokenizer = TokenizerRegistry.load_tokenizer("test_tokenizer", "abc")
99 assert isinstance(tokenizer, TestTokenizer)
100 assert tokenizer.path_or_repo_id == "abc"
101 assert tokenizer.bos_token_id == 0
102 assert tokenizer.eos_token_id == 1
103 assert tokenizer.pad_token_id == 2
104
105 tokenizer = get_tokenizer("abc", tokenizer_mode="test_tokenizer")
106 assert isinstance(tokenizer, TestTokenizer)
107 assert tokenizer.path_or_repo_id == "abc"
108 assert tokenizer.bos_token_id == 0
109 assert tokenizer.eos_token_id == 1
110 assert tokenizer.pad_token_id == 2
111
112
113 def test_cached_tokenizer_from_config_registers_local_config(tmp_path: Path):
114 (tmp_path / "config.json").write_text(
115 json.dumps({"model_type": "qwen3_5_moe"}),
116 encoding="utf-8",
117 )
118
119 model_config = SimpleNamespace(
120 skip_tokenizer_init=False,
121 tokenizer=str(tmp_path),
122 runner_type="generate",
123 tokenizer_mode="hf",
124 tokenizer_revision=None,
125 trust_remote_code=True,
126 hf_config=Qwen3_5MoeConfig(),
127 )
128
129 registered_config = CONFIG_MAPPING._extra_content.pop("qwen3_5_moe", None)
130 cached_get_tokenizer.cache_clear()
131 cached_resolve_tokenizer_args.cache_clear()
132
133 try:
134
135 def fake_from_pretrained(path_or_repo_id: str, *args, **kwargs):
136 passed_config = kwargs.pop("config")
137 assert isinstance(passed_config, Qwen3_5MoeConfig)
138 loaded_config = AutoConfig.from_pretrained(
139 path_or_repo_id,
140 trust_remote_code=False,
141 )
142 assert isinstance(loaded_config, Qwen3_5MoeConfig)
143 return SimpleNamespace(is_fast=True)
144
145 with (
146 patch(
147 "vllm.tokenizers.registry.logger.debug_once",
148 lambda *args, **kwargs: None,
149 ),
150 patch(
151 "vllm.tokenizers.hf.AutoTokenizer.from_pretrained",
152 side_effect=fake_from_pretrained,
153 ),
154 patch(
155 "vllm.tokenizers.hf.get_cached_tokenizer",
156 side_effect=lambda tokenizer: tokenizer,
157 ),
158 ):
159 tokenizer = cached_tokenizer_from_config(model_config)
160
161 assert tokenizer.is_fast is True
162 finally:
163 cached_get_tokenizer.cache_clear()
164 cached_resolve_tokenizer_args.cache_clear()
165 CONFIG_MAPPING._extra_content.pop("qwen3_5_moe", None)
166 if registered_config is not None:
167 CONFIG_MAPPING._extra_content["qwen3_5_moe"] = registered_config