fastembed-gpu 0.3.0__tar.gz → 0.3.2__tar.gz

{fastembed_gpu-0.3.0 → fastembed_gpu-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fastembed-gpu
-Version: 0.3.0
+Version: 0.3.2
 Summary: Fast, light, accurate library built for retrieval embedding generation
 Home-page: https://github.com/qdrant/fastembed
 License: Apache License
@@ -19,8 +19,8 @@ Requires-Dist: PyStemmer (>=2.2.0,<3.0.0)
 Requires-Dist: huggingface-hub (>=0.20,<1.0)
 Requires-Dist: loguru (>=0.7.2,<0.8.0)
 Requires-Dist: mmh3 (>=4.0,<5.0)
-Requires-Dist: numpy (>=1.21) ; python_version < "3.12"
-Requires-Dist: numpy (>=1.26) ; python_version >= "3.12"
+Requires-Dist: numpy (>=1.21,<2) ; python_version < "3.12"
+Requires-Dist: numpy (>=1.26,<2) ; python_version >= "3.12"
 Requires-Dist: onnxruntime-gpu (>=1.17.0,<2.0.0)
 Requires-Dist: pillow (>=10.3.0,<11.0.0)
 Requires-Dist: requests (>=2.31,<3.0)
@@ -50,11 +50,9 @@ To install the FastEmbed library, pip works best. You can install it with or wit
 
 ```bash
 pip install fastembed
-```
 
-### ⚡️ With GPU
+# or with GPU support
 
-```bash
 pip install fastembed-gpu
 ```
 
@@ -80,7 +78,99 @@ embeddings_list = list(embedding_model.embed(documents))
 len(embeddings_list[0]) # Vector of 384 dimensions
 ```
 
-### ⚡️ FastEmbed on a GPU
+Fastembed supports a variety of models for different tasks and modalities.
+The list of all the available models can be found [here](https://qdrant.github.io/fastembed/examples/Supported_Models/)
+### 🎒 Dense text embeddings
+
+```python
+from fastembed import TextEmbedding
+
+model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
+embeddings = list(embedding_model.embed(documents))
+
+# [
+#   array([-0.1115,  0.0097,  0.0052,  0.0195, ...], dtype=float32),
+#   array([-0.1019,  0.0635, -0.0332,  0.0522, ...], dtype=float32)
+# ]
+
+```
+
+
+
+### 🔱 Sparse text embeddings
+
+* SPLADE++
+
+```python
+from fastembed import SparseTextEmbedding
+
+model = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1")
+embeddings = list(embedding_model.embed(documents))
+
+# [
+#   SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
+#   SparseEmbedding(indices=[ 38,  12,  91, ... ], values=[0.11, 0.22, 0.39, ...])
+# ]
+```
+
+<!--
+* BM42 - ([link](ToDo))
+
+```
+from fastembed import SparseTextEmbedding
+
+model = SparseTextEmbedding(model_name="Qdrant/bm42-all-minilm-l6-v2-attentions")
+embeddings = list(embedding_model.embed(documents))
+
+# [
+#   SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
+#   SparseEmbedding(indices=[ 38,  12,  91, ... ], values=[0.11, 0.22, 0.39, ...])
+# ]
+```
+-->
+
+### 🦥 Late interaction models (aka ColBERT)
+
+
+```python
+from fastembed import LateInteractionTextEmbedding
+
+model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0")
+embeddings = list(embedding_model.embed(documents))
+
+# [
+#   array([
+#       [-0.1115,  0.0097,  0.0052,  0.0195, ...],
+#       [-0.1019,  0.0635, -0.0332,  0.0522, ...],
+#   ]),
+#   array([
+#       [-0.9019,  0.0335, -0.0032,  0.0991, ...],
+#       [-0.2115,  0.8097,  0.1052,  0.0195, ...],
+#   ]),  
+# ]
+```
+
+### 🖼️ Image embeddings
+
+```python
+from fastembed import ImageEmbedding
+
+images = [
+    "./path/to/image1.jpg",
+    "./path/to/image2.jpg",
+]
+
+model = ImageEmbedding(model_name="Qdrant/clip-ViT-B-32-vision")
+embeddings = list(embedding_model.embed(images))
+
+# [
+#   array([-0.1115,  0.0097,  0.0052,  0.0195, ...], dtype=float32),
+#   array([-0.1019,  0.0635, -0.0332,  0.0522, ...], dtype=float32)
+# ]
+```
+
+
+## ⚡️ FastEmbed on a GPU
 
 FastEmbed supports running on GPU devices.
 It requires installation of the `fastembed-gpu` package.

{fastembed_gpu-0.3.0 → fastembed_gpu-0.3.2}/README.md

@@ -18,11 +18,9 @@ To install the FastEmbed library, pip works best. You can install it with or wit
 
 ```bash
 pip install fastembed
-```
 
-### ⚡️ With GPU
+# or with GPU support
 
-```bash
 pip install fastembed-gpu
 ```
 
@@ -48,7 +46,99 @@ embeddings_list = list(embedding_model.embed(documents))
 len(embeddings_list[0]) # Vector of 384 dimensions
 ```
 
-### ⚡️ FastEmbed on a GPU
+Fastembed supports a variety of models for different tasks and modalities.
+The list of all the available models can be found [here](https://qdrant.github.io/fastembed/examples/Supported_Models/)
+### 🎒 Dense text embeddings
+
+```python
+from fastembed import TextEmbedding
+
+model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5")
+embeddings = list(embedding_model.embed(documents))
+
+# [
+#   array([-0.1115,  0.0097,  0.0052,  0.0195, ...], dtype=float32),
+#   array([-0.1019,  0.0635, -0.0332,  0.0522, ...], dtype=float32)
+# ]
+
+```
+
+
+
+### 🔱 Sparse text embeddings
+
+* SPLADE++
+
+```python
+from fastembed import SparseTextEmbedding
+
+model = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1")
+embeddings = list(embedding_model.embed(documents))
+
+# [
+#   SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
+#   SparseEmbedding(indices=[ 38,  12,  91, ... ], values=[0.11, 0.22, 0.39, ...])
+# ]
+```
+
+<!--
+* BM42 - ([link](ToDo))
+
+```
+from fastembed import SparseTextEmbedding
+
+model = SparseTextEmbedding(model_name="Qdrant/bm42-all-minilm-l6-v2-attentions")
+embeddings = list(embedding_model.embed(documents))
+
+# [
+#   SparseEmbedding(indices=[ 17, 123, 919, ... ], values=[0.71, 0.22, 0.39, ...]),
+#   SparseEmbedding(indices=[ 38,  12,  91, ... ], values=[0.11, 0.22, 0.39, ...])
+# ]
+```
+-->
+
+### 🦥 Late interaction models (aka ColBERT)
+
+
+```python
+from fastembed import LateInteractionTextEmbedding
+
+model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0")
+embeddings = list(embedding_model.embed(documents))
+
+# [
+#   array([
+#       [-0.1115,  0.0097,  0.0052,  0.0195, ...],
+#       [-0.1019,  0.0635, -0.0332,  0.0522, ...],
+#   ]),
+#   array([
+#       [-0.9019,  0.0335, -0.0032,  0.0991, ...],
+#       [-0.2115,  0.8097,  0.1052,  0.0195, ...],
+#   ]),  
+# ]
+```
+
+### 🖼️ Image embeddings
+
+```python
+from fastembed import ImageEmbedding
+
+images = [
+    "./path/to/image1.jpg",
+    "./path/to/image2.jpg",
+]
+
+model = ImageEmbedding(model_name="Qdrant/clip-ViT-B-32-vision")
+embeddings = list(embedding_model.embed(images))
+
+# [
+#   array([-0.1115,  0.0097,  0.0052,  0.0195, ...], dtype=float32),
+#   array([-0.1019,  0.0635, -0.0332,  0.0522, ...], dtype=float32)
+# ]
+```
+
+
+## ⚡️ FastEmbed on a GPU
 
 FastEmbed supports running on GPU devices.
 It requires installation of the `fastembed-gpu` package.

{fastembed_gpu-0.3.0 → fastembed_gpu-0.3.2}/fastembed/__init__.py

@@ -1,9 +1,9 @@
 import importlib.metadata
 
 from fastembed.image import ImageEmbedding
-from fastembed.text import TextEmbedding
-from fastembed.sparse import SparseTextEmbedding, SparseEmbedding
 from fastembed.late_interaction import LateInteractionTextEmbedding
+from fastembed.sparse import SparseEmbedding, SparseTextEmbedding
+from fastembed.text import TextEmbedding
 
 try:
     version = importlib.metadata.version("fastembed")

fastembed_gpu-0.3.2/fastembed/common/__init__.py

@@ -0,0 +1,3 @@
+from fastembed.common.types import ImageInput, OnnxProvider, PathInput
+
+__all__ = ["OnnxProvider", "ImageInput", "PathInput"]

{fastembed_gpu-0.3.0 → fastembed_gpu-0.3.2}/fastembed/common/model_management.py

@@ -2,13 +2,13 @@ import os
 import shutil
 import tarfile
 from pathlib import Path
-from typing import List, Optional, Dict, Any
+from typing import Any, Dict, List, Optional
 
 import requests
 from huggingface_hub import snapshot_download
 from huggingface_hub.utils import RepositoryNotFoundError
-from tqdm import tqdm
 from loguru import logger
+from tqdm import tqdm
 
 
 class ModelManagement:
@@ -42,7 +42,9 @@ class ModelManagement:
         raise ValueError(f"Model {model_name} is not supported in {cls.__name__}.")
 
     @classmethod
-    def download_file_from_gcs(cls, url: str, output_path: str, show_progress: bool = True) -> str:
+    def download_file_from_gcs(
+        cls, url: str, output_path: str, show_progress: bool = True
+    ) -> str:
         """
         Downloads a file from Google Cloud Storage.
 
@@ -71,12 +73,17 @@ class ModelManagement:
 
         # Warn if the total size is zero
         if total_size_in_bytes == 0:
-            print(f"Warning: Content-length header is missing or zero in the response from {url}.")
+            print(
+                f"Warning: Content-length header is missing or zero in the response from {url}."
+            )
 
         show_progress = total_size_in_bytes and show_progress
 
         with tqdm(
-            total=total_size_in_bytes, unit="iB", unit_scale=True, disable=not show_progress
+            total=total_size_in_bytes,
+            unit="iB",
+            unit_scale=True,
+            disable=not show_progress,
         ) as progress_bar:
             with open(output_path, "wb") as file:
                 for chunk in response.iter_content(chunk_size=1024):
@@ -156,7 +163,9 @@ class ModelManagement:
         return cache_dir
 
     @classmethod
-    def retrieve_model_gcs(cls, model_name: str, source_url: str, cache_dir: str) -> Path:
+    def retrieve_model_gcs(
+        cls, model_name: str, source_url: str, cache_dir: str
+    ) -> Path:
         fast_model_name = f"fast-{model_name.split('/')[-1]}"
 
         cache_tmp_dir = Path(cache_dir) / "tmp"
@@ -182,8 +191,12 @@ class ModelManagement:
             output_path=str(model_tar_gz),
         )
 
-        cls.decompress_to_cache(targz_path=str(model_tar_gz), cache_dir=str(cache_tmp_dir))
-        assert model_tmp_dir.exists(), f"Could not find {model_tmp_dir} in {cache_tmp_dir}"
+        cls.decompress_to_cache(
+            targz_path=str(model_tar_gz), cache_dir=str(cache_tmp_dir)
+        )
+        assert (
+            model_tmp_dir.exists()
+        ), f"Could not find {model_tmp_dir} in {cache_tmp_dir}"
 
         model_tar_gz.unlink()
         # Rename from tmp to final name is atomic

{fastembed_gpu-0.3.0 → fastembed_gpu-0.3.2}/fastembed/common/onnx_model.py

@@ -1,7 +1,17 @@
+import warnings
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Dict, Generic, Iterable, Optional, Tuple, Type, TypeVar, Sequence
-import warnings
+from typing import (
+    Any,
+    Dict,
+    Generic,
+    Iterable,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    TypeVar,
+)
 
 import numpy as np
 import onnxruntime as ort
@@ -9,7 +19,6 @@ import onnxruntime as ort
 from fastembed.common.types import OnnxProvider
 from fastembed.parallel_processor import Worker
 
-
 # Holds type of the embedding result
 T = TypeVar("T")
 
@@ -51,7 +60,9 @@ class OnnxModel(Generic[T]):
         model_path = model_dir / model_file
         # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
 
-        onnx_providers = ["CPUExecutionProvider"] if providers is None else list(providers)
+        onnx_providers = (
+            ["CPUExecutionProvider"] if providers is None else list(providers)
+        )
         available_providers = ort.get_available_providers()
         requested_provider_names = []
         for provider in onnx_providers:
@@ -92,6 +103,7 @@ class EmbeddingWorker(Worker):
         self,
         model_name: str,
         cache_dir: str,
+        **kwargs,
     ) -> OnnxModel:
         raise NotImplementedError()
 
@@ -99,15 +111,13 @@ class EmbeddingWorker(Worker):
         self,
         model_name: str,
         cache_dir: str,
+        **kwargs,
     ):
-        self.model = self.init_embedding(model_name, cache_dir)
+        self.model = self.init_embedding(model_name, cache_dir, **kwargs)
 
     @classmethod
     def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker":
-        return cls(
-            model_name=model_name,
-            cache_dir=cache_dir,
-        )
+        return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
 
     def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
         raise NotImplementedError("Subclasses must implement this method")

{fastembed_gpu-0.3.0 → fastembed_gpu-0.3.2}/fastembed/common/preprocessor_utils.py

@@ -2,7 +2,7 @@ import json
 from pathlib import Path
 from typing import Tuple
 
-from tokenizers import Tokenizer, AddedToken
+from tokenizers import AddedToken, Tokenizer
 
 from fastembed.image.transform.operators import Compose
 
@@ -40,7 +40,9 @@ def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tuple[Tokenizer, d
     tokens_map = load_special_tokens(model_dir)
 
     tokenizer = Tokenizer.from_file(str(tokenizer_path))
-    tokenizer.enable_truncation(max_length=min(tokenizer_config["model_max_length"], max_length))
+    tokenizer.enable_truncation(
+        max_length=min(tokenizer_config["model_max_length"], max_length)
+    )
     tokenizer.enable_padding(
         pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"]
     )

{fastembed_gpu-0.3.0 → fastembed_gpu-0.3.2}/fastembed/common/types.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from typing import Union, Iterable, Tuple, Dict, Any
+from typing import Any, Dict, Iterable, Tuple, Union
 
 if sys.version_info >= (3, 10):
     from typing import TypeAlias

{fastembed_gpu-0.3.0 → fastembed_gpu-0.3.2}/fastembed/common/utils.py

@@ -2,7 +2,7 @@ import os
 import tempfile
 from itertools import islice
 from pathlib import Path
-from typing import Union, Iterable, Generator, Optional
+from typing import Generator, Iterable, Optional, Union
 
 import numpy as np
 

{fastembed_gpu-0.3.0 → fastembed_gpu-0.3.2}/fastembed/image/__init__.py

@@ -1,4 +1,3 @@
 from fastembed.image.image_embedding import ImageEmbedding
 
-
 __all__ = ["ImageEmbedding"]

{fastembed_gpu-0.3.0 → fastembed_gpu-0.3.2}/fastembed/image/image_embedding.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Iterable, List, Optional, Type, Sequence
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Type
 
 import numpy as np
 
@@ -51,9 +51,16 @@ class ImageEmbedding(ImageEmbeddingBase):
 
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
             supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
-            if any(model_name.lower() == model["model"].lower() for model in supported_models):
+            if any(
+                model_name.lower() == model["model"].lower()
+                for model in supported_models
+            ):
                 self.model = EMBEDDING_MODEL_TYPE(
-                    model_name, cache_dir, threads, providers=providers, **kwargs
+                    model_name,
+                    cache_dir,
+                    threads=threads,
+                    providers=providers,
+                    **kwargs,
                 )
                 return
 

{fastembed_gpu-0.3.0 → fastembed_gpu-0.3.2}/fastembed/image/image_embedding_base.py

@@ -31,6 +31,11 @@ class ImageEmbeddingBase(ModelManagement):
 
         Args:
             images - The list of image paths to preprocess and embed.
+            batch_size: Batch size for encoding
+            parallel:
+                If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets.
+                If 0, use all available cores.
+                If None, don't use data-parallel processing, use default onnxruntime threading instead.
             **kwargs: Additional keyword argument to pass to the embed method.
 
         Yields:

{fastembed_gpu-0.3.0 → fastembed_gpu-0.3.2}/fastembed/image/onnx_embedding.py

@@ -1,12 +1,12 @@
-from typing import Dict, Optional, Iterable, Type, List, Any, Sequence
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Type
 
 import numpy as np
 
-from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.utils import normalize, define_cache_dir
 from fastembed.common import ImageInput, OnnxProvider
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.utils import define_cache_dir, normalize
 from fastembed.image.image_embedding_base import ImageEmbeddingBase
-from fastembed.image.onnx_image_model import OnnxImageModel, ImageEmbeddingWorker
+from fastembed.image.onnx_image_model import ImageEmbeddingWorker, OnnxImageModel
 
 supported_onnx_models = [
     {
@@ -56,9 +56,9 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
         super().__init__(model_name, cache_dir, threads, **kwargs)
 
         model_description = self._get_model_description(model_name)
-        cache_dir = define_cache_dir(cache_dir)
+        self.cache_dir = define_cache_dir(cache_dir)
         model_dir = self.download_model(
-            model_description, cache_dir, local_files_only=self._local_files_only
+            model_description, self.cache_dir, local_files_only=self._local_files_only
         )
 
         self.load_onnx_model(
@@ -106,6 +106,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
             images=images,
             batch_size=batch_size,
             parallel=parallel,
+            **kwargs,
         )
 
     @classmethod
@@ -126,9 +127,5 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
 
 
 class OnnxImageEmbeddingWorker(ImageEmbeddingWorker):
-    def init_embedding(
-        self,
-        model_name: str,
-        cache_dir: str,
-    ) -> OnnxImageEmbedding:
-        return OnnxImageEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1)
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxImageEmbedding:
+        return OnnxImageEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)

{fastembed_gpu-0.3.0 → fastembed_gpu-0.3.2}/fastembed/image/onnx_image_model.py

@@ -1,15 +1,15 @@
-import os
 import contextlib
+import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Sequence
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type
 
-from PIL import Image
 import numpy as np
+from PIL import Image
 
+from fastembed.common import ImageInput, OnnxProvider, PathInput
+from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_preprocessor
-from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker, T, OnnxOutputContext
-from fastembed.common import PathInput, ImageInput, OnnxProvider
 from fastembed.common.utils import iter_batch
 from fastembed.parallel_processor import ParallelWorkerPool
 
@@ -44,7 +44,10 @@ class OnnxImageModel(OnnxModel[T]):
         providers: Optional[Sequence[OnnxProvider]] = None,
     ) -> None:
         super().load_onnx_model(
-            model_dir=model_dir, model_file=model_file, threads=threads, providers=providers
+            model_dir=model_dir,
+            model_file=model_file,
+            threads=threads,
+            providers=providers,
         )
         self.processor = load_preprocessor(model_dir=model_dir)
 
@@ -59,9 +62,7 @@ class OnnxImageModel(OnnxModel[T]):
         onnx_input = self._preprocess_onnx_input(onnx_input)
         model_output = self.model.run(None, onnx_input)
         embeddings = model_output[0].reshape(len(images), -1)
-        return OnnxOutputContext(
-            model_output=embeddings
-        )
+        return OnnxOutputContext(model_output=embeddings)
 
     def _embed_images(
         self,
@@ -70,6 +71,7 @@ class OnnxImageModel(OnnxModel[T]):
         images: ImageInput,
         batch_size: int = 256,
         parallel: Optional[int] = None,
+        **kwargs,
     ) -> Iterable[T]:
         is_small = False
 
@@ -88,11 +90,10 @@ class OnnxImageModel(OnnxModel[T]):
             for batch in iter_batch(images, batch_size):
                 yield from self._post_process_onnx_output(self.onnx_embed(batch))
         else:
-            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
-            params = {
-                "model_name": model_name,
-                "cache_dir": cache_dir,
-            }
+            start_method = (
+                "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+            )
+            params = {"model_name": model_name, "cache_dir": cache_dir, **kwargs}
             pool = ParallelWorkerPool(
                 parallel, self._get_worker_class(), start_method=start_method
             )

{fastembed_gpu-0.3.0 → fastembed_gpu-0.3.2}/fastembed/image/transform/functional.py

@@ -1,8 +1,7 @@
-from typing import Union, Tuple, Sized
-
-from PIL import Image
+from typing import Sized, Tuple, Union
 
 import numpy as np
+from PIL import Image
 
 
 def convert_to_rgb(image: Image.Image) -> Image.Image: