[Frontend][3/N] Improve all pooling task | Support binary embedding response (vllm-project#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>

Author: 3 people

examples/online_serving/pooling/README.md

@@ -6,10 +6,16 @@
 python examples/online_serving/pooling/cohere_rerank_client.py
 ```
 
-## Embedding embed_dtype usage
+## Embedding requests base64 encoding_format usage
 
 ```bash
-python examples/online_serving/pooling/embedding_embed_dtype_client.py
+python examples/online_serving/pooling/embedding_requests_base64_client.py
+```
+
+## Embedding requests bytes encoding_format usage
+
+```bash
+python examples/online_serving/pooling/embedding_requests_bytes_client.py
 ```
 
 ## Jinaai rerank usage

examples/online_serving/pooling/embedding_embed_dtype_client.py renamed to examples/online_serving/pooling/embedding_requests_base64_client.py

@@ -12,7 +12,11 @@
 import requests
 import torch
 
-from vllm.entrypoints.openai.protocol import EMBED_DTYPE_TO_TORCH_DTYPE
+from vllm.utils.serial_utils import (
+    EMBED_DTYPE_TO_TORCH_DTYPE,
+    ENDIANNESS,
+    binary2tensor,
+)
 
 
 def post_http_request(prompt: dict, api_url: str) -> requests.Response:
@@ -34,24 +38,25 @@ def main(args):
     api_url = f"http://{args.host}:{args.port}/v1/embeddings"
     model_name = args.model
 
-    for embed_dtype, torch_dtype in EMBED_DTYPE_TO_TORCH_DTYPE.items():
-        prompt = {
-            "model": model_name,
-            "input": "vLLM is great!",
-            "encoding_format": "base64",
-            "embed_dtype": embed_dtype,
-        }
-        response = post_http_request(prompt=prompt, api_url=api_url)
-
-        embedding = []
-        for data in response.json()["data"]:
-            embedding.append(
-                torch.frombuffer(
-                    base64.b64decode(data["embedding"]), dtype=torch_dtype
-                ).to(torch.float32)
-            )
-        embedding = torch.cat(embedding)
-        print(embed_dtype, embedding.shape)
+    # The OpenAI client does not support the embed_dtype and endianness parameters.
+    for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE:
+        for endianness in ENDIANNESS:
+            prompt = {
+                "model": model_name,
+                "input": "vLLM is great!",
+                "encoding_format": "base64",
+                "embed_dtype": embed_dtype,
+                "endianness": endianness,
+            }
+            response = post_http_request(prompt=prompt, api_url=api_url)
+
+            embedding = []
+            for data in response.json()["data"]:
+                binary = base64.b64decode(data["embedding"])
+                tensor = binary2tensor(binary, (-1,), embed_dtype, endianness)
+                embedding.append(tensor.to(torch.float32))
+            embedding = torch.cat(embedding)
+            print(embed_dtype, endianness, embedding.shape)
 
 
 if __name__ == "__main__":

examples/online_serving/pooling/embedding_requests_bytes_client.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Example Python client for embedding API using vLLM API server
+NOTE:
+    start a supported embeddings model server with `vllm serve`, e.g.
+    vllm serve intfloat/e5-small
+"""
+
+import argparse
+import json
+
+import requests
+import torch
+
+from vllm.utils.serial_utils import (
+    EMBED_DTYPE_TO_TORCH_DTYPE,
+    ENDIANNESS,
+    MetadataItem,
+    decode_pooling_output,
+)
+
+
+def post_http_request(prompt: dict, api_url: str) -> requests.Response:
+    headers = {"User-Agent": "Test Client"}
+    response = requests.post(api_url, headers=headers, json=prompt)
+    return response
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--model", type=str, default="intfloat/e5-small")
+
+    return parser.parse_args()
+
+
+def main(args):
+    api_url = f"http://{args.host}:{args.port}/v1/embeddings"
+    model_name = args.model
+
+    # The OpenAI client does not support the bytes encoding_format.
+    # The OpenAI client does not support the embed_dtype and endianness parameters.
+    for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE:
+        for endianness in ENDIANNESS:
+            prompt = {
+                "model": model_name,
+                "input": "vLLM is great!",
+                "encoding_format": "bytes",
+                "embed_dtype": embed_dtype,
+                "endianness": endianness,
+            }
+            response = post_http_request(prompt=prompt, api_url=api_url)
+            metadata = json.loads(response.headers["metadata"])
+            body = response.content
+            items = [MetadataItem(**x) for x in metadata["data"]]
+
+            embedding = decode_pooling_output(items=items, body=body)
+            embedding = [x.to(torch.float32) for x in embedding]
+            embedding = torch.cat(embedding)
+            print(embed_dtype, endianness, embedding.shape)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

tests/entrypoints/pooling/openai/test_embedding.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import base64
+import json
 
 import numpy as np
 import openai
@@ -15,11 +16,17 @@
 from tests.models.utils import check_embeddings_close
 from tests.utils import RemoteOpenAIServer
 from vllm.entrypoints.openai.protocol import (
-    EMBED_DTYPE_TO_TORCH_DTYPE,
     EmbeddingResponse,
     PoolingResponse,
 )
 from vllm.transformers_utils.tokenizer import get_tokenizer
+from vllm.utils.serial_utils import (
+    EMBED_DTYPE_TO_TORCH_DTYPE,
+    ENDIANNESS,
+    MetadataItem,
+    binary2tensor,
+    decode_pooling_output,
+)
 
 MODEL_NAME = "intfloat/multilingual-e5-small"
 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(
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
-async def test_base64_embed_dtype(
-    hf_model, server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
+async def test_base64_embed_dtype_and_endianness(
+    server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
 ):
     input_texts = [
         "The best thing about vLLM is that it supports many different models",
@@ -262,59 +269,100 @@ async def test_base64_embed_dtype(
     )
     float_data = [d.embedding for d in responses_float.data]
 
-    for embed_dtype, torch_dtype in EMBED_DTYPE_TO_TORCH_DTYPE.items():
-        responses_base64 = requests.post(
-            server.url_for("/v1/embeddings"),
-            json={
-                "model": model_name,
-                "input": input_texts,
-                "encoding_format": "base64",
-                "embed_dtype": embed_dtype,
-            },
-        )
-
-        base64_data = []
-        for data in responses_base64.json()["data"]:
-            base64_data.append(
-                torch.frombuffer(base64.b64decode(data["embedding"]), dtype=torch_dtype)
-                .to(torch.float32)
-                .tolist()
+    for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE:
+        for endianness in ENDIANNESS:
+            responses_base64 = requests.post(
+                server.url_for("/v1/embeddings"),
+                json={
+                    "model": model_name,
+                    "input": input_texts,
+                    "encoding_format": "base64",
+                    "embed_dtype": embed_dtype,
+                    "endianness": endianness,
+                },
             )
 
-        check_embeddings_close(
-            embeddings_0_lst=float_data,
-            embeddings_1_lst=base64_data,
-            name_0="float_data",
-            name_1="base64_data",
-            tol=1e-2,
-        )
+            base64_data = []
+            for data in responses_base64.json()["data"]:
+                binary = base64.b64decode(data["embedding"])
+                tensor = binary2tensor(binary, (-1,), embed_dtype, endianness)
+                base64_data.append(tensor.to(torch.float32).tolist())
+
+            check_embeddings_close(
+                embeddings_0_lst=float_data,
+                embeddings_1_lst=base64_data,
+                name_0="float_data",
+                name_1="base64_data",
+                tol=1e-2,
+            )
 
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
-async def test_base64_embed_dtype_not_supported(
-    hf_model, server: RemoteOpenAIServer, model_name: str
+async def test_bytes_embed_dtype_and_endianness(
+    server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
 ):
     input_texts = [
         "The best thing about vLLM is that it supports many different models",
     ]
 
-    bad_embed_dtype = "bad_embed_dtype"
+    responses_float = await client.embeddings.create(
+        input=input_texts, model=model_name, encoding_format="float"
+    )
+    float_data = [d.embedding for d in responses_float.data]
+
+    for embed_dtype in list(EMBED_DTYPE_TO_TORCH_DTYPE.keys()):
+        for endianness in ENDIANNESS:
+            responses_bytes = requests.post(
+                server.url_for("/v1/embeddings"),
+                json={
+                    "model": model_name,
+                    "input": input_texts,
+                    "encoding_format": "bytes",
+                    "embed_dtype": embed_dtype,
+                    "endianness": endianness,
+                },
+            )
+
+            metadata = json.loads(responses_bytes.headers["metadata"])
+            body = responses_bytes.content
+            items = [MetadataItem(**x) for x in metadata["data"]]
+
+            bytes_data = decode_pooling_output(items=items, body=body)
+            bytes_data = [x.to(torch.float32).tolist() for x in bytes_data]
+
+            check_embeddings_close(
+                embeddings_0_lst=float_data,
+                embeddings_1_lst=bytes_data,
+                name_0="float_data",
+                name_1="bytes_data",
+                tol=1e-2,
+            )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
+@pytest.mark.parametrize("param_name", ["encoding_format", "embed_dtype", "endianness"])
+async def test_params_not_supported(
+    server: RemoteOpenAIServer, model_name: str, param_name: str
+):
+    input_texts = [
+        "The best thing about vLLM is that it supports many different models",
+    ]
 
     responses_base64 = requests.post(
         server.url_for("/v1/embeddings"),
         json={
             "model": model_name,
             "input": input_texts,
             "encoding_format": "base64",
-            "embed_dtype": bad_embed_dtype,
+            param_name: f"bad_{param_name}",
         },
     )
 
     assert responses_base64.status_code == 400
-    assert responses_base64.json()["error"]["message"].startswith(
-        f"embed_dtype={bad_embed_dtype!r} is not supported."
-    )
+    assert "literal_error" in responses_base64.json()["error"]["message"]
+    assert f"bad_{param_name}" in responses_base64.json()["error"]["message"]
 
 
 @pytest.mark.asyncio