Skip to content

Commit

Permalink
[Bugfix]: Use float32 for base64 embedding (vllm-project#7855)
Browse files Browse the repository at this point in the history
Signed-off-by: Hollow Man <[email protected]>
  • Loading branch information
HollowMan6 authored Aug 26, 2024
1 parent 8f141eb commit a8a817c
Show file tree
Hide file tree
Showing 3 changed files with 13 additions and 3 deletions.
1 change: 0 additions & 1 deletion examples/openai_embedding_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
"The best thing about vLLM is that it supports many different models"
],
model=model,
encoding_format="float",
)

for data in responses.data:
Expand Down
11 changes: 10 additions & 1 deletion tests/entrypoints/openai/test_embedding.py
Original file line number Diff line number Diff line change
Expand Up @@ -128,9 +128,18 @@ async def test_batch_base64_embedding(embedding_client: openai.AsyncOpenAI,
for data in responses_base64.data:
decoded_responses_base64_data.append(
np.frombuffer(base64.b64decode(data.embedding),
dtype="float").tolist())
dtype="float32").tolist())

assert responses_float.data[0].embedding == decoded_responses_base64_data[
0]
assert responses_float.data[1].embedding == decoded_responses_base64_data[
1]

# Default response is float32 decoded from base64 by OpenAI Client
responses_default = await embedding_client.embeddings.create(
input=input_texts, model=model_name)

assert responses_float.data[0].embedding == responses_default.data[
0].embedding
assert responses_float.data[1].embedding == responses_default.data[
1].embedding
4 changes: 3 additions & 1 deletion vllm/entrypoints/openai/serving_embedding.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,9 @@ def _get_embedding(
if encoding_format == "float":
return output.embedding
elif encoding_format == "base64":
embedding_bytes = np.array(output.embedding).tobytes()
# Force to use float32 for base64 encoding
# to match the OpenAI python client behavior
embedding_bytes = np.array(output.embedding, dtype="float32").tobytes()
return base64.b64encode(embedding_bytes).decode("utf-8")

assert_never(encoding_format)
Expand Down

0 comments on commit a8a817c

Please sign in to comment.