A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 07:15 UTC
[Misc] add return token strs for tokenize (#18941)
Signed-off-by: reidliu41 <reid201711@gmail.com> Co-authored-by: reidliu41 <reid201711@gmail.com>
Reid committed
Jun 1, 2025 at 02:00 UTC
20079c6e365752713c05de0eb719ac0307979c3b
3 files changed
+66
-15
tests/entrypoints/openai/test_tokenization.py
+50
-15
@@ -76,11 +76,11 @@ async def test_tokenize_completions(
76
})
77
response.raise_for_status()
78
79
- assert response.json() == {
80
- "tokens": tokens,
81
- "count": len(tokens),
82
- "max_model_len": 8192
83
- }
79
+ result = response.json()
80
+ assert result["tokens"] == tokens
81
+ assert result["count"] == len(tokens)
82
+ assert result["max_model_len"] == 8192
83
+ assert result["token_strs"] is None
84
85
86
@pytest.mark.asyncio
@@ -138,11 +138,11 @@ async def test_tokenize_chat(
138
})
139
response.raise_for_status()
140
141
- assert response.json() == {
142
- "tokens": tokens,
143
- "count": len(tokens),
144
- "max_model_len": 8192
145
- }
141
+ result = response.json()
142
+ assert result["tokens"] == tokens
143
+ assert result["count"] == len(tokens)
144
+ assert result["max_model_len"] == 8192
145
+ assert result["token_strs"] is None
146
147
148
@pytest.mark.asyncio
@@ -215,11 +215,46 @@ async def test_tokenize_chat_with_tools(
215
)
216
response.raise_for_status()
217
218
- assert response.json() == {
219
- "tokens": tokens,
220
- "count": len(tokens),
221
- "max_model_len": 8192,
222
- }
218
+ result = response.json()
219
+ assert result["tokens"] == tokens
220
+ assert result["count"] == len(tokens)
221
+ assert result["max_model_len"] == 8192
222
+ assert result["token_strs"] is None
223
+
224
+
225
+@pytest.mark.asyncio
226
+@pytest.mark.parametrize(
227
+ "model_name, tokenizer_name",
228
+ [(MODEL_NAME, MODEL_NAME), ("zephyr-lora2", "zephyr-lora2")],
229
+ indirect=["tokenizer_name"],
230
+)
231
+async def test_tokenize_with_return_token_strs(
232
+ server: RemoteOpenAIServer,
233
+ model_name: str,
234
+ tokenizer_name: str,
235
+):
236
+ tokenizer = get_tokenizer(tokenizer_name=tokenizer_name,
237
+ tokenizer_mode="fast")
238
+
239
+ prompt = "This is a token_strs test prompt! vllm1"
240
+ response = requests.post(
241
+ server.url_for("tokenize"),
242
+ json={
243
+ "prompt": prompt,
244
+ "model": model_name,
245
+ "return_token_strs": True
246
+ },
247
+ )
248
+ response.raise_for_status()
249
+
250
+ tokens = tokenizer.encode(prompt, add_special_tokens=True)
251
+ tokens_str = tokenizer.convert_ids_to_tokens(tokens)
252
+
253
+ result = response.json()
254
+ assert result["tokens"] == tokens
255
+ assert result["count"] == len(tokens)
256
+ assert result["max_model_len"] == 8192
257
+ assert result["token_strs"] == tokens_str
258
259
260
@pytest.mark.asyncio
vllm/entrypoints/openai/protocol.py
+11
@@ -1563,6 +1563,11 @@ class TokenizeCompletionRequest(OpenAIBaseModel):
1563
"If true (the default), special tokens (e.g. BOS) will be added to "
1564
"the prompt."),
1565
)
1566
+ return_token_strs: Optional[bool] = Field(
1567
+ default=False,
1568
+ description=("If true, also return the token strings "
1569
+ "corresponding to the token ids."),
1570
+ )
1571
1572
1573
class TokenizeChatRequest(OpenAIBaseModel):
@@ -1576,6 +1581,11 @@ class TokenizeChatRequest(OpenAIBaseModel):
1581
"This is a parameter used by chat template in tokenizer config of the "
1582
"model."),
1583
)
1584
+ return_token_strs: Optional[bool] = Field(
1585
+ default=False,
1586
+ description=("If true, also return the token strings "
1587
+ "corresponding to the token ids."),
1588
+ )
1589
continue_final_message: bool = Field(
1590
default=False,
1591
description=
@@ -1633,6 +1643,7 @@ class TokenizeResponse(OpenAIBaseModel):
1643
count: int
1644
max_model_len: int
1645
tokens: list[int]
1646
+ token_strs: Optional[list[str]] = None
1647
1648
1649
class DetokenizeRequest(OpenAIBaseModel):
vllm/entrypoints/openai/serving_tokenization.py
+5
@@ -110,7 +110,12 @@ class OpenAIServingTokenization(OpenAIServing):
110
dict) and "prompt_token_ids" in engine_prompt:
111
input_ids.extend(engine_prompt["prompt_token_ids"])
112
113
+ token_strs = None
114
+ if request.return_token_strs:
115
+ token_strs = tokenizer.convert_ids_to_tokens(input_ids)
116
+
117
return TokenizeResponse(tokens=input_ids,
118
+ token_strs=token_strs,
119
count=len(input_ids),
120
max_model_len=self.max_model_len)
121