Revert "[ENH] 1965 Split up embedding functions"

chromadb/api/types.py

@@ -1,7 +1,7 @@
 from typing import Optional, Union, TypeVar, List, Dict, Any, Tuple, cast
 from numpy.typing import NDArray
 import numpy as np
-from typing_extensions import Literal, TypedDict, Protocol, runtime_checkable
+from typing_extensions import Literal, TypedDict, Protocol
 import chromadb.errors as errors
 from chromadb.types import (
     Metadata,
@@ -56,7 +56,7 @@ def maybe_cast_one_to_many_ids(target: OneOrMany[ID]) -> IDs:
 
 
 def maybe_cast_one_to_many_embedding(
-    target: Union[OneOrMany[Embedding], OneOrMany[np.ndarray]]  # type: ignore[type-arg]
+    target: Union[OneOrMany[Embedding], OneOrMany[np.ndarray]]
 ) -> Embeddings:
     if isinstance(target, List):
         # One Embedding
@@ -101,7 +101,7 @@ def maybe_cast_one_to_many_document(target: OneOrMany[Document]) -> Documents:
 
 
 # Images
-ImageDType = Union[np.uint, np.int_, np.float_]  # type: ignore[name-defined]
+ImageDType = Union[np.uint, np.int_, np.float_]
 Image = NDArray[ImageDType]
 Images = List[Image]
 
@@ -184,7 +184,6 @@ class IndexMetadata(TypedDict):
     time_created: float
 
 
-@runtime_checkable
 class EmbeddingFunction(Protocol[D]):
     def __call__(self, input: D) -> Embeddings:
         ...
@@ -200,10 +199,8 @@ def __call__(self: EmbeddingFunction[D], input: D) -> Embeddings:
 
         setattr(cls, "__call__", __call__)
 
-    def embed_with_retries(
-        self, input: D, **retry_kwargs: Dict[str, Any]
-    ) -> Embeddings:
-        return cast(Embeddings, retry(**retry_kwargs)(self.__call__)(input))
+    def embed_with_retries(self, input: D, **retry_kwargs: Dict) -> Embeddings:
+        return retry(**retry_kwargs)(self.__call__)(input)
 
 
 def validate_embedding_function(

chromadb/test/ef/test_default_ef.py

@@ -7,10 +7,7 @@
 import pytest
 from hypothesis import given, settings
 
-from chromadb.utils.embedding_functions.onnx_mini_lm_l6_v2 import (
-    ONNXMiniLM_L6_V2,
-    _verify_sha256,
-)
+from chromadb.utils.embedding_functions import ONNXMiniLM_L6_V2, _verify_sha256
 
 
 def unique_by(x: Hashable) -> Hashable: