[ENH] Multimodal Embeddings (#1293)

## Description of changes This PR introduces multi-modal embeddings into Chroma. - It adds the generic `EmbeddingFunction` which can take various data types. Existing functions take the `Documents` type. - Adds `Images` as a type (numpy NDArray taking ints or floats) - Add `OpenCLIPEmbeddingFunction` which is an `EmbeddingFunction[Union[Documents, Images]]` ## Test Integration tests pass. A new test for multimodal embedding functions: [chromadb/test/ef/test_multimodal_ef.py](https://github.com/chroma-core/chroma/blob/86a9e2620352ee0b2844bc3233f9e001cc4aa3d9/chromadb/test/ef/test_multimodal_ef.py) ## Documentation See #1294 ## TODOs - [x] Tests - [x] ~Wiring through FastAPI~ Nothing to wire through - [x] Documentation - [x] Telemetry - [ ] ~JavaScript~
chroma-core · Nov 7, 2023 · 4db9955 · 4db9955
1 parent ae0206b
commit 4db9955
Show file tree

Hide file tree

Showing 12 changed files with 696 additions and 186 deletions.
diff --git a/chromadb/api/__init__.py b/chromadb/api/__init__.py
@@ -8,6 +8,7 @@
 from chromadb.api.types import (
     CollectionMetadata,
     Documents,
+    Embeddable,
     EmbeddingFunction,
     Embeddings,
     IDs,
@@ -58,7 +59,9 @@ def create_collection(
         self,
         name: str,
         metadata: Optional[CollectionMetadata] = None,
-        embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(),
+        embedding_function: Optional[
+            EmbeddingFunction[Embeddable]
+        ] = ef.DefaultEmbeddingFunction(),  # type: ignore
         get_or_create: bool = False,
     ) -> Collection:
         """Create a new collection with the given name and metadata.
@@ -90,9 +93,11 @@ def create_collection(
     @abstractmethod
     def get_collection(
         self,
-        name: Optional[str] = None,
+        name: str,
         id: Optional[UUID] = None,
-        embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(),
+        embedding_function: Optional[
+            EmbeddingFunction[Embeddable]
+        ] = ef.DefaultEmbeddingFunction(),  # type: ignore
     ) -> Collection:
         """Get a collection with the given name.
         Args:
@@ -119,7 +124,9 @@ def get_or_create_collection(
         self,
         name: str,
         metadata: Optional[CollectionMetadata] = None,
-        embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(),
+        embedding_function: Optional[
+            EmbeddingFunction[Embeddable]
+        ] = ef.DefaultEmbeddingFunction(),  # type: ignore
     ) -> Collection:
         """Get or create a collection with the given name and metadata.
         Args:
@@ -486,7 +493,9 @@ def create_collection(
         self,
         name: str,
         metadata: Optional[CollectionMetadata] = None,
-        embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(),
+        embedding_function: Optional[
+            EmbeddingFunction[Embeddable]
+        ] = ef.DefaultEmbeddingFunction(),  # type: ignore
         get_or_create: bool = False,
         tenant: str = DEFAULT_TENANT,
         database: str = DEFAULT_DATABASE,
@@ -497,9 +506,11 @@ def create_collection(
     @override
     def get_collection(
         self,
-        name: Optional[str] = None,
+        name: str,
         id: Optional[UUID] = None,
-        embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(),
+        embedding_function: Optional[
+            EmbeddingFunction[Embeddable]
+        ] = ef.DefaultEmbeddingFunction(),  # type: ignore
         tenant: str = DEFAULT_TENANT,
         database: str = DEFAULT_DATABASE,
     ) -> Collection:
@@ -511,7 +522,9 @@ def get_or_create_collection(
         self,
         name: str,
         metadata: Optional[CollectionMetadata] = None,
-        embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(),
+        embedding_function: Optional[
+            EmbeddingFunction[Embeddable]
+        ] = ef.DefaultEmbeddingFunction(),  # type: ignore
         tenant: str = DEFAULT_TENANT,
         database: str = DEFAULT_DATABASE,
     ) -> Collection:

diff --git a/chromadb/api/fastapi.py b/chromadb/api/fastapi.py
@@ -14,6 +14,7 @@
 from chromadb.api.models.Collection import Collection
 from chromadb.api.types import (
     Documents,
+    Embeddable,
     Embeddings,
     EmbeddingFunction,
     IDs,
@@ -219,7 +220,9 @@ def create_collection(
         self,
         name: str,
         metadata: Optional[CollectionMetadata] = None,
-        embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(),
+        embedding_function: Optional[
+            EmbeddingFunction[Embeddable]
+        ] = ef.DefaultEmbeddingFunction(),  # type: ignore
         get_or_create: bool = False,
         tenant: str = DEFAULT_TENANT,
         database: str = DEFAULT_DATABASE,
@@ -250,9 +253,9 @@ def create_collection(
     @override
     def get_collection(
         self,
-        name: Optional[str] = None,
+        name: str,
         id: Optional[UUID] = None,
-        embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(),
+        embedding_function: Optional[EmbeddingFunction[Embeddable]] = ef.DefaultEmbeddingFunction(),  # type: ignore
         tenant: str = DEFAULT_TENANT,
         database: str = DEFAULT_DATABASE,
     ) -> Collection:
@@ -284,17 +287,20 @@ def get_or_create_collection(
         self,
         name: str,
         metadata: Optional[CollectionMetadata] = None,
-        embedding_function: Optional[EmbeddingFunction] = ef.DefaultEmbeddingFunction(),
+        embedding_function: Optional[EmbeddingFunction[Embeddable]] = ef.DefaultEmbeddingFunction(),  # type: ignore
         tenant: str = DEFAULT_TENANT,
         database: str = DEFAULT_DATABASE,
     ) -> Collection:
-        return self.create_collection(
-            name,
-            metadata,
-            embedding_function,
-            get_or_create=True,
-            tenant=tenant,
-            database=database,
+        return cast(
+            Collection,
+            self.create_collection(
+                name,
+                metadata,
+                embedding_function,
+                get_or_create=True,
+                tenant=tenant,
+                database=database,
+            ),
         )
 
     @trace_method("FastAPI._modify", OpenTelemetryGranularity.OPERATION)
@@ -347,10 +353,13 @@ def _peek(
         collection_id: UUID,
         n: int = 10,
     ) -> GetResult:
-        return self._get(
-            collection_id,
-            limit=n,
-            include=["embeddings", "documents", "metadatas"],
+        return cast(
+            GetResult,
+            self._get(
+                collection_id,
+                limit=n,
+                include=["embeddings", "documents", "metadatas"],
+            ),
         )
 
     @trace_method("FastAPI._get", OpenTelemetryGranularity.OPERATION)