Skip to content

vllm.entrypoints.pooling.embed.protocol

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

Classes:

Functions:

EmbeddingBatchChatInputRequest

Bases: EmbeddingBatchChatRequest

OpenAI embeddings request with batched chat conversations in input.

Source code in vllm/entrypoints/pooling/embed/protocol.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
class EmbeddingBatchChatInputRequest(EmbeddingBatchChatRequest):
    """OpenAI embeddings request with batched chat conversations in ``input``."""

    input: Sequence[
        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

EmbeddingBatchChatRequest

Bases: 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.

Source code in vllm/entrypoints/pooling/embed/protocol.py 
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
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: Sequence[
        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,
        )

EmbeddingChatInputRequest

Bases: EmbeddingChatRequest

OpenAI embeddings request with one chat conversation in input.

Source code in vllm/entrypoints/pooling/embed/protocol.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
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

EmbeddingChatRequest

Bases: PoolingBasicRequestMixin, ChatRequestMixin, EmbedRequestMixin, EmbeddingTokenizeParamsMixin

OpenAI embeddings request with one top-level chat conversation.

Source code in vllm/entrypoints/pooling/embed/protocol.py 
69
70
71
72
73
74
75
76
77
78
79
80
81
82
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,
        )

_encode_base64_embeddings(float_embeddings)

Encode float embeddings as base64 (little-endian float32).

Source code in vllm/entrypoints/pooling/embed/protocol.py 
326
327
328
329
330
331
332
333
334
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

_pack_binary_embeddings(float_embeddings, signed)

Bit-pack float embeddings: positive -> 1, negative -> 0.

Bits are packed MSB-first, eight per byte.

Source code in vllm/entrypoints/pooling/embed/protocol.py 
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
def _pack_binary_embeddings(
    float_embeddings: list[list[float]],
    signed: bool,
) -> list[list[int]]:
    """Bit-pack float embeddings: positive -> 1, negative -> 0.

    Bits are packed MSB-first, eight per byte.
    """
    if not float_embeddings:
        return []

    array = np.asarray(float_embeddings, dtype=np.float64)
    if array.ndim != 2:
        raise ValueError(
            f"Expected a 2D batch of embeddings, but got {array.ndim}D input."
        )

    dim = array.shape[1]
    if dim % 8 != 0:
        raise ValueError(
            "Embedding dimension must be a multiple of 8 for binary "
            f"embedding types, but got {dim}."
        )

    packed = np.packbits(array >= 0, axis=-1)
    if signed:
        packed = packed.astype(np.int16) - _UNSIGNED_TO_SIGNED_DIFF
    return packed.tolist()

build_typed_embeddings(float_embeddings, embedding_types)

Convert float embeddings to all requested Cohere embedding types.

Source code in vllm/entrypoints/pooling/embed/protocol.py 
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
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