A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Frontend][3/N] Improve all pooling task | Support binary embedding response (#27066)
Signed-off-by: wang.yuqi <noooop@126.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
wang.yuqi committed
Oct 22, 2025 at 18:38 UTC
1f633b86329b1657b08798b4c3a8c6dd46334b74
12 files changed
+693
-232
examples/online_serving/pooling/README.md
+8
-2
@@ -6,10 +6,16 @@
6
python examples/online_serving/pooling/cohere_rerank_client.py
7
```
8
9
-## Embedding embed_dtype usage
9
+## Embedding requests base64 encoding_format usage
10
11
```bash
12
-python examples/online_serving/pooling/embedding_embed_dtype_client.py
12
+python examples/online_serving/pooling/embedding_requests_base64_client.py
13
+```
14
+
15
+## Embedding requests bytes encoding_format usage
16
+
17
+```bash
18
+python examples/online_serving/pooling/embedding_requests_bytes_client.py
19
```
20
21
## Jinaai rerank usage
examples/online_serving/pooling/embedding_requests_base64_client.py
renamed
+24
-19
@@ -12,7 +12,11 @@ import base64
12
import requests
13
import torch
14
15
-from vllm.entrypoints.openai.protocol import EMBED_DTYPE_TO_TORCH_DTYPE
15
+from vllm.utils.serial_utils import (
16
+ EMBED_DTYPE_TO_TORCH_DTYPE,
17
+ ENDIANNESS,
18
+ binary2tensor,
19
+)
20
21
22
def post_http_request(prompt: dict, api_url: str) -> requests.Response:
@@ -34,24 +38,25 @@ def main(args):
38
api_url = f"http://{args.host}:{args.port}/v1/embeddings"
39
model_name = args.model
40
37
- for embed_dtype, torch_dtype in EMBED_DTYPE_TO_TORCH_DTYPE.items():
38
- prompt = {
39
- "model": model_name,
40
- "input": "vLLM is great!",
41
- "encoding_format": "base64",
42
- "embed_dtype": embed_dtype,
43
- }
44
- response = post_http_request(prompt=prompt, api_url=api_url)
45
-
46
- embedding = []
47
- for data in response.json()["data"]:
48
- embedding.append(
49
- torch.frombuffer(
50
- base64.b64decode(data["embedding"]), dtype=torch_dtype
51
- ).to(torch.float32)
52
- )
53
- embedding = torch.cat(embedding)
54
- print(embed_dtype, embedding.shape)
41
+ # The OpenAI client does not support the embed_dtype and endianness parameters.
42
+ for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE:
43
+ for endianness in ENDIANNESS:
44
+ prompt = {
45
+ "model": model_name,
46
+ "input": "vLLM is great!",
47
+ "encoding_format": "base64",
48
+ "embed_dtype": embed_dtype,
49
+ "endianness": endianness,
50
+ }
51
+ response = post_http_request(prompt=prompt, api_url=api_url)
52
+
53
+ embedding = []
54
+ for data in response.json()["data"]:
55
+ binary = base64.b64decode(data["embedding"])
56
+ tensor = binary2tensor(binary, (-1,), embed_dtype, endianness)
57
+ embedding.append(tensor.to(torch.float32))
58
+ embedding = torch.cat(embedding)
59
+ print(embed_dtype, endianness, embedding.shape)
60
61
62
if __name__ == "__main__":
examples/online_serving/pooling/embedding_requests_bytes_client.py
new
+66
@@ -0,0 +1,66 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Example Python client for embedding API using vLLM API server
4
+NOTE:
5
+ start a supported embeddings model server with `vllm serve`, e.g.
6
+ vllm serve intfloat/e5-small
7
+"""
8
+
9
+import argparse
10
+import json
11
+
12
+import requests
13
+import torch
14
+
15
+from vllm.utils.serial_utils import (
16
+ EMBED_DTYPE_TO_TORCH_DTYPE,
17
+ ENDIANNESS,
18
+ MetadataItem,
19
+ decode_pooling_output,
20
+)
21
+
22
+
23
+def post_http_request(prompt: dict, api_url: str) -> requests.Response:
24
+ headers = {"User-Agent": "Test Client"}
25
+ response = requests.post(api_url, headers=headers, json=prompt)
26
+ return response
27
+
28
+
29
+def parse_args():
30
+ parser = argparse.ArgumentParser()
31
+ parser.add_argument("--host", type=str, default="localhost")
32
+ parser.add_argument("--port", type=int, default=8000)
33
+ parser.add_argument("--model", type=str, default="intfloat/e5-small")
34
+
35
+ return parser.parse_args()
36
+
37
+
38
+def main(args):
39
+ api_url = f"http://{args.host}:{args.port}/v1/embeddings"
40
+ model_name = args.model
41
+
42
+ # The OpenAI client does not support the bytes encoding_format.
43
+ # The OpenAI client does not support the embed_dtype and endianness parameters.
44
+ for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE:
45
+ for endianness in ENDIANNESS:
46
+ prompt = {
47
+ "model": model_name,
48
+ "input": "vLLM is great!",
49
+ "encoding_format": "bytes",
50
+ "embed_dtype": embed_dtype,
51
+ "endianness": endianness,
52
+ }
53
+ response = post_http_request(prompt=prompt, api_url=api_url)
54
+ metadata = json.loads(response.headers["metadata"])
55
+ body = response.content
56
+ items = [MetadataItem(**x) for x in metadata["data"]]
57
+
58
+ embedding = decode_pooling_output(items=items, body=body)
59
+ embedding = [x.to(torch.float32) for x in embedding]
60
+ embedding = torch.cat(embedding)
61
+ print(embed_dtype, endianness, embedding.shape)
62
+
63
+
64
+if __name__ == "__main__":
65
+ args = parse_args()
66
+ main(args)
tests/entrypoints/pooling/openai/test_embedding.py
+82
-34
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import base64
5
+import json
6
7
import numpy as np
8
import openai
@@ -15,11 +16,17 @@ from tests.models.language.pooling.embed_utils import run_embedding_correctness_
16
from tests.models.utils import check_embeddings_close
17
from tests.utils import RemoteOpenAIServer
18
from vllm.entrypoints.openai.protocol import (
18
- EMBED_DTYPE_TO_TORCH_DTYPE,
19
EmbeddingResponse,
20
PoolingResponse,
21
)
22
from vllm.transformers_utils.tokenizer import get_tokenizer
23
+from vllm.utils.serial_utils import (
24
+ EMBED_DTYPE_TO_TORCH_DTYPE,
25
+ ENDIANNESS,
26
+ MetadataItem,
27
+ binary2tensor,
28
+ decode_pooling_output,
29
+)
30
31
MODEL_NAME = "intfloat/multilingual-e5-small"
32
DUMMY_CHAT_TEMPLATE = """{% for message in messages %}{{message['role'] + ': ' + message['content'] + '\\n'}}{% endfor %}""" # noqa: E501
@@ -250,8 +257,8 @@ async def test_batch_base64_embedding(
257
258
@pytest.mark.asyncio
259
@pytest.mark.parametrize("model_name", [MODEL_NAME])
253
-async def test_base64_embed_dtype(
254
- hf_model, server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
260
+async def test_base64_embed_dtype_and_endianness(
261
+ server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
262
):
263
input_texts = [
264
"The best thing about vLLM is that it supports many different models",
@@ -262,44 +269,86 @@ async def test_base64_embed_dtype(
269
)
270
float_data = [d.embedding for d in responses_float.data]
271
265
- for embed_dtype, torch_dtype in EMBED_DTYPE_TO_TORCH_DTYPE.items():
266
- responses_base64 = requests.post(
267
- server.url_for("/v1/embeddings"),
268
- json={
269
- "model": model_name,
270
- "input": input_texts,
271
- "encoding_format": "base64",
272
- "embed_dtype": embed_dtype,
273
- },
274
- )
275
-
276
- base64_data = []
277
- for data in responses_base64.json()["data"]:
278
- base64_data.append(
279
- torch.frombuffer(base64.b64decode(data["embedding"]), dtype=torch_dtype)
280
- .to(torch.float32)
281
- .tolist()
272
+ for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE:
273
+ for endianness in ENDIANNESS:
274
+ responses_base64 = requests.post(
275
+ server.url_for("/v1/embeddings"),
276
+ json={
277
+ "model": model_name,
278
+ "input": input_texts,
279
+ "encoding_format": "base64",
280
+ "embed_dtype": embed_dtype,
281
+ "endianness": endianness,
282
+ },
283
)
284
284
- check_embeddings_close(
285
- embeddings_0_lst=float_data,
286
- embeddings_1_lst=base64_data,
287
- name_0="float_data",
288
- name_1="base64_data",
289
- tol=1e-2,
290
- )
285
+ base64_data = []
286
+ for data in responses_base64.json()["data"]:
287
+ binary = base64.b64decode(data["embedding"])
288
+ tensor = binary2tensor(binary, (-1,), embed_dtype, endianness)
289
+ base64_data.append(tensor.to(torch.float32).tolist())
290
+
291
+ check_embeddings_close(
292
+ embeddings_0_lst=float_data,
293
+ embeddings_1_lst=base64_data,
294
+ name_0="float_data",
295
+ name_1="base64_data",
296
+ tol=1e-2,
297
+ )
298
299
300
@pytest.mark.asyncio
301
@pytest.mark.parametrize("model_name", [MODEL_NAME])
295
-async def test_base64_embed_dtype_not_supported(
296
- hf_model, server: RemoteOpenAIServer, model_name: str
302
+async def test_bytes_embed_dtype_and_endianness(
303
+ server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
304
):
305
input_texts = [
306
"The best thing about vLLM is that it supports many different models",
307
]
308
302
- bad_embed_dtype = "bad_embed_dtype"
309
+ responses_float = await client.embeddings.create(
310
+ input=input_texts, model=model_name, encoding_format="float"
311
+ )
312
+ float_data = [d.embedding for d in responses_float.data]
313
+
314
+ for embed_dtype in list(EMBED_DTYPE_TO_TORCH_DTYPE.keys()):
315
+ for endianness in ENDIANNESS:
316
+ responses_bytes = requests.post(
317
+ server.url_for("/v1/embeddings"),
318
+ json={
319
+ "model": model_name,
320
+ "input": input_texts,
321
+ "encoding_format": "bytes",
322
+ "embed_dtype": embed_dtype,
323
+ "endianness": endianness,
324
+ },
325
+ )
326
+
327
+ metadata = json.loads(responses_bytes.headers["metadata"])
328
+ body = responses_bytes.content
329
+ items = [MetadataItem(**x) for x in metadata["data"]]
330
+
331
+ bytes_data = decode_pooling_output(items=items, body=body)
332
+ bytes_data = [x.to(torch.float32).tolist() for x in bytes_data]
333
+
334
+ check_embeddings_close(
335
+ embeddings_0_lst=float_data,
336
+ embeddings_1_lst=bytes_data,
337
+ name_0="float_data",
338
+ name_1="bytes_data",
339
+ tol=1e-2,
340
+ )
341
+
342
+
343
+@pytest.mark.asyncio
344
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
345
+@pytest.mark.parametrize("param_name", ["encoding_format", "embed_dtype", "endianness"])
346
+async def test_params_not_supported(
347
+ server: RemoteOpenAIServer, model_name: str, param_name: str
348
+):
349
+ input_texts = [
350
+ "The best thing about vLLM is that it supports many different models",
351
+ ]
352
353
responses_base64 = requests.post(
354
server.url_for("/v1/embeddings"),
@@ -307,14 +356,13 @@ async def test_base64_embed_dtype_not_supported(
356
"model": model_name,
357
"input": input_texts,
358
"encoding_format": "base64",
310
- "embed_dtype": bad_embed_dtype,
359
+ param_name: f"bad_{param_name}",
360
},
361
)
362
363
assert responses_base64.status_code == 400
315
- assert responses_base64.json()["error"]["message"].startswith(
316
- f"embed_dtype={bad_embed_dtype!r} is not supported."
317
- )
364
+ assert "literal_error" in responses_base64.json()["error"]["message"]
365
+ assert f"bad_{param_name}" in responses_base64.json()["error"]["message"]
366
367
368
@pytest.mark.asyncio
tests/entrypoints/pooling/openai/test_pooling.py
+90
-32
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import base64
5
+import json
6
7
import numpy as np
8
import pytest
@@ -10,8 +11,15 @@ import torch
11
12
from tests.models.utils import check_embeddings_close
13
from tests.utils import RemoteOpenAIServer
13
-from vllm.entrypoints.openai.protocol import EMBED_DTYPE_TO_TORCH_DTYPE, PoolingResponse
14
+from vllm.entrypoints.openai.protocol import PoolingResponse
15
from vllm.transformers_utils.tokenizer import get_tokenizer
16
+from vllm.utils.serial_utils import (
17
+ EMBED_DTYPE_TO_TORCH_DTYPE,
18
+ ENDIANNESS,
19
+ MetadataItem,
20
+ binary2tensor,
21
+ decode_pooling_output,
22
+)
23
24
MODEL_NAME = "internlm/internlm2-1_8b-reward"
25
DUMMY_CHAT_TEMPLATE = """{% for message in messages %}{{message['role'] + ': ' + message['content'] + '\\n'}}{% endfor %}""" # noqa: E501
@@ -251,7 +259,9 @@ async def test_batch_base64_pooling(server: RemoteOpenAIServer, model_name: str)
259
260
@pytest.mark.asyncio
261
@pytest.mark.parametrize("model_name", [MODEL_NAME])
254
-async def test_base64_embed_dtype(server: RemoteOpenAIServer, model_name: str):
262
+async def test_base64_embed_dtype_and_endianness(
263
+ server: RemoteOpenAIServer, model_name: str
264
+):
265
input_texts = [
266
"The best thing about vLLM is that it supports many different models",
267
]
@@ -268,44 +278,93 @@ async def test_base64_embed_dtype(server: RemoteOpenAIServer, model_name: str):
278
responses_float = PoolingResponse.model_validate(float_response.json())
279
float_data = [np.array(d.data).squeeze(-1).tolist() for d in responses_float.data]
280
271
- for embed_dtype, torch_dtype in EMBED_DTYPE_TO_TORCH_DTYPE.items():
272
- responses_base64 = requests.post(
273
- url,
274
- json={
275
- "model": model_name,
276
- "input": input_texts,
277
- "encoding_format": "base64",
278
- "embed_dtype": embed_dtype,
279
- },
280
- )
281
-
282
- base64_data = []
283
- for data in responses_base64.json()["data"]:
284
- base64_data.append(
285
- torch.frombuffer(base64.b64decode(data["data"]), dtype=torch_dtype)
286
- .to(torch.float32)
287
- .tolist()
281
+ for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE:
282
+ for endianness in ENDIANNESS:
283
+ responses_base64 = requests.post(
284
+ url,
285
+ json={
286
+ "model": model_name,
287
+ "input": input_texts,
288
+ "encoding_format": "base64",
289
+ "embed_dtype": embed_dtype,
290
+ "endianness": endianness,
291
+ },
292
)
293
290
- check_embeddings_close(
291
- embeddings_0_lst=float_data,
292
- embeddings_1_lst=base64_data,
293
- name_0="float_data",
294
- name_1="base64_data",
295
- tol=1e-2,
296
- )
294
+ base64_data = []
295
+ for data in responses_base64.json()["data"]:
296
+ binary = base64.b64decode(data["data"])
297
+ tensor = binary2tensor(binary, (-1,), embed_dtype, endianness)
298
+ base64_data.append(tensor.to(torch.float32).tolist())
299
+
300
+ check_embeddings_close(
301
+ embeddings_0_lst=float_data,
302
+ embeddings_1_lst=base64_data,
303
+ name_0="float_data",
304
+ name_1="base64_data",
305
+ tol=1e-2,
306
+ )
307
308
309
@pytest.mark.asyncio
310
@pytest.mark.parametrize("model_name", [MODEL_NAME])
301
-async def test_base64_embed_dtype_not_supported(
311
+async def test_bytes_embed_dtype_and_endianness(
312
server: RemoteOpenAIServer, model_name: str
313
):
314
input_texts = [
315
"The best thing about vLLM is that it supports many different models",
316
]
317
308
- bad_embed_dtype = "bad_embed_dtype"
318
+ url = server.url_for("pooling")
319
+ float_response = requests.post(
320
+ url,
321
+ json={
322
+ "model": model_name,
323
+ "input": input_texts,
324
+ "encoding_format": "float",
325
+ },
326
+ )
327
+ responses_float = PoolingResponse.model_validate(float_response.json())
328
+ float_data = [np.array(d.data).squeeze(-1).tolist() for d in responses_float.data]
329
+
330
+ for embed_dtype in list(EMBED_DTYPE_TO_TORCH_DTYPE.keys()):
331
+ for endianness in ENDIANNESS:
332
+ responses_bytes = requests.post(
333
+ url,
334
+ json={
335
+ "model": model_name,
336
+ "input": input_texts,
337
+ "encoding_format": "bytes",
338
+ "embed_dtype": embed_dtype,
339
+ "endianness": endianness,
340
+ },
341
+ )
342
+
343
+ metadata = json.loads(responses_bytes.headers["metadata"])
344
+ body = responses_bytes.content
345
+ items = [MetadataItem(**x) for x in metadata["data"]]
346
+
347
+ bytes_data = decode_pooling_output(items=items, body=body)
348
+ bytes_data = [x.to(torch.float32).view(-1).tolist() for x in bytes_data]
349
+
350
+ check_embeddings_close(
351
+ embeddings_0_lst=float_data,
352
+ embeddings_1_lst=bytes_data,
353
+ name_0="float_data",
354
+ name_1="bytes_data",
355
+ tol=1e-2,
356
+ )
357
+
358
+
359
+@pytest.mark.asyncio
360
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
361
+@pytest.mark.parametrize("param_name", ["encoding_format", "embed_dtype", "endianness"])
362
+async def test_params_not_supported(
363
+ server: RemoteOpenAIServer, model_name: str, param_name: str
364
+):
365
+ input_texts = [
366
+ "The best thing about vLLM is that it supports many different models",
367
+ ]
368
369
responses_base64 = requests.post(
370
server.url_for("pooling"),
@@ -313,14 +372,13 @@ async def test_base64_embed_dtype_not_supported(
372
"model": model_name,
373
"input": input_texts,
374
"encoding_format": "base64",
316
- "embed_dtype": bad_embed_dtype,
375
+ param_name: f"bad_{param_name}",
376
},
377
)
378
379
assert responses_base64.status_code == 400
321
- assert responses_base64.json()["error"]["message"].startswith(
322
- f"embed_dtype={bad_embed_dtype!r} is not supported."
323
- )
380
+ assert "literal_error" in responses_base64.json()["error"]["message"]
381
+ assert f"bad_{param_name}" in responses_base64.json()["error"]["message"]
382
383
384
@pytest.mark.asyncio
tests/utils_/test_serial_utils.py
new
+40
@@ -0,0 +1,40 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+import torch
5
+
6
+from tests.models.utils import check_embeddings_close
7
+from vllm.utils.serial_utils import (
8
+ EMBED_DTYPE_TO_TORCH_DTYPE,
9
+ ENDIANNESS,
10
+ binary2tensor,
11
+ tensor2binary,
12
+)
13
+
14
+
15
+@pytest.mark.parametrize("endianness", ENDIANNESS)
16
+@pytest.mark.parametrize("embed_dtype", EMBED_DTYPE_TO_TORCH_DTYPE.keys())
17
+@torch.inference_mode
18
+def test_encode_and_decode(embed_dtype: str, endianness: str):
19
+ for i in range(10):
20
+ tensor = torch.rand(2, 3, 5, 7, 11, 13, device="cpu", dtype=torch.float32)
21
+ shape = tensor.shape
22
+ binary = tensor2binary(tensor, embed_dtype, endianness)
23
+ new_tensor = binary2tensor(binary, shape, embed_dtype, endianness).to(
24
+ torch.float32
25
+ )
26
+
27
+ if embed_dtype in ["float32", "float16"]:
28
+ torch.testing.assert_close(tensor, new_tensor, atol=0.001, rtol=0.001)
29
+ elif embed_dtype == "bfloat16":
30
+ torch.testing.assert_close(tensor, new_tensor, atol=0.01, rtol=0.01)
31
+ else: # for fp8
32
+ torch.testing.assert_close(tensor, new_tensor, atol=0.1, rtol=0.1)
33
+
34
+ check_embeddings_close(
35
+ embeddings_0_lst=tensor.view(1, -1),
36
+ embeddings_1_lst=new_tensor.view(1, -1),
37
+ name_0="gt",
38
+ name_1="new",
39
+ tol=1e-2,
40
+ )
vllm/entrypoints/openai/api_server.py
+18
-1
@@ -58,12 +58,14 @@ from vllm.entrypoints.openai.protocol import (
58
CompletionResponse,
59
DetokenizeRequest,
60
DetokenizeResponse,
61
+ EmbeddingBytesResponse,
62
EmbeddingRequest,
63
EmbeddingResponse,
64
ErrorInfo,
65
ErrorResponse,
66
IOProcessorResponse,
67
LoadLoRAAdapterRequest,
68
+ PoolingBytesResponse,
69
PoolingRequest,
70
PoolingResponse,
71
RerankRequest,
@@ -681,7 +683,10 @@ async def create_completion(request: CompletionRequest, raw_request: Request):
683
)
684
@with_cancellation
685
@load_aware_call
684
-async def create_embedding(request: EmbeddingRequest, raw_request: Request):
686
+async def create_embedding(
687
+ request: EmbeddingRequest,
688
+ raw_request: Request,
689
+):
690
handler = embedding(raw_request)
691
if handler is None:
692
return base(raw_request).create_error_response(
@@ -701,6 +706,12 @@ async def create_embedding(request: EmbeddingRequest, raw_request: Request):
706
)
707
elif isinstance(generator, EmbeddingResponse):
708
return JSONResponse(content=generator.model_dump())
709
+ elif isinstance(generator, EmbeddingBytesResponse):
710
+ return StreamingResponse(
711
+ content=generator.body,
712
+ headers={"metadata": generator.metadata},
713
+ media_type=generator.media_type,
714
+ )
715
716
assert_never(generator)
717
@@ -733,6 +744,12 @@ async def create_pooling(request: PoolingRequest, raw_request: Request):
744
)
745
elif isinstance(generator, (PoolingResponse, IOProcessorResponse)):
746
return JSONResponse(content=generator.model_dump())
747
+ elif isinstance(generator, PoolingBytesResponse):
748
+ return StreamingResponse(
749
+ content=generator.body,
750
+ headers={"metadata": generator.metadata},
751
+ media_type=generator.media_type,
752
+ )
753
754
assert_never(generator)
755
vllm/entrypoints/openai/protocol.py
+57
-23
@@ -48,6 +48,12 @@ from openai.types.responses.response_reasoning_item import (
48
Content as ResponseReasoningTextContent,
49
)
50
51
+from vllm.utils.serial_utils import (
52
+ EmbedDType,
53
+ EncodingFormat,
54
+ Endianness,
55
+)
56
+
57
# Backward compatibility for OpenAI client versions
58
try: # For older openai versions (< 1.100.0)
59
from openai.types.responses import ResponseTextConfig
@@ -84,18 +90,6 @@ from vllm.sampling_params import (
90
from vllm.utils import random_uuid
91
from vllm.utils.import_utils import resolve_obj_by_qualname
92
87
-EMBED_DTYPE_TO_TORCH_DTYPE = {
88
- "float32": torch.float32,
89
- "float16": torch.float16,
90
- "bfloat16": torch.bfloat16,
91
- # I'm not sure if other platforms' CPUs support the fp8 data format.
92
- # EMBED_DTYPE only uses the fp8 data representation,
93
- # does not use fp8 computation, and only occurs on the CPU.
94
- # Apologize for any possible break.
95
- "fp8_e4m3": torch.float8_e4m3fn,
96
- "fp8_e5m2": torch.float8_e5m2,
97
-}
98
-
93
logger = init_logger(__name__)
94
95
_LONG_INFO = torch.iinfo(torch.long)
@@ -1531,7 +1525,7 @@ class EmbeddingCompletionRequest(OpenAIBaseModel):
1525
# https://platform.openai.com/docs/api-reference/embeddings
1526
model: str | None = None
1527
input: list[int] | list[list[int]] | str | list[str]
1534
- encoding_format: Literal["float", "base64"] = "float"
1528
+ encoding_format: EncodingFormat = "float"
1529
dimensions: int | None = None
1530
user: str | None = None
1531
truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
@@ -1564,11 +1558,20 @@ class EmbeddingCompletionRequest(OpenAIBaseModel):
1558
default=None,
1559
description="Whether to normalize the embeddings outputs. Default is True.",
1560
)
1567
- embed_dtype: str = Field(
1561
+ embed_dtype: EmbedDType = Field(
1562
default="float32",
1563
description=(
1570
- "What dtype to use for base64 encoding. Default to using "
1571
- "float32 for base64 encoding to match the OpenAI python client behavior."
1564
+ "What dtype to use for encoding. Default to using float32 for base64 "
1565
+ "encoding to match the OpenAI python client behavior. "
1566
+ "This parameter will affect base64 and binary_response."
1567
+ ),
1568
+ )
1569
+ endianness: Endianness = Field(
1570
+ default="native",
1571
+ description=(
1572
+ "What endianness to use for encoding. Default to using native for "
1573
+ "base64 encoding to match the OpenAI python client behavior."
1574
+ "This parameter will affect base64 and binary_response."
1575
),
1576
)
1577
# --8<-- [end:embedding-extra-params]
@@ -1585,7 +1588,7 @@ class EmbeddingChatRequest(OpenAIBaseModel):
1588
model: str | None = None
1589
messages: list[ChatCompletionMessageParam]
1590
1588
- encoding_format: Literal["float", "base64"] = "float"
1591
+ encoding_format: EncodingFormat = "float"
1592
dimensions: int | None = None
1593
user: str | None = None
1594
truncate_prompt_tokens: Annotated[int, Field(ge=-1)] | None = None
@@ -1650,11 +1653,20 @@ class EmbeddingChatRequest(OpenAIBaseModel):
1653
default=None,
1654
description="Whether to normalize the embeddings outputs. Default is True.",
1655
)
1653
- embed_dtype: str = Field(
1656
+ embed_dtype: EmbedDType = Field(
1657
default="float32",
1658
description=(
1656
- "Which dtype to use for base64 encoding. Defaults to float32 "
1657
- "to match OpenAI API."
1659
+ "What dtype to use for encoding. Default to using float32 for base64 "
1660
+ "encoding to match the OpenAI python client behavior. "
1661
+ "This parameter will affect base64 and binary_response."
1662
+ ),
1663
+ )
1664
+ endianness: Endianness = Field(
1665
+ default="native",
1666
+ description=(
1667
+ "What endianness to use for encoding. Default to using native for "
1668
+ "base64 encoding to match the OpenAI python client behavior."
1669
+ "This parameter will affect base64 and binary_response."
1670
),
1671
)
1672
# --8<-- [end:chat-embedding-extra-params]
@@ -1701,11 +1713,21 @@ class IOProcessorRequest(OpenAIBaseModel, Generic[T]):
1713
"""
1714
activation: bool = False
1715
1704
- embed_dtype: str = Field(
1716
+ encoding_format: EncodingFormat = "float"
1717
+ embed_dtype: EmbedDType = Field(
1718
default="float32",
1719
description=(
1707
- "What dtype to use for base64 encoding. Default to using "
1708
- "float32 for base64 encoding to match the OpenAI python client behavior."
1720
+ "What dtype to use for encoding. Default to using float32 for base64 "
1721
+ "encoding to match the OpenAI python client behavior. "
1722
+ "This parameter will affect base64 and binary_response."
1723
+ ),
1724
+ )
1725
+ endianness: Endianness = Field(
1726
+ default="native",
1727
+ description=(
1728
+ "What endianness to use for encoding. Default to using native for "
1729
+ "base64 encoding to match the OpenAI python client behavior."
1730
+ "This parameter will affect base64 and binary_response."
1731
),
1732
)
1733
@@ -1905,6 +1927,12 @@ class EmbeddingResponse(OpenAIBaseModel):
1927
usage: UsageInfo
1928
1929
1930
+class EmbeddingBytesResponse(OpenAIBaseModel):
1931
+ body: list[bytes]
1932
+ metadata: str
1933
+ media_type: str = "application/octet-stream"
1934
+
1935
+
1936
class PoolingResponseData(OpenAIBaseModel):
1937
index: int
1938
object: str = "pooling"
@@ -1920,6 +1948,12 @@ class PoolingResponse(OpenAIBaseModel):
1948
usage: UsageInfo
1949
1950
1951
+class PoolingBytesResponse(OpenAIBaseModel):
1952
+ body: list[bytes]
1953
+ metadata: str
1954
+ media_type: str = "application/octet-stream"
1955
+
1956
+
1957
class ScoreResponseData(OpenAIBaseModel):
1958
index: int
1959
object: str = "score"
vllm/entrypoints/openai/serving_embedding.py
+69
-34
@@ -1,18 +1,19 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
3
+import json
4
from collections.abc import AsyncGenerator, Mapping
5
from typing import Any, Final, cast
6
7
import torch
8
from fastapi import Request
9
-from typing_extensions import override
9
+from fastapi.responses import Response
10
+from typing_extensions import assert_never, override
11
12
from vllm.engine.protocol import EngineClient
13
from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
14
from vllm.entrypoints.logger import RequestLogger
15
from vllm.entrypoints.openai.protocol import (
15
- EMBED_DTYPE_TO_TORCH_DTYPE,
16
+ EmbeddingBytesResponse,
17
EmbeddingChatRequest,
18
EmbeddingCompletionRequest,
19
EmbeddingRequest,
@@ -28,7 +29,6 @@ from vllm.entrypoints.openai.serving_engine import (
29
TextTokensPrompt,
30
)
31
from vllm.entrypoints.openai.serving_models import OpenAIServingModels
31
-from vllm.entrypoints.openai.utils import encoding_pooling_output
32
from vllm.entrypoints.renderer import RenderConfig
33
from vllm.inputs.data import TokensPrompt as EngineTokensPrompt
34
from vllm.logger import init_logger
@@ -41,6 +41,13 @@ from vllm.outputs import (
41
from vllm.pooling_params import PoolingParams
42
from vllm.utils.async_utils import merge_async_iterators
43
from vllm.utils.collection_utils import chunk_list
44
+from vllm.utils.serial_utils import (
45
+ EmbedDType,
46
+ EncodingFormat,
47
+ Endianness,
48
+ encode_pooling_bytes,
49
+ encode_pooling_output,
50
+)
51
52
logger = init_logger(__name__)
53
@@ -68,12 +75,6 @@ class EmbeddingMixin(OpenAIServing):
75
) -> ErrorResponse | None:
76
ctx = cast(EmbeddingServeContext, ctx)
77
try:
71
- if ctx.request.embed_dtype not in EMBED_DTYPE_TO_TORCH_DTYPE:
72
- return self.create_error_response(
73
- f"embed_dtype={ctx.request.embed_dtype!r} is not supported. "
74
- f"Supported types: {EMBED_DTYPE_TO_TORCH_DTYPE.keys()}"
75
- )
76
-
78
ctx.lora_request = self._maybe_get_adapters(ctx.request)
79
80
tokenizer = await self.engine_client.get_tokenizer()
@@ -121,36 +122,70 @@ class EmbeddingMixin(OpenAIServing):
122
def _build_response(
123
self,
124
ctx: ServeContext,
124
- ) -> EmbeddingResponse | ErrorResponse:
125
- items: list[EmbeddingResponseData] = []
126
- num_prompt_tokens = 0
127
-
125
+ ) -> EmbeddingResponse | Response | ErrorResponse:
126
final_res_batch_checked = cast(list[PoolingRequestOutput], ctx.final_res_batch)
127
130
- for idx, final_res in enumerate(final_res_batch_checked):
131
- item = EmbeddingResponseData(
132
- index=idx,
133
- embedding=encoding_pooling_output(
134
- final_res, ctx.request.encoding_format, ctx.request.embed_dtype
135
- ),
128
+ encoding_format: EncodingFormat = ctx.request.encoding_format
129
+ embed_dtype: EmbedDType = ctx.request.embed_dtype
130
+ endianness: Endianness = ctx.request.endianness
131
+
132
+ def encode_float_base64():
133
+ items: list[EmbeddingResponseData] = []
134
+ num_prompt_tokens = 0
135
+
136
+ for idx, final_res in enumerate(final_res_batch_checked):
137
+ item = EmbeddingResponseData(
138
+ index=idx,
139
+ embedding=encode_pooling_output(
140
+ final_res,
141
+ encoding_format=encoding_format,
142
+ embed_dtype=embed_dtype,
143
+ endianness=endianness,
144
+ ),
145
+ )
146
+ prompt_token_ids = final_res.prompt_token_ids
147
+
148
+ items.append(item)
149
+ num_prompt_tokens += len(prompt_token_ids)
150
+
151
+ usage = UsageInfo(
152
+ prompt_tokens=num_prompt_tokens,
153
+ total_tokens=num_prompt_tokens,
154
)
137
- prompt_token_ids = final_res.prompt_token_ids
155
139
- items.append(item)
140
- num_prompt_tokens += len(prompt_token_ids)
156
+ return EmbeddingResponse(
157
+ id=ctx.request_id,
158
+ created=ctx.created_time,
159
+ model=ctx.model_name,
160
+ data=items,
161
+ usage=usage,
162
+ )
163
142
- usage = UsageInfo(
143
- prompt_tokens=num_prompt_tokens,
144
- total_tokens=num_prompt_tokens,
145
- )
164
+ def encode_bytes():
165
+ body, items, usage = encode_pooling_bytes(
166
+ pooling_outputs=final_res_batch_checked,
167
+ embed_dtype=embed_dtype,
168
+ endianness=endianness,
169
+ )
170
147
- return EmbeddingResponse(
148
- id=ctx.request_id,
149
- created=ctx.created_time,
150
- model=ctx.model_name,
151
- data=items,
152
- usage=usage,
153
- )
171
+ metadata = {
172
+ "id": ctx.request_id,
173
+ "created": ctx.created_time,
174
+ "model": ctx.model_name,
175
+ "data": items,
176
+ "usage": usage,
177
+ }
178
+ return EmbeddingBytesResponse(
179
+ body=body,
180
+ metadata=json.dumps(metadata),
181
+ )
182
+
183
+ if encoding_format == "float" or encoding_format == "base64":
184
+ return encode_float_base64()
185
+ elif encoding_format == "bytes":
186
+ return encode_bytes()
187
+ else:
188
+ assert_never(encoding_format)
189
190
def _get_max_position_embeddings(self) -> int:
191
"""Get the model's effective maximum sequence length for chunking."""
vllm/entrypoints/openai/serving_pooling.py
+70
-54
@@ -2,14 +2,12 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import asyncio
5
-import base64
5
+import json
6
import time
7
from collections.abc import AsyncGenerator
8
-from typing import Final, Literal, cast
8
+from typing import Final, cast
9
10
import jinja2
11
-import numpy as np
12
-import torch
11
from fastapi import Request
12
from typing_extensions import assert_never
13
@@ -17,10 +15,10 @@ from vllm.engine.protocol import EngineClient
15
from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
16
from vllm.entrypoints.logger import RequestLogger
17
from vllm.entrypoints.openai.protocol import (
20
- EMBED_DTYPE_TO_TORCH_DTYPE,
18
ErrorResponse,
19
IOProcessorRequest,
20
IOProcessorResponse,
21
+ PoolingBytesResponse,
22
PoolingChatRequest,
23
PoolingCompletionRequest,
24
PoolingRequest,
@@ -30,33 +28,23 @@ from vllm.entrypoints.openai.protocol import (
28
)
29
from vllm.entrypoints.openai.serving_engine import OpenAIServing
30
from vllm.entrypoints.openai.serving_models import OpenAIServingModels
33
-from vllm.entrypoints.openai.utils import encoding_pooling_output
31
from vllm.entrypoints.renderer import RenderConfig
32
from vllm.entrypoints.utils import _validate_truncation_size
33
from vllm.logger import init_logger
37
-from vllm.outputs import PoolingOutput, PoolingRequestOutput
34
+from vllm.outputs import PoolingRequestOutput
35
from vllm.tasks import SupportedTask
36
from vllm.utils.async_utils import merge_async_iterators
37
+from vllm.utils.serial_utils import (
38
+ EmbedDType,
39
+ EncodingFormat,
40
+ Endianness,
41
+ encode_pooling_bytes,
42
+ encode_pooling_output,
43
+)
44
45
logger = init_logger(__name__)
46
47
44
-def _get_data(
45
- output: PoolingOutput,
46
- encoding_format: Literal["float", "base64"],
47
-) -> list[float] | str:
48
- if encoding_format == "float":
49
- return output.data.tolist()
50
- elif encoding_format == "base64":
51
- # Force to use float32 for base64 encoding
52
- # to match the OpenAI python client behavior
53
- pt_float32 = output.data.to(dtype=torch.float32)
54
- pooling_bytes = np.array(pt_float32, dtype="float32").tobytes()
55
- return base64.b64encode(pooling_bytes).decode("utf-8")
56
-
57
- assert_never(encoding_format)
58
-
59
-
48
class OpenAIServingPooling(OpenAIServing):
49
def __init__(
50
self,
@@ -86,7 +74,7 @@ class OpenAIServingPooling(OpenAIServing):
74
self,
75
request: PoolingRequest,
76
raw_request: Request | None = None,
89
- ) -> PoolingResponse | IOProcessorResponse | ErrorResponse:
77
+ ) -> PoolingResponse | IOProcessorResponse | PoolingBytesResponse | ErrorResponse:
78
"""
79
See https://platform.openai.com/docs/api-reference/embeddings/create
80
for the API specification. This API mimics the OpenAI Embedding API.
@@ -95,12 +83,6 @@ class OpenAIServingPooling(OpenAIServing):
83
if error_check_ret is not None:
84
return error_check_ret
85
98
- if request.embed_dtype not in EMBED_DTYPE_TO_TORCH_DTYPE:
99
- return self.create_error_response(
100
- f"embed_dtype={request.embed_dtype!r} is not supported. "
101
- f"Supported types: {EMBED_DTYPE_TO_TORCH_DTYPE.keys()}"
102
- )
103
-
86
model_name = self.models.model_name()
87
88
request_id = f"pool-{self._base_request_id(raw_request)}"
@@ -256,6 +238,7 @@ class OpenAIServingPooling(OpenAIServing):
238
model_name,
239
request.encoding_format,
240
request.embed_dtype,
241
+ request.endianness,
242
)
243
except asyncio.CancelledError:
244
return self.create_error_response("Client disconnected")
@@ -271,34 +254,67 @@ class OpenAIServingPooling(OpenAIServing):
254
request_id: str,
255
created_time: int,
256
model_name: str,
274
- encoding_format: Literal["float", "base64"],
275
- embed_dtype: str,
276
- ) -> PoolingResponse:
277
- items: list[PoolingResponseData] = []
278
- num_prompt_tokens = 0
279
-
280
- for idx, final_res in enumerate(final_res_batch):
281
- item = PoolingResponseData(
282
- index=idx,
283
- data=encoding_pooling_output(final_res, encoding_format, embed_dtype),
257
+ encoding_format: EncodingFormat,
258
+ embed_dtype: EmbedDType,
259
+ endianness: Endianness,
260
+ ) -> PoolingResponse | PoolingBytesResponse:
261
+ def encode_float_base64():
262
+ items: list[PoolingResponseData] = []
263
+ num_prompt_tokens = 0
264
+
265
+ for idx, final_res in enumerate(final_res_batch):
266
+ item = PoolingResponseData(
267
+ index=idx,
268
+ data=encode_pooling_output(
269
+ final_res,
270
+ encoding_format=encoding_format,
271
+ embed_dtype=embed_dtype,
272
+ endianness=endianness,
273
+ ),
274
+ )
275
+ prompt_token_ids = final_res.prompt_token_ids
276
+
277
+ items.append(item)
278
+ num_prompt_tokens += len(prompt_token_ids)
279
+
280
+ usage = UsageInfo(
281
+ prompt_tokens=num_prompt_tokens,
282
+ total_tokens=num_prompt_tokens,
283
)
285
- prompt_token_ids = final_res.prompt_token_ids
284
287
- items.append(item)
288
- num_prompt_tokens += len(prompt_token_ids)
285
+ return PoolingResponse(
286
+ id=request_id,
287
+ created=created_time,
288
+ model=model_name,
289
+ data=items,
290
+ usage=usage,
291
+ )
292
290
- usage = UsageInfo(
291
- prompt_tokens=num_prompt_tokens,
292
- total_tokens=num_prompt_tokens,
293
- )
293
+ def encode_bytes():
294
+ body, items, usage = encode_pooling_bytes(
295
+ pooling_outputs=final_res_batch,
296
+ embed_dtype=embed_dtype,
297
+ endianness=endianness,
298
+ )
299
295
- return PoolingResponse(
296
- id=request_id,
297
- created=created_time,
298
- model=model_name,
299
- data=items,
300
- usage=usage,
301
- )
300
+ metadata = {
301
+ "id": request_id,
302
+ "created": created_time,
303
+ "model": model_name,
304
+ "data": items,
305
+ "usage": usage,
306
+ }
307
+ return PoolingBytesResponse(
308
+ body=body,
309
+ metadata=json.dumps(metadata),
310
+ )
311
+
312
+ if encoding_format == "float" or encoding_format == "base64":
313
+ return encode_float_base64()
314
+ elif encoding_format == "bytes":
315
+ return encode_bytes()
316
+ else:
317
+ assert_never(encoding_format)
318
319
def _build_render_config(self, request: PoolingCompletionRequest) -> RenderConfig:
320
return RenderConfig(
vllm/entrypoints/openai/utils.py
deleted
-33
@@ -1,33 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import base64
4
-from typing import Literal
5
-
6
-import torch
7
-from typing_extensions import assert_never
8
-
9
-from vllm import PoolingRequestOutput
10
-from vllm.entrypoints.openai.protocol import EMBED_DTYPE_TO_TORCH_DTYPE
11
-
12
-
13
-def encoding_pooling_output(
14
- output: PoolingRequestOutput,
15
- encoding_format: Literal["float", "base64"],
16
- embed_dtype: str,
17
-) -> list[float] | str:
18
- if encoding_format == "float":
19
- return output.outputs.data.tolist()
20
- elif encoding_format == "base64":
21
- assert embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE
22
- torch_dtype = EMBED_DTYPE_TO_TORCH_DTYPE[embed_dtype]
23
- embedding_bytes = (
24
- output.outputs.data.to(torch_dtype)
25
- .flatten()
26
- .contiguous()
27
- .view(torch.uint8)
28
- .numpy()
29
- .tobytes()
30
- )
31
- return base64.b64encode(embedding_bytes).decode("utf-8")
32
-
33
- assert_never(encoding_format)
vllm/utils/serial_utils.py
new
+169
@@ -0,0 +1,169 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import base64
4
+import sys
5
+from dataclasses import dataclass
6
+from typing import Literal
7
+
8
+import numpy as np
9
+import torch
10
+from typing_extensions import assert_never
11
+
12
+from vllm import PoolingRequestOutput
13
+
14
+sys_byteorder = sys.byteorder
15
+
16
+
17
+EMBED_DTYPE_TO_TORCH_DTYPE = {
18
+ "float32": torch.float32,
19
+ "float16": torch.float16,
20
+ "bfloat16": torch.bfloat16,
21
+ # I'm not sure if other platforms' CPUs support the fp8 data format.
22
+ # EMBED_DTYPE only uses the fp8 data representation,
23
+ # does not use fp8 computation, and only occurs on the CPU.
24
+ # Apologize for any possible break.
25
+ "fp8_e4m3": torch.float8_e4m3fn,
26
+ "fp8_e5m2": torch.float8_e5m2,
27
+}
28
+
29
+
30
+EMBED_DTYPE_TO_TORCH_DTYPE_VIEW = {
31
+ "float32": torch.float32,
32
+ "float16": torch.float16,
33
+ # numpy does not support bfloat16 and fp8
34
+ "bfloat16": torch.float16,
35
+ "fp8_e4m3": torch.uint8,
36
+ "fp8_e5m2": torch.uint8,
37
+}
38
+
39
+EMBED_DTYPE_TO_NUMPY_DTYPE_VIEW = {
40
+ "float32": np.float32,
41
+ "float16": np.float16,
42
+ # numpy does not support bfloat16 and fp8
43
+ "bfloat16": np.float16,
44
+ "fp8_e4m3": np.uint8,
45
+ "fp8_e5m2": np.uint8,
46
+}
47
+
48
+ENDIANNESS = ["native", "big", "little"]
49
+
50
+EmbedDType = Literal["float32", "float16", "bfloat16", "fp8_e4m3", "fp8_e5m2"]
51
+Endianness = Literal["native", "big", "little"]
52
+EncodingFormat = Literal["float", "base64", "bytes"]
53
+
54
+
55
+def tensor2binary(
56
+ tensor: torch.Tensor, embed_dtype: EmbedDType, endianness: Endianness
57
+) -> bytes:
58
+ assert isinstance(tensor, torch.Tensor)
59
+ assert embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE
60
+ assert endianness in ENDIANNESS
61
+
62
+ torch_dtype = EMBED_DTYPE_TO_TORCH_DTYPE[embed_dtype]
63
+ torch_view_dtype = EMBED_DTYPE_TO_TORCH_DTYPE_VIEW[embed_dtype]
64
+
65
+ np_array = (
66
+ tensor.to(torch_dtype).flatten().contiguous().view(torch_view_dtype).numpy()
67
+ )
68
+
69
+ if endianness != "native" and endianness != sys_byteorder:
70
+ np_array = np_array.byteswap()
71
+
72
+ return np_array.tobytes()
73
+
74
+
75
+def binary2tensor(
76
+ binary: bytes,
77
+ shape: tuple[int, ...],
78
+ embed_dtype: EmbedDType,
79
+ endianness: Endianness,
80
+) -> torch.Tensor:
81
+ assert embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE
82
+ assert embed_dtype in EMBED_DTYPE_TO_NUMPY_DTYPE_VIEW
83
+ assert endianness in ENDIANNESS
84
+
85
+ torch_dtype = EMBED_DTYPE_TO_TORCH_DTYPE[embed_dtype]
86
+ np_dtype = EMBED_DTYPE_TO_NUMPY_DTYPE_VIEW[embed_dtype]
87
+
88
+ np_array = np.frombuffer(binary, dtype=np_dtype).reshape(shape)
89
+
90
+ if endianness != "native" and endianness != sys_byteorder:
91
+ np_array = np_array.byteswap()
92
+
93
+ return torch.from_numpy(np_array).view(torch_dtype)
94
+
95
+
96
+def encode_pooling_output(
97
+ output: PoolingRequestOutput,
98
+ encoding_format: EncodingFormat,
99
+ embed_dtype: EmbedDType,
100
+ endianness: Endianness,
101
+) -> list[float] | str | bytes:
102
+ if encoding_format == "float":
103
+ return output.outputs.data.tolist()
104
+ elif encoding_format == "base64":
105
+ embedding_bytes = tensor2binary(output.outputs.data, embed_dtype, endianness)
106
+ return base64.b64encode(embedding_bytes).decode("utf-8")
107
+ elif encoding_format == "bytes":
108
+ return tensor2binary(output.outputs.data, embed_dtype, endianness)
109
+ assert_never(encoding_format)
110
+
111
+
112
+@dataclass
113
+class MetadataItem:
114
+ index: int
115
+ embed_dtype: EmbedDType
116
+ endianness: Endianness
117
+ start: int
118
+ end: int
119
+ shape: tuple[int, ...]
120
+
121
+
122
+def encode_pooling_bytes(
123
+ pooling_outputs: list[PoolingRequestOutput],
124
+ embed_dtype: EmbedDType,
125
+ endianness: Endianness,
126
+):
127
+ num_prompt_tokens = 0
128
+ items: list[dict[str, MetadataItem]] = []
129
+ body = []
130
+ offset = 0
131
+ for idx, output in enumerate(pooling_outputs):
132
+ binary = tensor2binary(
133
+ tensor=output.outputs.data,
134
+ embed_dtype=embed_dtype,
135
+ endianness=endianness,
136
+ )
137
+ size = len(binary)
138
+
139
+ item = {
140
+ "index": idx,
141
+ "embed_dtype": embed_dtype,
142
+ "endianness": endianness,
143
+ "start": offset,
144
+ "end": offset + size,
145
+ "shape": output.outputs.data.shape,
146
+ }
147
+
148
+ body.append(binary)
149
+ items.append(item)
150
+ prompt_token_ids = output.prompt_token_ids
151
+ num_prompt_tokens += len(prompt_token_ids)
152
+ offset += size
153
+
154
+ usage = {
155
+ "prompt_tokens": num_prompt_tokens,
156
+ "total_tokens": num_prompt_tokens,
157
+ }
158
+ return body, items, usage
159
+
160
+
161
+def decode_pooling_output(items: list[MetadataItem], body: bytes) -> list[torch.Tensor]:
162
+ items.sort(key=lambda x: x.index)
163
+
164
+ tensor_list: list[torch.Tensor] = []
165
+ for item in items:
166
+ binary = body[item.start : item.end]
167
+ tensor = binary2tensor(binary, item.shape, item.embed_dtype, item.endianness)
168
+ tensor_list.append(tensor)
169
+ return tensor_list