list_supported_models type fix

I8dNLo · I8dNLo · commit f4e29fb7e9c4 · 2025-02-11T19:00:46.000+01:00
diff --git a/fastembed/image/image_embedding.py b/fastembed/image/image_embedding.py
@@ -4,13 +4,14 @@
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.image.image_embedding_base import ImageEmbeddingBase
 from fastembed.image.onnx_embedding import OnnxImageEmbedding
+from fastembed.common.model_description import ModelDescription
 
 
 class ImageEmbedding(ImageEmbeddingBase):
     EMBEDDINGS_REGISTRY: list[Type[ImageEmbeddingBase]] = [OnnxImageEmbedding]
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def list_supported_models(cls) -> list[ModelDescription]:
         """
         Lists the supported models.
 
diff --git a/fastembed/image/onnx_embedding.py b/fastembed/image/onnx_embedding.py
@@ -137,7 +137,7 @@ def load_onnx_model(self) -> None:
         )
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def list_supported_models(cls) -> list[ModelDescription]:
         """
         Lists the supported models.
 
diff --git a/fastembed/late_interaction/colbert.py b/fastembed/late_interaction/colbert.py
@@ -108,7 +108,7 @@ def _tokenize_documents(self, documents: list[str]) -> list[Encoding]:
         return encoded
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def list_supported_models(cls) -> list[ModelDescription]:
         """Lists the supported models.
 
         Returns:
diff --git a/fastembed/late_interaction/jina_colbert.py b/fastembed/late_interaction/jina_colbert.py
@@ -29,7 +29,7 @@ def _get_worker_class(cls) -> Type[ColbertEmbeddingWorker]:
         return JinaColbertEmbeddingWorker
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def list_supported_models(cls) -> list[ModelDescription]:
         """Lists the supported models.
 
         Returns:
diff --git a/fastembed/late_interaction/late_interaction_text_embedding.py b/fastembed/late_interaction/late_interaction_text_embedding.py
@@ -7,13 +7,14 @@
 from fastembed.late_interaction.late_interaction_embedding_base import (
     LateInteractionTextEmbeddingBase,
 )
+from fastembed.common.model_description import ModelDescription
 
 
 class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
     EMBEDDINGS_REGISTRY: list[Type[LateInteractionTextEmbeddingBase]] = [Colbert, JinaColbert]
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def list_supported_models(cls) -> list[ModelDescription]:
         """
         Lists the supported models.
 
@@ -37,7 +38,7 @@ def list_supported_models(cls) -> list[dict[str, Any]]:
                 ]
                 ```
         """
-        result: list[dict[str, Any]] = []
+        result: list[ModelDescription] = []
         for embedding in cls.EMBEDDINGS_REGISTRY:
             result.extend(embedding.list_supported_models())
         return result
