[Frontend][1/N] Improve all pooling task | Support FP16 Embedding Base64 (Still uses fp32 by default). (#26414)

Signed-off-by: wang.yuqi <noooop@126.com>
Co-authored-by: Maximilien de Bayser <maxdebayser@gmail.com>
Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>

tests/entrypoints/pooling/openai/test_embedding.py

@@ -14,7 +14,10 @@ import torch.nn.functional as F
 from tests.models.language.pooling.embed_utils import run_embedding_correctness_test
 from tests.models.utils import check_embeddings_close
 from tests.utils import RemoteOpenAIServer
-from vllm.entrypoints.openai.protocol import EmbeddingResponse
+from vllm.entrypoints.openai.protocol import (
+    EMBED_DTYPE_TO_TORCH_DTYPE,
+    EmbeddingResponse,
+)
 from vllm.transformers_utils.tokenizer import get_tokenizer
 
 MODEL_NAME = "intfloat/multilingual-e5-small"
@@ -244,6 +247,75 @@ async def test_batch_base64_embedding(
     run_embedding_correctness_test(hf_model, input_texts, default_data)
 
 
+@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
+):
+    input_texts = [
+        "The best thing about vLLM is that it supports many different models",
+    ]
+
+    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, 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()
+            )
+
+        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
+):
+    input_texts = [
+        "The best thing about vLLM is that it supports many different models",
+    ]
+
+    bad_embed_dtype = "bad_embed_dtype"
+
+    responses_base64 = requests.post(
+        server.url_for("/v1/embeddings"),
+        json={
+            "model": model_name,
+            "input": input_texts,
+            "encoding_format": "base64",
+            "embed_dtype": bad_embed_dtype,
+        },
+    )
+
+    assert responses_base64.status_code == 400
+    assert responses_base64.json()["error"]["message"].startswith(
+        f"embed_dtype={bad_embed_dtype!r} is not supported."
+    )
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("model_name", [MODEL_NAME])
 async def test_single_embedding_truncation(client: openai.AsyncOpenAI, model_name: str):