preprocess.py 19.4 KB
Newer Older
1
import asyncio
2
from typing import List, Optional
3
4
5
6
7
8
9
10

from typing_extensions import assert_never

from vllm.config import ModelConfig
from vllm.logger import init_logger
from vllm.lora.request import LoRARequest
from vllm.prompt_adapter.request import PromptAdapterRequest
from vllm.transformers_utils.tokenizer_group import BaseTokenizerGroup
11
from vllm.utils import print_warning_once
12

13
14
from .data import (DecoderOnlyInputs, EncoderDecoderInputs, ProcessorInputs,
                   PromptType, SingletonInputs, SingletonPrompt, token_inputs)
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
from .parse import is_explicit_encoder_decoder_prompt, parse_singleton_prompt

logger = init_logger(__name__)


class InputPreprocessor:

    def __init__(
        self,
        model_config: ModelConfig,
        tokenizer: Optional[BaseTokenizerGroup],
    ) -> None:
        super().__init__()

        self.model_config = model_config
        self.tokenizer = tokenizer

    def get_tokenizer_group(self) -> BaseTokenizerGroup:
        if self.tokenizer is None:
            raise ValueError("You cannot pass text prompts when "
                             "`skip_tokenizer_init` is True")

        return self.tokenizer

    def get_bos_token_id(self,
                         lora_request: Optional[LoRARequest] = None
                         ) -> Optional[int]:
        if self.tokenizer is None:
            logger.warning("Using None for BOS token id because tokenizer "
                           "is not initialized")
            return None

        return self.tokenizer.get_lora_tokenizer(lora_request).bos_token_id

    def get_eos_token_id(self,
                         lora_request: Optional[LoRARequest] = None
                         ) -> Optional[int]:
        if self.tokenizer is None:
            logger.warning("Using None for EOS token id because tokenizer "
                           "is not initialized")
            return None

        return self.tokenizer.get_lora_tokenizer(lora_request).eos_token_id

    def get_decoder_start_token_id(self) -> Optional[int]:
        '''
        Obtain the decoder start token id employed by an encoder/decoder
        model. Returns None for non-encoder/decoder models or if the
        model config is unavailable.
        '''

        if not self.is_encoder_decoder_model():
67
68
            print_warning_once("Using None for decoder start token id because "
                               "this is not an encoder/decoder model.")
69
70
71
            return None

        if (self.model_config is None or self.model_config.hf_config is None):
72
73
            print_warning_once("Using None for decoder start token id because "
                               "model config is not available.")
74
75
76
77
78
            return None

        dec_start_token_id = getattr(self.model_config.hf_config,
                                     'decoder_start_token_id', None)
        if dec_start_token_id is None:
79
80
81
            print_warning_once("Falling back on <BOS> for decoder start token "
                               "id because decoder start token id is not "
                               "available.")
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
            dec_start_token_id = self.get_bos_token_id()

        return dec_start_token_id

    def _get_default_enc_dec_decoder_prompt(self) -> List[int]:
        '''
        Specifically for encoder/decoder models:
        generate a default decoder prompt for when
        the user specifies only the encoder prompt.

        Encoder/decoder models utilize the decoder
        prompt in different ways; as new models are
        added, it is intended that this function
        will be extended to produce differing
        default decoder prompts, depending on the
        model variety.

        Absent a special case, the default behavior
        of this method is to mirror the behavior of
        the HuggingFace (HF) GenerationMixin for a None
        decoder prompt, which is to employ a logit processor
        setting to force the first decoded token to be <BOS>.
        Here, this behavior is approximated by having the
        "default" decoder prompt be <BOS>.

        However, it is possible that in the future
108
        other models may have different or more
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
        complex logic for the default decoder prompt.
        This motivates having a special helper method
        for default decoder prompts.

        Returns:

        * prompt_token_ids
        '''

        bos_token_id = self.get_bos_token_id()
        assert bos_token_id is not None
        return [bos_token_id]

    def _prepare_decoder_input_ids_for_generation(
        self,
        decoder_input_ids: Optional[List[int]],
    ) -> List[int]:
        """
        Prepares `decoder_input_ids` for generation with encoder-decoder models.

        Based on

        https://github.com/huggingface/transformers/blob/
        4037a2b5b1278736e566aec12e169100275545ea/
        src/transformers/generation/utils.py

        specifically GenerationMixin._prepare_decoder_input_ids_for_generation()

        Arguments:

        * decoder_input_ids: input token ids to preprocess

        Returns:

        * Processed token list
        """

        decoder_start_token_id = self.get_decoder_start_token_id()
        assert decoder_start_token_id is not None

        if decoder_input_ids is None:
            # no decoder prompt input ->
            # use decoder_start_token_id as decoder_input_ids
            decoder_input_ids = self._get_default_enc_dec_decoder_prompt()

