data.py 11.2 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from collections.abc import Iterable
from typing import TYPE_CHECKING, Any, Generic, Literal, Optional, Union, cast
5

6
import torch
7
from typing_extensions import NotRequired, TypedDict, TypeIs, TypeVar
8
9

if TYPE_CHECKING:
10
11
    from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalInputs,
                                        MultiModalUUIDDict)
12
13
14
15
16
17
18
19


class TextPrompt(TypedDict):
    """Schema for a text prompt."""

    prompt: str
    """The input text to be tokenized before passing to the model."""

20
    multi_modal_data: NotRequired["MultiModalDataDict"]
21
22
23
24
25
    """
    Optional multi-modal data to pass to the model,
    if the model supports it.
    """

26
    mm_processor_kwargs: NotRequired[dict[str, Any]]
27
28
29
30
31
32
33
    """
    Optional multi-modal processor kwargs to be forwarded to the
    multimodal input mapper & processor. Note that if multiple modalities
    have registered mappers etc for the model being considered, we attempt
    to pass the mm_processor_kwargs to each of them.
    """

34
35
36
37
38
39
40
41
42
    multi_modal_uuids: NotRequired["MultiModalUUIDDict"]
    """
    Optional user-specified UUIDs for multimodal items, mapped by modality.
    Lists must match the number of items per modality and may contain `None`.
    For `None` entries, the hasher will compute IDs automatically; non-None
    entries override the default hashes for caching, and MUST be unique per
    multimodal item.
    """

43
44
45
46
47
    cache_salt: NotRequired[str]
    """
    Optional cache salt to be used for prefix caching.
    """

48
49
50
51

class TokensPrompt(TypedDict):
    """Schema for a tokenized prompt."""

52
    prompt_token_ids: list[int]
53
54
    """A list of token IDs to pass to the model."""

55
    token_type_ids: NotRequired[list[int]]
56
57
    """A list of token type IDs to pass to the cross encoder model."""

58
    multi_modal_data: NotRequired["MultiModalDataDict"]
59
    """
60
    Optional multi-modal data to pass to the model,
61
62
63
    if the model supports it.
    """

64
    mm_processor_kwargs: NotRequired[dict[str, Any]]
65
    """
66
    Optional multi-modal processor kwargs to be forwarded to the
67
68
69
70
71
    multimodal input mapper & processor. Note that if multiple modalities
    have registered mappers etc for the model being considered, we attempt
    to pass the mm_processor_kwargs to each of them.
    """

72
73
74
75
76
77
78
79
    multi_modal_uuids: NotRequired["MultiModalUUIDDict"]
    """
    Optional user-specified UUIDs for multimodal items, mapped by modality.
    Lists must match the number of items per modality and may contain `None`.
    For `None` entries, the hasher will compute IDs automatically; non-None
    entries override the default hashes for caching.
    """

80
81
82
83
84
    cache_salt: NotRequired[str]
    """
    Optional cache salt to be used for prefix caching.
    """

85

86
87
88
89
90
91
class EmbedsPrompt(TypedDict):
    """Schema for a prompt provided via token embeddings."""

    prompt_embeds: torch.Tensor
    """The embeddings of the prompt."""

92
93
94
95
96
    cache_salt: NotRequired[str]
    """
    Optional cache salt to be used for prefix caching.
    """

97
98

SingletonPrompt = Union[str, TextPrompt, TokensPrompt, EmbedsPrompt]
99
"""
100
Set of possible schemas for a single prompt:
101

102
103
104
- A text prompt ([`str`][] or [`TextPrompt`][vllm.inputs.data.TextPrompt])
- A tokenized prompt ([`TokensPrompt`][vllm.inputs.data.TokensPrompt])
- An embeddings prompt ([`EmbedsPrompt`][vllm.inputs.data.EmbedsPrompt])
105
106
107
108
109

Note that "singleton" is as opposed to a data structure
which encapsulates multiple prompts, i.e. of the sort
which may be utilized for encoder/decoder models when
the user desires to express both the encoder & decoder
110
111
prompts explicitly, i.e. 
[`ExplicitEncoderDecoderPrompt`][vllm.inputs.data.ExplicitEncoderDecoderPrompt]
112

113
114
A prompt of type [`SingletonPrompt`][vllm.inputs.data.SingletonPrompt] may be 
employed as (1) input to a decoder-only model, (2) input to
115
116
117
the encoder of an encoder/decoder model, in the scenario
where the decoder-prompt is not specified explicitly, or
(3) as a member of a larger data structure encapsulating
118
119
more than one prompt, i.e. 
[`ExplicitEncoderDecoderPrompt`][vllm.inputs.data.ExplicitEncoderDecoderPrompt]
120
121
"""

