| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import importlib.util |
| 5 | import json |
| 6 | import math |
| 7 | from collections.abc import Callable, Sequence |
| 8 | from dataclasses import dataclass |
| 9 | from functools import lru_cache, partial |
| 10 | from typing import Any, Literal, cast |
| 11 | |
| 12 | import pybase64 |
| 13 | import torch |
| 14 | from fastapi.responses import JSONResponse, StreamingResponse |
| 15 | |
| 16 | from vllm.config import ModelConfig |
| 17 | from vllm.entrypoints.serve.engine.protocol import UsageInfo |
| 18 | from vllm.logger import init_logger |
| 19 | from vllm.outputs import PoolingRequestOutput |
| 20 | from vllm.tasks import SupportedTask |
| 21 | from vllm.utils.serial_utils import ( |
| 22 | EMBED_DTYPES, |
| 23 | EmbedDType, |
| 24 | Endianness, |
| 25 | binary2tensor, |
| 26 | tensor2binary, |
| 27 | ) |
| 28 | |
| 29 | logger = init_logger(__name__) |
| 30 | |
| 31 | JsonEncodingFormat = Literal["float", "base64"] |
| 32 | BytesEncodingFormat = Literal["bytes", "bytes_only"] |
| 33 | FloatEncodedPoolingOutput = list[float] | list[list[float]] |
| 34 | JsonEncodedPoolingOutput = FloatEncodedPoolingOutput | str |
| 35 | |
| 36 | |
| 37 | @dataclass |
| 38 | class MetadataItem: |
| 39 | index: int |
| 40 | embed_dtype: EmbedDType |
| 41 | endianness: Endianness |
| 42 | start: int |
| 43 | end: int |
| 44 | shape: tuple[int, ...] |
| 45 | |
| 46 | |
| 47 | def build_metadata_items( |
| 48 | embed_dtype: EmbedDType, |
| 49 | endianness: Endianness, |
| 50 | shape: tuple[int, ...], |
| 51 | n_request: int, |
| 52 | ) -> list[MetadataItem]: |
| 53 | n_bytes = EMBED_DTYPES[embed_dtype].nbytes |
| 54 | size = math.prod(shape) |
| 55 | |
| 56 | return [ |
| 57 | MetadataItem( |
| 58 | index=i, |
| 59 | embed_dtype=embed_dtype, |
| 60 | endianness=endianness, |
| 61 | start=i * size * n_bytes, |
| 62 | end=(i + 1) * size * n_bytes, |
| 63 | shape=shape, |
| 64 | ) |
| 65 | for i in range(n_request) |
| 66 | ] |
| 67 | |
| 68 | |
| 69 | def encode_pooling_output_float( |
| 70 | output: PoolingRequestOutput, |
| 71 | ) -> FloatEncodedPoolingOutput: |
| 72 | return output.outputs.data.tolist() |
| 73 | |
| 74 | |
| 75 | def encode_pooling_output_float_or_ndarray(output: PoolingRequestOutput) -> Any: |
| 76 | """Return an ndarray when the response renderer can serialize NumPy.""" |
| 77 | try: |
| 78 | data = output.outputs.data |
| 79 | if not data.is_contiguous(): |
| 80 | data = data.contiguous() |
| 81 | return data.numpy() |
| 82 | except (RuntimeError, TypeError): |
| 83 | return output.outputs.data.tolist() |
| 84 | |
| 85 | |
| 86 | def encode_pooling_output_base64( |
| 87 | output: PoolingRequestOutput, |
| 88 | embed_dtype: EmbedDType, |
| 89 | endianness: Endianness, |
| 90 | ) -> str: |
| 91 | embedding_bytes = tensor2binary(output.outputs.data, embed_dtype, endianness) |
| 92 | return pybase64.b64encode(embedding_bytes).decode("utf-8") |
| 93 | |
| 94 | |
| 95 | def encode_pooling_bytes( |
| 96 | pooling_outputs: list[PoolingRequestOutput], |
| 97 | embed_dtype: EmbedDType, |
| 98 | endianness: Endianness, |
| 99 | ) -> tuple[list[bytes], list[dict[str, Any]]]: |
| 100 | items: list[dict[str, Any]] = [] |
| 101 | body: list[bytes] = [] |
| 102 | offset = 0 |
| 103 | for idx, output in enumerate(pooling_outputs): |
| 104 | binary = tensor2binary( |
| 105 | tensor=output.outputs.data, |
| 106 | embed_dtype=embed_dtype, |
| 107 | endianness=endianness, |
| 108 | ) |
| 109 | size = len(binary) |
| 110 | |
| 111 | # Dictionary form of MetadataItem |
| 112 | item = dict( |
| 113 | index=idx, |
| 114 | embed_dtype=embed_dtype, |
| 115 | endianness=endianness, |
| 116 | start=offset, |
| 117 | end=offset + size, |
| 118 | shape=output.outputs.data.shape, |
| 119 | ) |
| 120 | |
| 121 | body.append(binary) |
| 122 | items.append(item) |
| 123 | offset += size |
| 124 | |
| 125 | return body, items |
| 126 | |
| 127 | |
| 128 | def get_pooling_output_encoder( |
| 129 | encoding_format: JsonEncodingFormat, |
| 130 | embed_dtype: EmbedDType, |
| 131 | endianness: Endianness, |
| 132 | ) -> Callable[[PoolingRequestOutput], JsonEncodedPoolingOutput]: |
| 133 | return cast( |
| 134 | Callable[[PoolingRequestOutput], JsonEncodedPoolingOutput], |
| 135 | ( |
| 136 | encode_pooling_output_float |
| 137 | if encoding_format == "float" |
| 138 | else partial( |
| 139 | encode_pooling_output_base64, |
| 140 | embed_dtype=embed_dtype, |
| 141 | endianness=endianness, |
| 142 | ) |
| 143 | ), |
| 144 | ) |
| 145 | |
| 146 | |
| 147 | def get_pooling_usage( |
| 148 | pooling_outputs: Sequence[PoolingRequestOutput], |
| 149 | ) -> UsageInfo: |
| 150 | num_prompt_tokens = sum( |
| 151 | len(output.prompt_token_ids) if output.prompt_token_ids is not None else 0 |
| 152 | for output in pooling_outputs |
| 153 | ) |
| 154 | return UsageInfo( |
| 155 | prompt_tokens=num_prompt_tokens, |
| 156 | total_tokens=num_prompt_tokens, |
| 157 | ) |
| 158 | |
| 159 | |
| 160 | def get_pooling_usage_payload( |
| 161 | pooling_outputs: Sequence[PoolingRequestOutput], |
| 162 | ) -> dict[str, int]: |
| 163 | usage = get_pooling_usage(pooling_outputs) |
| 164 | return { |
| 165 | "prompt_tokens": usage.prompt_tokens, |
| 166 | "total_tokens": usage.total_tokens, |
| 167 | } |
| 168 | |
| 169 | |
| 170 | def build_pooling_bytes_streaming_response( |
| 171 | pooling_outputs: list[PoolingRequestOutput], |
| 172 | request_id: str, |
| 173 | created_time: int, |
| 174 | model_name: str, |
| 175 | encoding_format: BytesEncodingFormat, |
| 176 | embed_dtype: EmbedDType, |
| 177 | endianness: Endianness, |
| 178 | ) -> StreamingResponse: |
| 179 | content, items = encode_pooling_bytes( |
| 180 | pooling_outputs=pooling_outputs, |
| 181 | embed_dtype=embed_dtype, |
| 182 | endianness=endianness, |
| 183 | ) |
| 184 | |
| 185 | headers = ( |
| 186 | None |
| 187 | if encoding_format == "bytes_only" |
| 188 | else { |
| 189 | "metadata": json.dumps( |
| 190 | { |
| 191 | "id": request_id, |
| 192 | "created": created_time, |
| 193 | "model": model_name, |
| 194 | "data": items, |
| 195 | "usage": get_pooling_usage_payload(pooling_outputs), |
| 196 | } |
| 197 | ) |
| 198 | } |
| 199 | ) |
| 200 | |
| 201 | return StreamingResponse( |
| 202 | content=content, |
| 203 | headers=headers, |
| 204 | media_type="application/octet-stream", |
| 205 | ) |
| 206 | |
| 207 | |
| 208 | def decode_pooling_output(items: list[MetadataItem], body: bytes) -> list[torch.Tensor]: |
| 209 | return [ |
| 210 | binary2tensor( |
| 211 | body[item.start : item.end], |
| 212 | item.shape, |
| 213 | item.embed_dtype, |
| 214 | item.endianness, |
| 215 | ) |
| 216 | for item in sorted(items, key=lambda x: x.index) |
| 217 | ] |
| 218 | |
| 219 | |
| 220 | @lru_cache(maxsize=1) |
| 221 | def get_json_response_cls() -> type[JSONResponse]: |
| 222 | if importlib.util.find_spec("orjson") is not None: |
| 223 | from fastapi.responses import ORJSONResponse |
| 224 | |
| 225 | return ORJSONResponse |
| 226 | logger.warning_once( |
| 227 | "To make v1/embeddings API fast, please install orjson by `pip install orjson`" |
| 228 | ) |
| 229 | return JSONResponse |
| 230 | |
| 231 | |
| 232 | def enable_scoring_api( |
| 233 | supported_tasks: tuple["SupportedTask", ...], |
| 234 | model_config: ModelConfig | None = None, |
| 235 | ) -> bool: |
| 236 | if model_config is None: |
| 237 | return False |
| 238 | |
| 239 | pooling_task = model_config.get_pooling_task(supported_tasks) |
| 240 | if pooling_task in ("embed", "token_embed"): |
| 241 | return True |
| 242 | |
| 243 | if pooling_task == "classify": |
| 244 | num_labels = getattr(model_config.hf_config, "num_labels", 0) |
| 245 | if num_labels != 1: |
| 246 | logger.debug_once("Scoring API is only enabled for num_labels == 1.") |
| 247 | return False |
| 248 | return True |
| 249 | |
| 250 | return False |