main
py 366 lines 13.2 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5 from typing import Annotated, Any, Literal
6
7 from pydantic import Field, model_validator
8
9 from vllm.config import ModelConfig
10 from vllm.entrypoints.chat_utils import (
11 ChatCompletionMessageParam,
12 ChatTemplateContentFormatOption,
13 )
14 from vllm.entrypoints.generate.base.protocol import validate_cache_salt
15 from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel
16 from vllm.exceptions import VLLMValidationError
17 from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs
18 from vllm.tasks import check_removed_pooling_task
19 from vllm.utils import random_uuid
20 from vllm.utils.serial_utils import EmbedDType, EncodingFormat, Endianness
21
22
23 def reject_removed_pooling_parameters(data):
24 if not isinstance(data, dict):
25 return data
26 if "normalize" in data:
27 raise VLLMValidationError(
28 "Parameter `normalize` was removed; use `use_activation` instead.",
29 parameter="normalize",
30 )
31 check_removed_pooling_task(data.get("task"))
32 return data
33
34
35 class PoolingBasicRequestMixin(OpenAIBaseModel):
36 # --8<-- [start:pooling-common-params]
37 model: str | None = None
38 user: str | None = None
39 # --8<-- [end:pooling-common-params]
40
41 # --8<-- [start:pooling-common-extra-params]
42 truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
43 padding: Literal["max_length", "do_not_pad"] | None = Field(
44 default=None,
45 description=(
46 "Whether to pad the prompt, using the same names as the "
47 "Transformers tokenizer. 'max_length' pads to the maximum input "
48 "length, 'do_not_pad' leaves the prompt as is. Models trained "
49 "with a fixed sequence length and no attention mask, such as "
50 "SigLIP, need 'max_length' to match training, otherwise their "
51 "embeddings are not comparable."
52 ),
53 )
54 truncation_side: Literal["left", "right"] | None = Field(
55 default=None,
56 description=(
57 "Which side to truncate from when truncate_prompt_tokens is active. "
58 "'right' keeps the first N tokens. "
59 "'left' keeps the last N tokens."
60 ),
61 )
62 request_id: str = Field(
63 default_factory=random_uuid,
64 description=(
65 "The request_id related to this request. If the caller does "
66 "not set it, a random_uuid will be generated. This id is used "
67 "through out the inference process and return in response."
68 ),
69 )
70 priority: int = Field(
71 default=0,
72 ge=-(2**63),
73 le=2**63 - 1,
74 description=(
75 "The priority of the request (lower means earlier handling; "
76 "default: 0). Any priority other than 0 will raise an error "
77 "if the served model does not use priority scheduling."
78 ),
79 )
80 mm_processor_kwargs: dict[str, Any] | None = Field(
81 default=None,
82 description="Additional kwargs to pass to the HF processor.",
83 )
84 cache_salt: str | None = Field(
85 default=None,
86 min_length=1,
87 max_length=1024,
88 description=(
89 "If specified, the prefix cache will be salted with the provided "
90 "string to prevent an attacker to guess prompts in multi-user "
91 "environments. The salt should be random, protected from "
92 "access by 3rd parties, and long enough to be "
93 "unpredictable (e.g., 43 characters base64-encoded, corresponding "
94 "to 256 bit)."
95 ),
96 )
97 # --8<-- [end:pooling-common-extra-params]
98
99 @model_validator(mode="before")
100 @classmethod
101 def check_cache_salt_support(cls, data: Any) -> Any:
102 if not isinstance(data, dict):
103 return data
104
105 validate_cache_salt(data.get("cache_salt"))
106 return data
107
108 def _build_pooling_tok_params(
109 self,
110 model_config: ModelConfig,
111 *,
112 add_special_tokens: bool,
113 max_total_tokens: int | None,
114 max_output_tokens: int,
115 max_total_tokens_param: str = "max_model_len",
116 max_output_tokens_param: str | None = None,
117 ) -> TokenizeParams:
118 encoder_config = model_config.encoder_config or {}
119 if max_output_tokens_param is None:
120 tok_params = TokenizeParams(
121 max_total_tokens=max_total_tokens,
122 max_output_tokens=max_output_tokens,
123 truncate_prompt_tokens=self.truncate_prompt_tokens,
124 truncation_side=self.truncation_side,
125 do_lower_case=encoder_config.get("do_lower_case", False),
126 add_special_tokens=add_special_tokens,
127 max_total_tokens_param=max_total_tokens_param,
128 )
129 else:
130 tok_params = TokenizeParams(
131 max_total_tokens=max_total_tokens,
132 max_output_tokens=max_output_tokens,
133 truncate_prompt_tokens=self.truncate_prompt_tokens,
134 truncation_side=self.truncation_side,
135 do_lower_case=encoder_config.get("do_lower_case", False),
136 add_special_tokens=add_special_tokens,
137 max_total_tokens_param=max_total_tokens_param,
138 max_output_tokens_param=max_output_tokens_param,
139 )
140
141 if self.padding is None:
142 return tok_params
143
144 return tok_params.with_kwargs(padding=self.padding)
145
146
147 class PoolingTokenizeParamsMixin:
148 add_special_tokens: bool
149
150 def _build_pooling_tok_params(
151 self,
152 model_config: ModelConfig,
153 *,
154 add_special_tokens: bool,
155 max_total_tokens: int | None,
156 max_output_tokens: int,
157 max_total_tokens_param: str = "max_model_len",
158 max_output_tokens_param: str | None = None,
159 ) -> TokenizeParams:
160 raise NotImplementedError
161
162
163 class FixedMaxLenTokenizeParamsMixin(PoolingTokenizeParamsMixin):
164 def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
165 return self._build_pooling_tok_params(
166 model_config,
167 add_special_tokens=self.add_special_tokens,
168 max_total_tokens=model_config.max_model_len,
169 max_output_tokens=0,
170 )
171
172
173 class EmbeddingTokenizeParamsMixin(PoolingTokenizeParamsMixin):
174 padding: Literal["max_length", "do_not_pad"] | None
175
176 def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
177 default_max_total_tokens = model_config.max_model_len
178 max_total_tokens: int | None = default_max_total_tokens
179 max_output_tokens = 0
180
181 pooler_config = model_config.pooler_config
182 if pooler_config is not None:
183 if pooler_config.enable_chunked_processing:
184 max_total_tokens = pooler_config.max_embed_len
185 else:
186 max_embed_len = pooler_config.max_embed_len or default_max_total_tokens
187 max_output_tokens = default_max_total_tokens - max_embed_len
188
189 tok_params = self._build_pooling_tok_params(
190 model_config,
191 add_special_tokens=self.add_special_tokens,
192 max_total_tokens=max_total_tokens,
193 max_output_tokens=max_output_tokens,
194 max_output_tokens_param="max_model_len - max_embed_len",
195 )
196
197 if (
198 self.padding == "max_length"
199 and pooler_config is not None
200 and pooler_config.enable_chunked_processing
201 ):
202 max_padding_len = pooler_config.max_embed_len
203 if max_padding_len is None:
204 max_padding_len = default_max_total_tokens
205 return tok_params.with_kwargs(
206 pad_prompt_tokens=min(default_max_total_tokens, max_padding_len)
207 )
208
209 return tok_params
210
211
212 class CompletionRequestMixin(OpenAIBaseModel):
213 # --8<-- [start:completion-params]
214 input: list[int] | list[list[int]] | str | list[str]
215 # --8<-- [end:completion-params]
216
217 # --8<-- [start:completion-extra-params]
218 add_special_tokens: bool = Field(
219 default=True,
220 description=(
221 "If true (the default), special tokens (e.g. BOS) will be added to "
222 "the prompt."
223 ),
224 )
225 # --8<-- [end:completion-extra-params]
226
227
228 class ChatRequestOptionsMixin(OpenAIBaseModel):
229 # --8<-- [start:chat-extra-params]
230 add_generation_prompt: bool = Field(
231 default=False,
232 description=(
233 "If true, the generation prompt will be added to the chat template. "
234 "This is a parameter used by chat template in tokenizer config of the "
235 "model."
236 ),
237 )
238 continue_final_message: bool = Field(
239 default=False,
240 description=(
241 "If this is set, the chat will be formatted so that the final "
242 "message in the chat is open-ended, without any EOS tokens. The "
243 "model will continue this message rather than starting a new one. "
244 'This allows you to "prefill" part of the model\'s response for it. '
245 "Cannot be used at the same time as `add_generation_prompt`."
246 ),
247 )
248 add_special_tokens: bool = Field(
249 default=False,
250 description=(
251 "If true, special tokens (e.g. BOS) will be added to the prompt "
252 "on top of what is added by the chat template. "
253 "For most models, the chat template takes care of adding the "
254 "special tokens so this should be set to false (as is the "
255 "default)."
256 ),
257 )
258 chat_template: str | None = Field(
259 default=None,
260 description=(
261 "A Jinja template to use for this conversion. "
262 "As of transformers v4.44, default chat template is no longer "
263 "allowed, so you must provide a chat template if the tokenizer "
264 "does not define one."
265 ),
266 )
267 chat_template_kwargs: dict[str, Any] | None = Field(
268 default=None,
269 description=(
270 "Additional keyword args to pass to the template renderer. "
271 "Will be accessible by the chat template."
272 ),
273 )
274 media_io_kwargs: dict[str, dict[str, Any]] | None = Field(
275 default=None,
276 description=(
277 "Additional kwargs to pass to the media IO connectors, "
278 "keyed by modality. Merged with engine-level media_io_kwargs."
279 ),
280 )
281 # --8<-- [end:chat-extra-params]
282
283 @model_validator(mode="before")
284 @classmethod
285 def check_generation_prompt(cls, data):
286 if not isinstance(data, dict):
287 return data
288 if data.get("continue_final_message") and data.get("add_generation_prompt"):
289 raise VLLMValidationError(
290 "Cannot set both `continue_final_message` and "
291 "`add_generation_prompt` to True.",
292 )
293 return data
294
295 def build_chat_params(
296 self,
297 default_template: str | None,
298 default_template_content_format: ChatTemplateContentFormatOption,
299 ) -> ChatParams:
300 return ChatParams(
301 chat_template=self.chat_template or default_template,
302 chat_template_content_format=default_template_content_format,
303 chat_template_kwargs=merge_kwargs(
304 self.chat_template_kwargs,
305 dict(
306 add_generation_prompt=self.add_generation_prompt,
307 continue_final_message=self.continue_final_message,
308 ),
309 ),
310 media_io_kwargs=self.media_io_kwargs,
311 )
312
313
314 class ChatRequestMixin(ChatRequestOptionsMixin):
315 # --8<-- [start:chat-params]
316 messages: list[ChatCompletionMessageParam]
317 # --8<-- [end:chat-params]
318
319
320 class EncodingRequestMixin(OpenAIBaseModel):
321 # --8<-- [start:encoding-params]
322 encoding_format: EncodingFormat = "float"
323 # --8<-- [end:encoding-params]
324
325 # --8<-- [start:encoding-extra-params]
326 embed_dtype: EmbedDType = Field(
327 default="float32",
328 description=(
329 "What dtype to use for encoding. Default to using float32 for base64 "
330 "encoding to match the OpenAI python client behavior. "
331 "This parameter will affect base64 and binary_response."
332 ),
333 )
334 endianness: Endianness = Field(
335 default="native",
336 description=(
337 "What endianness to use for encoding. Default to using native for "
338 "base64 encoding to match the OpenAI python client behavior."
339 "This parameter will affect base64 and binary_response."
340 ),
341 )
342 # --8<-- [end:encoding-extra-params]
343
344
345 class EmbedRequestMixin(EncodingRequestMixin):
346 # --8<-- [start:embed-params]
347 dimensions: int | None = None
348 # --8<-- [end:embed-params]
349
350 # --8<-- [start:embed-extra-params]
351 use_activation: bool | None = Field(
352 default=None,
353 description="Whether to use activation for the pooler outputs. "
354 "`None` uses the pooler's default, which is `True` in most cases.",
355 )
356 # --8<-- [end:embed-extra-params]
357
358
359 class ClassifyRequestMixin(OpenAIBaseModel):
360 # --8<-- [start:classify-extra-params]
361 use_activation: bool | None = Field(
362 default=None,
363 description="Whether to use activation for the pooler outputs. "
364 "`None` uses the pooler's default, which is `True` in most cases.",
365 )
366 # --8<-- [end:classify-extra-params]