122
123
124
125
126
127
128
129
130
131
132

def is_tokens_prompt(prompt: SingletonPrompt) -> TypeIs[TokensPrompt]:
    return (isinstance(prompt, dict) and "prompt_token_ids" in prompt
            and "prompt_embeds" not in prompt)


def is_embeds_prompt(prompt: SingletonPrompt) -> TypeIs[EmbedsPrompt]:
    return (isinstance(prompt, dict) and "prompt_token_ids" not in prompt
            and "prompt_embeds" in prompt)


133
_T1_co = TypeVar("_T1_co",
134
135
                 bound=SingletonPrompt,
                 default=SingletonPrompt,
136
137
                 covariant=True)
_T2_co = TypeVar("_T2_co",
138
139
                 bound=SingletonPrompt,
                 default=SingletonPrompt,
140
                 covariant=True)
141

142
143
144

# TODO: Make fields ReadOnly once mypy supports it
class ExplicitEncoderDecoderPrompt(TypedDict, Generic[_T1_co, _T2_co]):
145
146
147
    """
    Represents an encoder/decoder model input prompt,
    comprising an explicit encoder prompt and a decoder prompt.
148

149
    The encoder and decoder prompts, respectively, may be formatted
150
151
    according to any of the
    [`SingletonPrompt`][vllm.inputs.data.SingletonPrompt] schemas,
152
    and are not required to have the same schema.
153

154
155
156
    Only the encoder prompt may have multi-modal data. mm_processor_kwargs
    should be at the top-level, and should not be set in the encoder/decoder
    prompts, since they are agnostic to the encoder/decoder.
157

158
159
160
    Note that an
    [`ExplicitEncoderDecoderPrompt`][vllm.inputs.data.ExplicitEncoderDecoderPrompt]
    may not be used as an input to a decoder-only model,
161
    and that the `encoder_prompt` and `decoder_prompt`
162
    fields of this data structure themselves must be
163
    [`SingletonPrompt`][vllm.inputs.data.SingletonPrompt] instances.
164
165
    """

166
    encoder_prompt: _T1_co
167

168
    decoder_prompt: Optional[_T2_co]
169

170
    mm_processor_kwargs: NotRequired[dict[str, Any]]
171

172

173
PromptType = Union[SingletonPrompt, ExplicitEncoderDecoderPrompt]
174
175
176
177
"""
Set of possible schemas for an LLM input, including
both decoder-only and encoder/decoder input types:

178
179
180
- A text prompt ([`str`][] or [`TextPrompt`][vllm.inputs.data.TextPrompt])
- A tokenized prompt ([`TokensPrompt`][vllm.inputs.data.TokensPrompt])
- An embeddings prompt ([`EmbedsPrompt`][vllm.inputs.data.EmbedsPrompt])
181
- A single data structure containing both an encoder and a decoder prompt
182
  ([`ExplicitEncoderDecoderPrompt`][vllm.inputs.data.ExplicitEncoderDecoderPrompt])
183
184
185
"""


186
187
class TokenInputs(TypedDict):
    """Represents token-based inputs."""
188
189
190
191

    type: Literal["token"]
    """The type of inputs."""

192
    prompt_token_ids: list[int]
193
194
    """The token IDs of the prompt."""

195
    prompt: NotRequired[str]
196
197
198
199
    """
    The original prompt text corresponding to the token IDs, if available.
    """

200
201
202
203
204
    cache_salt: NotRequired[str]
    """
    Optional cache salt to be used for prefix caching.
    """

205

206
def token_inputs(
207
    prompt_token_ids: list[int],
208
    prompt: Optional[str] = None,
209
    cache_salt: Optional[str] = None,
210
) -> TokenInputs:
211
212
    """Construct [`TokenInputs`][vllm.inputs.data.TokenInputs] from optional
    values."""
213
    inputs = TokenInputs(type="token", prompt_token_ids=prompt_token_ids)
214
215
216

    if prompt is not None:
        inputs["prompt"] = prompt
217
218
    if cache_salt is not None:
        inputs["cache_salt"] = cache_salt
219
220
221
222

    return inputs


223
224
225
226
227
228
229
230
231
class EmbedsInputs(TypedDict):
    """Represents embeddings-based inputs."""

    type: Literal["embeds"]
    """The type of inputs."""

    prompt_embeds: torch.Tensor
    """The embeddings of the prompt."""

232
233
234
235
236
    cache_salt: NotRequired[str]
    """
    Optional cache salt to be used for prefix caching.
    """

237

238
239
240
241
def embeds_inputs(
    prompt_embeds: torch.Tensor,
    cache_salt: Optional[str] = None,
) -> EmbedsInputs:
242
243
    """Construct [`EmbedsInputs`][vllm.inputs.data.EmbedsInputs] from optional
    values."""