154
155
        if (len(decoder_input_ids) == 0
                or decoder_input_ids[0] != decoder_start_token_id):
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
            decoder_input_ids = [decoder_start_token_id] + decoder_input_ids

        return decoder_input_ids

    def _apply_prompt_adapter(
        self,
        prompt_token_ids: List[int],
        prompt_adapter_request: Optional[PromptAdapterRequest],
    ) -> List[int]:
        if prompt_adapter_request:
            prompt_token_ids = (
                [0] * prompt_adapter_request.prompt_adapter_num_virtual_tokens
                + prompt_token_ids)

        return prompt_token_ids

    def _tokenize_prompt(
        self,
        prompt: str,
        request_id: str,
        lora_request: Optional[LoRARequest],
    ) -> List[int]:
        """
        Apply the model's tokenizer to a text prompt, returning the
        corresponding token IDs.
        """
        tokenizer = self.get_tokenizer_group()

        return tokenizer.encode(request_id=request_id,
                                prompt=prompt,
                                lora_request=lora_request)

    async def _tokenize_prompt_async(
        self,
        prompt: str,
        request_id: str,
        lora_request: Optional[LoRARequest],
    ) -> List[int]:
        """Async version of :meth:`_tokenize_prompt`."""
        tokenizer = self.get_tokenizer_group()

        return await tokenizer.encode_async(request_id=request_id,
                                            prompt=prompt,
                                            lora_request=lora_request)

201
    def _prompt_to_llm_inputs(
202
        self,
203
        prompt: SingletonPrompt,
204
205
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
206
    ) -> SingletonInputs:
207
208
209
210
211
212
        '''
        Extract the components of any single encoder or decoder input prompt.

        Arguments:

        * request_id
213
        * prompt: single encoder or decoder input prompt
214
215
216
217
218
219
220
        * lora_request: this is only valid for decoder prompts

        Returns:

        * prompt
        * prompt_token_ids
        * multi_modal_data
221
        * mm_processor_kwargs (request-level input processor/mapper overrides)
222
223
        '''

224
        parsed = parse_singleton_prompt(prompt)
225
226

        if parsed["type"] == "str":
227
            prompt_text = parsed["content"]
228
            prompt_token_ids = self._tokenize_prompt(
229
                prompt_text,
230
231
232
                request_id=request_id,
                lora_request=lora_request,
            )
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255

            return token_inputs(
                prompt=prompt_text,
                prompt_token_ids=prompt_token_ids,
            )

        if parsed["type"] == "tokens":
            tokens_content = parsed["content"]

            prompt_token_ids = tokens_content["prompt_token_ids"]
            multi_modal_data = tokens_content.get("multi_modal_data")
            mm_processor_kwargs = tokens_content.get("mm_processor_kwargs")

            return token_inputs(
                prompt_token_ids=prompt_token_ids,
                multi_modal_data=multi_modal_data,
                mm_processor_kwargs=mm_processor_kwargs,
            )

        if parsed["type"] == "text":
            text_content = parsed["content"]

            prompt_text = text_content["prompt"]
256
            prompt_token_ids = self._tokenize_prompt(
257
                prompt_text,
258
259
260
                request_id=request_id,
                lora_request=lora_request,
            )
261
262
263
264
265
266
267
268
269
            multi_modal_data = text_content.get("multi_modal_data")
            mm_processor_kwargs = text_content.get("mm_processor_kwargs")

            return token_inputs(
                prompt=prompt_text,
                prompt_token_ids=prompt_token_ids,
                multi_modal_data=multi_modal_data,
                mm_processor_kwargs=mm_processor_kwargs,
            )
270

271
        assert_never(parsed)
272

273
    async def _prompt_to_llm_inputs_async(
274
        self,
275
        prompt: SingletonPrompt,
276
277
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
278
    ) -> SingletonInputs:
279
        """Async version of :meth:`_extract_prompt_components`."""
280
        parsed = parse_singleton_prompt(prompt)
281
282

        if parsed["type"] == "str":
283
            prompt_text = parsed["content"]
284
            prompt_token_ids = await self._tokenize_prompt_async(
285
                prompt_text,
286
287
288
                request_id=request_id,
                lora_request=lora_request,
            )
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311

            return token_inputs(
                prompt=prompt_text,
                prompt_token_ids=prompt_token_ids,
            )

        if parsed["type"] == "tokens":
            tokens_content = parsed["content"]

            prompt_token_ids = tokens_content["prompt_token_ids"]
            multi_modal_data = tokens_content.get("multi_modal_data")
            mm_processor_kwargs = tokens_content.get("mm_processor_kwargs")

            return token_inputs(
                prompt_token_ids=prompt_token_ids,
                multi_modal_data=multi_modal_data,
                mm_processor_kwargs=mm_processor_kwargs,
            )

        if parsed["type"] == "text":
            text_content = parsed["content"]

            prompt_text = text_content["prompt"]