diff --git a/fastembed/late_interaction_multimodal/colpali.py b/fastembed/late_interaction_multimodal/colpali.py
@@ -108,7 +108,7 @@ def __init__(
             self.load_onnx_model()
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def list_supported_models(cls) -> list[ModelDescription]:
         """Lists the supported models.
 
         Returns:
diff --git a/fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py b/fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py
@@ -7,13 +7,14 @@
 from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import (
     LateInteractionMultimodalEmbeddingBase,
 )
+from fastembed.common.model_description import ModelDescription
 
 
 class LateInteractionMultimodalEmbedding(LateInteractionMultimodalEmbeddingBase):
     EMBEDDINGS_REGISTRY: list[Type[LateInteractionMultimodalEmbeddingBase]] = [ColPali]
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def list_supported_models(cls) -> list[ModelDescription]:
         """
         Lists the supported models.
 
diff --git a/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py b/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py
@@ -72,7 +72,7 @@
 
 class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def list_supported_models(cls) -> list[ModelDescription]:
         """Lists the supported models.
 
         Returns:
diff --git a/fastembed/rerank/cross_encoder/text_cross_encoder.py b/fastembed/rerank/cross_encoder/text_cross_encoder.py
@@ -3,6 +3,7 @@
 from fastembed.common import OnnxProvider
 from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
 from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase
+from fastembed.common.model_description import ModelDescription
 
 
 class TextCrossEncoder(TextCrossEncoderBase):
@@ -11,7 +12,7 @@ class TextCrossEncoder(TextCrossEncoderBase):
     ]
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def list_supported_models(cls) -> list[ModelDescription]:
         """Lists the supported models.
 
         Returns:
diff --git a/fastembed/sparse/bm25.py b/fastembed/sparse/bm25.py
@@ -147,7 +147,7 @@ def __init__(
         self.tokenizer = SimpleTokenizer
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def list_supported_models(cls) -> list[SparseModelDescription]:
         """Lists the supported models.
 
         Returns:
diff --git a/fastembed/sparse/bm42.py b/fastembed/sparse/bm42.py
@@ -251,7 +251,7 @@ def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[Spars
             yield SparseEmbedding.from_dict(rescored)
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def list_supported_models(cls) -> list[SparseModelDescription]:
         """Lists the supported models.
 
         Returns:
diff --git a/fastembed/sparse/sparse_text_embedding.py b/fastembed/sparse/sparse_text_embedding.py
@@ -9,13 +9,14 @@
 )
 from fastembed.sparse.splade_pp import SpladePP
 import warnings
+from fastembed.common.model_description import SparseModelDescription
 
 
 class SparseTextEmbedding(SparseTextEmbeddingBase):
     EMBEDDINGS_REGISTRY: list[Type[SparseTextEmbeddingBase]] = [SpladePP, Bm42, Bm25]
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def list_supported_models(cls) -> list[SparseModelDescription]:
         """
         Lists the supported models.
 
diff --git a/fastembed/sparse/splade_pp.py b/fastembed/sparse/splade_pp.py
@@ -52,7 +52,7 @@ def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[Spars
             yield SparseEmbedding(values=scores, indices=indices)
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def list_supported_models(cls) -> list[SparseModelDescription]:
         """Lists the supported models.
 
         Returns:
diff --git a/fastembed/text/clip_embedding.py b/fastembed/text/clip_embedding.py
@@ -27,7 +27,7 @@ def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
         return CLIPEmbeddingWorker
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def list_supported_models(cls) -> list[ModelDescription]:
         """Lists the supported models.
 
         Returns:
diff --git a/fastembed/text/multitask_embedding.py b/fastembed/text/multitask_embedding.py
@@ -53,7 +53,7 @@ def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
         return JinaEmbeddingV3Worker
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def list_supported_models(cls) -> list[ModelDescription]:
         return supported_multitask_models
 
     def _preprocess_onnx_input(
diff --git a/fastembed/text/onnx_embedding.py b/fastembed/text/onnx_embedding.py
@@ -196,7 +196,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
     """Implementation of the Flag Embedding model."""
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def list_supported_models(cls) -> list[ModelDescription]:
         """
         Lists the supported models.
 
diff --git a/fastembed/text/pooled_embedding.py b/fastembed/text/pooled_embedding.py
@@ -105,7 +105,7 @@ def mean_pooling(cls, model_output: NumpyArray, attention_mask: NumpyArray) -> N
         return pooled_embeddings
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def list_supported_models(cls) -> list[ModelDescription]:
         """Lists the supported models.
 
         Returns:
diff --git a/fastembed/text/pooled_normalized_embedding.py b/fastembed/text/pooled_normalized_embedding.py
@@ -118,7 +118,7 @@ def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
         return PooledNormalizedEmbeddingWorker
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def list_supported_models(cls) -> list[ModelDescription]:
         """Lists the supported models.
 
         Returns:
diff --git a/fastembed/text/text_embedding.py b/fastembed/text/text_embedding.py
@@ -8,6 +8,7 @@
 from fastembed.text.multitask_embedding import JinaEmbeddingV3
 from fastembed.text.onnx_embedding import OnnxTextEmbedding
 from fastembed.text.text_embedding_base import TextEmbeddingBase
+from fastembed.common.model_description import ModelDescription
 
 
 class TextEmbedding(TextEmbeddingBase):
@@ -20,7 +21,7 @@ class TextEmbedding(TextEmbeddingBase):
     ]
 
     @classmethod
-    def list_supported_models(cls) -> list[dict[str, Any]]:
+    def list_supported_models(cls) -> list[ModelDescription]:
         """
         Lists the supported models.
 

Original file line number	Diff line number	Diff line change
`@@ -137,7 +137,7 @@ def load_onnx_model(self) -> None:`
`137`	`137`	`)`
`138`	`138`
`139`	`139`	`@classmethod`
`140`		`- def list_supported_models(cls) -> list[dict[str, Any]]:`
	`140`	`+ def list_supported_models(cls) -> list[ModelDescription]:`
`141`	`141`	`"""`
`142`	`142`	`Lists the supported models.`
`143`	`143`