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

4
import itertools
5
from collections.abc import Sequence
6
from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast
7

8
import cloudpickle
9
import torch.nn as nn
10
from pydantic import ValidationError
11
from tqdm.auto import tqdm
12
from typing_extensions import TypeVar
13

14
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
from vllm.beam_search import (
    BeamSearchInstance,
    BeamSearchOutput,
    BeamSearchSequence,
    create_sort_beams_key_function,
)
from vllm.config import (
    CompilationConfig,
    ModelDType,
    StructuredOutputsConfig,
    TokenizerMode,
    is_init_field,
)
from vllm.engine.arg_utils import (
    ConvertOption,
    EngineArgs,
    HfOverrides,
    PoolerConfig,
    RunnerOption,
)
from vllm.entrypoints.chat_utils import (
    ChatCompletionMessageParam,
    ChatTemplateContentFormatOption,
    apply_hf_chat_template,
    apply_mistral_chat_template,
    parse_chat_messages,
    resolve_chat_template_content_format,
)
from vllm.entrypoints.score_utils import (
    ScoreContentPartParam,
    ScoreMultiModalParam,
    _cosine_similarity,
    _validate_score_input_lens,
    compress_token_type_ids,
    get_score_prompt,
)
from vllm.entrypoints.utils import _validate_truncation_size, log_non_default_args
from vllm.inputs import (
    DataPrompt,
    PromptType,
    SingletonPrompt,
    TextPrompt,
    TokensPrompt,
)
58
from vllm.inputs.parse import get_prompt_components
59
from vllm.logger import init_logger
60
from vllm.lora.request import LoRARequest
61
from vllm.model_executor.layers.quantization import QuantizationMethods
62
63
64
65
66
67
68
from vllm.outputs import (
    ClassificationRequestOutput,
    EmbeddingRequestOutput,
    PoolingRequestOutput,
    RequestOutput,
    ScoringRequestOutput,
)
69
from vllm.plugins.io_processors import get_io_processor
70
from vllm.pooling_params import PoolingParams
71
from vllm.sampling_params import BeamSearchParams, RequestOutputKind, SamplingParams
72
from vllm.tasks import PoolingTask
73
74
75
76
77
from vllm.transformers_utils.tokenizer import (
    AnyTokenizer,
    MistralTokenizer,
    get_cached_tokenizer,
)
yhu422's avatar
yhu422 committed
78
from vllm.usage.usage_lib import UsageContext
79
from vllm.utils import Counter, Device, as_iter, is_list_of
80
from vllm.v1.engine import EngineCoreRequest
81
from vllm.v1.engine.llm_engine import LLMEngine
82
from vllm.v1.engine.processor import Processor
83
from vllm.v1.sample.logits_processor import LogitsProcessor
84

85
86
87
if TYPE_CHECKING:
    from vllm.v1.metrics.reader import Metric

88
89
logger = init_logger(__name__)

90
91
_R = TypeVar("_R", default=Any)

92
93

class LLM:
Woosuk Kwon's avatar
Woosuk Kwon committed
94
95
96
97
98
99
100
101
102
103
    """An LLM for generating texts from given prompts and sampling parameters.

    This class includes a tokenizer, a language model (possibly distributed
    across multiple GPUs), and GPU memory space allocated for intermediate
    states (aka KV cache). Given a batch of prompts and sampling parameters,
    this class generates texts from the model, using an intelligent batching
    mechanism and efficient memory management.

    Args:
        model: The name or path of a HuggingFace Transformers model.
104
        tokenizer: The name or path of a HuggingFace Transformers tokenizer.
105
106
        tokenizer_mode: The tokenizer mode. "auto" will use the fast tokenizer
            if available, and "slow" will always use the slow tokenizer.
107
108
109
        skip_tokenizer_init: If true, skip initialization of tokenizer and
            detokenizer. Expect valid prompt_token_ids and None for prompt
            from the input.
110
111
        trust_remote_code: Trust remote code (e.g., from HuggingFace) when
            downloading the model and tokenizer.
112
113
114
115
        allowed_local_media_path: Allowing API requests to read local images
            or videos from directories specified by the server file system.
            This is a security risk. Should only be enabled in trusted
            environments.
116
        allowed_media_domains: If set, only media URLs that belong to this
117
            domain can be used for multi-modal inputs.
Woosuk Kwon's avatar
Woosuk Kwon committed
118
119
120
        tensor_parallel_size: The number of GPUs to use for distributed
            execution with tensor parallelism.
        dtype: The data type for the model weights and activations. Currently,
Woosuk Kwon's avatar
Woosuk Kwon committed
121
122
123
124
            we support `float32`, `float16`, and `bfloat16`. If `auto`, we use
            the `torch_dtype` attribute specified in the model config file.
            However, if the `torch_dtype` in the config is `float32`, we will
            use `float16` instead.
125
        quantization: The method used to quantize the model weights. Currently,
126
            we support "awq", "gptq", and "fp8" (experimental).
127
128
129
130
            If None, we first check the `quantization_config` attribute in the
            model config file. If that is None, we assume the model weights are
            not quantized and use `dtype` to determine the data type of
            the weights.
Jasmond L's avatar
Jasmond L committed
131
132
        revision: The specific model version to use. It can be a branch name,
            a tag name, or a commit id.
133
134
        tokenizer_revision: The specific tokenizer version to use. It can be a
            branch name, a tag name, or a commit id.
135
136
137
138
139
140
        seed: The seed to initialize the random number generator for sampling.
        gpu_memory_utilization: The ratio (between 0 and 1) of GPU memory to
            reserve for the model weights, activations, and KV cache. Higher
            values will increase the KV cache size and thus improve the model's
            throughput. However, if the value is too high, it may cause out-of-
            memory (OOM) errors.
141
142
143
144
145
146
147
148
        kv_cache_memory_bytes: Size of KV Cache per GPU in bytes. By default,
            this is set to None and vllm can automatically infer the kv cache
            size based on gpu_memory_utilization. However, users may want to
            manually specify the kv cache memory size. kv_cache_memory_bytes
            allows more fine-grain control of how much memory gets used when
            compared with using gpu_memory_memory_utilization. Note that
            kv_cache_memory_bytes (when not-None) ignores
            gpu_memory_utilization
149
        swap_space: The size (GiB) of CPU memory per GPU to use as swap space.
150
151
152
153
154
            This can be used for temporarily storing the states of the requests
            when their `best_of` sampling parameters are larger than 1. If all
            requests will have `best_of=1`, you can safely set this to 0.
            Noting that `best_of` is only supported in V0. Otherwise, too small
            values may cause out-of-memory (OOM) errors.
155
156
157
158
        cpu_offload_gb: The size (GiB) of CPU memory to use for offloading
            the model weights. This virtually increases the GPU memory space
            you can use to hold the model weights, at the cost of CPU-GPU data
            transfer for every forward pass.
159
160
161
        enforce_eager: Whether to enforce eager execution. If True, we will
            disable CUDA graph and always execute the model in eager mode.
            If False, we will use CUDA graph and eager execution in hybrid.
162
163
        disable_custom_all_reduce: See
            [ParallelConfig][vllm.config.ParallelConfig].
164
        hf_token: The token to use as HTTP bearer authorization for remote files
165
            . If `True`, will use the token generated when running
166
            `huggingface-cli login` (stored in `~/.huggingface`).
167
168
169
        hf_overrides: If a dictionary, contains arguments to be forwarded to the
            HuggingFace config. If a callable, it is called to update the
            HuggingFace config.
170
171
172
173
174
        mm_processor_kwargs: Arguments to be forwarded to the model's processor
            for multi-modal data, e.g., image processor. Overrides for the
            multi-modal processor obtained from `AutoProcessor.from_pretrained`.
            The available overrides depend on the model that is being run.
            For example, for Phi-3-Vision: `{"num_crops": 4}`.
175
176
177
178
179
        pooler_config: Initialize non-default pooling config for the pooling
            model. e.g. `PoolerConfig(pooling_type="mean", normalize=False)`.
        override_pooler_config: [DEPRECATED] Use `pooler_config` instead. This
            argument is deprecated and will be removed in v0.12.0 or v1.0.0,
            whichever is sooner.
180
181
182
        compilation_config: Either an integer or a dictionary. If it is an
            integer, it is used as the level of compilation optimization. If it
            is a dictionary, it can specify the full compilation configuration.
183
        **kwargs: Arguments for [`EngineArgs`][vllm.EngineArgs].
nunjunj's avatar
nunjunj committed
184

185
186
    Note:
        This class is intended to be used for offline inference. For online
187
        serving, use the [AsyncLLMEngine][vllm.AsyncLLMEngine] class instead.
188
    """
