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 
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
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

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 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
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,
        )

EmbeddingChatInputRequest

Bases: EmbeddingChatRequest

OpenAI embeddings request with one chat conversation in input.

Source code in vllm/entrypoints/pooling/embed/protocol.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
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 
67
68
69
70
71
72
73
74
75
76
77
78
79
80
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.

Each bit is shifted left by 7 - idx%8, and every 8 bits are packed into one byte.

Source code in vllm/entrypoints/pooling/embed/protocol.py 
293
294
295
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.

    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

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