chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,354 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Embedding API protocol models for OpenAI and Cohere formats.
|
||||
|
||||
OpenAI: https://platform.openai.com/docs/api-reference/embeddings
|
||||
Cohere: https://docs.cohere.com/reference/embed
|
||||
"""
|
||||
|
||||
import builtins
|
||||
import struct
|
||||
import time
|
||||
from collections.abc import Sequence
|
||||
from typing import Annotated, Any, Literal, TypeAlias
|
||||
|
||||
import pybase64 as base64
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
|
||||
from vllm import PoolingParams
|
||||
from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
|
||||
from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo
|
||||
from vllm.utils import random_uuid
|
||||
|
||||
from ..base.protocol import (
|
||||
ChatRequestMixin,
|
||||
ChatRequestOptionsMixin,
|
||||
CompletionRequestMixin,
|
||||
EmbeddingTokenizeParamsMixin,
|
||||
EmbedRequestMixin,
|
||||
PoolingBasicRequestMixin,
|
||||
)
|
||||
|
||||
|
||||
class EmbeddingCompletionRequest(
|
||||
PoolingBasicRequestMixin,
|
||||
CompletionRequestMixin,
|
||||
EmbedRequestMixin,
|
||||
EmbeddingTokenizeParamsMixin,
|
||||
):
|
||||
def to_pooling_params(self):
|
||||
return PoolingParams(
|
||||
task="embed",
|
||||
dimensions=self.dimensions,
|
||||
use_activation=self.use_activation,
|
||||
)
|
||||
|
||||
|
||||
def _is_chat_message(value: Any) -> bool:
|
||||
return isinstance(value, dict) and isinstance(value.get("role"), str)
|
||||
|
||||
|
||||
def _is_chat_messages(value: Any) -> bool:
|
||||
return (
|
||||
isinstance(value, list)
|
||||
and bool(value)
|
||||
and all(_is_chat_message(item) for item in value)
|
||||
)
|
||||
|
||||
|
||||
def _is_batched_chat_messages(value: Any) -> bool:
|
||||
return (
|
||||
isinstance(value, list)
|
||||
and bool(value)
|
||||
and all(_is_chat_messages(item) for item in value)
|
||||
)
|
||||
|
||||
|
||||
class EmbeddingChatRequest(
|
||||
PoolingBasicRequestMixin,
|
||||
ChatRequestMixin,
|
||||
EmbedRequestMixin,
|
||||
EmbeddingTokenizeParamsMixin,
|
||||
):
|
||||
"""OpenAI embeddings request with one top-level chat conversation."""
|
||||
|
||||
def to_pooling_params(self):
|
||||
return PoolingParams(
|
||||
task="embed",
|
||||
dimensions=self.dimensions,
|
||||
use_activation=self.use_activation,
|
||||
)
|
||||
|
||||
|
||||
class EmbeddingBatchChatRequest(
|
||||
PoolingBasicRequestMixin,
|
||||
ChatRequestOptionsMixin,
|
||||
EmbedRequestMixin,
|
||||
EmbeddingTokenizeParamsMixin,
|
||||
):
|
||||
"""OpenAI embeddings request with batched top-level chat conversations.
|
||||
|
||||
Mirrors ``BatchChatCompletionRequest`` by keeping batched conversations in
|
||||
``messages`` instead of introducing a separate batch-specific field.
|
||||
"""
|
||||
|
||||
messages: list[Annotated[list[ChatCompletionMessageParam], Field(min_length=1)]] = (
|
||||
Field(..., min_length=1)
|
||||
)
|
||||
|
||||
def to_pooling_params(self):
|
||||
return PoolingParams(
|
||||
task="embed",
|
||||
dimensions=self.dimensions,
|
||||
use_activation=self.use_activation,
|
||||
)
|
||||
|
||||
|
||||
class EmbeddingChatInputRequest(
|
||||
EmbeddingChatRequest,
|
||||
):
|
||||
"""OpenAI embeddings request with one chat conversation in ``input``."""
|
||||
|
||||
input: list[ChatCompletionMessageParam]
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def normalize_input_messages(cls, data):
|
||||
if not isinstance(data, dict):
|
||||
return data
|
||||
|
||||
if "messages" in data or "input" not in data:
|
||||
return data
|
||||
|
||||
input_data = data["input"]
|
||||
if not _is_chat_messages(input_data):
|
||||
return data
|
||||
|
||||
normalized = dict(data)
|
||||
normalized["messages"] = input_data
|
||||
return normalized
|
||||
|
||||
|
||||
class EmbeddingBatchChatInputRequest(EmbeddingBatchChatRequest):
|
||||
"""OpenAI embeddings request with batched chat conversations in ``input``."""
|
||||
|
||||
input: list[Annotated[list[ChatCompletionMessageParam], Field(min_length=1)]] = (
|
||||
Field(..., min_length=1)
|
||||
)
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def normalize_input_messages(cls, data):
|
||||
if not isinstance(data, dict):
|
||||
return data
|
||||
|
||||
if "messages" in data or "input" not in data:
|
||||
return data
|
||||
|
||||
input_data = data["input"]
|
||||
if not _is_batched_chat_messages(input_data):
|
||||
return data
|
||||
|
||||
normalized = dict(data)
|
||||
normalized["messages"] = input_data
|
||||
return normalized
|
||||
|
||||
|
||||
EmbeddingRequest: TypeAlias = (
|
||||
EmbeddingCompletionRequest
|
||||
| EmbeddingChatRequest
|
||||
| EmbeddingBatchChatRequest
|
||||
| EmbeddingChatInputRequest
|
||||
| EmbeddingBatchChatInputRequest
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# OpenAI /v1/embeddings — response models
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class EmbeddingResponseData(OpenAIBaseModel):
|
||||
index: int
|
||||
object: str = "embedding"
|
||||
embedding: list[float] | str
|
||||
|
||||
|
||||
class EmbeddingResponse(OpenAIBaseModel):
|
||||
id: str = Field(default_factory=lambda: f"embd-{random_uuid()}")
|
||||
object: str = "list"
|
||||
created: int = Field(default_factory=lambda: int(time.time()))
|
||||
model: str | None = None
|
||||
data: list[EmbeddingResponseData]
|
||||
usage: UsageInfo
|
||||
|
||||
|
||||
class EmbeddingBytesResponse(OpenAIBaseModel):
|
||||
content: list[bytes]
|
||||
headers: dict[str, str] | None = None
|
||||
media_type: str = "application/octet-stream"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cohere /v2/embed — request models
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
CohereEmbeddingType = Literal[
|
||||
"float",
|
||||
"binary",
|
||||
"ubinary",
|
||||
"base64",
|
||||
]
|
||||
CohereTruncate = Literal["NONE", "START", "END"]
|
||||
|
||||
|
||||
class CohereEmbedContent(BaseModel):
|
||||
type: Literal["text", "image_url"]
|
||||
text: str | None = None
|
||||
image_url: dict[str, str] | None = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_content_payload(self):
|
||||
if self.type == "text":
|
||||
if self.text is None:
|
||||
raise ValueError("CohereEmbedContent with type='text' requires text")
|
||||
elif not self.image_url or not self.image_url.get("url"):
|
||||
raise ValueError(
|
||||
"CohereEmbedContent with type='image_url' requires image_url.url"
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
class CohereEmbedInput(BaseModel):
|
||||
content: list[CohereEmbedContent]
|
||||
|
||||
|
||||
class CohereEmbedRequest(BaseModel):
|
||||
model: str | None = None
|
||||
input_type: str | None = None
|
||||
texts: list[str] | None = None
|
||||
images: list[str] | None = None
|
||||
inputs: list[CohereEmbedInput] | None = None
|
||||
output_dimension: int | None = None
|
||||
embedding_types: list[CohereEmbeddingType] | None = None
|
||||
truncate: CohereTruncate = "END"
|
||||
max_tokens: int | None = None
|
||||
priority: int = 0
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_input_fields(self):
|
||||
input_fields = (self.texts, self.images, self.inputs)
|
||||
provided_fields = [field for field in input_fields if field is not None]
|
||||
if len(provided_fields) != 1 or not provided_fields[0]:
|
||||
raise ValueError(
|
||||
"Exactly one of texts, images, or inputs must be provided, "
|
||||
"and it must be non-empty"
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cohere /v2/embed — response models
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class CohereApiVersion(BaseModel):
|
||||
version: str = "2"
|
||||
|
||||
|
||||
class CohereBilledUnits(BaseModel):
|
||||
input_tokens: int | None = None
|
||||
image_tokens: int | None = None
|
||||
|
||||
|
||||
class CohereMeta(BaseModel):
|
||||
api_version: CohereApiVersion = Field(default_factory=CohereApiVersion)
|
||||
billed_units: CohereBilledUnits | None = None
|
||||
|
||||
|
||||
class CohereEmbedByTypeEmbeddings(BaseModel):
|
||||
# The field name ``float`` shadows the builtin type, so the annotation
|
||||
# must use ``builtins.float`` to avoid a self-referential type error.
|
||||
float: list[list[builtins.float]] | None = None
|
||||
binary: list[list[int]] | None = None
|
||||
ubinary: list[list[int]] | None = None
|
||||
base64: list[str] | None = None
|
||||
|
||||
|
||||
class CohereEmbedResponse(BaseModel):
|
||||
id: str = Field(default_factory=lambda: f"embd-{random_uuid()}")
|
||||
embeddings: CohereEmbedByTypeEmbeddings
|
||||
texts: list[str] | None = None
|
||||
meta: CohereMeta | None = None
|
||||
response_type: Literal["embeddings_by_type"] = "embeddings_by_type"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cohere embedding type conversion helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_UNSIGNED_TO_SIGNED_DIFF = 1 << 7 # 128
|
||||
|
||||
|
||||
def _pack_binary_embeddings(
|
||||
float_embeddings: list[list[float]],
|
||||
signed: bool,
|
||||
) -> list[list[int]]:
|
||||
"""Bit-pack float embeddings: positive -> 1, negative -> 0.
|
||||
|
||||
Each bit is shifted left by ``7 - idx%8``, and every 8 bits are packed
|
||||
into one byte.
|
||||
"""
|
||||
result: list[list[int]] = []
|
||||
for embedding in float_embeddings:
|
||||
dim = len(embedding)
|
||||
if dim % 8 != 0:
|
||||
raise ValueError(
|
||||
"Embedding dimension must be a multiple of 8 for binary "
|
||||
f"embedding types, but got {dim}."
|
||||
)
|
||||
packed_len = dim // 8
|
||||
packed: list[int] = []
|
||||
byte_val = 0
|
||||
for idx, value in enumerate(embedding):
|
||||
bit = 1 if value >= 0 else 0
|
||||
byte_val += bit << (7 - idx % 8)
|
||||
if (idx + 1) % 8 == 0:
|
||||
if signed:
|
||||
byte_val -= _UNSIGNED_TO_SIGNED_DIFF
|
||||
packed.append(byte_val)
|
||||
byte_val = 0
|
||||
assert len(packed) == packed_len
|
||||
result.append(packed)
|
||||
return result
|
||||
|
||||
|
||||
def _encode_base64_embeddings(
|
||||
float_embeddings: list[list[float]],
|
||||
) -> list[str]:
|
||||
"""Encode float embeddings as base64 (little-endian float32)."""
|
||||
result: list[str] = []
|
||||
for embedding in float_embeddings:
|
||||
buf = struct.pack(f"<{len(embedding)}f", *embedding)
|
||||
result.append(base64.b64encode(buf).decode("utf-8"))
|
||||
return result
|
||||
|
||||
|
||||
def build_typed_embeddings(
|
||||
float_embeddings: list[list[float]],
|
||||
embedding_types: Sequence[str],
|
||||
) -> CohereEmbedByTypeEmbeddings:
|
||||
"""Convert float embeddings to all requested Cohere embedding types."""
|
||||
result = CohereEmbedByTypeEmbeddings()
|
||||
|
||||
for emb_type in embedding_types:
|
||||
if emb_type == "float":
|
||||
result.float = float_embeddings
|
||||
elif emb_type == "binary":
|
||||
result.binary = _pack_binary_embeddings(float_embeddings, signed=True)
|
||||
elif emb_type == "ubinary":
|
||||
result.ubinary = _pack_binary_embeddings(float_embeddings, signed=False)
|
||||
elif emb_type == "base64":
|
||||
result.base64 = _encode_base64_embeddings(float_embeddings)
|
||||
|
||||
return result
|
||||
Reference in New Issue
Block a user