main
py 243 lines 7.5 KB
Raw
1 """Schema and utilities for input prompts to the LLM API."""
2
3 # SPDX-License-Identifier: Apache-2.0
4 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
5 from collections.abc import Mapping, Sequence
6 from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar, final
7
8 from typing_extensions import NotRequired, TypedDict
9
10 if TYPE_CHECKING:
11 import torch
12
13 from vllm.multimodal.inputs import AudioItem, ImageItem, VideoItem, VisionChunk
14
15
16 _T = TypeVar("_T")
17
18 ModalityData: TypeAlias = _T | list[_T | None] | None
19 """
20 Either a single data item, or a list of data items. Can only be None if UUID
21 is provided.
22
23 The number of data items allowed per modality is restricted by
24 `--limit-mm-per-prompt`.
25 """
26
27
28 @final
29 class MultiModalDataBuiltins(TypedDict, total=False):
30 """Type annotations for modality types predefined by vLLM."""
31
32 image: ModalityData["ImageItem"]
33 """The input image(s)."""
34
35 video: ModalityData["VideoItem"]
36 """The input video(s)."""
37
38 audio: ModalityData["AudioItem"]
39 """The input audio(s)."""
40
41 vision_chunk: ModalityData["VisionChunk"]
42 """The input visual atom(s) - unified modality for images and video chunks."""
43
44
45 MultiModalDataDict: TypeAlias = Mapping[str, ModalityData[Any]]
46 """
47 A dictionary containing an entry for each modality type to input.
48
49 The built-in modalities are defined by
50 [`MultiModalDataBuiltins`][vllm.inputs.llm.MultiModalDataBuiltins].
51 """
52
53 MultiModalUUIDDict: TypeAlias = Mapping[str, Sequence[str | None] | str]
54 """
55 A dictionary containing user-provided UUIDs for items in each modality.
56 If a UUID for an item is not provided, its entry will be `None` and
57 MultiModalHasher will compute a hash for the item.
58
59 The UUID will be used to identify the item for all caching purposes
60 (input processing caching, embedding caching, prefix caching, etc).
61 """
62
63
64 class _PromptOptions(TypedDict):
65 """Additional options available to all
66 [`SingletonPrompt`][vllm.inputs.llm.SingletonPrompt] types.
67 """
68
69 multi_modal_data: NotRequired[MultiModalDataDict | None]
70 """
71 Optional multi-modal data to pass to the model,
72 if the model supports it.
73 """
74
75 media_io_kwargs: NotRequired[dict[str, dict[str, Any]] | None]
76 """
77 Optional per-modality media loading and decoding arguments.
78 Used only for hash derivation when using multi-modal UUIDs.
79 """
80
81 mm_processor_kwargs: NotRequired[dict[str, Any] | None]
82 """
83 Optional multi-modal processor kwargs to be forwarded to the
84 multimodal input mapper & processor. Note that if multiple modalities
85 have registered mappers etc for the model being considered, we attempt
86 to pass the mm_processor_kwargs to each of them.
87 """
88
89 multi_modal_uuids: NotRequired[MultiModalUUIDDict]
90 """
91 Optional user-specified UUIDs for multimodal items, mapped by modality.
92 Lists must match the number of items per modality and may contain `None`.
93 For `None` entries, the hasher will compute IDs automatically; non-None
94 entries override the default hashes for caching, and MUST be unique per
95 multimodal item.
96 """
97
98 cache_salt: NotRequired[str]
99 """
100 Optional cache salt to be used for prefix caching.
101 """
102
103
104 class TextPrompt(_PromptOptions):
105 """Schema for a text prompt."""
106
107 prompt: str
108 """The input text to be tokenized before passing to the model."""
109
110
111 class TokensPrompt(_PromptOptions):
112 """Schema for a tokenized prompt."""
113
114 prompt_token_ids: list[int]
115 """A list of token IDs to pass to the model."""
116
117 prompt: NotRequired[str]
118 """The prompt text corresponding to the token IDs, if available."""
119
120 token_type_ids: NotRequired[list[int]]
121 """A list of token type IDs to pass to the cross encoder model."""
122
123 prompt_token_offsets: NotRequired[list[tuple[int, int]] | None]
124 """Char-level (start, end) offsets per token, relative to the
125 tokenized source string. Present only when offsets were requested
126 AND a Fast (Rust-backed) tokenizer was used AND no multimodal data
127 was present. The list length equals the length of `prompt_token_ids`."""
128
129
130 class EmbedsPrompt(_PromptOptions):
131 """Schema for a prompt provided via token embeddings."""
132
133 prompt_embeds: "torch.Tensor"
134 """The embeddings of the prompt."""
135
136 prompt: NotRequired[str]
137 """The prompt text corresponding to the token embeddings, if available."""
138
139 prompt_token_ids: NotRequired[list[int]]
140 """Token IDs for mixed-mode inputs (chat completion with
141 `prompt_embeds` content parts). The tokens at positions where
142 `prompt_is_token_ids` is `False` are placeholder tokens that
143 get replaced by entries from `prompt_embeds` in the forward pass."""
144
145 prompt_is_token_ids: NotRequired[list[bool]]
146 """Per-position mask, `True` uses the real token ID, `False` uses
147 the corresponding entry from `prompt_embeds`.
148 Must be the same length as `prompt_token_ids` when both are set."""
149
150
151 DecoderOnlyPrompt: TypeAlias = (
152 str | TextPrompt | list[int] | TokensPrompt | EmbedsPrompt
153 )
154 """
155 Schema of a prompt for a decoder-only model:
156
157 - A text prompt (string or [`TextPrompt`][vllm.inputs.llm.TextPrompt])
158 - A tokenized prompt (list of token IDs, or
159 [`TokensPrompt`][vllm.inputs.llm.TokensPrompt])
160 - An embeddings prompt ([`EmbedsPrompt`][vllm.inputs.llm.EmbedsPrompt])
161
162 For encoder-decoder models, passing a singleton prompt is shorthand for passing
163 `ExplicitEncoderDecoderPrompt(encoder_prompt=prompt, decoder_prompt=None)`.
164 """
165
166
167 EncoderPrompt: TypeAlias = str | TextPrompt | list[int] | TokensPrompt
168 """
169 Schema of a prompt for the encoder part of a encoder-decoder model:
170
171 - A text prompt (string or [`TextPrompt`][vllm.inputs.llm.TextPrompt])
172 - A tokenized prompt (list of token IDs, or
173 [`TokensPrompt`][vllm.inputs.llm.TokensPrompt])
174 """
175
176
177 DecoderPrompt: TypeAlias = str | TextPrompt | list[int] | TokensPrompt
178 """
179 Schema of a prompt for the decoder part of an encoder-decoder model:
180
181 - A text prompt (string or [`TextPrompt`][vllm.inputs.llm.TextPrompt])
182 - A tokenized prompt (list of token IDs, or
183 [`TokensPrompt`][vllm.inputs.llm.TokensPrompt])
184
185 Note:
186 Multi-modal inputs are not supported for decoder prompts.
187 """
188
189
190 class ExplicitEncoderDecoderPrompt(TypedDict):
191 """Schema for a pair of encoder and decoder singleton prompts.
192
193 Note:
194 This schema is not valid for decoder-only models.
195
196 """
197
198 encoder_prompt: EncoderPrompt
199 """The prompt for the encoder part of the model."""
200
201 decoder_prompt: DecoderPrompt | None
202 """
203 The prompt for the decoder part of the model.
204
205 Passing `None` will cause the prompt to be inferred automatically.
206 """
207
208
209 EncoderDecoderPrompt: TypeAlias = EncoderPrompt | ExplicitEncoderDecoderPrompt
210 """
211 Schema for a prompt for an encoder-decoder model.
212
213 You can pass a singleton encoder prompt, in which case the decoder prompt is
214 considered to be `None` (i.e., infer automatically).
215 """
216
217
218 SingletonPrompt: TypeAlias = DecoderOnlyPrompt | EncoderPrompt | DecoderPrompt
219 """
220 Schema for a single prompt. This is as opposed to a data structure
221 which encapsulates multiple prompts, such as
222 [`ExplicitEncoderDecoderPrompt`][vllm.inputs.llm.ExplicitEncoderDecoderPrompt].
223 """
224
225
226 PromptType: TypeAlias = DecoderOnlyPrompt | EncoderDecoderPrompt
227 """
228 Schema for any prompt, regardless of model type.
229
230 This is the input format accepted by most [`LLM`][vllm.entrypoints.llm.LLM] APIs.
231 """
232
233
234 class DataPrompt(_PromptOptions):
235 """Represents generic inputs that are converted to
236 [`PromptType`][vllm.inputs.llm.PromptType] by IO processor plugins.
237 """
238
239 data: Any
240 """The input data."""
241
242 data_format: str
243 """The input data format."""