Skip to content

Commit 0814b40

Browse files
atroynAnush008
authored andcommitted
[BUG] 1965 Split up embedding functions - Redux (chroma-core#2395)
## Description of changes The original attempt to split up the embedding functions failed because of python 3.9 and 3.10 incompatibilities with `issubtype`. Original PR here: chroma-core#2034 Failing tests here: https://github.com/chroma-core/chroma/actions/runs/9605053108/job/26491923410 The fix is changing `issubtype` to `isinstance`, which has the same functionality. ## Test plan Along with CI, tested locally with python 3.9 and 3.10 and confirmed passing. ## Documentation Changes N/A
1 parent db099c2 commit 0814b40

19 files changed

Lines changed: 1240 additions & 1035 deletions

chromadb/api/types.py

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from typing import Optional, Union, TypeVar, List, Dict, Any, Tuple, cast
22
from numpy.typing import NDArray
33
import numpy as np
4-
from typing_extensions import Literal, TypedDict, Protocol
4+
from typing_extensions import Literal, TypedDict, Protocol, runtime_checkable
55
import chromadb.errors as errors
66
from chromadb.types import (
77
Metadata,
@@ -56,7 +56,7 @@ def maybe_cast_one_to_many_ids(target: OneOrMany[ID]) -> IDs:
5656

5757

5858
def maybe_cast_one_to_many_embedding(
59-
target: Union[OneOrMany[Embedding], OneOrMany[np.ndarray]]
59+
target: Union[OneOrMany[Embedding], OneOrMany[np.ndarray]] # type: ignore[type-arg]
6060
) -> Embeddings:
6161
if isinstance(target, List):
6262
# One Embedding
@@ -101,7 +101,7 @@ def maybe_cast_one_to_many_document(target: OneOrMany[Document]) -> Documents:
101101

102102

103103
# Images
104-
ImageDType = Union[np.uint, np.int_, np.float_]
104+
ImageDType = Union[np.uint, np.int_, np.float_] # type: ignore[name-defined]
105105
Image = NDArray[ImageDType]
106106
Images = List[Image]
107107

@@ -184,6 +184,7 @@ class IndexMetadata(TypedDict):
184184
time_created: float
185185

186186

187+
@runtime_checkable
187188
class EmbeddingFunction(Protocol[D]):
188189
def __call__(self, input: D) -> Embeddings:
189190
...
@@ -199,8 +200,10 @@ def __call__(self: EmbeddingFunction[D], input: D) -> Embeddings:
199200

200201
setattr(cls, "__call__", __call__)
201202

202-
def embed_with_retries(self, input: D, **retry_kwargs: Dict) -> Embeddings:
203-
return retry(**retry_kwargs)(self.__call__)(input)
203+
def embed_with_retries(
204+
self, input: D, **retry_kwargs: Dict[str, Any]
205+
) -> Embeddings:
206+
return cast(Embeddings, retry(**retry_kwargs)(self.__call__)(input))
204207

205208

206209
def validate_embedding_function(

chromadb/test/ef/test_default_ef.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,10 @@
77
import pytest
88
from hypothesis import given, settings
99

10-
from chromadb.utils.embedding_functions import ONNXMiniLM_L6_V2, _verify_sha256
10+
from chromadb.utils.embedding_functions.onnx_mini_lm_l6_v2 import (
11+
ONNXMiniLM_L6_V2,
12+
_verify_sha256,
13+
)
1114

1215

1316
def unique_by(x: Hashable) -> Hashable:

chromadb/test/ef/test_ef.py

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
from chromadb.utils import embedding_functions
2+
from chromadb.api.types import EmbeddingFunction
3+
4+
5+
def test_get_builtins_holds() -> None:
6+
"""
7+
Ensure that `get_builtins` is consistent after the ef migration.
8+
9+
This test is intended to be temporary until the ef migration is complete as
10+
these expected builtins are likely to grow as long as users add new
11+
embedding functions.
12+
13+
REMOVE ME ON THE NEXT EF ADDITION
14+
"""
15+
expected_builtins = {
16+
"AmazonBedrockEmbeddingFunction",
17+
"CohereEmbeddingFunction",
18+
"GoogleGenerativeAiEmbeddingFunction",
19+
"GooglePalmEmbeddingFunction",
20+
"GoogleVertexEmbeddingFunction",
21+
"HuggingFaceEmbeddingFunction",
22+
"HuggingFaceEmbeddingServer",
23+
"InstructorEmbeddingFunction",
24+
"JinaEmbeddingFunction",
25+
"ONNXMiniLM_L6_V2",
26+
"OllamaEmbeddingFunction",
27+
"OpenAIEmbeddingFunction",
28+
"OpenCLIPEmbeddingFunction",
29+
"RoboflowEmbeddingFunction",
30+
"SentenceTransformerEmbeddingFunction",
31+
"Text2VecEmbeddingFunction",
32+
"ChromaLangchainEmbeddingFunction",
33+
}
34+
35+
assert expected_builtins == embedding_functions.get_builtins()
36+
37+
38+
def test_default_ef_exists() -> None:
39+
assert hasattr(embedding_functions, "DefaultEmbeddingFunction")
40+
default_ef = embedding_functions.DefaultEmbeddingFunction()
41+
42+
assert default_ef is not None
43+
assert isinstance(default_ef, EmbeddingFunction)
44+
45+
46+
def test_ef_imports() -> None:
47+
for ef in embedding_functions.get_builtins():
48+
# Langchain embedding function is a special snowflake
49+
if ef == "ChromaLangchainEmbeddingFunction":
50+
continue
51+
assert hasattr(embedding_functions, ef)
52+
assert isinstance(getattr(embedding_functions, ef), type)
53+
assert issubclass(getattr(embedding_functions, ef), EmbeddingFunction)

0 commit comments

Comments
 (0)