feat: add pooling to HuggingFaceEmbedder, add CompositeEmbedder

Author: nnethercott

meilisearch/index.py

@@ -25,6 +25,7 @@
 from meilisearch.errors import version_error_hint_message
 from meilisearch.models.document import Document, DocumentsResults
 from meilisearch.models.index import (
+    CompositeEmbedder,
     Embedders,
     Faceting,
     HuggingFaceEmbedder,
@@ -961,6 +962,7 @@ def get_settings(self) -> Dict[str, Any]:
                 | HuggingFaceEmbedder
                 | OllamaEmbedder
                 | RestEmbedder
+                | CompositeEmbedder
                 | UserProvidedEmbedder,
             ] = {}
             for k, v in settings["embedders"].items():
@@ -972,6 +974,8 @@ def get_settings(self) -> Dict[str, Any]:
                     embedders[k] = HuggingFaceEmbedder(**v)
                 elif v.get("source") == "rest":
                     embedders[k] = RestEmbedder(**v)
+                elif v.get("source") == "composite":
+                    embedders[k] = CompositeEmbedder(**v)
                 else:
                     embedders[k] = UserProvidedEmbedder(**v)
 
@@ -1900,6 +1904,7 @@ def get_embedders(self) -> Embedders | None:
             | HuggingFaceEmbedder
             | OllamaEmbedder
             | RestEmbedder
+            | CompositeEmbedder
             | UserProvidedEmbedder,
         ] = {}
         for k, v in response.items():
@@ -1911,6 +1916,8 @@ def get_embedders(self) -> Embedders | None:
                 embedders[k] = HuggingFaceEmbedder(**v)
             elif v.get("source") == "rest":
                 embedders[k] = RestEmbedder(**v)
+            elif v.get("source") == "composite":
+                embedders[k] = CompositeEmbedder(**v)
             else:
                 embedders[k] = UserProvidedEmbedder(**v)
 

meilisearch/models/index.py

@@ -64,6 +64,12 @@ class LocalizedAttributes(CamelBase):
     locales: List[str]
 
 
+class PoolingOpt(str, Enum):
+    USE_MODEL = "useModel"
+    FORCE_MEAN = "forceMean"
+    FORCE_CLS = "forceCls"
+
+
 class OpenAiEmbedder(CamelBase):
     source: str = "openAi"
     url: Optional[str] = None
@@ -84,6 +90,7 @@ class HuggingFaceEmbedder(CamelBase):
     document_template_max_bytes: Optional[int] = None  # Default to 400
     distribution: Optional[EmbedderDistribution] = None
     binary_quantized: Optional[bool] = None
+    pooling: Optional[PoolingOpt] = None
 
 
 class OllamaEmbedder(CamelBase):
@@ -117,10 +124,33 @@ class UserProvidedEmbedder(CamelBase):
     binary_quantized: Optional[bool] = None
 
 
+class CompositeEmbedder(CamelBase):
+    source: str = "composite"
+    search_embedder: Union[
+        OpenAiEmbedder,
+        HuggingFaceEmbedder,
+        OllamaEmbedder,
+        RestEmbedder,
+        UserProvidedEmbedder
+    ]
+    indexing_embedder: Union[
+        OpenAiEmbedder,
+        HuggingFaceEmbedder,
+        OllamaEmbedder,
+        RestEmbedder,
+        UserProvidedEmbedder
+    ]
+
+
 class Embedders(CamelBase):
     embedders: Dict[
         str,
         Union[
-            OpenAiEmbedder, HuggingFaceEmbedder, OllamaEmbedder, RestEmbedder, UserProvidedEmbedder
+            OpenAiEmbedder,
+            HuggingFaceEmbedder,
+            OllamaEmbedder,
+            RestEmbedder,
+            UserProvidedEmbedder,
+            CompositeEmbedder,
         ],
     ]