189
190
191
192

    def __init__(
        self,
        model: str,
193
        *,
194
195
        runner: RunnerOption = "auto",
        convert: ConvertOption = "auto",
196
        tokenizer: Optional[str] = None,
197
        tokenizer_mode: TokenizerMode = "auto",
198
        skip_tokenizer_init: bool = False,
199
        trust_remote_code: bool = False,
200
        allowed_local_media_path: str = "",
201
        allowed_media_domains: Optional[list[str]] = None,
202
        tensor_parallel_size: int = 1,
203
204
        dtype: ModelDType = "auto",
        quantization: Optional[QuantizationMethods] = None,
205
        revision: Optional[str] = None,
206
        tokenizer_revision: Optional[str] = None,
207
        seed: Optional[int] = None,
208
        gpu_memory_utilization: float = 0.9,
209
        swap_space: float = 4,
210
        cpu_offload_gb: float = 0,
211
        enforce_eager: bool = False,
212
        disable_custom_all_reduce: bool = False,
213
        hf_token: Optional[Union[bool, str]] = None,
214
        hf_overrides: Optional[HfOverrides] = None,
215
        mm_processor_kwargs: Optional[dict[str, Any]] = None,
216
        pooler_config: Optional[PoolerConfig] = None,
217
        override_pooler_config: Optional[PoolerConfig] = None,
218
219
220
        structured_outputs_config: Optional[
            Union[dict[str, Any], StructuredOutputsConfig]
        ] = None,
221
        kv_cache_memory_bytes: Optional[int] = None,
222
223
224
225
        compilation_config: Optional[
            Union[int, dict[str, Any], CompilationConfig]
        ] = None,
        logits_processors: Optional[list[Union[str, type[LogitsProcessor]]]] = None,
226
        **kwargs: Any,
227
    ) -> None:
228
        """LLM constructor."""
229

230
231
        if "disable_log_stats" not in kwargs:
            kwargs["disable_log_stats"] = True
232

233
234
235
236
237
238
239
        if "worker_cls" in kwargs:
            worker_cls = kwargs["worker_cls"]
            # if the worker_cls is not qualified string name,
            # we serialize it using cloudpickle to avoid pickling issues
            if isinstance(worker_cls, type):
                kwargs["worker_cls"] = cloudpickle.dumps(worker_cls)

240
        if "kv_transfer_config" in kwargs and isinstance(
241
242
            kwargs["kv_transfer_config"], dict
        ):
243
            from vllm.config.kv_transfer import KVTransferConfig
244

245
246
            raw_config_dict = kwargs["kv_transfer_config"]
            try:
247
                kwargs["kv_transfer_config"] = KVTransferConfig(**raw_config_dict)
248
249
250
251
            except ValidationError as e:
                logger.error(
                    "Failed to convert 'kv_transfer_config' dict to "
                    "KVTransferConfig object. Dict: %s. Error: %s",
252
253
254
                    raw_config_dict,
                    e,
                )
255
256
                # Consider re-raising a more specific vLLM error or ValueError
                # to provide better context to the user.
257
                raise ValueError(f"Invalid 'kv_transfer_config' provided: {e}") from e
258

259
260
261
        if hf_overrides is None:
            hf_overrides = {}

262
        if compilation_config is not None:
263
264
            if isinstance(compilation_config, int):
                compilation_config_instance = CompilationConfig(
265
266
                    level=compilation_config
                )
267
268
            elif isinstance(compilation_config, dict):
                compilation_config_instance = CompilationConfig(
269
270
271
272
                    **{
                        k: v
                        for k, v in compilation_config.items()
                        if is_init_field(CompilationConfig, k)
273
274
                    }
                )
275
276
            else:
                compilation_config_instance = compilation_config
277
        else:
278
            compilation_config_instance = CompilationConfig()
279

280
281
282
283
284
285
286
        if structured_outputs_config is not None:
            if isinstance(structured_outputs_config, dict):
                structured_outputs_instance = StructuredOutputsConfig(
                    **{
                        k: v
                        for k, v in structured_outputs_config.items()
                        if is_init_field(StructuredOutputsConfig, k)
287
288
                    }
                )
289
290
291
292
293
            else:
                structured_outputs_instance = structured_outputs_config
        else:
            structured_outputs_instance = StructuredOutputsConfig()

Zhuohan Li's avatar
Zhuohan Li committed
294
        engine_args = EngineArgs(
295
            model=model,
296
297
            runner=runner,
            convert=convert,
298
            tokenizer=tokenizer,
299
            tokenizer_mode=tokenizer_mode,
300
            skip_tokenizer_init=skip_tokenizer_init,
301
            trust_remote_code=trust_remote_code,
302
            allowed_local_media_path=allowed_local_media_path,
303
            allowed_media_domains=allowed_media_domains,
304
305
            tensor_parallel_size=tensor_parallel_size,
            dtype=dtype,
306
            quantization=quantization,
307
            revision=revision,
308
            tokenizer_revision=tokenizer_revision,
309
310
            seed=seed,
            gpu_memory_utilization=gpu_memory_utilization,
311
            kv_cache_memory_bytes=kv_cache_memory_bytes,
312
            swap_space=swap_space,
313
            cpu_offload_gb=cpu_offload_gb,
314
            enforce_eager=enforce_eager,
315
            disable_custom_all_reduce=disable_custom_all_reduce,
316
            hf_token=hf_token,
317
            hf_overrides=hf_overrides,
318
            mm_processor_kwargs=mm_processor_kwargs,
319
            pooler_config=pooler_config,
320
            override_pooler_config=override_pooler_config,
321
            structured_outputs_config=structured_outputs_instance,
322
            compilation_config=compilation_config_instance,
323
            logits_processors=logits_processors,
324
325
            **kwargs,
        )
326

327
328
        log_non_default_args(engine_args)

329
330
        # Create the Engine (autoselects V0 vs V1)
        self.llm_engine = LLMEngine.from_engine_args(
331
332
            engine_args=engine_args, usage_context=UsageContext.LLM_CLASS
        )
333
        self.engine_class = type(self.llm_engine)
334

335
        self.request_counter = Counter()
336
        self.default_sampling_params: Union[dict[str, Any], None] = None
337

338
        supported_tasks = self.llm_engine.get_supported_tasks()  # type: ignore
339
340
341
342
343

        logger.info("Supported_tasks: %s", supported_tasks)

        self.supported_tasks = supported_tasks

344
345
        # Load the Input/Output processor plugin if any
        io_processor_plugin = self.llm_engine.model_config.io_processor_plugin
346
347
348
        self.io_processor = get_io_processor(
            self.llm_engine.vllm_config, io_processor_plugin
        )
349

350
351
352
353
    @property
    def model_config(self):
        return self.llm_engine.model_config

354
355
    def get_tokenizer(self) -> AnyTokenizer:
        return self.llm_engine.get_tokenizer()
356
357

    def set_tokenizer(self, tokenizer: AnyTokenizer) -> None:
358
359
360
361
        # While CachedTokenizer is dynamic, have no choice but
        # compare class name. Misjudgment will arise from
        # user-defined tokenizer started with 'Cached'
        if tokenizer.__class__.__name__.startswith("Cached"):
362
            self.llm_engine.tokenizer = tokenizer
363
        else:
364
            self.llm_engine.tokenizer = get_cached_tokenizer(tokenizer)
365

366
367
368
    def _get_processor(self) -> Processor:
        if not hasattr(self, "_processor"):
            vllm_config = self.llm_engine.vllm_config
369
370
            self._processor = Processor(vllm_config)

371
372
        return self._processor

373
    def get_default_sampling_params(self) -> SamplingParams:
374
375
        if self.default_sampling_params is None:
            self.default_sampling_params = (
376
377
                self.llm_engine.model_config.get_diff_sampling_param()
            )
378
379
        if self.default_sampling_params:
            return SamplingParams.from_optional(**self.default_sampling_params)
380
381
        return SamplingParams()

382
383
384
    def generate(
        self,
        prompts: Union[PromptType, Sequence[PromptType]],
385
386
387
        sampling_params: Optional[
            Union[SamplingParams, Sequence[SamplingParams]]
        ] = None,
388
        *,
389
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
390
391
392
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
        priority: Optional[list[int]] = None,
    ) -> list[RequestOutput]:
