wip: dataclass idea, small fixes (#475)

* wip: dataclass idea, small fixes

* fix: fix exception message in base model description

Author: joein

fastembed/common/model_description.py

@@ -1,5 +1,5 @@
-from dataclasses import dataclass, field, InitVar
-from typing import Optional, List, Dict
+from dataclasses import dataclass, field
+from typing import Optional, Any
 
 
 @dataclass(frozen=True)
@@ -15,67 +15,54 @@ def __post_init__(self) -> None:
 
 
 @dataclass(frozen=True)
-class ModelDescription:
+class BaseModelDescription:
     model: str
     sources: ModelSource
     model_file: str
-    dim: Optional[int]
+    description: str = ""
+    license: str = ""
+    size_in_GB: Optional[float] = None
+    additional_files: list[str] = field(default_factory=list)
 
-    description: str
-    license: str
-    size_in_GB: Optional[float]
-    additional_files: List[str] = field(default_factory=list)
-    tasks: Dict[str, int] = field(default_factory=dict)
+    def validate_info(self) -> None:
+        if self.license == "":
+            raise ValueError("license is required in builtin model description")
+
+        if self.description == "":
+            raise ValueError("description is required in builtin model description")
+
+        if self.size_in_GB is None:
+            raise ValueError("size_in_GB is required in builtin model description")
+
+    def __post_init__(self) -> None:
+        self.validate_info()
 
 
 @dataclass(frozen=True)
-class MultimodalModelDescription(ModelDescription):
-    dim: int
+class DenseModelDescription(BaseModelDescription):
+    dim: Optional[int] = None
+    tasks: Optional[dict[str, Any]] = None
+
+    def __post_init__(self) -> None:
+        assert self.dim is not None, "dim is required for dense model description"
+        self.validate_info()
 
 
 @dataclass(frozen=True)
-class SparseModelDescription(ModelDescription):
-    _vocab_size: InitVar[Optional[int]] = None
-    _requires_idf: InitVar[Optional[bool]] = None
-
-    vocab_size: int = field(init=False)
-    requires_idf: Optional[bool] = field(init=False, default=None)
-    dim: Optional[int] = field(default=None, init=False)
-
-    def __init__(
-        self,
-        *,
-        model: str,
-        sources: ModelSource,
-        model_file: str,
-        description: str,
-        license: str,
-        size_in_GB: Optional[float],
-        dim: Optional[int] = None,
-        additional_files: Optional[List[str]] = None,
-        tasks: Optional[Dict[str, int]] = None,
-        vocab_size: int,
-        requires_idf: Optional[bool] = None,
-    ):
-        # Call the parent initializer with the fields it needs.
-        object.__setattr__(self, "model", model)
-        object.__setattr__(self, "sources", sources)
-        object.__setattr__(self, "model_file", model_file)
-        object.__setattr__(self, "dim", dim if dim else None)
-        object.__setattr__(self, "description", description)
-        object.__setattr__(self, "license", license)
-        object.__setattr__(self, "size_in_GB", size_in_GB)
-        object.__setattr__(
-            self, "additional_files", additional_files if additional_files is not None else []
-        )
-        object.__setattr__(self, "tasks", tasks if tasks is not None else {})
-        # Set new fields.
-        object.__setattr__(self, "vocab_size", vocab_size)
-        object.__setattr__(self, "requires_idf", requires_idf)
+class SparseModelDescription(BaseModelDescription):
+    requires_idf: Optional[bool] = None
+    vocab_size: Optional[int] = None
 
 
 @dataclass(frozen=True)
-class CustomModelDescription(ModelDescription):
-    description: str = ""
-    license: str = ""
-    size_in_GB: Optional[float] = None
+class CustomDenseModelDescription(DenseModelDescription):
+    def __post_init__(self) -> None:
+        if self.dim is None:
+            raise ValueError("dim is required for custom dense model description")
+        # disable self.validate_info
+
+
+@dataclass(frozen=True)
+class CustomSparseModelDescription(SparseModelDescription):
+    def __post_init__(self) -> None:
+        pass  # disable self.validate_info

fastembed/common/model_management.py

@@ -4,7 +4,7 @@
 import shutil
 import tarfile
 from pathlib import Path
-from typing import Any, Optional, Union, Sequence, TypeVar, Generic
+from typing import Any, Optional, Union, TypeVar, Generic
 
 import requests
 from huggingface_hub import snapshot_download, model_info, list_repo_tree
@@ -16,16 +16,16 @@
 )
 from loguru import logger
 from tqdm import tqdm
-from fastembed.common.model_description import ModelDescription
+from fastembed.common.model_description import BaseModelDescription
 