244
245
246
247
    inputs = EmbedsInputs(type="embeds", prompt_embeds=prompt_embeds)

    if cache_salt is not None:
        inputs["cache_salt"] = cache_salt
248
249
250
251
252

    return inputs


DecoderOnlyInputs = Union[TokenInputs, EmbedsInputs, "MultiModalInputs"]
253
"""
254
The inputs in [`LLMEngine`][vllm.engine.llm_engine.LLMEngine] before they are
255
256
257
258
259
passed to the model executor.
This specifies the data required for decoder-only models.
"""


260
class EncoderDecoderInputs(TypedDict):
261
    """
262
263
    The inputs in [`LLMEngine`][vllm.engine.llm_engine.LLMEngine] before they
    are passed to the model executor.
264
265
266

    This specifies the required data for encoder-decoder models.
    """
267

268
    encoder: Union[TokenInputs, "MultiModalInputs"]
269
    """The inputs for the encoder portion."""
270

271
    decoder: Union[TokenInputs, "MultiModalInputs"]
272
    """The inputs for the decoder portion."""
273

274

275
SingletonInputs = Union[TokenInputs, EmbedsInputs, "MultiModalInputs"]
276
"""
277
278
A processed [`SingletonPrompt`][vllm.inputs.data.SingletonPrompt] which can be 
passed to [`vllm.sequence.Sequence`][].
279
280
281
282
"""

ProcessorInputs = Union[DecoderOnlyInputs, EncoderDecoderInputs]
"""
283
The outputs from [`vllm.inputs.preprocess.InputPreprocessor`][].
284
"""
285

286
287
_T1 = TypeVar("_T1", bound=SingletonPrompt, default=SingletonPrompt)
_T2 = TypeVar("_T2", bound=SingletonPrompt, default=SingletonPrompt)
288
289


290
291
292
def build_explicit_enc_dec_prompt(
    encoder_prompt: _T1,
    decoder_prompt: Optional[_T2],
293
    mm_processor_kwargs: Optional[dict[str, Any]] = None,
294
) -> ExplicitEncoderDecoderPrompt[_T1, _T2]:
295
296
297
298
299
    if mm_processor_kwargs is None:
        mm_processor_kwargs = {}
    return ExplicitEncoderDecoderPrompt(
        encoder_prompt=encoder_prompt,
        decoder_prompt=decoder_prompt,
300
301
        mm_processor_kwargs=mm_processor_kwargs,
    )
302
303
304
305
306


def zip_enc_dec_prompts(
    enc_prompts: Iterable[_T1],
    dec_prompts: Iterable[Optional[_T2]],
307
308
309
    mm_processor_kwargs: Optional[Union[Iterable[dict[str, Any]],
                                        dict[str, Any]]] = None,
) -> list[ExplicitEncoderDecoderPrompt[_T1, _T2]]:
310
    """
311
    Zip encoder and decoder prompts together into a list of
312
313
    [`ExplicitEncoderDecoderPrompt`][vllm.inputs.data.ExplicitEncoderDecoderPrompt]
    instances.
314

315
316
317
    ``mm_processor_kwargs`` may also be provided; if a dict is passed, the same
    dictionary will be used for every encoder/decoder prompt. If an iterable is
    provided, it will be zipped with the encoder/decoder prompts.
318
319
    """
    if mm_processor_kwargs is None:
320
        mm_processor_kwargs = cast(dict[str, Any], {})
321
    if isinstance(mm_processor_kwargs, dict):
322
        return [
323
            build_explicit_enc_dec_prompt(
324
325
326
327
328
                encoder_prompt,
                decoder_prompt,
                cast(dict[str, Any], mm_processor_kwargs),
            ) for (encoder_prompt,
                   decoder_prompt) in zip(enc_prompts, dec_prompts)
329
        ]
330
    return [
331
332
333
334
        build_explicit_enc_dec_prompt(encoder_prompt, decoder_prompt,
                                      mm_proc_kwargs)
        for (encoder_prompt, decoder_prompt, mm_proc_kwargs
             ) in zip(enc_prompts, dec_prompts, mm_processor_kwargs)
335
336
    ]

337

338
339
def to_enc_dec_tuple_list(
    enc_dec_prompts: Iterable[ExplicitEncoderDecoderPrompt[_T1, _T2]],
340
) -> list[tuple[_T1, Optional[_T2]]]:
341
342
343
    return [(enc_dec_prompt["encoder_prompt"],
             enc_dec_prompt["decoder_prompt"])
            for enc_dec_prompt in enc_dec_prompts]