Woosuk Kwon's avatar
Woosuk Kwon committed
393
394
        """Generates the completions for the input prompts.

395
        This class automatically batches the given prompts, considering
Woosuk Kwon's avatar
Woosuk Kwon committed
396
397
398
399
        the memory constraint. For the best performance, put all of your prompts
        into a single list and pass it to this method.

        Args:
400
            prompts: The prompts to the LLM. You may pass a sequence of prompts
401
                for batch inference. See [PromptType][vllm.inputs.PromptType]
402
                for more details about the format of each prompt.
Woosuk Kwon's avatar
Woosuk Kwon committed
403
            sampling_params: The sampling parameters for text generation. If
nunjunj's avatar
nunjunj committed
404
405
406
                None, we use the default sampling parameters.
                When it is a single value, it is applied to every prompt.
                When it is a list, the list must have the same length as the
407
                prompts and it is paired one by one with the prompt.
408
409
410
411
            use_tqdm: If `True`, shows a tqdm progress bar.
                If a callable (e.g., `functools.partial(tqdm, leave=False)`),
                it is used to create the progress bar.
                If `False`, no progress bar is created.
412
            lora_request: LoRA request to use for generation, if any.
413
414
            priority: The priority of the requests, if any.
                Only applicable when priority scheduling policy is enabled.
Woosuk Kwon's avatar
Woosuk Kwon committed
415
416

        Returns:
417
            A list of `RequestOutput` objects containing the
418
            generated completions in the same order as the input prompts.
419

420
421
422
423
        Note:
            Using `prompts` and `prompt_token_ids` as keyword parameters is
            considered legacy and may be deprecated in the future. You should
            instead pass them via the `inputs` parameter.
424
        """
425
426
427
        model_config = self.llm_engine.model_config
        runner_type = model_config.runner_type
        if runner_type != "generate":
428
429
430
            raise ValueError(
                "LLM.generate() is only supported for generative models. "
                "Try passing `--runner generate` to use the model as a "
431
432
                "generative model."
            )
433

434
435
        if sampling_params is None:
            # Use default sampling params.
436
            sampling_params = self.get_default_sampling_params()
437

438
        # Add any modality specific loras to the corresponding prompts
439
        lora_request = self._get_modality_specific_lora_reqs(prompts, lora_request)
440

441
        self._validate_and_add_requests(
442
            prompts=prompts,
443
            params=sampling_params,
444
            use_tqdm=use_tqdm,
445
            lora_request=lora_request,
446
447
            priority=priority,
        )
448

449
        outputs = self._run_engine(use_tqdm=use_tqdm)
Joe Runde's avatar
Joe Runde committed
450
        return self.engine_class.validate_outputs(outputs, RequestOutput)
451

452
    def _get_modality_specific_lora_reqs(
453
454
455
456
        self,
        prompts: Union[PromptType, Sequence[PromptType]],
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]],
    ):
457
458
459
460
461
462
        # Grab the lora config off the vllm config on the engine,
        # since this is the same for both v0 & v1.
        lora_config = self.llm_engine.vllm_config.lora_config

        # If there's no lora config / default_mm_loras, or the model
        # isn't multimodal, leave the lora as is.
463
464
465
466
467
        if (
            lora_config is None
            or not self.llm_engine.model_config.is_multimodal_model
            or (lora_config and lora_config.default_mm_loras is None)
        ):
468
469
            return lora_request

470
471
        if not isinstance(prompts, Sequence):
            prompts = [prompts]
472

473
474
475
476
477
        optional_loras = (
            [lora_request] * len(prompts)
            if not isinstance(lora_request, Sequence)
            else lora_request
        )
478
479
480

        return [
            self._resolve_single_prompt_mm_lora(
481
                prompt,
482
483
                opt_lora_req,
                lora_config.default_mm_loras,
484
485
            )
            for prompt, opt_lora_req in zip(prompts, optional_loras)
486
487
        ]

488
489
490
491
492
493
494
495
496
497
498
    def _resolve_single_prompt_mm_lora(
        self,
        prompt: PromptType,
        lora_request: Optional[LoRARequest],
        default_mm_loras: Optional[dict[str, str]],
    ):
        if (
            not default_mm_loras
            or not isinstance(prompt, dict)
            or "multi_modal_data" not in prompt
        ):
499
500
            return lora_request

501
        prompt = cast(Union[TextPrompt, TokensPrompt], prompt)
502

503
504
505
        intersection = set(prompt["multi_modal_data"].keys()).intersection(
            default_mm_loras.keys()
        )
506
507
508
509
510
511
512
513
514
        if not intersection:
            return lora_request
        if len(intersection) > 1:
            # TODO: Would be nice to be able to have multiple loras per prompt
            logger.warning(
                "Multiple modality specific loras were registered and would be"
                " used by a single prompt consuming several modalities; "
                " currently we only support one lora per request; as such,"
                " lora(s) registered with modalities: %s"
515
516
517
                " will be skipped",
                intersection,
            )
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
            return lora_request

        # Build the LoRA request; the ID of the default mm lora is the
        # index of the modality name sorted alphabetically + 1.
        modality_name = intersection.pop()
        modality_lora_path = default_mm_loras[modality_name]
        modality_lora_id = sorted(default_mm_loras).index(modality_name) + 1

        # If we have a collision, warn if there is a collision,
        # but always send the explicitly provided request.
        if lora_request:
            if lora_request.lora_int_id != modality_lora_id:
                logger.warning(
                    "A modality with a registered lora and a lora_request "
                    "with a different ID were provided; falling back to the "
533
534
                    "lora_request as we only apply one LoRARequest per prompt"
                )
535
536
537
538
539
540
541
542
            return lora_request

        return LoRARequest(
            modality_name,
            modality_lora_id,
            modality_lora_path,
        )

543
544
545
546
547
548
549
    def collective_rpc(
        self,
        method: Union[str, Callable[..., _R]],
        timeout: Optional[float] = None,
        args: tuple = (),
        kwargs: Optional[dict[str, Any]] = None,
    ) -> list[_R]:
550
551
552
553
554
555
556
557
558
559
560
        """
        Execute an RPC call on all workers.

        Args:
            method: Name of the worker method to execute, or a callable that
                is serialized and sent to all workers to execute.

                If the method is a callable, it should accept an additional
                `self` argument, in addition to the arguments passed in `args`
                and `kwargs`. The `self` argument will be the worker object.
            timeout: Maximum time in seconds to wait for execution. Raises a
561
                [`TimeoutError`][] on timeout. `None` means wait indefinitely.
562
563
564
565
566
            args: Positional arguments to pass to the worker method.
            kwargs: Keyword arguments to pass to the worker method.

        Returns:
            A list containing the results from each worker.
567

568
569
570
        Note:
            It is recommended to use this API to only pass control messages,
            and set up data-plane communication to pass data.
571
        """
572
573

        return self.llm_engine.collective_rpc(method, timeout, args, kwargs)
574
575

    def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
576
        """
577
578
        Run a function directly on the model inside each worker,
        returning the result for each of them.
579
580
581
582
583
584

        !!! warning
            To reduce the overhead of data transfer, avoid returning large
            arrays or tensors from this method. If you must return them,
            make sure you move them to CPU first to avoid taking up additional
            VRAM!
585
        """
586
        return self.llm_engine.apply_model(func)
587

588
589
590
591
592
593
    def _get_beam_search_lora_requests(
        self,
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]],
        prompts: list[Union[TokensPrompt, TextPrompt]],
    ) -> list[Optional[LoRARequest]]:
        """Get the optional lora request corresponding to each prompt."""
594
        if isinstance(lora_request, Sequence) and len(lora_request) != len(prompts):
595
            raise ValueError(
596
597
                "Lora request list should be the same length as the prompts"
            )
598
599
600
601
602
603

        if lora_request is None or isinstance(lora_request, LoRARequest):
            return [lora_request] * len(prompts)

        raise TypeError(f"Invalid lora_request type {type(lora_request)}")

604
605
    def beam_search(
        self,
606
        prompts: list[Union[TokensPrompt, TextPrompt]],
607
        params: BeamSearchParams,
608
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
609
        use_tqdm: bool = False,
610
        concurrency_limit: Optional[int] = None,
611
    ) -> list[BeamSearchOutput]:
612
613
614
615
616
617
        """
        Generate sequences using beam search.

        Args:
            prompts: A list of prompts. Each prompt can be a string or a list
                of token IDs.
618
            params: The beam search parameters.
619
            lora_request: LoRA request to use for generation, if any.
620
            use_tqdm: Whether to use tqdm to display the progress bar.
621
622
            concurrency_limit: The maximum number of concurrent requests.
                If None, the number of concurrent requests is unlimited.
623
        """
624
625
        # TODO: how does beam search work together with length penalty,
        # frequency, penalty, and stopping criteria, etc.?
626
627
628
629
        beam_width = params.beam_width
        max_tokens = params.max_tokens
        temperature = params.temperature
        ignore_eos = params.ignore_eos
630
631
        length_penalty = params.length_penalty

632
        lora_requests = self._get_beam_search_lora_requests(lora_request, prompts)
633

634
635
636
637
638
        tokenizer = self.get_tokenizer()
        sort_beams_key = create_sort_beams_key_function(
            tokenizer.eos_token_id,
            length_penalty,
        )
639

640
641
642
        if use_tqdm and concurrency_limit is not None:
            logger.warning(
                "Progress bar is not supported when using concurrency_limit. "
643
644
                "Disabling progress bar."
            )
645
646
647
648
649
            use_tqdm = False

        if concurrency_limit is None:
            concurrency_limit = len(prompts)

