| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Callable, Sequence |
| 5 | from typing import Any |
| 6 | |
| 7 | from tqdm.auto import tqdm |
| 8 | |
| 9 | from vllm.entrypoints.chat_utils import ChatTemplateConfig |
| 10 | from vllm.entrypoints.offline_utils import OfflineInferenceMixin |
| 11 | from vllm.inputs import DataPrompt, PromptType |
| 12 | from vllm.logger import init_logger |
| 13 | from vllm.lora.request import LoRARequest |
| 14 | from vllm.outputs import ( |
| 15 | ClassificationRequestOutput, |
| 16 | EmbeddingRequestOutput, |
| 17 | PoolingRequestOutput, |
| 18 | ScoringRequestOutput, |
| 19 | ) |
| 20 | from vllm.pooling_params import PoolingParams |
| 21 | from vllm.tasks import SCORE_TYPE_MAP, PoolingTask, SupportedTask |
| 22 | |
| 23 | from .base.io_processor import PoolingIOProcessor |
| 24 | from .factories import init_pooling_io_processors |
| 25 | from .scoring.io_processor import ScoringIOProcessor |
| 26 | from .scoring.typing import ScoreInput |
| 27 | from .typing import ( |
| 28 | AnyOfflineInputsContext, |
| 29 | OfflineEncodeInputsContext, |
| 30 | OfflineOutputsContext, |
| 31 | OfflinePluginInputsContext, |
| 32 | OfflineScoringInputsContext, |
| 33 | RequestFactory, |
| 34 | ) |
| 35 | |
| 36 | logger = init_logger(__name__) |
| 37 | |
| 38 | |
| 39 | class PoolingOfflineMixin(OfflineInferenceMixin): |
| 40 | """Offline inference for pooling models.""" |
| 41 | |
| 42 | runner_type: str |
| 43 | chat_template: str | None |
| 44 | supported_tasks: tuple[SupportedTask, ...] |
| 45 | |
| 46 | def __init__(self): |
| 47 | self.pooling_task = self.model_config.get_pooling_task(self.supported_tasks) |
| 48 | if self.pooling_task is not None: |
| 49 | logger.info("Supported pooling task: %s", self.pooling_task) |
| 50 | |
| 51 | self.chat_template_config = ChatTemplateConfig(chat_template=self.chat_template) |
| 52 | self.pooling_io_processors = init_pooling_io_processors( |
| 53 | supported_tasks=self.supported_tasks, |
| 54 | vllm_config=self.llm_engine.vllm_config, |
| 55 | renderer=self.renderer, |
| 56 | chat_template_config=self.chat_template_config, |
| 57 | ) |
| 58 | |
| 59 | # Use thread pool executor to accelerate preprocessing. |
| 60 | self._executor = self.renderer._executor |
| 61 | |
| 62 | def encode( |
| 63 | self, |
| 64 | prompts: PromptType | Sequence[PromptType] | DataPrompt, |
| 65 | pooling_params: PoolingParams | Sequence[PoolingParams] | None = None, |
| 66 | *, |
| 67 | use_tqdm: bool | Callable[..., tqdm] = True, |
| 68 | lora_request: list[LoRARequest] | LoRARequest | None = None, |
| 69 | pooling_task: PoolingTask | None = None, |
| 70 | tokenization_kwargs: dict[str, Any] | None = None, |
| 71 | ) -> list[PoolingRequestOutput]: |
| 72 | """Apply pooling to the hidden states corresponding to the input |
| 73 | prompts. |
| 74 | |
| 75 | This class automatically batches the given prompts, considering |
| 76 | the memory constraint. For the best performance, put all of your prompts |
| 77 | into a single list and pass it to this method. |
| 78 | |
| 79 | Args: |
| 80 | prompts: The prompts to the LLM. You may pass a sequence of prompts |
| 81 | for batch inference. See [PromptType][vllm.inputs.PromptType] |
| 82 | for more details about the format of each prompt. |
| 83 | pooling_params: The pooling parameters for pooling. If None, we |
| 84 | use the default pooling parameters. |
| 85 | use_tqdm: If `True`, shows a tqdm progress bar. |
| 86 | If a callable (e.g., `functools.partial(tqdm, leave=False)`), |
| 87 | it is used to create the progress bar. |
| 88 | If `False`, no progress bar is created. |
| 89 | lora_request: LoRA request to use for generation, if any. |
| 90 | pooling_task: Override the pooling task to use. |
| 91 | tokenization_kwargs: Overrides for `tokenizer.encode`. |
| 92 | |
| 93 | Returns: |
| 94 | A list of `PoolingRequestOutput` objects containing the |
| 95 | pooled hidden states in the same order as the input prompts. |
| 96 | |
| 97 | """ |
| 98 | if isinstance(prompts, dict) and "data" in prompts and pooling_task != "plugin": |
| 99 | raise ValueError( |
| 100 | "The 'data' field is only supported for the 'plugin' pooling task." |
| 101 | ) |
| 102 | self._verify_pooling_task(pooling_task) |
| 103 | assert pooling_task is not None and pooling_task in self.pooling_io_processors |
| 104 | |
| 105 | io_processor = self.pooling_io_processors[pooling_task] |
| 106 | |
| 107 | ctx: AnyOfflineInputsContext |
| 108 | if isinstance(prompts, dict) and "data" in prompts: |
| 109 | ctx = OfflinePluginInputsContext( |
| 110 | pooling_task=pooling_task, |
| 111 | prompts=prompts, # type: ignore[arg-type] |
| 112 | tokenization_kwargs=tokenization_kwargs, |
| 113 | pooling_params=pooling_params, |
| 114 | lora_request=lora_request, |
| 115 | priorities=None, |
| 116 | ) |
| 117 | else: |
| 118 | ctx = OfflineEncodeInputsContext( |
| 119 | pooling_task=pooling_task, |
| 120 | prompts=prompts, |
| 121 | tokenization_kwargs=tokenization_kwargs, |
| 122 | pooling_params=pooling_params, |
| 123 | lora_request=lora_request, |
| 124 | priorities=None, |
| 125 | ) |
| 126 | |
| 127 | request_factory, num_requests = io_processor.get_request_factory_offline(ctx) |
| 128 | outputs = self._run_tiling_engine( |
| 129 | io_processor, request_factory, num_requests, use_tqdm=use_tqdm |
| 130 | ) |
| 131 | outputs = io_processor.post_process_offline( |
| 132 | ctx=OfflineOutputsContext(outputs=outputs) |
| 133 | ) |
| 134 | return outputs |
| 135 | |
| 136 | def _verify_pooling_task(self, pooling_task: PoolingTask | None): |
| 137 | if self.runner_type != "pooling": |
| 138 | raise ValueError( |
| 139 | "LLM.encode() is only supported for pooling models. " |
| 140 | "Try passing `--runner pooling` to use the model as a " |
| 141 | "pooling model." |
| 142 | ) |
| 143 | |
| 144 | if pooling_task is None: |
| 145 | raise ValueError( |
| 146 | """ |
| 147 | pooling_task required for `LLM.encode`. |
| 148 | Please use one of the more specific methods or set the pooling_task when using `LLM.encode`: |
| 149 | - For embeddings, use `LLM.embed(...)` or `pooling_task="embed"`. |
| 150 | - For classification logits, use `LLM.classify(...)` or `pooling_task="classify"`. |
| 151 | - For similarity scores, use `LLM.score(...)`. |
| 152 | - For rewards, `pooling_task="classify"` or `pooling_task="token_classify"`. |
| 153 | - For token classification, use `pooling_task="token_classify"`. |
| 154 | - For multi-vector retrieval, use `pooling_task="token_embed"`. |
| 155 | """ # noqa: E501 |
| 156 | ) |
| 157 | |
| 158 | if ( |
| 159 | pooling_task in ("embed", "token_embed") |
| 160 | and pooling_task not in self.supported_tasks |
| 161 | ): |
| 162 | raise ValueError( |
| 163 | "Embedding API is not supported by this model. " |
| 164 | "Try converting the model using `--convert embed`." |
| 165 | ) |
| 166 | |
| 167 | if ( |
| 168 | pooling_task in ("classify", "token_classify") |
| 169 | and pooling_task not in self.supported_tasks |
| 170 | ): |
| 171 | raise ValueError( |
| 172 | "Classification API is not supported by this model. " |
| 173 | "Try converting the model using `--convert classify`." |
| 174 | ) |
| 175 | |
| 176 | # plugin task uses io_processor.parse_data to verify inputs |
| 177 | if pooling_task != "plugin" and pooling_task != self.pooling_task: |
| 178 | if pooling_task not in self.supported_tasks: |
| 179 | raise ValueError( |
| 180 | f"Unsupported task: {pooling_task!r} " |
| 181 | f"Supported tasks: {self.supported_tasks}" |
| 182 | ) |
| 183 | else: |
| 184 | raise ValueError( |
| 185 | f"Try switching the model's pooling_task " |
| 186 | f'via `PoolerConfig(task="{pooling_task}")`' |
| 187 | ) |
| 188 | |
| 189 | if pooling_task == "plugin" and "plugin" not in self.pooling_io_processors: |
| 190 | raise ValueError( |
| 191 | "No IOProcessor plugin installed. Please refer " |
| 192 | "to the documentation and to the " |
| 193 | "'prithvi_geospatial_mae_io_processor' " |
| 194 | "offline inference example for more details." |
| 195 | ) |
| 196 | |
| 197 | def embed( |
| 198 | self, |
| 199 | prompts: PromptType | Sequence[PromptType], |
| 200 | *, |
| 201 | use_tqdm: bool | Callable[..., tqdm] = True, |
| 202 | pooling_params: PoolingParams | Sequence[PoolingParams] | None = None, |
| 203 | lora_request: list[LoRARequest] | LoRARequest | None = None, |
| 204 | tokenization_kwargs: dict[str, Any] | None = None, |
| 205 | ) -> list[EmbeddingRequestOutput]: |
| 206 | """Generate an embedding vector for each prompt. |
| 207 | |
| 208 | This class automatically batches the given prompts, considering |
| 209 | the memory constraint. For the best performance, put all of your prompts |
| 210 | into a single list and pass it to this method. |
| 211 | |
| 212 | Args: |
| 213 | prompts: The prompts to the LLM. You may pass a sequence of prompts |
| 214 | for batch inference. See [PromptType][vllm.inputs.PromptType] |
| 215 | for more details about the format of each prompt. |
| 216 | pooling_params: The pooling parameters for pooling. If None, we |
| 217 | use the default pooling parameters. |
| 218 | use_tqdm: If `True`, shows a tqdm progress bar. |
| 219 | If a callable (e.g., `functools.partial(tqdm, leave=False)`), |
| 220 | it is used to create the progress bar. |
| 221 | If `False`, no progress bar is created. |
| 222 | lora_request: LoRA request to use for generation, if any. |
| 223 | tokenization_kwargs: Overrides for `tokenizer.encode`. |
| 224 | |
| 225 | Returns: |
| 226 | A list of `EmbeddingRequestOutput` objects containing the |
| 227 | embedding vectors in the same order as the input prompts. |
| 228 | |
| 229 | """ |
| 230 | items = self.encode( |
| 231 | prompts, |
| 232 | use_tqdm=use_tqdm, |
| 233 | pooling_params=pooling_params, |
| 234 | lora_request=lora_request, |
| 235 | pooling_task="embed", |
| 236 | tokenization_kwargs=tokenization_kwargs, |
| 237 | ) |
| 238 | |
| 239 | return [EmbeddingRequestOutput.from_base(item) for item in items] |
| 240 | |
| 241 | def classify( |
| 242 | self, |
| 243 | prompts: PromptType | Sequence[PromptType], |
| 244 | *, |
| 245 | pooling_params: PoolingParams | Sequence[PoolingParams] | None = None, |
| 246 | use_tqdm: bool | Callable[..., tqdm] = True, |
| 247 | lora_request: list[LoRARequest] | LoRARequest | None = None, |
| 248 | tokenization_kwargs: dict[str, Any] | None = None, |
| 249 | ) -> list[ClassificationRequestOutput]: |
| 250 | """Generate class logits for each prompt. |
| 251 | |
| 252 | This class automatically batches the given prompts, considering |
| 253 | the memory constraint. For the best performance, put all of your prompts |
| 254 | into a single list and pass it to this method. |
| 255 | |
| 256 | Args: |
| 257 | prompts: The prompts to the LLM. You may pass a sequence of prompts |
| 258 | for batch inference. See [PromptType][vllm.inputs.PromptType] |
| 259 | for more details about the format of each prompt. |
| 260 | pooling_params: The pooling parameters for pooling. If None, we |
| 261 | use the default pooling parameters. |
| 262 | use_tqdm: If `True`, shows a tqdm progress bar. |
| 263 | If a callable (e.g., `functools.partial(tqdm, leave=False)`), |
| 264 | it is used to create the progress bar. |
| 265 | If `False`, no progress bar is created. |
| 266 | lora_request: LoRA request to use for generation, if any. |
| 267 | tokenization_kwargs: Overrides for `tokenizer.encode`. |
| 268 | |
| 269 | Returns: |
| 270 | A list of `ClassificationRequestOutput` objects containing the |
| 271 | embedding vectors in the same order as the input prompts. |
| 272 | |
| 273 | """ |
| 274 | items = self.encode( |
| 275 | prompts, |
| 276 | use_tqdm=use_tqdm, |
| 277 | pooling_params=pooling_params, |
| 278 | lora_request=lora_request, |
| 279 | pooling_task="classify", |
| 280 | tokenization_kwargs=tokenization_kwargs, |
| 281 | ) |
| 282 | |
| 283 | return [ClassificationRequestOutput.from_base(item) for item in items] |
| 284 | |
| 285 | def score( |
| 286 | self, |
| 287 | data_1: ScoreInput | list[ScoreInput], |
| 288 | data_2: ScoreInput | list[ScoreInput], |
| 289 | /, |
| 290 | *, |
| 291 | use_tqdm: bool | Callable[..., tqdm] = True, |
| 292 | pooling_params: PoolingParams | None = None, |
| 293 | lora_request: list[LoRARequest] | LoRARequest | None = None, |
| 294 | tokenization_kwargs: dict[str, Any] | None = None, |
| 295 | chat_template: str | None = None, |
| 296 | ) -> list[ScoringRequestOutput]: |
| 297 | """Generate similarity scores for all pairs `<text,text_pair>` or |
| 298 | `<multi-modal data, multi-modal data pair>`. |
| 299 | |
| 300 | The inputs can be `1 -> 1`, `1 -> N` or `N -> N`. |
| 301 | In the `1 - N` case the `data_1` input will be replicated `N` |
| 302 | times to pair with the `data_2` inputs. |
| 303 | The input pairs are used to build a list of prompts for the |
| 304 | cross encoder model. This class automatically batches the prompts, |
| 305 | considering the memory constraint. For the best performance, put all |
| 306 | of your inputs into a single list and pass it to this method. |
| 307 | |
| 308 | Supports both text and multi-modal data (images, etc.) when used with |
| 309 | appropriate multi-modal models. For multi-modal inputs, ensure the |
| 310 | prompt structure matches the model's expected input format. |
| 311 | |
| 312 | Args: |
| 313 | data_1: Can be a single prompt, a list of prompts or |
| 314 | `ScoreMultiModalParam`, which can contain either text or |
| 315 | multi-modal data. When a list, it must have the same length as |
| 316 | the `data_2` list. |
| 317 | data_2: The data to pair with the query to form the input to |
| 318 | the LLM. Can be text or multi-modal data. See [PromptType] |
| 319 | [vllm.inputs.PromptType] for more details about the format of |
| 320 | each prompt. |
| 321 | pooling_params: The pooling parameters for pooling. If None, we |
| 322 | use the default pooling parameters. |
| 323 | use_tqdm: If `True`, shows a tqdm progress bar. |
| 324 | If a callable (e.g., `functools.partial(tqdm, leave=False)`), |
| 325 | it is used to create the progress bar. |
| 326 | If `False`, no progress bar is created. |
| 327 | lora_request: LoRA request to use for generation, if any. |
| 328 | chat_template: The chat template to use for the scoring. If None, we |
| 329 | use the model's default chat template. |
| 330 | tokenization_kwargs: Overrides for `tokenizer.encode`. |
| 331 | |
| 332 | Returns: |
| 333 | A list of `ScoringRequestOutput` objects containing the |
| 334 | generated scores in the same order as the input prompts. |
| 335 | |
| 336 | """ |
| 337 | if self.runner_type != "pooling": |
| 338 | raise ValueError( |
| 339 | "LLM.score() is only supported for pooling models. " |
| 340 | "Try passing `--runner pooling` to use the model as a " |
| 341 | "pooling model." |
| 342 | ) |
| 343 | |
| 344 | score_type: str | None = SCORE_TYPE_MAP.get(self.pooling_task, None) # type: ignore[arg-type] |
| 345 | if ( |
| 346 | score_type == "cross-encoder" |
| 347 | and getattr(self.model_config.hf_config, "num_labels", 0) != 1 |
| 348 | ): |
| 349 | raise ValueError("Scoring API is only enabled for num_labels == 1.") |
| 350 | |
| 351 | if score_type is None or score_type not in self.pooling_io_processors: |
| 352 | raise ValueError("This model does not support the Scoring API.") |
| 353 | |
| 354 | io_processor = self.pooling_io_processors[score_type] |
| 355 | assert isinstance(io_processor, ScoringIOProcessor) |
| 356 | |
| 357 | scoring_data = io_processor.valid_inputs(data_1, data_2) |
| 358 | n_queries = len(scoring_data.data_1) |
| 359 | |
| 360 | if pooling_params is None: |
| 361 | pooling_params = PoolingParams() |
| 362 | |
| 363 | assert isinstance(pooling_params, PoolingParams) |
| 364 | pooling_task = io_processor.pooling_task |
| 365 | if pooling_params.task is None: |
| 366 | pooling_params.task = pooling_task |
| 367 | elif pooling_params.task != pooling_task: |
| 368 | msg = ( |
| 369 | f"You cannot overwrite {pooling_params.task=!r} with {pooling_task=!r}!" |
| 370 | ) |
| 371 | raise ValueError(msg) |
| 372 | |
| 373 | ctx = OfflineScoringInputsContext( |
| 374 | pooling_task=pooling_task, |
| 375 | scoring_data=scoring_data, |
| 376 | pooling_params=pooling_params, |
| 377 | tokenization_kwargs=tokenization_kwargs, |
| 378 | lora_request=lora_request, |
| 379 | chat_template=chat_template, |
| 380 | priorities=None, |
| 381 | ) |
| 382 | |
| 383 | request_factory, num_requests = io_processor.get_request_factory_offline(ctx) |
| 384 | outputs = self._run_tiling_engine( |
| 385 | io_processor, request_factory, num_requests, use_tqdm=use_tqdm |
| 386 | ) |
| 387 | |
| 388 | outputs = io_processor.post_process_offline( |
| 389 | ctx=OfflineOutputsContext(outputs=outputs, n_queries=n_queries), |
| 390 | ) |
| 391 | |
| 392 | return [ScoringRequestOutput.from_base(item) for item in outputs] |
| 393 | |
| 394 | def _run_tiling_engine( |
| 395 | self, |
| 396 | io_processor: PoolingIOProcessor, |
| 397 | request_factory: RequestFactory, |
| 398 | num_requests: int, |
| 399 | use_tqdm: bool | Callable[..., tqdm] = True, |
| 400 | ): |
| 401 | if num_requests == 0: |
| 402 | raise ValueError("You must pass at least one prompt") |
| 403 | |
| 404 | # Keeping max_num_seqs * 2 requests in the core can already saturate the core. |
| 405 | # Therefore, keep most requests waiting outside the core. |
| 406 | max_requests_in_core = ( |
| 407 | self.llm_engine.vllm_config.scheduler_config.max_num_seqs * 2 |
| 408 | ) |
| 409 | num_requests_in_core = 0 |
| 410 | num_waited_requests = num_requests |
| 411 | |
| 412 | if use_tqdm: |
| 413 | tqdm_func = use_tqdm if callable(use_tqdm) else tqdm |
| 414 | pbar = tqdm_func( |
| 415 | total=num_requests, |
| 416 | desc="Processed prompts", |
| 417 | dynamic_ncols=True, |
| 418 | postfix=f"est. speed input: {0:.2f} toks/s, output: {0:.2f} toks/s", |
| 419 | ) |
| 420 | |
| 421 | outputs: list[PoolingRequestOutput] = [] |
| 422 | added_request_ids: set[str] = set() |
| 423 | |
| 424 | it = self._executor.map(io_processor.render, request_factory()) |
| 425 | |
| 426 | try: |
| 427 | while num_waited_requests or self.llm_engine.has_unfinished_requests(): |
| 428 | requests = [] |
| 429 | for _ in range(max_requests_in_core - num_requests_in_core): |
| 430 | if num_waited_requests == 0: |
| 431 | break |
| 432 | try: |
| 433 | request = next(it) |
| 434 | requests.append(request) |
| 435 | except StopIteration: |
| 436 | num_waited_requests = 0 |
| 437 | break |
| 438 | |
| 439 | num_waited_requests -= 1 |
| 440 | num_requests_in_core += 1 |
| 441 | |
| 442 | if requests: |
| 443 | request_ids = self._render_and_add_requests( |
| 444 | prompts=[x["prompts"] for x in requests], |
| 445 | params=[x["params"] for x in requests], |
| 446 | lora_requests=[x["lora_requests"] for x in requests], |
| 447 | priorities=[x["priorities"] for x in requests], |
| 448 | ) |
| 449 | |
| 450 | for request_id in request_ids: |
| 451 | # undo assign_request_id |
| 452 | request_id = request_id.split("-", 1)[0] |
| 453 | added_request_ids.add(request_id) |
| 454 | |
| 455 | step_outputs = self.llm_engine.step() |
| 456 | for output in step_outputs: |
| 457 | assert isinstance(output, PoolingRequestOutput) |
| 458 | assert output.finished |
| 459 | outputs.append(output) |
| 460 | added_request_ids.discard(output.request_id) |
| 461 | num_requests_in_core -= 1 |
| 462 | |
| 463 | if use_tqdm: |
| 464 | pbar.update(1) |
| 465 | if pbar.n == num_requests: |
| 466 | pbar.refresh() |
| 467 | |
| 468 | except Exception: |
| 469 | if added_request_ids: |
| 470 | self.llm_engine.abort_request(list(added_request_ids)) |
| 471 | raise |
| 472 | |
| 473 | finally: |
| 474 | if use_tqdm: |
| 475 | pbar.close() |
| 476 | |
| 477 | return sorted(outputs, key=lambda x: int(x.request_id)) |