-T = TypeVar("T", bound=ModelDescription)
+T = TypeVar("T", bound=BaseModelDescription)
 
 
 class ModelManagement(Generic[T]):
     METADATA_FILE = "files_metadata.json"
 
     @classmethod
-    def list_supported_models(cls) -> Sequence[T]:
+    def list_supported_models(cls) -> list[T]:
         """Lists the supported models.
 
         Returns:

fastembed/image/image_embedding.py

@@ -4,14 +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
+from fastembed.common.model_description import DenseModelDescription
 
 
 class ImageEmbedding(ImageEmbeddingBase):
     EMBEDDINGS_REGISTRY: list[Type[ImageEmbeddingBase]] = [OnnxImageEmbedding]
 
     @classmethod
-    def list_supported_models(cls) -> list[ModelDescription]:
+    def list_supported_models(cls) -> list[DenseModelDescription]:
         """
         Lists the supported models.
 
@@ -35,7 +35,7 @@ def list_supported_models(cls) -> list[ModelDescription]:
                 ]
                 ```
         """
-        result: list[ModelDescription] = []
+        result: list[DenseModelDescription] = []
         for embedding in cls.EMBEDDINGS_REGISTRY:
             result.extend(embedding.list_supported_models())
         return result

fastembed/image/image_embedding_base.py

@@ -1,11 +1,12 @@
 from typing import Iterable, Optional, Any, Union
 
+from fastembed.common.model_description import DenseModelDescription
 from fastembed.common.types import NumpyArray
 from fastembed.common.model_management import ModelManagement
 from fastembed.common.types import ImageInput
 
 
-class ImageEmbeddingBase(ModelManagement):
+class ImageEmbeddingBase(ModelManagement[DenseModelDescription]):
     def __init__(
         self,
         model_name: str,

fastembed/image/onnx_embedding.py

@@ -9,10 +9,10 @@
 from fastembed.image.image_embedding_base import ImageEmbeddingBase
 from fastembed.image.onnx_image_model import ImageEmbeddingWorker, OnnxImageModel
 
-from fastembed.common.model_description import ModelDescription, ModelSource
+from fastembed.common.model_description import DenseModelDescription, ModelSource
 
-supported_onnx_models: list[ModelDescription] = [
-    ModelDescription(
+supported_onnx_models: list[DenseModelDescription] = [
+    DenseModelDescription(
         model="Qdrant/clip-ViT-B-32-vision",
         dim=512,
         description="Image embeddings, Multimodal (text&image), 2021 year",
@@ -21,7 +21,7 @@
         sources=ModelSource(hf="Qdrant/clip-ViT-B-32-vision"),
         model_file="model.onnx",
     ),
-    ModelDescription(
+    DenseModelDescription(
         model="Qdrant/resnet50-onnx",
         dim=2048,
         description="Image embeddings, Unimodal (image), 2016 year",
@@ -30,7 +30,7 @@
         sources=ModelSource(hf="Qdrant/resnet50-onnx"),
         model_file="model.onnx",
     ),
-    ModelDescription(
+    DenseModelDescription(
         model="Qdrant/Unicom-ViT-B-16",
         dim=768,
         description="Image embeddings (more detailed than Unicom-ViT-B-32), Multimodal (text&image), 2023 year",
@@ -39,7 +39,7 @@
         sources=ModelSource(hf="Qdrant/Unicom-ViT-B-16"),
         model_file="model.onnx",
     ),
-    ModelDescription(
+    DenseModelDescription(
         model="Qdrant/Unicom-ViT-B-32",
         dim=512,
         description="Image embeddings, Multimodal (text&image), 2023 year",
@@ -48,7 +48,7 @@
         sources=ModelSource(hf="Qdrant/Unicom-ViT-B-32"),
         model_file="model.onnx",
     ),
-    ModelDescription(
+    DenseModelDescription(
         model="jinaai/jina-clip-v1",
         dim=768,
         description="Image embeddings, Multimodal (text&image), 2024 year",
@@ -137,7 +137,7 @@ def load_onnx_model(self) -> None:
         )
 
     @classmethod
-    def list_supported_models(cls) -> list[ModelDescription]:
+    def list_supported_models(cls) -> list[DenseModelDescription]:
         """
         Lists the supported models.
 

fastembed/late_interaction/colbert.py

@@ -12,10 +12,10 @@
     LateInteractionTextEmbeddingBase,
 )
 from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
-from fastembed.common.model_description import ModelDescription, ModelSource
+from fastembed.common.model_description import DenseModelDescription, ModelSource
 
-supported_colbert_models: list[ModelDescription] = [
-    ModelDescription(
+supported_colbert_models: list[DenseModelDescription] = [
+    DenseModelDescription(
         model="colbert-ir/colbertv2.0",
         dim=128,
         description="Late interaction model",
@@ -24,7 +24,7 @@
         sources=ModelSource(hf="colbert-ir/colbertv2.0"),
         model_file="model.onnx",
     ),
-    ModelDescription(
+    DenseModelDescription(
         model="answerdotai/answerai-colbert-small-v1",
         dim=96,
         description="Text embeddings, Unimodal (text), Multilingual (~100 languages), 512 input tokens truncation, 2024 year",
@@ -108,7 +108,7 @@ def _tokenize_documents(self, documents: list[str]) -> list[Encoding]:
         return encoded
 
     @classmethod
-    def list_supported_models(cls) -> list[ModelDescription]:
+    def list_supported_models(cls) -> list[DenseModelDescription]:
         """Lists the supported models.
 
         Returns:

fastembed/late_interaction/jina_colbert.py

@@ -2,10 +2,10 @@
 
 from fastembed.common.types import NumpyArray
 from fastembed.late_interaction.colbert import Colbert, ColbertEmbeddingWorker
-from fastembed.common.model_description import ModelDescription, ModelSource
+from fastembed.common.model_description import DenseModelDescription, ModelSource
 
-supported_jina_colbert_models: list[ModelDescription] = [
-    ModelDescription(
+supported_jina_colbert_models: list[DenseModelDescription] = [
+    DenseModelDescription(
         model="jinaai/jina-colbert-v2",
         dim=128,
         description="New model that expands capabilities of colbert-v1 with multilingual and context length of 8192, 2024 year",
@@ -29,7 +29,7 @@ def _get_worker_class(cls) -> Type[ColbertEmbeddingWorker]:
         return JinaColbertEmbeddingWorker
 
     @classmethod
-    def list_supported_models(cls) -> list[ModelDescription]:
+    def list_supported_models(cls) -> list[DenseModelDescription]:
         """Lists the supported models.
 
         Returns:

fastembed/late_interaction/late_interaction_embedding_base.py

@@ -1,10 +1,11 @@
 from typing import Iterable, Optional, Union, Any
 
+from fastembed.common.model_description import DenseModelDescription
 from fastembed.common.types import NumpyArray
 from fastembed.common.model_management import ModelManagement
 
 
-class LateInteractionTextEmbeddingBase(ModelManagement):
+class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
     def __init__(
         self,
         model_name: str,

fastembed/late_interaction/late_interaction_text_embedding.py

@@ -7,19 +7,19 @@
 from fastembed.late_interaction.late_interaction_embedding_base import (
     LateInteractionTextEmbeddingBase,
 )
-from fastembed.common.model_description import ModelDescription
+from fastembed.common.model_description import DenseModelDescription
 
 
 class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
     EMBEDDINGS_REGISTRY: list[Type[LateInteractionTextEmbeddingBase]] = [Colbert, JinaColbert]
 
     @classmethod
-    def list_supported_models(cls) -> list[ModelDescription]:
+    def list_supported_models(cls) -> list[DenseModelDescription]:
         """
         Lists the supported models.
 
         Returns:
-            list[ModelDescription]: A list of dictionaries containing the model information.
+            list[DenseModelDescription]: A list of dictionaries containing the model information.
 
             Example:
                 ```