650
651
        def create_tokens_prompt_from_beam(beam: BeamSearchSequence) -> TokensPrompt:
            token_prompt_kwargs: TokensPrompt = {"prompt_token_ids": beam.tokens}
652
653
654
655
            if beam.multi_modal_data is not None:
                token_prompt_kwargs["multi_modal_data"] = beam.multi_modal_data

            if beam.mm_processor_kwargs is not None:
656
                token_prompt_kwargs["mm_processor_kwargs"] = beam.mm_processor_kwargs
657
            return TokensPrompt(**token_prompt_kwargs)
658

659
660
661
        # generate 2 * beam_width candidates at each step
        # following the huggingface transformers implementation
        # at https://github.com/huggingface/transformers/blob/e15687fffe5c9d20598a19aeab721ae0a7580f8a/src/transformers/generation/beam_search.py#L534 # noqa
662
663
664
        beam_search_params = SamplingParams(
            logprobs=2 * beam_width, max_tokens=1, temperature=temperature
        )
665
        instances: list[BeamSearchInstance] = []
666

667
        for lora_req, prompt in zip(lora_requests, prompts):
668
669
670
671
672
            # Add multimodal processor kwargs & data
            mm_kwargs = {}
            if "multi_modal_data" in prompt:
                mm_kwargs["multi_modal_data"] = prompt["multi_modal_data"]
            if "mm_processor_kwargs" in prompt:
673
                mm_kwargs["mm_processor_kwargs"] = prompt["mm_processor_kwargs"]
674

675
676
            if "prompt_token_ids" in prompt:
                prompt = cast(TokensPrompt, prompt)  # Needed for mypy
677
678
679
                prompt_tokens = prompt["prompt_token_ids"]
            else:
                prompt_tokens = tokenizer.encode(prompt["prompt"])
680

681
            instances.append(
682
683
684
685
686
                BeamSearchInstance(
                    prompt_tokens,
                    lora_request=lora_req,
                    logprobs=None,
                    **mm_kwargs,
687
688
                ),
            )
689

690
        for prompt_start in range(0, len(prompts), concurrency_limit):
691
            instances_batch = instances[prompt_start : prompt_start + concurrency_limit]
692
693
694

            token_iter = range(max_tokens)
            if use_tqdm:
695
696
697
                token_iter = tqdm(
                    token_iter, desc="Beam search", unit="token", unit_scale=False
                )
698
699
700
                logger.warning(
                    "The progress bar shows the upper bound on token steps and "
                    "may finish early due to stopping conditions. It does not "
701
702
                    "reflect instance-level progress."
                )
703
704
            for _ in token_iter:
                all_beams: list[BeamSearchSequence] = list(
705
706
                    sum((instance.beams for instance in instances_batch), [])
                )
707
708
                pos = [0] + list(
                    itertools.accumulate(
709
710
711
                        len(instance.beams) for instance in instances_batch
                    )
                )
712
                instance_start_and_end: list[tuple[int, int]] = list(
713
714
                    zip(pos[:-1], pos[1:])
                )
715
716
717
718
719
720

                if len(all_beams) == 0:
                    break

                # create corresponding batch entries for prompt & optional lora
                prompts_batch, lora_req_batch = zip(
721
722
723
724
725
                    *[
                        (create_tokens_prompt_from_beam(beam), beam.lora_request)
                        for beam in all_beams
                    ]
                )
726
727
728

                # only runs for one step
                # we don't need to use tqdm here
729
730
731
732
733
734
                output = self.generate(
                    prompts_batch,
                    sampling_params=beam_search_params,
                    use_tqdm=False,
                    lora_request=lora_req_batch,
                )
735

736
737
738
                for (start, end), instance in zip(
                    instance_start_and_end, instances_batch
                ):
739
740
741
742
743
744
745
746
747
748
749
750
751
752
                    instance_new_beams = []
                    for i in range(start, end):
                        current_beam = all_beams[i]
                        result = output[i]

                        if result.outputs[0].logprobs is not None:
                            # if `result.outputs[0].logprobs` is None, it means
                            # the sequence is completed because of the
                            # max-model-len or abortion. we don't need to add
                            # it to the new beams.
                            logprobs = result.outputs[0].logprobs[0]
                            for token_id, logprob_obj in logprobs.items():
                                new_beam = BeamSearchSequence(
                                    tokens=current_beam.tokens + [token_id],
753
                                    logprobs=current_beam.logprobs + [logprobs],
754
                                    lora_request=current_beam.lora_request,
755
756
757
758
759
760
761
762
763
764
                                    cum_logprob=current_beam.cum_logprob
                                    + logprob_obj.logprob,
                                    multi_modal_data=current_beam.multi_modal_data,
                                    mm_processor_kwargs=current_beam.mm_processor_kwargs,
                                )

                                if (
                                    token_id == tokenizer.eos_token_id
                                    and not ignore_eos
                                ):
765
766
767
                                    instance.completed.append(new_beam)
                                else:
                                    instance_new_beams.append(new_beam)
768
769
770
                    sorted_beams = sorted(
                        instance_new_beams, key=sort_beams_key, reverse=True
                    )
771
                    instance.beams = sorted_beams[:beam_width]
772
773
774
775

        outputs = []
        for instance in instances:
            instance.completed.extend(instance.beams)
776
777
778
            sorted_completed = sorted(
                instance.completed, key=sort_beams_key, reverse=True
            )
779
780
781
782
783
784
785
786
            best_beams = sorted_completed[:beam_width]

            for beam in best_beams:
                beam.text = tokenizer.decode(beam.tokens)
            outputs.append(BeamSearchOutput(sequences=best_beams))

        return outputs

787
    def preprocess_chat(
nunjunj's avatar
nunjunj committed
788
        self,
789
790
791
        messages: Union[
            list[ChatCompletionMessageParam], list[list[ChatCompletionMessageParam]]
        ],
nunjunj's avatar
nunjunj committed
792
        chat_template: Optional[str] = None,
793
        chat_template_content_format: ChatTemplateContentFormatOption = "auto",
794
        add_generation_prompt: bool = True,
795
        continue_final_message: bool = False,
796
        tools: Optional[list[dict[str, Any]]] = None,
797
        chat_template_kwargs: Optional[dict[str, Any]] = None,
798
        mm_processor_kwargs: Optional[dict[str, Any]] = None,
799
    ) -> list[TokensPrompt]:
nunjunj's avatar
nunjunj committed
800
        """
801
802
        Generate prompt for a chat conversation. The pre-processed
        prompt can then be used as input for the other LLM methods.
nunjunj's avatar
nunjunj committed
803

804
        Refer to `chat` for a complete description of the arguments.
nunjunj's avatar
nunjunj committed
805
        Returns:
806
807
808
            A list of `TokensPrompts` objects containing the tokenized
            prompt after chat template interpolation, and the
            pre-processed multi-modal inputs.
nunjunj's avatar
nunjunj committed
809
        """
810
        list_of_messages: list[list[ChatCompletionMessageParam]]
nunjunj's avatar
nunjunj committed
811

812
813
        # Handle multi and single conversations
        if is_list_of(messages, list):
814
            # messages is list[list[...]]
815
            list_of_messages = cast(list[list[ChatCompletionMessageParam]], messages)
816
        else:
817
            # messages is list[...]
818
            list_of_messages = [cast(list[ChatCompletionMessageParam], messages)]
819

820
        tokenizer = self.get_tokenizer()
821
822
823
        model_config = self.llm_engine.get_model_config()
        resolved_content_format = resolve_chat_template_content_format(
            chat_template,
824
            tools,
825
826
            chat_template_content_format,
            tokenizer,
827
            model_config=model_config,
828
829
        )

830
831
832
833
834
835
836
837
        _chat_template_kwargs: dict[str, Any] = dict(
            chat_template=chat_template,
            add_generation_prompt=add_generation_prompt,
            continue_final_message=continue_final_message,
            tools=tools,
        )
        _chat_template_kwargs.update(chat_template_kwargs or {})

838
        prompts: list[TokensPrompt] = []
839
840

        for msgs in list_of_messages:
841
842
843
            # NOTE: _parse_chat_message_content_parts() currently doesn't
            # handle mm_processor_kwargs, since there is no implementation in
            # the chat message parsing for it.
844
            conversation, mm_data, mm_uuids = parse_chat_messages(
845
846
847
848
849
                msgs,
                model_config,
                tokenizer,
                content_format=resolved_content_format,
            )
850
851

            if isinstance(tokenizer, MistralTokenizer):
852
                prompt_token_ids = apply_mistral_chat_template(
853
854
                    tokenizer,
                    messages=msgs,
855
                    **_chat_template_kwargs,
856
857
                )
            else:
858
                prompt_str = apply_hf_chat_template(
859
                    tokenizer=tokenizer,
860
                    conversation=conversation,
861
                    model_config=model_config,
862
                    **_chat_template_kwargs,
863
                )
