Skip to content

Commit

Permalink
feat(embeddings): use stdlib array type for improved performance (#2060)
Browse files Browse the repository at this point in the history
  • Loading branch information
tonybaloney authored Feb 6, 2025
1 parent e0ca9f0 commit 9a95db9
Showing 1 changed file with 17 additions and 12 deletions.
29 changes: 17 additions & 12 deletions src/openai/resources/embeddings.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

from __future__ import annotations

import array
import base64
from typing import List, Union, Iterable, cast
from typing_extensions import Literal
Expand Down Expand Up @@ -102,7 +103,7 @@ def create(
"dimensions": dimensions,
"encoding_format": encoding_format,
}
if not is_given(encoding_format) and has_numpy():
if not is_given(encoding_format):
params["encoding_format"] = "base64"

def parser(obj: CreateEmbeddingResponse) -> CreateEmbeddingResponse:
Expand All @@ -113,12 +114,14 @@ def parser(obj: CreateEmbeddingResponse) -> CreateEmbeddingResponse:
for embedding in obj.data:
data = cast(object, embedding.embedding)
if not isinstance(data, str):
# numpy is not installed / base64 optimisation isn't enabled for this model yet
continue

embedding.embedding = np.frombuffer( # type: ignore[no-untyped-call]
base64.b64decode(data), dtype="float32"
).tolist()
if not has_numpy():
# use array for base64 optimisation
embedding.embedding = array.array("f", base64.b64decode(data)).tolist()
else:
embedding.embedding = np.frombuffer( # type: ignore[no-untyped-call]
base64.b64decode(data), dtype="float32"
).tolist()

return obj

Expand Down Expand Up @@ -215,7 +218,7 @@ async def create(
"dimensions": dimensions,
"encoding_format": encoding_format,
}
if not is_given(encoding_format) and has_numpy():
if not is_given(encoding_format):
params["encoding_format"] = "base64"

def parser(obj: CreateEmbeddingResponse) -> CreateEmbeddingResponse:
Expand All @@ -226,12 +229,14 @@ def parser(obj: CreateEmbeddingResponse) -> CreateEmbeddingResponse:
for embedding in obj.data:
data = cast(object, embedding.embedding)
if not isinstance(data, str):
# numpy is not installed / base64 optimisation isn't enabled for this model yet
continue

embedding.embedding = np.frombuffer( # type: ignore[no-untyped-call]
base64.b64decode(data), dtype="float32"
).tolist()
if not has_numpy():
# use array for base64 optimisation
embedding.embedding = array.array("f", base64.b64decode(data)).tolist()
else:
embedding.embedding = np.frombuffer( # type: ignore[no-untyped-call]
base64.b64decode(data), dtype="float32"
).tolist()

return obj

Expand Down

0 comments on commit 9a95db9

Please sign in to comment.