fastembed-gpu 0.7.0__tar.gz → 0.7.1__tar.gz

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: fastembed-gpu
-Version: 0.7.0
+Version: 0.7.1
 Summary: Fast, light, accurate library built for retrieval embedding generation
 License: Apache License
 Keywords: vector,embedding,neural,search,qdrant,sentence-transformers

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/image/image_embedding.py

@@ -77,6 +77,40 @@ class ImageEmbedding(ImageEmbeddingBase):
             "Please check the supported models using `ImageEmbedding.list_supported_models()`"
         )
 
+    @property
+    def embedding_size(self) -> int:
+        """Get the embedding size of the current model"""
+        if self._embedding_size is None:
+            self._embedding_size = self.get_embedding_size(self.model_name)
+        return self._embedding_size
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Get the embedding size of the passed model
+
+        Args:
+            model_name (str): The name of the model to get embedding size for.
+
+        Returns:
+            int: The size of the embedding.
+
+        Raises:
+            ValueError: If the model name is not found in the supported models.
+        """
+        descriptions = cls._list_supported_models()
+        embedding_size: Optional[int] = None
+        for description in descriptions:
+            if description.model.lower() == model_name.lower():
+                embedding_size = description.dim
+                break
+        if embedding_size is None:
+            model_names = [description.model for description in descriptions]
+            raise ValueError(
+                f"Embedding size for model {model_name} was None. "
+                f"Available model names: {model_names}"
+            )
+        return embedding_size
+
     def embed(
         self,
         images: Union[ImageInput, Iterable[ImageInput]],

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/image/image_embedding_base.py

@@ -18,6 +18,7 @@ class ImageEmbeddingBase(ModelManagement[DenseModelDescription]):
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
+        self._embedding_size: Optional[int] = None
 
     def embed(
         self,
@@ -42,3 +43,13 @@ class ImageEmbeddingBase(ModelManagement[DenseModelDescription]):
             Iterable[NdArray]: The embeddings.
         """
         raise NotImplementedError()
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Returns embedding size of the chosen model."""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @property
+    def embedding_size(self) -> int:
+        """Returns embedding size for the current model"""
+        raise NotImplementedError("Subclasses must implement this method")

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/image/onnx_embedding.py

@@ -112,11 +112,12 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
 
         self.model_description = self._get_model_description(model_name)
         self.cache_dir = str(define_cache_dir(cache_dir))
+        self._specific_model_path = specific_model_path
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
 
         if not self.lazy_load:
@@ -177,6 +178,8 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
             providers=self.providers,
             cuda=self.cuda,
             device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
             **kwargs,
         )
 

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/image/onnx_image_model.py

@@ -97,6 +97,8 @@ class OnnxImageModel(OnnxModel[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
+        local_files_only: bool = False,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -123,6 +125,8 @@ class OnnxImageModel(OnnxModel[T]):
                 "model_name": model_name,
                 "cache_dir": cache_dir,
                 "providers": providers,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
                 **kwargs,
             }
 

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/late_interaction/colbert.py

@@ -169,11 +169,12 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         self.model_description = self._get_model_description(model_name)
         self.cache_dir = str(define_cache_dir(cache_dir))
 
+        self._specific_model_path = specific_model_path
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
         self.mask_token_id: Optional[int] = None
         self.pad_token_id: Optional[int] = None
@@ -233,6 +234,8 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
             providers=self.providers,
             cuda=self.cuda,
             device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
             **kwargs,
         )
 

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/late_interaction/late_interaction_embedding_base.py

@@ -17,6 +17,7 @@ class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
+        self._embedding_size: Optional[int] = None
 
     def embed(
         self,
@@ -58,3 +59,13 @@ class LateInteractionTextEmbeddingBase(ModelManagement[DenseModelDescription]):
             yield from self.embed([query], **kwargs)
         else:
             yield from self.embed(query, **kwargs)
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Returns embedding size of the chosen model."""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @property
+    def embedding_size(self) -> int:
+        """Returns embedding size for the current model"""
+        raise NotImplementedError("Subclasses must implement this method")

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/late_interaction/late_interaction_text_embedding.py

@@ -80,6 +80,40 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
             "Please check the supported models using `LateInteractionTextEmbedding.list_supported_models()`"
         )
 
