io_processor.py 8.68 KB
Newer Older
1
2
3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4
from collections.abc import Sequence
5
6
from typing import Any, Final

7
8
from vllm import PoolingParams, PoolingRequestOutput, PromptType
from vllm.config import VllmConfig
9
10
11
12
13
14
15
from vllm.entrypoints.chat_utils import (
    ChatCompletionMessageParam,
    ChatTemplateConfig,
    ChatTemplateContentFormatOption,
    ConversationMessage,
)
from vllm.entrypoints.openai.engine.serving import RendererChatRequest, RendererRequest
16
from vllm.entrypoints.pooling.scoring.typing import ScoringData
17
from vllm.entrypoints.pooling.typing import (
18
19
    OfflineInputsContext,
    OfflineOutputsContext,
20
21
22
23
    PoolingChatLikeRequest,
    PoolingCompletionLikeRequest,
    PoolingServeContext,
)
24
from vllm.inputs import EngineInput, SingletonPrompt
25
from vllm.renderers import BaseRenderer, TokenizeParams, merge_kwargs
26
27
28
29
30
31
from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq
from vllm.tool_parsers import ToolParser
from vllm.utils.mistral import is_mistral_tokenizer


class PoolingIOProcessor:
32
33
    name: str

34
35
    def __init__(
        self,
36
        vllm_config: VllmConfig,
37
38
39
        renderer: BaseRenderer,
        chat_template_config: ChatTemplateConfig,
    ):
40
41
        self.vllm_config = vllm_config
        self.model_config = vllm_config.model_config
42
43
44
45
46
47
48
49
50
51
        self.renderer = renderer

        self.chat_template = chat_template_config.chat_template
        self.chat_template_content_format: Final = (
            chat_template_config.chat_template_content_format
        )
        self.trust_request_chat_template = (
            chat_template_config.trust_request_chat_template
        )

52
53
    #######################################
    # online APIs
54

55
56
57
    def create_pooling_params(self, request):
        return request.to_pooling_params()

58
59
60
61
62
63
64
65
66
    def pre_process_online(self, ctx: PoolingServeContext):
        request = ctx.request

        if isinstance(ctx.request, PoolingChatLikeRequest):
            self._validate_chat_template(
                request_chat_template=request.chat_template,
                chat_template_kwargs=request.chat_template_kwargs,
                trust_request_chat_template=self.trust_request_chat_template,
            )
67
            _, engine_inputs = self._preprocess_chat_online(
68
69
70
71
72
73
74
                request,
                request.messages,
                default_template=self.chat_template,
                default_template_content_format=self.chat_template_content_format,
                default_template_kwargs=None,
            )
        elif isinstance(request, PoolingCompletionLikeRequest):
75
            engine_inputs = self._preprocess_completion_online(
76
77
78
79
80
81
82
                request,
                prompt_input=request.input,
                prompt_embeds=None,
            )
        else:
            raise ValueError(f"Invalid {self.name} request type")

83
        ctx.engine_inputs = engine_inputs
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

    async def pre_process_online_async(self, ctx: PoolingServeContext):
        self.pre_process_online(ctx)

    def post_process_online(
        self,
        ctx: PoolingServeContext,
    ):
        pass

    async def post_process_online_async(
        self,
        ctx: PoolingServeContext,
    ):
        self.post_process_online(ctx)

    #######################################
    # offline APIs

103
    def pre_process_offline(self, ctx: OfflineInputsContext) -> Sequence[EngineInput]:
104
105
106
107
108
        assert not isinstance(ctx.prompts, ScoringData) and not (
            isinstance(ctx.prompts, dict) and "data" in ctx.prompts
        )

        prompts_seq = prompt_to_seq(ctx.prompts)
109
110
111
        tok_params = self.renderer.default_cmpl_tok_params.with_kwargs(
            **(ctx.tokenization_kwargs or {})
        )
112
        return self._preprocess_completion_offline(
113
            prompts=prompts_seq, tok_params=tok_params
114
        )
115

116
117
    async def pre_process_offline_async(self, ctx: OfflineInputsContext):
        return self.pre_process_offline(ctx)
118

119
120
    def post_process_offline(
        self,
121
        ctx: OfflineOutputsContext,
122
    ) -> list[PoolingRequestOutput]:
123
        return ctx.outputs
124

125
126
    async def post_process_offline_async(
        self,
127
        ctx: OfflineOutputsContext,
128
    ) -> list[PoolingRequestOutput]:
