typing fix for colpali

Author: I8dNLo

fastembed/common/model_description.py

@@ -28,6 +28,11 @@ class ModelDescription:
     tasks: Dict[str, int] = field(default_factory=dict)
 
 
+@dataclass(frozen=True)
+class MultimodalModelDescription(ModelDescription):
+    dim: int
+
+
 @dataclass(frozen=True)
 class SparseModelDescription(ModelDescription):
     _vocab_size: InitVar[Optional[int]] = None

fastembed/late_interaction_multimodal/colpali.py

@@ -15,10 +15,10 @@
     TextEmbeddingWorker,
     ImageEmbeddingWorker,
 )
-from fastembed.common.model_description import ModelDescription, ModelSource
+from fastembed.common.model_description import MultimodalModelDescription, ModelSource
 
-supported_colpali_models: list[ModelDescription] = [
-    ModelDescription(
+supported_colpali_models: list[MultimodalModelDescription] = [
+    MultimodalModelDescription(
         model="Qdrant/colpali-v1.3-fp16",
         dim=128,
         description="Text embeddings, Multimodal (text&image), English, 50 tokens query length truncation, 2024.",
@@ -108,7 +108,7 @@ def __init__(
             self.load_onnx_model()
 
     @classmethod
-    def list_supported_models(cls) -> list[ModelDescription]:
+    def list_supported_models(cls) -> list[MultimodalModelDescription]:
         """Lists the supported models.
 
         Returns:

fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py

@@ -1,4 +1,4 @@
-from typing import Any, Iterable, Optional, Sequence, Type, Union
+from typing import Any, Iterable, Optional, Sequence, Type, Union, cast
 
 from fastembed.common import OnnxProvider, ImageInput
 from fastembed.common.types import NumpyArray
@@ -7,14 +7,14 @@
 from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import (
     LateInteractionMultimodalEmbeddingBase,
 )
-from fastembed.common.model_description import ModelDescription
+from fastembed.common.model_description import MultimodalModelDescription
 
 
 class LateInteractionMultimodalEmbedding(LateInteractionMultimodalEmbeddingBase):
     EMBEDDINGS_REGISTRY: list[Type[LateInteractionMultimodalEmbeddingBase]] = [ColPali]
 
     @classmethod
-    def list_supported_models(cls) -> list[ModelDescription]:
+    def list_supported_models(cls) -> list[MultimodalModelDescription]:
         """
         Lists the supported models.
 
@@ -41,9 +41,9 @@ def list_supported_models(cls) -> list[ModelDescription]:
                 ]
                 ```
         """
-        result: list[ModelDescription] = []
+        result: list[MultimodalModelDescription] = []
         for embedding in cls.EMBEDDINGS_REGISTRY:
-            result.extend(embedding.list_supported_models())
+            result.extend(cast(MultimodalModelDescription, embedding.list_supported_models()))
         return result
 
     def __init__(