+    @property
+    def embedding_size(self) -> int:
+        """Get the embedding size of the current model"""
+        if self._embedding_size is None:
+            self._embedding_size = self.get_embedding_size(self.model_name)
+        return self._embedding_size
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Get the embedding size of the passed model
+
+        Args:
+            model_name (str): The name of the model to get embedding size for.
+
+        Returns:
+            int: The size of the embedding.
+
+        Raises:
+            ValueError: If the model name is not found in the supported models.
+        """
+        descriptions = cls._list_supported_models()
+        embedding_size: Optional[int] = None
+        for description in descriptions:
+            if description.model.lower() == model_name.lower():
+                embedding_size = description.dim
+                break
+        if embedding_size is None:
+            model_names = [description.model for description in descriptions]
+            raise ValueError(
+                f"Embedding size for model {model_name} was None. "
+                f"Available model names: {model_names}"
+            )
+        return embedding_size
+
     def embed(
         self,
         documents: Union[str, Iterable[str]],

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/late_interaction/token_embeddings.py

@@ -9,7 +9,7 @@ from fastembed.late_interaction.late_interaction_embedding_base import (
 )
 from fastembed.text.onnx_embedding import OnnxTextEmbedding
 from fastembed.text.onnx_text_model import TextEmbeddingWorker
-import numpy as np
+
 
 supported_token_embeddings_models = [
     DenseModelDescription(

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/late_interaction_multimodal/colpali.py

@@ -95,11 +95,12 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
         self.model_description = self._get_model_description(model_name)
         self.cache_dir = str(define_cache_dir(cache_dir))
 
+        self._specific_model_path = specific_model_path
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
         self.mask_token_id = None
         self.pad_token_id = None
@@ -174,7 +175,7 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
     ) -> dict[str, NumpyArray]:
         onnx_input["input_ids"] = np.array(
             [
-                self.QUERY_MARKER_TOKEN_ID + input_ids[2:].tolist()
+                self.QUERY_MARKER_TOKEN_ID + input_ids[2:].tolist()  # type: ignore[index]
                 for input_ids in onnx_input["input_ids"]
             ]
         )
@@ -235,6 +236,8 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
             providers=self.providers,
             cuda=self.cuda,
             device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
             **kwargs,
         )
 
@@ -268,6 +271,8 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
             providers=self.providers,
             cuda=self.cuda,
             device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
             **kwargs,
         )
 

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding.py

@@ -83,6 +83,40 @@ class LateInteractionMultimodalEmbedding(LateInteractionMultimodalEmbeddingBase)
             "Please check the supported models using `LateInteractionMultimodalEmbedding.list_supported_models()`"
         )
 
+    @property
+    def embedding_size(self) -> int:
+        """Get the embedding size of the current model"""
+        if self._embedding_size is None:
+            self._embedding_size = self.get_embedding_size(self.model_name)
+        return self._embedding_size
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Get the embedding size of the passed model
+
+        Args:
+            model_name (str): The name of the model to get embedding size for.
+
+        Returns:
+            int: The size of the embedding.
+
+        Raises:
+            ValueError: If the model name is not found in the supported models.
+        """
+        descriptions = cls._list_supported_models()
+        embedding_size: Optional[int] = None
+        for description in descriptions:
+            if description.model.lower() == model_name.lower():
+                embedding_size = description.dim
+                break
+        if embedding_size is None:
+            model_names = [description.model for description in descriptions]
+            raise ValueError(
+                f"Embedding size for model {model_name} was None. "
+                f"Available model names: {model_names}"
+            )
+        return embedding_size
+
     def embed_text(
         self,
         documents: Union[str, Iterable[str]],

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/late_interaction_multimodal/late_interaction_multimodal_embedding_base.py

@@ -19,6 +19,7 @@ class LateInteractionMultimodalEmbeddingBase(ModelManagement[DenseModelDescripti
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
+        self._embedding_size: Optional[int] = None
 
     def embed_text(
         self,
@@ -65,3 +66,13 @@ class LateInteractionMultimodalEmbeddingBase(ModelManagement[DenseModelDescripti
             List of embeddings, one per image
         """
         raise NotImplementedError()
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Returns embedding size of the chosen model."""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @property
+    def embedding_size(self) -> int:
+        """Returns embedding size for the current model"""
+        raise NotImplementedError("Subclasses must implement this method")

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/late_interaction_multimodal/onnx_multimodal_model.py

@@ -120,6 +120,8 @@ class OnnxMultimodalModel(OnnxModel[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
+        local_files_only: bool = False,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -146,6 +148,8 @@ class OnnxMultimodalModel(OnnxModel[T]):
                 "model_name": model_name,
                 "cache_dir": cache_dir,
                 "providers": providers,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
                 **kwargs,
             }
 
@@ -183,6 +187,8 @@ class OnnxMultimodalModel(OnnxModel[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
+        local_files_only: bool = False,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -209,6 +215,8 @@ class OnnxMultimodalModel(OnnxModel[T]):
                 "model_name": model_name,
                 "cache_dir": cache_dir,
                 "providers": providers,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
                 **kwargs,
             }
 

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py

@@ -131,11 +131,12 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
 
         self.model_description = self._get_model_description(model_name)
         self.cache_dir = str(define_cache_dir(cache_dir))
+        self._specific_model_path = specific_model_path
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
 
         if not self.lazy_load:
@@ -189,6 +190,8 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
             providers=self.providers,
             cuda=self.cuda,
             device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
             **kwargs,
         )
 

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -94,6 +94,8 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
+        local_files_only: bool = False,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ) -> Iterable[float]:
         is_small = False
@@ -120,6 +122,8 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
                 "model_name": model_name,
                 "cache_dir": cache_dir,
                 "providers": providers,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
                 **kwargs,
             }
 

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/sparse/bm25.py

@@ -115,11 +115,12 @@ class Bm25(SparseTextEmbeddingBase):
         model_description = self._get_model_description(model_name)
         self.cache_dir = str(define_cache_dir(cache_dir))
 
+        self._specific_model_path = specific_model_path
         self._model_dir = self.download_model(
             model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
 
         self.token_max_length = token_max_length
@@ -160,6 +161,8 @@ class Bm25(SparseTextEmbeddingBase):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
+        local_files_only: bool = False,
+        specific_model_path: Optional[str] = None,
     ) -> Iterable[SparseEmbedding]:
         is_small = False
 
@@ -188,6 +191,8 @@ class Bm25(SparseTextEmbeddingBase):
                 "language": self.language,
                 "token_max_length": self.token_max_length,
                 "disable_stemmer": self.disable_stemmer,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
             }
             pool = ParallelWorkerPool(
                 num_workers=parallel or 1,
@@ -226,6 +231,8 @@ class Bm25(SparseTextEmbeddingBase):
             documents=documents,
             batch_size=batch_size,
             parallel=parallel,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
         )
 
     def _stem(self, tokens: list[str]) -> list[str]:

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/sparse/bm42.py

@@ -110,11 +110,12 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         self.model_description = self._get_model_description(model_name)
         self.cache_dir = str(define_cache_dir(cache_dir))
 
+        self._specific_model_path = specific_model_path
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
 
         self.invert_vocab: dict[int, str] = {}
@@ -301,6 +302,8 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             cuda=self.cuda,
             device_ids=self.device_ids,
             alpha=self.alpha,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
         )
 
     @classmethod

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/sparse/minicoil.py

@@ -127,11 +127,12 @@ class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
 
         self.model_description = self._get_model_description(model_name)
         self.cache_dir = str(define_cache_dir(cache_dir))
+        self._specific_model_path = specific_model_path
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
 
         if not self.lazy_load:
@@ -211,6 +212,9 @@ class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             b=self.b,
             avg_len=self.avg_len,
             is_query=False,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+            **kwargs,
         )
 
     def query_embed(
@@ -230,6 +234,9 @@ class MiniCOIL(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             b=self.b,
             avg_len=self.avg_len,
             is_query=True,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
+            **kwargs,
         )
 
     @classmethod

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/sparse/sparse_text_embedding.py

@@ -62,7 +62,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
         **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
-        if model_name == "prithvida/Splade_PP_en_v1":
+        if model_name.lower() == "prithvida/Splade_PP_en_v1".lower():
             warnings.warn(
                 "The right spelling is prithivida/Splade_PP_en_v1. "
                 "Support of this name will be removed soon, please fix the model_name",

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/sparse/splade_pp.py

@@ -114,11 +114,12 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
         self.model_description = self._get_model_description(model_name)
         self.cache_dir = str(define_cache_dir(cache_dir))
 
+        self._specific_model_path = specific_model_path
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
 
         if not self.lazy_load:
@@ -165,6 +166,8 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
             providers=self.providers,
             cuda=self.cuda,
             device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
             **kwargs,
         )
 

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/text/custom_text_embedding.py

@@ -77,6 +77,11 @@ class CustomTextEmbedding(OnnxTextEmbedding):
         if self._pooling == PoolingType.DISABLED:
             return embeddings
 
+        raise ValueError(
+            f"Unsupported pooling type {self._pooling}. "
+            f"Supported types are: {PoolingType.CLS}, {PoolingType.MEAN}, {PoolingType.DISABLED}."
+        )
+
     def _normalize(self, embeddings: NumpyArray) -> NumpyArray:
         return normalize(embeddings) if self._normalization else embeddings
 

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/text/onnx_embedding.py

@@ -247,11 +247,12 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
 
         self.model_description = self._get_model_description(model_name)
         self.cache_dir = str(define_cache_dir(cache_dir))
+        self._specific_model_path = specific_model_path
         self._model_dir = self.download_model(
             self.model_description,
             self.cache_dir,
             local_files_only=self._local_files_only,
-            specific_model_path=specific_model_path,
+            specific_model_path=self._specific_model_path,
         )
 
         if not self.lazy_load:
@@ -288,6 +289,8 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
             providers=self.providers,
             cuda=self.cuda,
             device_ids=self.device_ids,
+            local_files_only=self._local_files_only,
+            specific_model_path=self._specific_model_path,
             **kwargs,
         )
 

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/text/onnx_text_model.py

@@ -108,6 +108,8 @@ class OnnxTextModel(OnnxModel[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
         device_ids: Optional[list[int]] = None,
+        local_files_only: bool = False,
+        specific_model_path: Optional[str] = None,
         **kwargs: Any,
     ) -> Iterable[T]:
         is_small = False
@@ -136,6 +138,8 @@ class OnnxTextModel(OnnxModel[T]):
                 "model_name": model_name,
                 "cache_dir": cache_dir,
                 "providers": providers,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
                 **kwargs,
             }
 

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/text/text_embedding.py

@@ -55,7 +55,7 @@ class TextEmbedding(TextEmbeddingBase):
     ) -> None:
         registered_models = cls._list_supported_models()
         for registered_model in registered_models:
-            if model == registered_model.model:
+            if model.lower() == registered_model.model.lower():
                 raise ValueError(
                     f"Model {model} is already registered in TextEmbedding, if you still want to add this model, "
                     f"please use another model name"
@@ -88,18 +88,18 @@ class TextEmbedding(TextEmbeddingBase):
         **kwargs: Any,
     ):
         super().__init__(model_name, cache_dir, threads, **kwargs)
-        if model_name == "nomic-ai/nomic-embed-text-v1.5-Q":
+        if model_name.lower() == "nomic-ai/nomic-embed-text-v1.5-Q".lower():
             warnings.warn(
                 "The model 'nomic-ai/nomic-embed-text-v1.5-Q' has been updated on HuggingFace. Please review "
                 "the latest documentation on HF and release notes to ensure compatibility with your workflow. ",
                 UserWarning,
                 stacklevel=2,
             )
-        if model_name in {
-            "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
-            "thenlper/gte-large",
-            "intfloat/multilingual-e5-large",
-            "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
+        if model_name.lower() in {
+            "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2".lower(),
+            "thenlper/gte-large".lower(),
+            "intfloat/multilingual-e5-large".lower(),
+            "sentence-transformers/paraphrase-multilingual-mpnet-base-v2".lower(),
         }:
             warnings.warn(
                 f"The model {model_name} now uses mean pooling instead of CLS embedding. "
@@ -128,6 +128,40 @@ class TextEmbedding(TextEmbeddingBase):
             "Please check the supported models using `TextEmbedding.list_supported_models()`"
         )
 
+    @property
+    def embedding_size(self) -> int:
+        """Get the embedding size of the current model"""
+        if self._embedding_size is None:
+            self._embedding_size = self.get_embedding_size(self.model_name)
+        return self._embedding_size
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Get the embedding size of the passed model
+
+        Args:
+            model_name (str): The name of the model to get embedding size for.
+
+        Returns:
+            int: The size of the embedding.
+
+        Raises:
+            ValueError: If the model name is not found in the supported models.
+        """
+        descriptions = cls._list_supported_models()
+        embedding_size: Optional[int] = None
+        for description in descriptions:
+            if description.model.lower() == model_name.lower():
+                embedding_size = description.dim
+                break
+        if embedding_size is None:
+            model_names = [description.model for description in descriptions]
+            raise ValueError(
+                f"Embedding size for model {model_name} was None. "
+                f"Available model names: {model_names}"
+            )
+        return embedding_size
+
     def embed(
         self,
         documents: Union[str, Iterable[str]],

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/fastembed/text/text_embedding_base.py

@@ -17,6 +17,7 @@ class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
         self.cache_dir = cache_dir
         self.threads = threads
         self._local_files_only = kwargs.pop("local_files_only", False)
+        self._embedding_size: Optional[int] = None
 
     def embed(
         self,
@@ -58,3 +59,13 @@ class TextEmbeddingBase(ModelManagement[DenseModelDescription]):
             yield from self.embed([query], **kwargs)
         else:
             yield from self.embed(query, **kwargs)
+
+    @classmethod
+    def get_embedding_size(cls, model_name: str) -> int:
+        """Returns embedding size of the passed model."""
+        raise NotImplementedError("Subclasses must implement this method")
+
+    @property
+    def embedding_size(self) -> int:
+        """Returns embedding size for the current model"""
+        raise NotImplementedError("Subclasses must implement this method")

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastembed-gpu"
-version = "0.7.0"
+version = "0.7.1"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"