main
py 189 lines 6.33 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, TypeAlias
6
7 from pydantic import ConfigDict, 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.openai.chat_completion.protocol import ChatCompletionToolsParam
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
19
20 class TokenizeCompletionRequest(OpenAIBaseModel):
21 model: str | None = None
22 prompt: str
23
24 add_special_tokens: bool = Field(
25 default=True,
26 description=(
27 "If true (the default), special tokens (e.g. BOS) will be added to "
28 "the prompt."
29 ),
30 )
31 return_token_strs: bool | None = Field(
32 default=False,
33 description=(
34 "If true, also return the token strings corresponding to the token ids."
35 ),
36 )
37
38 def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
39 return TokenizeParams(
40 max_total_tokens=None,
41 max_output_tokens=0,
42 add_special_tokens=self.add_special_tokens,
43 )
44
45
46 class TokenizeChatRequest(OpenAIBaseModel):
47 model: str | None = None
48 messages: list[ChatCompletionMessageParam]
49
50 add_generation_prompt: bool = Field(
51 default=True,
52 description=(
53 "If true, the generation prompt will be added to the chat template. "
54 "This is a parameter used by chat template in tokenizer config of the "
55 "model."
56 ),
57 )
58 return_token_strs: bool | None = Field(
59 default=False,
60 description=(
61 "If true, also return the token strings corresponding to the token ids."
62 ),
63 )
64 continue_final_message: bool = Field(
65 default=False,
66 description=(
67 "If this is set, the chat will be formatted so that the final "
68 "message in the chat is open-ended, without any EOS tokens. The "
69 "model will continue this message rather than starting a new one. "
70 'This allows you to "prefill" part of the model\'s response for it. '
71 "Cannot be used at the same time as `add_generation_prompt`."
72 ),
73 )
74 add_special_tokens: bool = Field(
75 default=False,
76 description=(
77 "If true, special tokens (e.g. BOS) will be added to the prompt "
78 "on top of what is added by the chat template. "
79 "For most models, the chat template takes care of adding the "
80 "special tokens so this should be set to false (as is the "
81 "default)."
82 ),
83 )
84 chat_template: str | None = Field(
85 default=None,
86 description=(
87 "A Jinja template to use for this conversion. "
88 "As of transformers v4.44, default chat template is no longer "
89 "allowed, so you must provide a chat template if the tokenizer "
90 "does not define one."
91 ),
92 )
93 chat_template_kwargs: dict[str, Any] | None = Field(
94 default=None,
95 description=(
96 "Additional keyword args to pass to the template renderer. "
97 "Will be accessible by the chat template."
98 ),
99 )
100 media_io_kwargs: dict[str, dict[str, Any]] | None = Field(
101 default=None,
102 description=(
103 "Additional kwargs to pass to the media IO connectors, "
104 "keyed by modality. Merged with engine-level media_io_kwargs."
105 ),
106 )
107 mm_processor_kwargs: dict[str, Any] | None = Field(
108 default=None,
109 description="Additional kwargs to pass to the HF processor.",
110 )
111 tools: list[ChatCompletionToolsParam] | None = Field(
112 default=None,
113 description="A list of tools the model may call.",
114 )
115
116 @model_validator(mode="before")
117 @classmethod
118 def check_generation_prompt(cls, data):
119 if not isinstance(data, dict):
120 return data
121 if data.get("continue_final_message") and data.get("add_generation_prompt"):
122 raise VLLMValidationError(
123 "Cannot set both `continue_final_message` and "
124 "`add_generation_prompt` to True.",
125 )
126 return data
127
128 def build_chat_params(
129 self,
130 default_template: str | None,
131 default_template_content_format: ChatTemplateContentFormatOption,
132 ) -> ChatParams:
133 return ChatParams(
134 chat_template=self.chat_template or default_template,
135 chat_template_content_format=default_template_content_format,
136 chat_template_kwargs=merge_kwargs(
137 self.chat_template_kwargs,
138 dict(
139 add_generation_prompt=self.add_generation_prompt,
140 continue_final_message=self.continue_final_message,
141 ),
142 ),
143 media_io_kwargs=self.media_io_kwargs,
144 )
145
146 def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
147 return TokenizeParams(
148 max_total_tokens=None,
149 max_output_tokens=0,
150 add_special_tokens=self.add_special_tokens,
151 )
152
153
154 TokenizeRequest: TypeAlias = TokenizeCompletionRequest | TokenizeChatRequest
155
156
157 class TokenizeResponse(OpenAIBaseModel):
158 count: int
159 max_model_len: int
160 tokens: list[int]
161 token_strs: list[str] | None = None
162
163
164 class DetokenizeRequest(OpenAIBaseModel):
165 model: str | None = None
166 # TODO: Factor `torch.iinfo` out. `torch.iinfo` pulls torch into a
167 # Pydantic protocol file that currently has no torch dependency.
168 # See: https://github.com/vllm-project/vllm/pull/34468#discussion_r2801173630
169 tokens: list[Annotated[int, Field(ge=0, le=2**63 - 1)]]
170
171 def build_tok_params(self, model_config: ModelConfig) -> TokenizeParams:
172 return TokenizeParams(
173 max_total_tokens=None,
174 max_output_tokens=0,
175 needs_detokenization=True,
176 )
177
178
179 class DetokenizeResponse(OpenAIBaseModel):
180 prompt: str
181
182
183 class TokenizerInfoResponse(OpenAIBaseModel):
184 """Response containing tokenizer configuration
185 equivalent to tokenizer_config.json
186 """
187
188 model_config = ConfigDict(extra="allow")
189 tokenizer_class: str