draft: move custom models logic to text embedding

Author: joein

fastembed/text/clip_embedding.py

@@ -22,8 +22,6 @@
 
 
 class CLIPOnnxEmbedding(OnnxTextEmbedding):
-    supported_models = supported_clip_models
-
     @classmethod
     def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
         return CLIPEmbeddingWorker
@@ -37,10 +35,6 @@ def list_supported_models(cls) -> list[dict[str, Any]]:
         """
         return cls.supported_models
 
-    @classmethod
-    def add_custom_model(cls, model_info: dict[str, Any]):
-        cls.supported_models.append(model_info)
-
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
         return output.model_output
 

fastembed/text/multitask_embedding.py

@@ -55,10 +55,6 @@ def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
     def list_supported_models(cls) -> list[dict[str, Any]]:
         return cls.supported_models
 
-    @classmethod
-    def add_custom_model(cls, model_info: dict[str, Any]):
-        cls.supported_models.append(model_info)
-
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, np.ndarray], **kwargs
     ) -> dict[str, np.ndarray]:

fastembed/text/pooled_embedding.py

@@ -79,8 +79,6 @@
 
 
 class PooledEmbedding(OnnxTextEmbedding):
-    supported_models = supported_pooled_models
-
     @classmethod
     def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
         return PooledEmbeddingWorker
@@ -105,10 +103,6 @@ def list_supported_models(cls) -> list[dict[str, Any]]:
         """
         return cls.supported_models
 
-    @classmethod
-    def add_custom_model(cls, model_info: dict[str, Any]):
-        cls.supported_models.append(model_info)
-
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
         if output.attention_mask is None:
             raise ValueError("attention_mask must be provided for document post-processing")

fastembed/text/pooled_normalized_embedding.py

@@ -88,8 +88,6 @@
 
 
 class PooledNormalizedEmbedding(PooledEmbedding):
-    supported_models = supported_pooled_normalized_models
-
     @classmethod
     def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
         return PooledNormalizedEmbeddingWorker
@@ -103,10 +101,6 @@ def list_supported_models(cls) -> list[dict[str, Any]]:
         """
         return cls.supported_models
 
-    @classmethod
-    def add_custom_model(cls, model_info: dict[str, Any]):
-        cls.supported_models.append(model_info)
-
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
         if output.attention_mask is None:
             raise ValueError("attention_mask must be provided for document post-processing")

fastembed/text/text_embedding.py

@@ -1,4 +1,5 @@
 import warnings
+from collections import defaultdict
 from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
@@ -19,6 +20,9 @@ class TextEmbedding(TextEmbeddingBase):
         PooledEmbedding,
         JinaEmbeddingV3,
     ]
+    CUSTOM_EMBEDDINGS_REGISTRY: dict[Type[TextEmbeddingBase], list[dict[str, Any]]] = defaultdict(
+        list
+    )
 
     @classmethod
     def list_supported_models(cls) -> list[dict[str, Any]]:
@@ -48,6 +52,9 @@ def list_supported_models(cls) -> list[dict[str, Any]]:
         result = []
         for embedding in cls.EMBEDDINGS_REGISTRY:
             result.extend(embedding.list_supported_models())
+        for embedding, models in cls.CUSTOM_EMBEDDINGS_REGISTRY.items():
+            for model in models:
+                result.append(model)
         return result
 
     @classmethod
@@ -74,13 +81,13 @@ def add_custom_model(
             None
         """
         if mean_pooling and not normalization:
-            PooledEmbedding.add_custom_model(model_info)
+            cls.CUSTOM_EMBEDDINGS_REGISTRY[PooledEmbedding].append(model_info)
         elif mean_pooling and normalization:
-            PooledNormalizedEmbedding.add_custom_model(model_info)
-        elif "clip" in model_info["model"].lower():
-            CLIPOnnxEmbedding.add_custom_model(model_info)
+            cls.CUSTOM_EMBEDDINGS_REGISTRY[PooledNormalizedEmbedding].append(model_info)
+        elif not mean_pooling and not normalization:
+            cls.CUSTOM_EMBEDDINGS_REGISTRY[OnnxTextEmbedding].append(model_info)
         else:
-            OnnxTextEmbedding.add_custom_model(model_info)
+            cls.CUSTOM_EMBEDDINGS_REGISTRY[PooledNormalizedEmbedding].append(model_info)
 
     def __init__(
         self,