Mypy fixes

Author: I8dNLo

fastembed/common/model_description.py

@@ -7,7 +7,7 @@ class ModelSource:
     hf: Optional[str] = None
     url: Optional[str] = None
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         if self.hf is None and self.url is None:
             raise ValueError(
                 f"At least one source should be set, current sources: hf={self.hf}, url={self.url}"
@@ -23,19 +23,11 @@ class ModelDescription:
 
     description: str
     license: str
-    size_in_GB: float
+    size_in_GB: Optional[float]
     additional_files: List[str] = field(default_factory=list)
     tasks: Dict[str, int] = field(default_factory=dict)
 
 
-# @dataclass(frozen=True, kw_only=True)
-# class SparseModelDescription(ModelDescription):
-#     vocab_size: int
-#     requires_idf: Optional[bool] = None
-#     # For sparse models, override dim to always be None.
-#     dim: Optional[int] = None
-
-
 @dataclass(frozen=True)
 class SparseModelDescription(ModelDescription):
     _vocab_size: InitVar[Optional[int]] = None
@@ -53,7 +45,7 @@ def __init__(
         model_file: str,
         description: str,
         license: str,
-        size_in_GB: float,
+        size_in_GB: Optional[float],
         dim: Optional[int] = None,
         additional_files: Optional[List[str]] = None,
         tasks: Optional[Dict[str, int]] = None,

fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py

@@ -41,7 +41,7 @@ def list_supported_models(cls) -> list[ModelDescription]:
                 ]
                 ```
         """
-        result: list[dict[str, Any]] = []
+        result: list[ModelDescription] = []
         for embedding in cls.EMBEDDINGS_REGISTRY:
             result.extend(embedding.list_supported_models())
         return result

fastembed/sparse/sparse_text_embedding.py

@@ -9,14 +9,14 @@
 )
 from fastembed.sparse.splade_pp import SpladePP
 import warnings
-from fastembed.common.model_description import SparseModelDescription
+from fastembed.common.model_description import ModelDescription
 
 
 class SparseTextEmbedding(SparseTextEmbeddingBase):
     EMBEDDINGS_REGISTRY: list[Type[SparseTextEmbeddingBase]] = [SpladePP, Bm42, Bm25]
 
     @classmethod
-    def list_supported_models(cls) -> list[SparseModelDescription]:
+    def list_supported_models(cls) -> list[ModelDescription]:
         """
         Lists the supported models.
 
@@ -39,7 +39,7 @@ def list_supported_models(cls) -> list[SparseModelDescription]:
                 ]
                 ```
         """
-        result: list[SparseModelDescription] = []
+        result: list[ModelDescription] = []
         for embedding in cls.EMBEDDINGS_REGISTRY:
             result.extend(embedding.list_supported_models())
         return result

fastembed/text/text_embedding.py

@@ -45,7 +45,7 @@ def list_supported_models(cls) -> list[ModelDescription]:
                 ]
                 ```
         """
-        result: list[dict[str, Any]] = []
+        result: list[ModelDescription] = []
         for embedding in cls.EMBEDDINGS_REGISTRY:
             result.extend(embedding.list_supported_models())
         return result