864
865
                # Special tokens are already included in chat templates so
                # should not be added by the tokenizer in this case.
866
867
868
                prompt_token_ids = tokenizer.encode(
                    prompt_str, add_special_tokens=False
                )
869

870
            prompt = TokensPrompt(prompt_token_ids=prompt_token_ids)
871
872
873
874

            if mm_data is not None:
                prompt["multi_modal_data"] = mm_data

875
876
877
            if mm_uuids is not None:
                prompt["multi_modal_uuids"] = mm_uuids

878
879
880
            if mm_processor_kwargs is not None:
                prompt["mm_processor_kwargs"] = mm_processor_kwargs

881
            prompts.append(prompt)
882

883
884
885
886
        return prompts

    def chat(
        self,
887
888
889
890
        messages: Union[
            list[ChatCompletionMessageParam], list[list[ChatCompletionMessageParam]]
        ],
        sampling_params: Optional[Union[SamplingParams, list[SamplingParams]]] = None,
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
        lora_request: Optional[LoRARequest] = None,
        chat_template: Optional[str] = None,
        chat_template_content_format: ChatTemplateContentFormatOption = "auto",
        add_generation_prompt: bool = True,
        continue_final_message: bool = False,
        tools: Optional[list[dict[str, Any]]] = None,
        chat_template_kwargs: Optional[dict[str, Any]] = None,
        mm_processor_kwargs: Optional[dict[str, Any]] = None,
    ) -> list[RequestOutput]:
        """
        Generate responses for a chat conversation.

        The chat conversation is converted into a text prompt using the
        tokenizer and calls the [generate][vllm.LLM.generate] method to generate
        the responses.

        Multi-modal inputs can be passed in the same way you would pass them
        to the OpenAI API.

        Args:
            messages: A list of conversations or a single conversation.

                - Each conversation is represented as a list of messages.
                - Each message is a dictionary with 'role' and 'content' keys.

            sampling_params: The sampling parameters for text generation.
                If None, we use the default sampling parameters. When it
                is a single value, it is applied to every prompt. When it
                is a list, the list must have the same length as the
                prompts and it is paired one by one with the prompt.
            use_tqdm: If `True`, shows a tqdm progress bar.
                If a callable (e.g., `functools.partial(tqdm, leave=False)`),
                it is used to create the progress bar.
                If `False`, no progress bar is created.
            lora_request: LoRA request to use for generation, if any.
            chat_template: The template to use for structuring the chat.
                If not provided, the model's default chat template will be used.
            chat_template_content_format: The format to render message content.

                - "string" will render the content as a string.
                  Example: `"Who are you?"`
                - "openai" will render the content as a list of dictionaries,
                  similar to OpenAI schema.
                  Example: `[{"type": "text", "text": "Who are you?"}]`

            add_generation_prompt: If True, adds a generation template
                to each message.
            continue_final_message: If True, continues the final message in
                the conversation instead of starting a new one. Cannot be
                `True` if `add_generation_prompt` is also `True`.
            chat_template_kwargs: Additional kwargs to pass to the chat
                template.
            mm_processor_kwargs: Multimodal processor kwarg overrides for this
                chat request. Only used for offline requests.

        Returns:
            A list of `RequestOutput` objects containing the generated
            responses in the same order as the input messages.
        """

        prompts = self.preprocess_chat(
            messages=messages,
            chat_template=chat_template,
            chat_template_content_format=chat_template_content_format,
            add_generation_prompt=add_generation_prompt,
            continue_final_message=continue_final_message,
            tools=tools,
            chat_template_kwargs=chat_template_kwargs,
            mm_processor_kwargs=mm_processor_kwargs,
        )

nunjunj's avatar
nunjunj committed
963
        return self.generate(
964
            prompts,
965
            sampling_params=sampling_params,
nunjunj's avatar
nunjunj committed
966
967
968
969
            use_tqdm=use_tqdm,
            lora_request=lora_request,
        )

970
971
    def encode(
        self,
972
        prompts: Union[PromptType, Sequence[PromptType], DataPrompt],
973
        pooling_params: Optional[Union[PoolingParams, Sequence[PoolingParams]]] = None,
974
        *,
975
        truncate_prompt_tokens: Optional[int] = None,
976
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
977
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
978
        pooling_task: PoolingTask = "encode",
979
        tokenization_kwargs: Optional[dict[str, Any]] = None,
980
    ) -> list[PoolingRequestOutput]:
981
982
        """Apply pooling to the hidden states corresponding to the input
        prompts.
983

984
        This class automatically batches the given prompts, considering
985
986
987
988
        the memory constraint. For the best performance, put all of your prompts
        into a single list and pass it to this method.

        Args:
989
            prompts: The prompts to the LLM. You may pass a sequence of prompts
990
                for batch inference. See [PromptType][vllm.inputs.PromptType]
991
                for more details about the format of each prompt.
992
993
            pooling_params: The pooling parameters for pooling. If None, we
                use the default pooling parameters.
994
995
996
997
            use_tqdm: If `True`, shows a tqdm progress bar.
                If a callable (e.g., `functools.partial(tqdm, leave=False)`),
                it is used to create the progress bar.
                If `False`, no progress bar is created.
998
            lora_request: LoRA request to use for generation, if any.
999
            pooling_task: Override the pooling task to use.
1000
1001
            tokenization_kwargs: overrides tokenization_kwargs set in
                pooling_params
1002
1003

        Returns:
1004
            A list of `PoolingRequestOutput` objects containing the
1005
            pooled hidden states in the same order as the input prompts.
1006

1007
1008
1009
1010
        Note:
            Using `prompts` and `prompt_token_ids` as keyword parameters is
            considered legacy and may be deprecated in the future. You should
            instead pass them via the `inputs` parameter.
1011
        """
1012
1013
1014
1015

        if self.supported_tasks == ["encode"] and pooling_task is None:
            pooling_task = "encode"

1016
        if pooling_task is None:
1017
            pooling_task = "embed" if "embed" in self.supported_tasks else "encode"
1018
1019
1020
1021
1022
1023

            logger.warning_once(
                "`LLM.encode` is currently using `pooling_task = %s`.\n"
                "Please use one of the more specific methods or set the "
                "task directly when using `LLM.encode`:\n"
                "  - For embeddings, use `LLM.embed(...)` "
1024
                'or `pooling_task="embed"`.\n'
1025
                "  - For classification logits, use `LLM.classify(...)` "
1026
                'or `pooling_task="classify"`.\n'
1027
                "  - For rewards, use `LLM.reward(...)` "
1028
                'or `pooling_task="reward"`\n'
1029
                "  - For similarity scores, use `LLM.score(...)`.",
1030
1031
                pooling_task,
            )
1032

1033
1034
        model_config = self.llm_engine.model_config
        runner_type = model_config.runner_type
1035
        if runner_type != "pooling":
1036
1037
1038
            raise ValueError(
                "LLM.encode() is only supported for pooling models. "
                "Try passing `--runner pooling` to use the model as a "
1039
1040
                "pooling model."
            )
1041

1042
        if pooling_task not in self.supported_tasks:
1043
            raise ValueError(f"pooling_task must be one of {self.supported_tasks}.")
1044

1045
1046
1047
        if pooling_params is None:
            # Use default pooling params.
            pooling_params = PoolingParams()
1048

1049
1050
1051
1052
1053
        for param in as_iter(pooling_params):
            param.verify(pooling_task, model_config)
            # for backwards compatibility
            if truncate_prompt_tokens is not None:
                param.truncate_prompt_tokens = truncate_prompt_tokens
1054

1055
1056
1057
1058
1059
1060
1061
1062
        io_processor_prompt = False
        if isinstance(prompts, dict) and "data" in prompts:
            io_processor_prompt = True
            if self.io_processor is None:
                raise ValueError(
                    "No IOProcessor plugin installed. Please refer "
                    "to the documentation and to the "
                    "'prithvi_geospatial_mae_io_processor' "
1063
1064
                    "offline inference example for more details."
                )
1065
1066
1067
1068
1069
1070
1071

            # Validate the request data is valid for the loaded plugin
            validated_prompt = self.io_processor.parse_request(prompts)

            # obtain the actual model prompts from the pre-processor
            prompts = self.io_processor.pre_process(prompt=validated_prompt)

1072
        self._validate_and_add_requests(
1073
            prompts=prompts,
1074
            params=pooling_params,
1075
            use_tqdm=use_tqdm,
1076
            lora_request=lora_request,
1077
1078
        )

1079
        outputs = self._run_engine(use_tqdm=use_tqdm)
1080
1081

        model_outputs = self.engine_class.validate_outputs(
1082
1083
            outputs, PoolingRequestOutput
        )
1084
1085
1086
1087
1088

        if io_processor_prompt:
            # get the post-processed model outputs
            assert self.io_processor is not None
            processed_outputs = self.io_processor.post_process(
1089
1090
                model_output=model_outputs
            )