129
        return self.post_process_offline(ctx)
130

131
132
    #######################################
    # helpers
133
134
135
136
137
138

    def _preprocess_completion_online(
        self,
        request: RendererRequest,
        prompt_input: str | list[str] | list[int] | list[list[int]] | None,
        prompt_embeds: bytes | list[bytes] | None,
139
    ) -> list[EngineInput]:
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
        renderer = self.renderer
        model_config = self.model_config

        prompts = list[SingletonPrompt | bytes]()
        if prompt_embeds is not None:  # embeds take higher priority
            prompts.extend(prompt_to_seq(prompt_embeds))
        if prompt_input is not None:
            prompts.extend(prompt_to_seq(prompt_input))

        parsed_prompts = [
            (
                prompt
                if isinstance(prompt, bytes)
                else parse_model_prompt(model_config, prompt)
            )
            for prompt in prompts
        ]
        tok_params = request.build_tok_params(model_config)

        return renderer.render_cmpl(
            parsed_prompts,
            tok_params,
            prompt_extras={
                k: v
                for k in ("mm_processor_kwargs", "cache_salt")
                if (v := getattr(request, k, None)) is not None
            },
        )

    def _preprocess_chat_online(
        self,
        request: RendererChatRequest,
        messages: list[ChatCompletionMessageParam],
        default_template: str | None,
        default_template_content_format: ChatTemplateContentFormatOption,
        default_template_kwargs: dict[str, Any] | None,
        tool_dicts: list[dict[str, Any]] | None = None,
177
        tool_parser: type[ToolParser] | None = None,
178
    ) -> tuple[list[ConversationMessage], list[EngineInput]]:
179
180
181
182
183
184
185
186
187
188
        renderer = self.renderer

        default_template_kwargs = merge_kwargs(
            default_template_kwargs,
            dict(
                tools=tool_dicts,
                tokenize=is_mistral_tokenizer(renderer.tokenizer),
            ),
        )

189
190
        mm_config = self.model_config.multimodal_config

191
192
193
        tok_params = request.build_tok_params(self.model_config)
        chat_params = request.build_chat_params(
            default_template, default_template_content_format
194
195
196
197
        ).with_defaults(
            default_template_kwargs,
            default_media_io_kwargs=(mm_config.media_io_kwargs if mm_config else None),
        )
198

199
        (conversation,), (engine_input,) = renderer.render_chat(
200
201
202
203
204
205
206
207
208
209
            [messages],
            chat_params,
            tok_params,
            prompt_extras={
                k: v
                for k in ("mm_processor_kwargs", "cache_salt")
                if (v := getattr(request, k, None)) is not None
            },
        )

210
        return conversation, [engine_input]
211
212
213
214

    def _preprocess_completion_offline(
        self,
        prompts: PromptType | Sequence[PromptType],
215
216
        tok_params: TokenizeParams,
        prompt_extras: dict[str, Any] | None = None,
217
    ) -> Sequence[EngineInput]:
218
219
220
221
222
        prompts = prompt_to_seq(prompts)
        parsed_prompts = [
            (
                prompt
                if isinstance(prompt, bytes)
223
                else parse_model_prompt(self.model_config, prompt)
224
225
226
227
            )
            for prompt in prompts
        ]

228
229
        return self.renderer.render_cmpl(
            parsed_prompts, tok_params, prompt_extras=prompt_extras
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
        )

    def _validate_chat_template(
        self,
        request_chat_template: str | None,
        chat_template_kwargs: dict[str, Any] | None,
        trust_request_chat_template: bool,
    ):
        if not trust_request_chat_template and (
            request_chat_template is not None
            or (
                chat_template_kwargs
                and chat_template_kwargs.get("chat_template") is not None
            )
        ):
            raise ValueError(
                "Chat template is passed with request, but "
                "--trust-request-chat-template is not set. "
                "Refused request with untrusted chat template."
            )
        return None
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266

    def _params_to_seq(
        self,
        params: PoolingParams | Sequence[PoolingParams],
        num_requests: int,
    ) -> Sequence[PoolingParams]:
        if isinstance(params, Sequence):
            if len(params) != num_requests:
                raise ValueError(
                    f"The lengths of prompts ({num_requests}) "
                    f"and params ({len(params)}) must be the same."
                )

            return params

        return [params] * num_requests