@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 354 lines 10.4 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Embedding API protocol models for OpenAI and Cohere formats.
4
5 OpenAI: https://platform.openai.com/docs/api-reference/embeddings
6 Cohere: https://docs.cohere.com/reference/embed
7 """
8
9 import builtins
10 import struct
11 import time
12 from collections.abc import Sequence
13 from typing import Annotated, Any, Literal, TypeAlias
14
15 import numpy as np
16 import pybase64 as base64
17 from pydantic import BaseModel, BeforeValidator, Field, model_validator
18
19 from vllm import PoolingParams
20 from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
21 from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel, UsageInfo
22 from vllm.utils import random_uuid
23
24 from ..base.protocol import (
25 ChatRequestMixin,
26 ChatRequestOptionsMixin,
27 CompletionRequestMixin,
28 EmbeddingTokenizeParamsMixin,
29 EmbedRequestMixin,
30 PoolingBasicRequestMixin,
31 reject_removed_pooling_parameters,
32 )
33
34
35 class EmbeddingCompletionRequest(
36 PoolingBasicRequestMixin,
37 CompletionRequestMixin,
38 EmbedRequestMixin,
39 EmbeddingTokenizeParamsMixin,
40 ):
41 def to_pooling_params(self):
42 return PoolingParams(
43 task="embed",
44 dimensions=self.dimensions,
45 use_activation=self.use_activation,
46 )
47
48
49 def _is_chat_message(value: Any) -> bool:
50 return isinstance(value, dict) and isinstance(value.get("role"), str)
51
52
53 def _is_chat_messages(value: Any) -> bool:
54 return (
55 isinstance(value, list)
56 and bool(value)
57 and all(_is_chat_message(item) for item in value)
58 )
59
60
61 def _is_batched_chat_messages(value: Any) -> bool:
62 return (
63 isinstance(value, list)
64 and bool(value)
65 and all(_is_chat_messages(item) for item in value)
66 )
67
68
69 class EmbeddingChatRequest(
70 PoolingBasicRequestMixin,
71 ChatRequestMixin,
72 EmbedRequestMixin,
73 EmbeddingTokenizeParamsMixin,
74 ):
75 """OpenAI embeddings request with one top-level chat conversation."""
76
77 def to_pooling_params(self):
78 return PoolingParams(
79 task="embed",
80 dimensions=self.dimensions,
81 use_activation=self.use_activation,
82 )
83
84
85 class EmbeddingBatchChatRequest(
86 PoolingBasicRequestMixin,
87 ChatRequestOptionsMixin,
88 EmbedRequestMixin,
89 EmbeddingTokenizeParamsMixin,
90 ):
91 """OpenAI embeddings request with batched top-level chat conversations.
92
93 Mirrors ``BatchChatCompletionRequest`` by keeping batched conversations in
94 ``messages`` instead of introducing a separate batch-specific field.
95 """
96
97 messages: Sequence[
98 Annotated[list[ChatCompletionMessageParam], Field(min_length=1)]
99 ] = Field(..., min_length=1)
100
101 def to_pooling_params(self):
102 return PoolingParams(
103 task="embed",
104 dimensions=self.dimensions,
105 use_activation=self.use_activation,
106 )
107
108
109 class EmbeddingChatInputRequest(
110 EmbeddingChatRequest,
111 ):
112 """OpenAI embeddings request with one chat conversation in ``input``."""
113
114 input: list[ChatCompletionMessageParam]
115
116 @model_validator(mode="before")
117 @classmethod
118 def normalize_input_messages(cls, data):
119 if not isinstance(data, dict):
120 return data
121
122 if "messages" in data or "input" not in data:
123 return data
124
125 input_data = data["input"]
126 if not _is_chat_messages(input_data):
127 return data
128
129 normalized = dict(data)
130 normalized["messages"] = input_data
131 return normalized
132
133
134 class EmbeddingBatchChatInputRequest(EmbeddingBatchChatRequest):
135 """OpenAI embeddings request with batched chat conversations in ``input``."""
136
137 input: Sequence[
138 Annotated[list[ChatCompletionMessageParam], Field(min_length=1)]
139 ] = Field(..., min_length=1)
140
141 @model_validator(mode="before")
142 @classmethod
143 def normalize_input_messages(cls, data):
144 if not isinstance(data, dict):
145 return data
146
147 if "messages" in data or "input" not in data:
148 return data
149
150 input_data = data["input"]
151 if not _is_batched_chat_messages(input_data):
152 return data
153
154 normalized = dict(data)
155 normalized["messages"] = input_data
156 return normalized
157
158
159 EmbeddingRequest: TypeAlias = Annotated[
160 EmbeddingCompletionRequest
161 | EmbeddingChatRequest
162 | EmbeddingBatchChatRequest
163 | EmbeddingChatInputRequest
164 | EmbeddingBatchChatInputRequest,
165 BeforeValidator(reject_removed_pooling_parameters),
166 ]
167
168
169 # ---------------------------------------------------------------------------
170 # OpenAI /v1/embeddings — response models
171 # ---------------------------------------------------------------------------
172
173
174 class EmbeddingResponseData(OpenAIBaseModel):
175 index: int
176 object: str = "embedding"
177 embedding: list[float] | str
178
179
180 class EmbeddingResponse(OpenAIBaseModel):
181 id: str = Field(default_factory=lambda: f"embd-{random_uuid()}")
182 object: str = "list"
183 created: int = Field(default_factory=lambda: int(time.time()))
184 model: str | None = None
185 data: list[EmbeddingResponseData]
186 usage: UsageInfo
187
188
189 class EmbeddingBytesResponse(OpenAIBaseModel):
190 content: list[bytes]
191 headers: dict[str, str] | None = None
192 media_type: str = "application/octet-stream"
193
194
195 # ---------------------------------------------------------------------------
196 # Cohere /v2/embed — request models
197 # ---------------------------------------------------------------------------
198
199 CohereEmbeddingType = Literal[
200 "float",
201 "binary",
202 "ubinary",
203 "base64",
204 ]
205 CohereTruncate = Literal["NONE", "START", "END"]
206
207
208 class CohereEmbedContent(BaseModel):
209 type: Literal["text", "image_url"]
210 text: str | None = None
211 image_url: dict[str, str] | None = None
212
213 @model_validator(mode="after")
214 def validate_content_payload(self):
215 if self.type == "text":
216 if self.text is None:
217 raise ValueError("CohereEmbedContent with type='text' requires text")
218 elif not self.image_url or not self.image_url.get("url"):
219 raise ValueError(
220 "CohereEmbedContent with type='image_url' requires image_url.url"
221 )
222 return self
223
224
225 class CohereEmbedInput(BaseModel):
226 content: list[CohereEmbedContent]
227
228
229 class CohereEmbedRequest(BaseModel):
230 model: str | None = None
231 input_type: str | None = None
232 texts: list[str] | None = None
233 images: list[str] | None = None
234 inputs: list[CohereEmbedInput] | None = None
235 output_dimension: int | None = None
236 embedding_types: list[CohereEmbeddingType] | None = None
237 truncate: CohereTruncate = "END"
238 max_tokens: int | None = None
239 priority: int = Field(default=0, ge=-(2**63), le=2**63 - 1)
240
241 @model_validator(mode="after")
242 def validate_input_fields(self):
243 input_fields = (self.texts, self.images, self.inputs)
244 provided_fields = [field for field in input_fields if field is not None]
245 if len(provided_fields) != 1 or not provided_fields[0]:
246 raise ValueError(
247 "Exactly one of texts, images, or inputs must be provided, "
248 "and it must be non-empty"
249 )
250 return self
251
252
253 # ---------------------------------------------------------------------------
254 # Cohere /v2/embed — response models
255 # ---------------------------------------------------------------------------
256
257
258 class CohereApiVersion(BaseModel):
259 version: str = "2"
260
261
262 class CohereBilledUnits(BaseModel):
263 input_tokens: int | None = None
264 image_tokens: int | None = None
265
266
267 class CohereMeta(BaseModel):
268 api_version: CohereApiVersion = Field(default_factory=CohereApiVersion)
269 billed_units: CohereBilledUnits | None = None
270
271
272 class CohereEmbedByTypeEmbeddings(BaseModel):
273 # The field name ``float`` shadows the builtin type, so the annotation
274 # must use ``builtins.float`` to avoid a self-referential type error.
275 float: list[list[builtins.float]] | None = None
276 binary: list[list[int]] | None = None
277 ubinary: list[list[int]] | None = None
278 base64: list[str] | None = None
279
280
281 class CohereEmbedResponse(BaseModel):
282 id: str = Field(default_factory=lambda: f"embd-{random_uuid()}")
283 embeddings: CohereEmbedByTypeEmbeddings
284 texts: list[str] | None = None
285 meta: CohereMeta | None = None
286 response_type: Literal["embeddings_by_type"] = "embeddings_by_type"
287
288
289 # ---------------------------------------------------------------------------
290 # Cohere embedding type conversion helpers
291 # ---------------------------------------------------------------------------
292
293 _UNSIGNED_TO_SIGNED_DIFF = 1 << 7 # 128
294
295
296 def _pack_binary_embeddings(
297 float_embeddings: list[list[float]],
298 signed: bool,
299 ) -> list[list[int]]:
300 """Bit-pack float embeddings: positive -> 1, negative -> 0.
301
302 Bits are packed MSB-first, eight per byte.
303 """
304 if not float_embeddings:
305 return []
306
307 array = np.asarray(float_embeddings, dtype=np.float64)
308 if array.ndim != 2:
309 raise ValueError(
310 f"Expected a 2D batch of embeddings, but got {array.ndim}D input."
311 )
312
313 dim = array.shape[1]
314 if dim % 8 != 0:
315 raise ValueError(
316 "Embedding dimension must be a multiple of 8 for binary "
317 f"embedding types, but got {dim}."
318 )
319
320 packed = np.packbits(array >= 0, axis=-1)
321 if signed:
322 packed = packed.astype(np.int16) - _UNSIGNED_TO_SIGNED_DIFF
323 return packed.tolist()
324
325
326 def _encode_base64_embeddings(
327 float_embeddings: list[list[float]],
328 ) -> list[str]:
329 """Encode float embeddings as base64 (little-endian float32)."""
330 result: list[str] = []
331 for embedding in float_embeddings:
332 buf = struct.pack(f"<{len(embedding)}f", *embedding)
333 result.append(base64.b64encode(buf).decode("utf-8"))
334 return result
335
336
337 def build_typed_embeddings(
338 float_embeddings: list[list[float]],
339 embedding_types: Sequence[str],
340 ) -> CohereEmbedByTypeEmbeddings:
341 """Convert float embeddings to all requested Cohere embedding types."""
342 result = CohereEmbedByTypeEmbeddings()
343
344 for emb_type in embedding_types:
345 if emb_type == "float":
346 result.float = float_embeddings
347 elif emb_type == "binary":
348 result.binary = _pack_binary_embeddings(float_embeddings, signed=True)
349 elif emb_type == "ubinary":
350 result.ubinary = _pack_binary_embeddings(float_embeddings, signed=False)
351 elif emb_type == "base64":
352 result.base64 = _encode_base64_embeddings(float_embeddings)
353
354 return result