1091
1092

            return [
1093
1094
1095
1096
1097
1098
                PoolingRequestOutput[Any](
                    request_id="",
                    outputs=processed_outputs,
                    prompt_token_ids=[],
                    finished=True,
                )
1099
1100
1101
            ]
        else:
            return model_outputs
1102

1103
1104
1105
1106
    def embed(
        self,
        prompts: Union[PromptType, Sequence[PromptType]],
        *,
1107
        truncate_prompt_tokens: Optional[int] = None,
1108
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
1109
        pooling_params: Optional[Union[PoolingParams, Sequence[PoolingParams]]] = None,
1110
1111
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
    ) -> list[EmbeddingRequestOutput]:
1112
1113
1114
1115
1116
1117
1118
1119
1120
        """
        Generate an embedding vector for each prompt.

        This class automatically batches the given prompts, considering
        the memory constraint. For the best performance, put all of your prompts
        into a single list and pass it to this method.

        Args:
            prompts: The prompts to the LLM. You may pass a sequence of prompts
1121
                for batch inference. See [PromptType][vllm.inputs.PromptType]
1122
                for more details about the format of each prompt.
1123
1124
            pooling_params: The pooling parameters for pooling. If None, we
                use the default pooling parameters.
1125
1126
1127
1128
            use_tqdm: If `True`, shows a tqdm progress bar.
                If a callable (e.g., `functools.partial(tqdm, leave=False)`),
                it is used to create the progress bar.
                If `False`, no progress bar is created.
1129
1130
1131
            lora_request: LoRA request to use for generation, if any.

        Returns:
1132
            A list of `EmbeddingRequestOutput` objects containing the
1133
1134
            embedding vectors in the same order as the input prompts.
        """
1135
        if "embed" not in self.supported_tasks:
1136
1137
            raise ValueError(
                "Embedding API is not supported by this model. "
1138
1139
                "Try converting the model using `--convert embed`."
            )
1140

1141
1142
1143
1144
1145
1146
1147
1148
        items = self.encode(
            prompts,
            truncate_prompt_tokens=truncate_prompt_tokens,
            use_tqdm=use_tqdm,
            pooling_params=pooling_params,
            lora_request=lora_request,
            pooling_task="embed",
        )
1149
1150
1151
1152
1153
1154
1155

        return [EmbeddingRequestOutput.from_base(item) for item in items]

    def classify(
        self,
        prompts: Union[PromptType, Sequence[PromptType]],
        *,
1156
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
1157
        pooling_params: Optional[Union[PoolingParams, Sequence[PoolingParams]]] = None,
1158
1159
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
    ) -> list[ClassificationRequestOutput]:
1160
1161
1162
1163
1164
1165
1166
1167
1168
        """
        Generate class logits for each prompt.

        This class automatically batches the given prompts, considering
        the memory constraint. For the best performance, put all of your prompts
        into a single list and pass it to this method.

        Args:
            prompts: The prompts to the LLM. You may pass a sequence of prompts
1169
                for batch inference. See [PromptType][vllm.inputs.PromptType]
1170
                for more details about the format of each prompt.
1171
1172
1173
1174
            use_tqdm: If `True`, shows a tqdm progress bar.
                If a callable (e.g., `functools.partial(tqdm, leave=False)`),
                it is used to create the progress bar.
                If `False`, no progress bar is created.
1175
            lora_request: LoRA request to use for generation, if any.
1176
1177
            pooling_params: The pooling parameters for pooling. If None, we
                use the default pooling parameters.
1178
        Returns:
1179
            A list of `ClassificationRequestOutput` objects containing the
1180
1181
            embedding vectors in the same order as the input prompts.
        """
1182
        if "classify" not in self.supported_tasks:
1183
            raise ValueError(
1184
                "Classification API is not supported by this model. "
1185
1186
                "Try converting the model using `--convert classify`."
            )
1187

1188
1189
1190
        items = self.encode(
            prompts,
            use_tqdm=use_tqdm,
1191
            pooling_params=pooling_params,
1192
1193
1194
            lora_request=lora_request,
            pooling_task="classify",
        )
1195
1196
1197

        return [ClassificationRequestOutput.from_base(item) for item in items]

1198
1199
1200
1201
1202
1203
1204
    def reward(
        self,
        prompts: Union[PromptType, Sequence[PromptType]],
        /,
        *,
        truncate_prompt_tokens: Optional[int] = None,
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
1205
        pooling_params: Optional[Union[PoolingParams, Sequence[PoolingParams]]] = None,
1206
1207
1208
1209
1210
1211
1212
1213
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
    ) -> list[PoolingRequestOutput]:
        """
        Generate rewards for each prompt.

        Args:
            prompts: The prompts to the LLM. You may pass a sequence of prompts
                for batch inference. See [PromptType][vllm.inputs.PromptType]
1214
                for more details about the format of each prompt.
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
            use_tqdm: If `True`, shows a tqdm progress bar.
                If a callable (e.g., `functools.partial(tqdm, leave=False)`),
                it is used to create the progress bar.
                If `False`, no progress bar is created.
            lora_request: LoRA request to use for generation, if any.
            pooling_params: The pooling parameters for pooling. If None, we
                use the default pooling parameters.
        Returns:
            A list of `PoolingRequestOutput` objects containing the
            pooled hidden states in the same order as the input prompts.
        """

        return self.encode(
            prompts,
            use_tqdm=use_tqdm,
            lora_request=lora_request,
            pooling_params=pooling_params,
            truncate_prompt_tokens=truncate_prompt_tokens,
            pooling_task="encode",
        )

1236
1237
1238
    def _embedding_score(
        self,
        tokenizer: AnyTokenizer,
1239
1240
        text_1: list[Union[str, TextPrompt, TokensPrompt]],
        text_2: list[Union[str, TextPrompt, TokensPrompt]],
1241
        truncate_prompt_tokens: Optional[int] = None,
1242
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
1243
        pooling_params: Optional[PoolingParams] = None,
1244
1245
1246
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
    ) -> list[ScoringRequestOutput]:
        encoded_output: list[PoolingRequestOutput] = self.encode(
1247
            text_1 + text_2,
1248
            truncate_prompt_tokens=truncate_prompt_tokens,
1249
1250
            use_tqdm=use_tqdm,
            lora_request=lora_request,
1251
            pooling_params=pooling_params,
1252
1253
            pooling_task="embed",
        )
1254

1255
1256
        encoded_output_1: list[PoolingRequestOutput] = encoded_output[0 : len(text_1)]
        encoded_output_2: list[PoolingRequestOutput] = encoded_output[len(text_1) :]
1257
1258
1259
1260

        if len(encoded_output_1) == 1:
            encoded_output_1 = encoded_output_1 * len(encoded_output_2)

1261
1262
1263
        scores = _cosine_similarity(
            tokenizer=tokenizer, embed_1=encoded_output_1, embed_2=encoded_output_2
        )
1264

1265
        items = self.engine_class.validate_outputs(scores, PoolingRequestOutput)
1266
1267
1268
1269
        return [ScoringRequestOutput.from_base(item) for item in items]

    def _cross_encoding_score(
        self,
1270
        tokenizer: AnyTokenizer,
1271
1272
        data_1: Union[list[str], list[ScoreContentPartParam]],
        data_2: Union[list[str], list[ScoreContentPartParam]],
1273
        truncate_prompt_tokens: Optional[int] = None,
1274
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
1275
        pooling_params: Optional[PoolingParams] = None,
1276
1277
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
    ) -> list[ScoringRequestOutput]:
1278
        model_config = self.llm_engine.model_config
1279
1280

        if isinstance(tokenizer, MistralTokenizer):
1281
            raise ValueError("Score API is not supported for Mistral tokenizer")
1282

1283
1284
        if len(data_1) == 1:
            data_1 = data_1 * len(data_2)
1285

1286
1287
1288
1289
1290
        if pooling_params is None:
            pooling_params = PoolingParams(task="score")

        model_config = self.llm_engine.model_config
        pooling_params.verify("score", model_config)
1291
        pooling_params_list = list[PoolingParams]()
1292

1293
        tokenization_kwargs: dict[str, Any] = {}
1294

1295
1296
1297
        _validate_truncation_size(
            model_config.max_model_len, truncate_prompt_tokens, tokenization_kwargs
        )
1298

1299
        prompts = list[PromptType]()
1300

1301
1302
        input_pairs = [(t1, t2) for t1, t2 in zip(data_1, data_2)]

1303
        model_config = self.llm_engine.model_config
1304

1305
1306
1307
1308
1309
1310
1311
1312
1313
        for q, d in input_pairs:
            _, engine_prompt = get_score_prompt(
                model_config=model_config,
                data_1=q,
                data_2=d,
                tokenizer=tokenizer,
                tokenization_kwargs=tokenization_kwargs,
            )

1314
            if token_type_ids := engine_prompt.pop("token_type_ids", None):
