feat: Support numpy-based Chroma embeddings (#177)

Closes #176

Author: tazarov

chroma_dp/__init__.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 try:
     import chromadb  # noqa: F401
 except ImportError:
@@ -19,7 +21,7 @@
 )
 
 from chromadb.api.types import Embedding
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, ConfigDict
 
 C = TypeVar("C")
 
@@ -31,17 +33,22 @@ class ResourceFeature(BaseModel, Generic[C]):
 
 Metadata = Dict[str, Union[str, int, float, bool]]
 
+EmbeddingWrapper = Union[Embedding, np.ndarray]
+
 
 class EmbeddableResource(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
     id: Optional[str] = Field(None, description="Document ID")
     metadata: Optional[Metadata] = Field(None, description="Document metadata")
-    embedding: Optional[Embedding] = Field(None, description="Document embedding")
+    embedding: Optional[EmbeddingWrapper] = Field(
+        None, description="Document embedding"
+    )
 
     @staticmethod
     def resource_features() -> Sequence[ResourceFeature]:
         return [
-            ResourceFeature[Embedding](
-                feature_name="embedding", feature_type=Embedding
+            ResourceFeature[EmbeddingWrapper](
+                feature_name="embedding", feature_type=EmbeddingWrapper
             ),
             ResourceFeature[Metadata](feature_name="metadata", eature_type=Metadata),
             ResourceFeature[str](feature_name="id", feature_type=str),
@@ -58,6 +65,13 @@ def resource_features() -> Sequence[ResourceFeature]:
             *super().resource_features(),
         ]
 
+    def model_dump(self, **kwargs):
+        # Convert NumPy arrays to lists before dumping
+        data = super().model_dump(**kwargs)
+        if isinstance(data["embedding"], np.ndarray):
+            data["embedding"] = data["embedding"].tolist()
+        return data
+
 
 D = TypeVar("D", bound=EmbeddableResource, contravariant=True)
 

chroma_dp/chroma/chroma_import.py

@@ -1,3 +1,4 @@
+import numpy as np
 import orjson as json
 import sys
 import uuid
@@ -28,13 +29,19 @@ def add_to_col(
     ef: EmbeddingFunction = None,
 ) -> None:
     try:
+        if "embeddings" in batch and len(batch["embeddings"]) > 0:
+            batch["embeddings"] = [
+                e.tolist() if isinstance(e, np.ndarray) else e
+                for e in batch["embeddings"]
+            ]
         if ef:
             batch["embeddings"] = ef(batch["documents"])
         if upsert:
             col.upsert(**batch)
         else:
             col.add(**batch)
     except Exception as e:
+        print(e, file=sys.stderr)
         raise e
 
 

chroma_dp/utils/chroma.py

@@ -3,6 +3,7 @@
 from typing import Optional, Dict, Any, List, cast
 from urllib.parse import urlparse, parse_qs
 import chromadb
+import numpy as np
 from chromadb import ClientAPI, GetResult
 from chromadb.api.models.Collection import Collection
 
@@ -223,6 +224,8 @@ def remap_features(
     _doc = in_dict[doc_feature]
     _embed = in_dict[embed_feature] if embed_feature else None
     _id = in_dict[id_feature] if id_feature else None
+    if not isinstance(_embed, np.ndarray):
+        _embed = np.array(_embed)
     return EmbeddableTextResource(
         text_chunk=_doc, embedding=_embed, metadata=_meta, id=_id
     )

docs/index.md

@@ -67,7 +67,7 @@ cdp ds-get "hf://tazarov/chroma-qna?split=train" | cdp import "http://localhost:
 **Importing from a directory with PDF files into Local Persisted Chroma DB:**
 
 ```bash
-cdp ds-get sample-data/papers/ | grep "2401.02412.pdf" | head -1 | cdp chunk -s 500 | cdp embed --ef default | cdp import "file://chroma-data/my-pdfs" --upsert --create
+cdp imp pdf sample-data/papers/ | grep "2401.02412.pdf" | head -1 | cdp chunk -s 500 | cdp embed --ef default | cdp import "file://chroma-data/my-pdfs" --upsert --create
 ```
 
 !!! note