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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: fastembed-gpu
-Version: 0.7.0
+Version: 0.7.2
 Summary: Fast, light, accurate library built for retrieval embedding generation
 License: Apache License
 Keywords: vector,embedding,neural,search,qdrant,sentence-transformers
@@ -281,37 +281,37 @@ pip install qdrant-client[fastembed-gpu]
 You might have to use quotes ```pip install 'qdrant-client[fastembed]'``` on zsh.
 
 ```python
-from qdrant_client import QdrantClient
+from qdrant_client import QdrantClient, models
 
 # Initialize the client
 client = QdrantClient("localhost", port=6333) # For production
-# client = QdrantClient(":memory:") # For small experiments
+# client = QdrantClient(":memory:") # For experimentation
 
-# Prepare your documents, metadata, and IDs
-docs = ["Qdrant has Langchain integrations", "Qdrant also has Llama Index integrations"]
-metadata = [
-    {"source": "Langchain-docs"},
-    {"source": "Llama-index-docs"},
+model_name = "sentence-transformers/all-MiniLM-L6-v2"
+payload = [
+    {"document": "Qdrant has Langchain integrations", "source": "Langchain-docs", },
+    {"document": "Qdrant also has Llama Index integrations", "source": "LlamaIndex-docs"},
 ]
+docs = [models.Document(text=data["document"], model=model_name) for data in payload]
 ids = [42, 2]
 
-# If you want to change the model:
-# client.set_model("sentence-transformers/all-MiniLM-L6-v2")
-# List of supported models: https://qdrant.github.io/fastembed/examples/Supported_Models
+client.create_collection(
+    "demo_collection",
+    vectors_config=models.VectorParams(
+        size=client.get_embedding_size(model_name), distance=models.Distance.COSINE)
+)
 
-# Use the new add() instead of upsert()
-# This internally calls embed() of the configured embedding model
-client.add(
+client.upload_collection(
     collection_name="demo_collection",
-    documents=docs,
-    metadata=metadata,
-    ids=ids
+    vectors=docs,
+    ids=ids,
+    payload=payload,
 )
 
-search_result = client.query(
+search_result = client.query_points(
     collection_name="demo_collection",
-    query_text="This is a query document"
-)
+    query=models.Document(text="This is a query document", model=model_name)
+).points
 print(search_result)
 ```
 

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.2}/README.md

@@ -246,36 +246,36 @@ pip install qdrant-client[fastembed-gpu]
 You might have to use quotes ```pip install 'qdrant-client[fastembed]'``` on zsh.
 
 ```python
-from qdrant_client import QdrantClient
+from qdrant_client import QdrantClient, models
 
 # Initialize the client
 client = QdrantClient("localhost", port=6333) # For production
-# client = QdrantClient(":memory:") # For small experiments
+# client = QdrantClient(":memory:") # For experimentation
 
-# Prepare your documents, metadata, and IDs
-docs = ["Qdrant has Langchain integrations", "Qdrant also has Llama Index integrations"]
-metadata = [
-    {"source": "Langchain-docs"},
-    {"source": "Llama-index-docs"},
+model_name = "sentence-transformers/all-MiniLM-L6-v2"
+payload = [
+    {"document": "Qdrant has Langchain integrations", "source": "Langchain-docs", },
+    {"document": "Qdrant also has Llama Index integrations", "source": "LlamaIndex-docs"},
 ]
+docs = [models.Document(text=data["document"], model=model_name) for data in payload]
 ids = [42, 2]
 
-# If you want to change the model:
-# client.set_model("sentence-transformers/all-MiniLM-L6-v2")
-# List of supported models: https://qdrant.github.io/fastembed/examples/Supported_Models
+client.create_collection(
+    "demo_collection",
+    vectors_config=models.VectorParams(
+        size=client.get_embedding_size(model_name), distance=models.Distance.COSINE)
+)
 
-# Use the new add() instead of upsert()
-# This internally calls embed() of the configured embedding model
-client.add(
+client.upload_collection(
     collection_name="demo_collection",
-    documents=docs,
-    metadata=metadata,
-    ids=ids
+    vectors=docs,
+    ids=ids,
+    payload=payload,
 )
 
-search_result = client.query(
+search_result = client.query_points(
     collection_name="demo_collection",
-    query_text="This is a query document"
-)
+    query=models.Document(text="This is a query document", model=model_name)
+).points
 print(search_result)
 ```

{fastembed_gpu-0.7.0 → fastembed_gpu-0.7.2}/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.2}/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.2}/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.2}/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.2}/fastembed/late_interaction/colbert.py

@@ -46,7 +46,8 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         self, output: OnnxOutputContext, is_doc: bool = True, **kwargs: Any
     ) -> Iterable[NumpyArray]:
         if not is_doc:
-            return output.model_output
+            for embedding in output.model_output:
+                yield embedding
 
         if output.input_ids is None or output.attention_mask is None:
             raise ValueError(
@@ -62,7 +63,9 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         norm = np.linalg.norm(output.model_output, ord=2, axis=2, keepdims=True)
         norm_clamped = np.maximum(norm, 1e-12)
         output.model_output /= norm_clamped
-        return output.model_output
+
+        for embedding, attention_mask in zip(output.model_output, output.attention_mask):
+            yield embedding[attention_mask == 1]
 
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], is_doc: bool = True, **kwargs: Any
@@ -169,11 +172,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 +237,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.2}/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.2}/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.2}/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.2}/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.2}/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.2}/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.2}/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.2/fastembed/postprocess/__init__.py

@@ -0,0 +1,3 @@
+from fastembed.postprocess.muvera import Muvera
+
+__all__ = ["Muvera"]