@@ -38,7 +38,7 @@ def list_supported_models(cls) -> list[ModelDescription]:
                 ]
                 ```
         """
-        result: list[ModelDescription] = []
+        result: list[DenseModelDescription] = []
         for embedding in cls.EMBEDDINGS_REGISTRY:
             result.extend(embedding.list_supported_models())
         return result

fastembed/late_interaction_multimodal/colpali.py

@@ -15,10 +15,10 @@
     TextEmbeddingWorker,
     ImageEmbeddingWorker,
 )
-from fastembed.common.model_description import MultimodalModelDescription, ModelSource
+from fastembed.common.model_description import DenseModelDescription, ModelSource
 
-supported_colpali_models: list[MultimodalModelDescription] = [
-    MultimodalModelDescription(
+supported_colpali_models: list[DenseModelDescription] = [
+    DenseModelDescription(
         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[MultimodalModelDescription]:
+    def list_supported_models(cls) -> list[DenseModelDescription]:
         """Lists the supported models.
 
         Returns:
@@ -139,6 +139,7 @@ def _post_process_onnx_image_output(
         Returns:
             Iterable[NumpyArray]: Post-processed output as NumPy arrays.
         """
+        assert self.model_description.dim is not None, "Model dim is not defined"
         return output.model_output.reshape(
             output.model_output.shape[0], -1, self.model_description.dim
         ).astype(np.float32)