312
            prompt_token_ids = await self._tokenize_prompt_async(
313
                prompt_text,
314
315
316
                request_id=request_id,
                lora_request=lora_request,
            )
317
318
319
320
321
322
323
324
325
            multi_modal_data = text_content.get("multi_modal_data")
            mm_processor_kwargs = text_content.get("mm_processor_kwargs")

            return token_inputs(
                prompt=prompt_text,
                prompt_token_ids=prompt_token_ids,
                multi_modal_data=multi_modal_data,
                mm_processor_kwargs=mm_processor_kwargs,
            )
326

327
        assert_never(parsed)
328
329
330

    def _build_enc_dec_llm_inputs(
        self,
331
332
        encoder_inputs: SingletonInputs,
        decoder_inputs: Optional[SingletonInputs],
333
    ) -> EncoderDecoderInputs:
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
        if encoder_inputs["type"] == "token":
            pass
        else:
            assert_never(encoder_inputs)

        if decoder_inputs is None:
            dec_token_ids = self._prepare_decoder_input_ids_for_generation(
                None)
            decoder_inputs = token_inputs(dec_token_ids)
        elif decoder_inputs["type"] == "token":
            dec_token_ids = self._prepare_decoder_input_ids_for_generation(
                decoder_inputs["prompt_token_ids"])
            decoder_inputs["prompt_token_ids"] = dec_token_ids

            if "multi_modal_data" in decoder_inputs:
                raise ValueError("Multi-modal decoder inputs of encoder-"
                                 "decoder models are not supported yet")
        else:
            assert_never(encoder_inputs)
353

354
        return EncoderDecoderInputs(
355
356
            encoder=encoder_inputs,
            decoder=decoder_inputs,
357
358
359
360
        )

    def _process_encoder_decoder_prompt(
        self,
361
        prompt: PromptType,
362
        request_id: str,
363
    ) -> EncoderDecoderInputs:
364
365
        '''
        For encoder/decoder models only:
366
        Process an input prompt into an :class:`EncoderDecoderInputs` instance.
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384

        There are two types of input prompts:
        singleton prompts which carry only the
        encoder prompt, and explicit encoder/decoder
        prompts which carry both the encoder and the
        decoder prompts as member variables.

        This function handles the following scenarios:
        * Singleton encoder prompt: extract encoder prompt
          token ids & infer default decoder prompt token ids
        * Explicit encoder/decoder prompt: extract encoder
          and decoder prompt token ids

        Note that for Explicit encoder/decoder prompts,
        each sub-prompt (encoder or decoder prompt) can
        have any possible singleton type; thus this
        method relies on helper functions to obtain
        token ids for the sub-prompts.
385

386
387
        Arguments:

388
        * prompt: an input prompt
389
390
391
392
        * request_id

        Returns:

393
        * :class:`EncoderDecoderInputs` instance
394
395
        '''

396
397
        encoder_inputs: SingletonInputs
        decoder_inputs: Optional[SingletonInputs]
398

399
        if is_explicit_encoder_decoder_prompt(prompt):
400
            encoder_inputs = self._prompt_to_llm_inputs(
401
                prompt["encoder_prompt"],
402
403
404
                request_id=request_id,
            )

405
            if (decoder_input := prompt["decoder_prompt"]) is None:
406
                decoder_inputs = None
407
            else:
408
                decoder_inputs = self._prompt_to_llm_inputs(
409
410
411
412
                    decoder_input,
                    request_id=request_id,
                )
        else:
413
            encoder_inputs = self._prompt_to_llm_inputs(
414
                prompt,
415
416
                request_id=request_id,
            )
417
418
419
420

            decoder_inputs = None

        return self._build_enc_dec_llm_inputs(encoder_inputs, decoder_inputs)
421
422
423

    async def _process_encoder_decoder_prompt_async(
        self,
424
        prompt: PromptType,
425
        request_id: str,
426
    ) -> EncoderDecoderInputs:
427
        """Async version of :meth:`_process_encoder_decoder_prompt`."""
428
429
        encoder_inputs: SingletonInputs
        decoder_inputs: Optional[SingletonInputs]
430

431
        if is_explicit_encoder_decoder_prompt(prompt):
432
            encoder_task = self._prompt_to_llm_inputs_async(
433
                prompt["encoder_prompt"],
434
435
436
                request_id=request_id,
            )