1315
1316
1317
1318
1319
1320
1321
                params = pooling_params.clone()
                compressed = compress_token_type_ids(token_type_ids)
                params.extra_kwargs = {"compressed_token_type_ids": compressed}
                pooling_params_list.append(params)
            else:
                pooling_params_list.append(pooling_params)

1322
            prompts.append(engine_prompt)
1323
1324

        self._validate_and_add_requests(
1325
            prompts=prompts,
1326
            params=pooling_params_list,
1327
            use_tqdm=use_tqdm,
1328
1329
1330
1331
            lora_request=lora_request,
        )

        outputs = self._run_engine(use_tqdm=use_tqdm)
1332
        items = self.engine_class.validate_outputs(outputs, PoolingRequestOutput)
1333
1334
1335

        return [ScoringRequestOutput.from_base(item) for item in items]

1336
1337
    def score(
        self,
1338
1339
        data_1: Union[SingletonPrompt, Sequence[SingletonPrompt], ScoreMultiModalParam],
        data_2: Union[SingletonPrompt, Sequence[SingletonPrompt], ScoreMultiModalParam],
1340
        /,
1341
        *,
1342
        truncate_prompt_tokens: Optional[int] = None,
1343
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
1344
        pooling_params: Optional[PoolingParams] = None,
1345
1346
        lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
    ) -> list[ScoringRequestOutput]:
1347
1348
        """Generate similarity scores for all pairs `<text,text_pair>` or
          `<multi-modal data, multi-modal data pair>`.
1349

1350
        The inputs can be `1 -> 1`, `1 -> N` or `N -> N`.
1351
1352
        In the `1 - N` case the `data_1` input will be replicated `N`
        times to pair with the `data_2` inputs.
1353
        The input pairs are used to build a list of prompts for the
1354
1355
        cross encoder model. This class automatically batches the prompts,
        considering the memory constraint. For the best performance, put all
1356
1357
1358
        of your inputs into a single list and pass it to this method.

        Supports both text and multi-modal data (images, etc.) when used with
1359
        appropriate multi-modal models. For multi-modal inputs, ensure the
1360
        prompt structure matches the model's expected input format.
1361
1362

        Args:
1363
1364
1365
            data_1: Can be a single prompt, a list of prompts or
                `ScoreMultiModalParam`, which can contain either text or
                multi-modal data. When a list, it must have the same length as
1366
                the `data_2` list.
1367
            data_2: The data to pair with the query to form the input to
1368
                the LLM. Can be text or multi-modal data. See [PromptType]
1369
                [vllm.inputs.PromptType] for more details about the format of
1370
                each prompt.
1371
1372
1373
1374
            use_tqdm: If `True`, shows a tqdm progress bar.
                If a callable (e.g., `functools.partial(tqdm, leave=False)`),
                it is used to create the progress bar.
                If `False`, no progress bar is created.
1375
            lora_request: LoRA request to use for generation, if any.
1376
1377
            pooling_params: The pooling parameters for pooling. If None, we
                use the default pooling parameters.
1378
        Returns:
1379
            A list of `ScoringRequestOutput` objects containing the
1380
1381
            generated scores in the same order as the input prompts.
        """
1382
1383
        model_config = self.llm_engine.model_config
        runner_type = model_config.runner_type
1384
        if runner_type != "pooling":
1385
1386
1387
            raise ValueError(
                "LLM.score() is only supported for pooling models. "
                "Try passing `--runner pooling` to use the model as a "
1388
1389
                "pooling model."
            )
1390

1391
1392
        supported_tasks = self.supported_tasks
        if all(t not in supported_tasks for t in ("embed", "classify")):
1393
1394
1395
1396
1397
            raise ValueError(
                "Score API is not supported by this model. "
                "Try converting the model using "
                "`--convert embed` or `--convert classify`."
            )
1398

1399
1400
1401
1402
        if (
            model_config.is_cross_encoder
            and getattr(model_config.hf_config, "num_labels", 0) != 1
        ):
1403
            raise ValueError("Score API is only enabled for num_labels == 1.")
1404
1405
1406
1407

        # the tokenizer for models such as
        # "cross-encoder/ms-marco-MiniLM-L-6-v2" doesn't support passing
        # lists of tokens to the `text` and `text_pair` kwargs
1408
        tokenizer = self.get_tokenizer()
1409

1410
        if not model_config.is_multimodal_model:
1411

1412
1413
1414
1415
1416
            def check_data_type(
                data: Union[
                    SingletonPrompt, Sequence[SingletonPrompt], ScoreMultiModalParam
                ],
            ):
1417
                if isinstance(data, dict) and "content" in data:
1418
1419
1420
1421
                    raise ValueError(
                        "ScoreMultiModalParam is not supported "
                        f"for {model_config.architecture}"
                    )
1422
1423
1424
1425
1426
1427
1428

            check_data_type(data_1)
            check_data_type(data_2)

            def ensure_str(prompt: SingletonPrompt):
                if isinstance(prompt, dict):
                    if "multi_modal_data" in prompt:
1429
1430
1431
                        raise ValueError(
                            "Multi-modal prompt is not supported for scoring"
                        )
1432
1433
                    elif "prompt_token_ids" in prompt:
                        prompt = tokenizer.decode(
1434
1435
                            cast(TokensPrompt, prompt)["prompt_token_ids"]
                        )
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
                    elif "prompt" in prompt:
                        prompt = cast(TextPrompt, prompt)["prompt"]
                assert type(prompt) is str
                return prompt

            if isinstance(data_1, (str, dict)):
                # Convert a single prompt to a list.
                data_1 = [data_1]  # type: ignore[list-item]

            data_1 = [ensure_str(t) for t in data_1]

            if isinstance(data_2, (str, dict)):
                # Convert a single prompt to a list.
                data_2 = [data_2]  # type: ignore[list-item]

            data_2 = [ensure_str(t) for t in data_2]

        if isinstance(data_1, dict) and "content" in data_1:
            data_1 = data_1.get("content")  # type: ignore[assignment]
        elif isinstance(data_1, str):
            data_1 = [data_1]

        if isinstance(data_2, dict) and "content" in data_2:
            data_2 = data_2.get("content")  # type: ignore[assignment]
        elif isinstance(data_2, str):
            data_2 = [data_2]

        _validate_score_input_lens(data_1, data_2)  # type: ignore[arg-type]
1464

1465
        if model_config.is_cross_encoder:
1466
1467
1468
1469
1470
1471
            return self._cross_encoding_score(
                tokenizer,
                data_1,  # type: ignore[arg-type]
                data_2,  # type: ignore[arg-type]
                truncate_prompt_tokens,
                use_tqdm,
1472
                pooling_params,
1473
1474
                lora_request,
            )
1475
        else:
1476
1477
            return self._embedding_score(
                tokenizer,
1478
1479
                data_1,  # type: ignore[arg-type]
                data_2,  # type: ignore[arg-type]
1480
1481
                truncate_prompt_tokens,
                use_tqdm,
1482
                pooling_params,
1483
1484
                lora_request,
            )
1485

1486
1487
1488
1489
1490
1491
    def start_profile(self) -> None:
        self.llm_engine.start_profile()

    def stop_profile(self) -> None:
        self.llm_engine.stop_profile()

1492
1493
    def reset_prefix_cache(self, device: Optional[Device] = None) -> bool:
        return self.llm_engine.reset_prefix_cache(device)
1494

1495
1496
1497
1498
1499
1500
    def sleep(self, level: int = 1):
        """
        Put the engine to sleep. The engine should not process any requests.
        The caller should guarantee that no requests are being processed
        during the sleep period, before `wake_up` is called.

1501
        Args:
1502
1503
            level: The sleep level. Level 1 sleep will offload the model
                weights and discard the kv cache. The content of kv cache
1504
                is forgotten. Level 1 sleep is good for sleeping and waking
1505
1506
1507
1508
1509
                up the engine to run the same model again. The model weights
                are backed up in CPU memory. Please make sure there's enough
                CPU memory to store the model weights. Level 2 sleep will
                discard both the model weights and the kv cache. The content
                of both the model weights and kv cache is forgotten. Level 2
1510
                sleep is good for sleeping and waking up the engine to run a
1511
                different model or update the model, where previous model
1512
                weights are not needed. It reduces CPU memory pressure.
1513
        """
1514
        self.reset_prefix_cache()
1515
1516
        self.llm_engine.sleep(level=level)

1517
    def wake_up(self, tags: Optional[list[str]] = None):
1518
        """
1519
1520
        Wake up the engine from sleep mode. See the [sleep][vllm.LLM.sleep]
        method for more details.
1521

1522
        Args:
1523
1524
            tags: An optional list of tags to reallocate the engine memory
                for specific memory allocations. Values must be in
1525
                `("weights", "kv_cache")`. If None, all memory is reallocated.
1526
                wake_up should be called with all tags (or None) before the
1527
1528
1529
                engine is used again.
        """
        self.llm_engine.wake_up(tags)