437
            if (decoder_input := prompt["decoder_prompt"]) is None:
438
439
                encoder_inputs = await encoder_task
                decoder_inputs = None
440
            else:
441
                decoder_task = self._prompt_to_llm_inputs_async(
442
443
444
445
                    decoder_input,
                    request_id=request_id,
                )

446
                encoder_inputs, decoder_inputs = await asyncio.gather(
447
448
                    encoder_task, decoder_task)
        else:
449
            encoder_inputs = await self._prompt_to_llm_inputs_async(
450
                prompt,
451
452
                request_id=request_id,
            )
453
454
455
456

            decoder_inputs = None

        return self._build_enc_dec_llm_inputs(encoder_inputs, decoder_inputs)
457
458
459

    def _build_decoder_only_llm_inputs(
        self,
460
        prompt_inputs: DecoderOnlyInputs,
461
        prompt_adapter_request: Optional[PromptAdapterRequest],
462
    ) -> DecoderOnlyInputs:
463
464
465
466
467
468
469
        if prompt_inputs["type"] == "token":
            prompt_inputs["prompt_token_ids"] = self._apply_prompt_adapter(
                prompt_inputs["prompt_token_ids"],
                prompt_adapter_request=prompt_adapter_request,
            )
        else:
            assert_never(prompt_inputs)
470

471
        return prompt_inputs
472
473
474

    def _process_decoder_only_prompt(
        self,
475
        prompt: SingletonPrompt,
476
477
478
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
479
    ) -> DecoderOnlyInputs:
480
481
        '''
        For decoder-only models:
482
        Process an input prompt into an :class:`DecoderOnlyInputs` instance.
483
484
485

        Arguments:

486
        * prompt: input prompt
487
488
489
490
491
492
        * request_id
        * lora_request
        * prompt_adapter_request

        Returns:

493
        * :class:`DecoderOnlyInputs` instance
494
495
        '''

496
        prompt_comps = self._prompt_to_llm_inputs(
497
            prompt,
498
499
500
501
502
503
504
505
506
507
508
            request_id=request_id,
            lora_request=lora_request,
        )

        return self._build_decoder_only_llm_inputs(
            prompt_comps,
            prompt_adapter_request=prompt_adapter_request,
        )

    async def _process_decoder_only_prompt_async(
        self,
509
        prompt: SingletonPrompt,
510
511
512
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
513
    ) -> DecoderOnlyInputs:
514
        """Async version of :meth:`_process_decoder_only_prompt`."""
515
        prompt_comps = await self._prompt_to_llm_inputs_async(
516
            prompt,
517
518
519
520
521
522
523
524
525
526
527
            request_id=request_id,
            lora_request=lora_request,
        )

        return self._build_decoder_only_llm_inputs(
            prompt_comps,
            prompt_adapter_request=prompt_adapter_request,
        )

    def preprocess(
        self,
528
        prompt: PromptType,
529
530
531
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
532
    ) -> ProcessorInputs:
533
534
535
536
537
        """Preprocess the input prompt."""
        if self.is_encoder_decoder_model():
            # Encoder-decoder model requires special mapping of
            # input prompts to encoder & decoder
            return self._process_encoder_decoder_prompt(
538
                prompt,
539
540
541
                request_id=request_id,
            )

542
        if is_explicit_encoder_decoder_prompt(prompt):
543
544
545
546
547
            raise ValueError("Cannot pass encoder-decoder prompt "
                             "to decoder-only models")

        # Decoder-only operation
        return self._process_decoder_only_prompt(
548
            prompt,
549
550
551
552
553
554
555
            request_id=request_id,
            lora_request=lora_request,
            prompt_adapter_request=prompt_adapter_request,
        )

    async def preprocess_async(
        self,
556
        prompt: PromptType,
557
558
559
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
560
    ) -> ProcessorInputs:
561
562
563
564
565
        """Async version of :meth:`preprocess`."""
        if self.is_encoder_decoder_model():
            # Encoder-decoder model requires special mapping of
            # input prompts to encoder & decoder
            return await self._process_encoder_decoder_prompt_async(
566
                prompt,
567
568
569
                request_id=request_id,
            )

570
        if is_explicit_encoder_decoder_prompt(prompt):
571
572
573
574
575
            raise ValueError("Cannot pass encoder-decoder prompt "
                             "to decoder-only models")

        # Decoder-only operation
        return await self._process_decoder_only_prompt_async(
576
            prompt,
577
578
579
580
581
582
            request_id=request_id,
            lora_request=lora_request,
            prompt_adapter_request=prompt_adapter_request,
        )

    def is_encoder_decoder_model(self):
583
        return self.model_config.is_encoder_decoder