1530

1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
    def get_metrics(self) -> list["Metric"]:
        """Return a snapshot of aggregated metrics from Prometheus.

        Returns:
            A ``MetricSnapshot`` instance capturing the current state
            of all aggregated metrics from Prometheus.

        Note:
            This method is only available with the V1 LLM engine.
        """
        return self.llm_engine.get_metrics()

1543
1544
    def _validate_and_add_requests(
        self,
1545
        prompts: Union[PromptType, Sequence[PromptType], DataPrompt],
1546
1547
1548
1549
1550
1551
        params: Union[
            SamplingParams,
            Sequence[SamplingParams],
            PoolingParams,
            Sequence[PoolingParams],
        ],
1552
        *,
1553
        use_tqdm: Union[bool, Callable[..., tqdm]] = True,
1554
        lora_request: Optional[Union[Sequence[LoRARequest], LoRARequest]],
1555
        priority: Optional[list[int]] = None,
1556
    ) -> None:
1557
        if isinstance(prompts, (str, dict)):
1558
            # Convert a single prompt to a list.
1559
            prompts = [prompts]  # type: ignore[list-item]
1560

1561
        num_requests = len(prompts)
1562
        if isinstance(params, Sequence) and len(params) != num_requests:
1563
1564
1565
1566
1567
1568
1569
            raise ValueError("The lengths of prompts and params must be the same.")
        if isinstance(lora_request, Sequence) and len(lora_request) != num_requests:
            raise ValueError(
                "The lengths of prompts and lora_request must be the same."
            )

        for sp in params if isinstance(params, Sequence) else (params,):
1570
1571
1572
            if isinstance(sp, SamplingParams):
                # We only care about the final output
                sp.output_kind = RequestOutputKind.FINAL_ONLY
1573

Zhuohan Li's avatar
Zhuohan Li committed
1574
        # Add requests to the engine.
1575
1576
        it = prompts
        if use_tqdm:
1577
1578
            tqdm_func = use_tqdm if callable(use_tqdm) else tqdm
            it = tqdm_func(it, desc="Adding requests")
1579
1580

        for i, prompt in enumerate(it):
1581
1582
            if isinstance(prompt, dict):
                self._validate_mm_data_and_uuids(
1583
1584
                    prompt.get("multi_modal_data"), prompt.get("multi_modal_uuids")
                )
1585

1586
            self._add_request(
1587
                prompt,
1588
                params[i] if isinstance(params, Sequence) else params,
1589
1590
1591
                lora_request=lora_request[i]
                if isinstance(lora_request, Sequence)
                else lora_request,
1592
                priority=priority[i] if priority else 0,
nunjunj's avatar
nunjunj committed
1593
            )
1594

1595
    def _validate_mm_data_and_uuids(
1596
1597
1598
        self,
        multi_modal_data: Optional[Any],  # MultiModalDataDict
        multi_modal_uuids: Optional[Any],  # MultiModalUUIDDict
1599
1600
1601
    ):
        """
        Validate that if any multi-modal data is skipped (i.e. None),
1602
        then its corresponding UUID must be set.
1603
1604
1605
1606
1607
1608
1609
1610
        """
        if multi_modal_data is None:
            return

        for modality, data in multi_modal_data.items():
            if isinstance(data, list):
                for i, d in enumerate(data):
                    if d is None:
1611
1612
1613
1614
1615
1616
1617
1618
                        if (
                            multi_modal_uuids is None
                            or modality not in multi_modal_uuids
                            or multi_modal_uuids[  # noqa: E501
                                modality
                            ]
                            is None
                        ):
1619
1620
                            raise ValueError(
                                f"Multi-modal data for {modality} is None "
1621
1622
                                f"but UUID is not provided"
                            )
1623
                        else:
1624
1625
1626
1627
                            if (
                                len(multi_modal_uuids[modality]) <= i
                                or multi_modal_uuids[modality][i] is None
                            ):
1628
1629
                                raise ValueError(
                                    f"Multi-modal data for {modality} is None "
1630
1631
                                    f"but UUID is not provided"
                                )
1632
            else:
1633
1634
1635
1636
1637
1638
1639
1640
1641
                if data is None and (
                    multi_modal_uuids is None
                    or modality not in multi_modal_uuids
                    or multi_modal_uuids[modality] is None
                ):
                    raise ValueError(
                        f"Multi-modal data for {modality} is None"
                        f" but UUID is not provided"
                    )
1642

1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
    def _process_inputs(
        self,
        request_id: str,
        engine_prompt: PromptType,
        params: Union[SamplingParams, PoolingParams],
        *,
        lora_request: Optional[LoRARequest],
        priority: int,
    ) -> tuple[EngineCoreRequest, dict[str, Any]]:
        """Use the Processor to process inputs for LLMEngine."""
        tokenization_kwargs: dict[str, Any] = {}
1654
1655
1656
1657
1658
        _validate_truncation_size(
            self.model_config.max_model_len,
            params.truncate_prompt_tokens,
            tokenization_kwargs,
        )
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670

        processor = self._get_processor()
        engine_request = processor.process_inputs(
            request_id,
            engine_prompt,
            params,
            lora_request=lora_request,
            tokenization_kwargs=tokenization_kwargs,
            priority=priority,
        )
        return engine_request, tokenization_kwargs

1671
    def _add_request(
nunjunj's avatar
nunjunj committed
1672
        self,
1673
        prompt: PromptType,
nunjunj's avatar
nunjunj committed
1674
        params: Union[SamplingParams, PoolingParams],
1675
        lora_request: Optional[LoRARequest] = None,
1676
        priority: int = 0,
1677
    ) -> None:
1678
        prompt_text, _, _ = get_prompt_components(prompt)
1679
        request_id = str(next(self.request_counter))
1680
1681

        engine_request, tokenization_kwargs = self._process_inputs(
1682
            request_id,
1683
            prompt,
1684
1685
            params,
            lora_request=lora_request,
1686
1687
1688
1689
1690
1691
1692
1693
            priority=priority,
        )

        self.llm_engine.add_request(
            request_id,
            engine_request,
            params,
            lora_request=lora_request,
1694
            tokenization_kwargs=tokenization_kwargs,
1695
            priority=priority,
1696
            prompt_text=prompt_text,
nunjunj's avatar
nunjunj committed
1697
        )
1698

1699
    def _run_engine(
1700
        self, *, use_tqdm: Union[bool, Callable[..., tqdm]] = True
1701
    ) -> list[Union[RequestOutput, PoolingRequestOutput]]:
1702
1703
        # Initialize tqdm.
        if use_tqdm:
Zhuohan Li's avatar
Zhuohan Li committed
1704
            num_requests = self.llm_engine.get_num_unfinished_requests()
1705
1706
            tqdm_func = use_tqdm if callable(use_tqdm) else tqdm
            pbar = tqdm_func(
1707
1708
1709
                total=num_requests,
                desc="Processed prompts",
                dynamic_ncols=True,
1710
                postfix=(f"est. speed input: {0:.2f} toks/s, output: {0:.2f} toks/s"),
1711
            )
1712

Zhuohan Li's avatar
Zhuohan Li committed
1713
        # Run the engine.
1714
        outputs: list[Union[RequestOutput, PoolingRequestOutput]] = []
1715
1716
        total_in_toks = 0
        total_out_toks = 0
Zhuohan Li's avatar
Zhuohan Li committed
1717
1718
        while self.llm_engine.has_unfinished_requests():
            step_outputs = self.llm_engine.step()
1719
            for output in step_outputs:
1720
                if output.finished:
1721
1722
                    outputs.append(output)
                    if use_tqdm:
1723
1724
                        if isinstance(output, RequestOutput):
                            # Calculate tokens only for RequestOutput
1725
                            n = len(output.outputs)
1726
                            assert output.prompt_token_ids is not None
1727
                            total_in_toks += len(output.prompt_token_ids) * n
1728
1729
                            in_spd = total_in_toks / pbar.format_dict["elapsed"]
                            total_out_toks += sum(
1730
1731
1732
                                len(stp.token_ids) for stp in output.outputs
                            )
                            out_spd = total_out_toks / pbar.format_dict["elapsed"]
1733
1734
                            pbar.postfix = (
                                f"est. speed input: {in_spd:.2f} toks/s, "
1735
1736
                                f"output: {out_spd:.2f} toks/s"
                            )
1737
                            pbar.update(n)
1738
1739
                        else:
                            pbar.update(1)
1740
1741
                        if pbar.n == num_requests:
                            pbar.refresh()
1742

1743
1744
        if use_tqdm:
            pbar.close()
1745
1746
1747
        # Sort the outputs by request ID.
        # This is necessary because some requests may be finished earlier than
        # its previous requests.
1748
        return sorted(outputs, key=lambda x: int(x.request_id))