fastembed-gpu 0.4.2__tar.gz → 0.5.0__tar.gz

{fastembed_gpu-0.4.2 → fastembed_gpu-0.5.0}/PKG-INFO

@@ -1,27 +1,31 @@
 Metadata-Version: 2.1
 Name: fastembed-gpu
-Version: 0.4.2
+Version: 0.5.0
 Summary: Fast, light, accurate library built for retrieval embedding generation
 Home-page: https://github.com/qdrant/fastembed
 License: Apache License
 Keywords: vector,embedding,neural,search,qdrant,sentence-transformers
 Author: Qdrant Team
 Author-email: info@qdrant.tech
-Requires-Python: >=3.8.0,<3.13
+Requires-Python: >=3.9.0
 Classifier: License :: Other/Proprietary License
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: huggingface-hub (>=0.20,<1.0)
 Requires-Dist: loguru (>=0.7.2,<0.8.0)
 Requires-Dist: mmh3 (>=4.1.0,<5.0.0)
-Requires-Dist: numpy (>=1.21) ; python_version < "3.12"
-Requires-Dist: numpy (>=1.26) ; python_version >= "3.12"
-Requires-Dist: onnx (>=1.15.0,<2.0.0)
-Requires-Dist: onnxruntime-gpu (>=1.17.0,<1.20.0)
+Requires-Dist: numpy (>=1.21) ; python_version >= "3.10" and python_version < "3.12"
+Requires-Dist: numpy (>=1.21,<2.1.0) ; python_version < "3.10"
+Requires-Dist: numpy (>=1.26) ; python_version >= "3.12" and python_version < "3.13"
+Requires-Dist: numpy (>=2.1.0) ; python_version >= "3.13"
+Requires-Dist: onnx (>=1.15.0)
+Requires-Dist: onnxruntime-gpu (>1.20.0) ; python_version >= "3.13"
+Requires-Dist: onnxruntime-gpu (>=1.17.0,!=1.20.0) ; python_version >= "3.10" and python_version < "3.13"
+Requires-Dist: onnxruntime-gpu (>=1.17.0,<1.20.0) ; python_version < "3.10"
 Requires-Dist: pillow (>=10.3.0,<11.0.0)
 Requires-Dist: py-rust-stemmers (>=0.1.0,<0.2.0)
 Requires-Dist: requests (>=2.31,<3.0)
@@ -60,10 +64,10 @@ pip install fastembed-gpu
 
 ```python
 from fastembed import TextEmbedding
-from typing import List
+
 
 # Example list of documents
-documents: List[str] = [
+documents: list[str] = [
     "This is built to be faster and lighter than other embedding libraries e.g. Transformers, Sentence-Transformers, etc.",
     "fastembed is supported by and maintained by Qdrant.",
 ]
@@ -171,11 +175,10 @@ embeddings = list(model.embed(images))
 
 ### 🔄 Rerankers
 ```python
-from typing import List
 from fastembed.rerank.cross_encoder import TextCrossEncoder
 
 query = "Who is maintaining Qdrant?"
-documents: List[str] = [
+documents: list[str] = [
     "This is built to be faster and lighter than other embedding libraries e.g. Transformers, Sentence-Transformers, etc.",
     "fastembed is supported by and maintained by Qdrant.",
 ]

{fastembed_gpu-0.4.2 → fastembed_gpu-0.5.0}/README.md

@@ -28,10 +28,10 @@ pip install fastembed-gpu
 
 ```python
 from fastembed import TextEmbedding
-from typing import List
+
 
 # Example list of documents
-documents: List[str] = [
+documents: list[str] = [
     "This is built to be faster and lighter than other embedding libraries e.g. Transformers, Sentence-Transformers, etc.",
     "fastembed is supported by and maintained by Qdrant.",
 ]
@@ -139,11 +139,10 @@ embeddings = list(model.embed(images))
 
 ### 🔄 Rerankers
 ```python
-from typing import List
 from fastembed.rerank.cross_encoder import TextCrossEncoder
 
 query = "Who is maintaining Qdrant?"
-documents: List[str] = [
+documents: list[str] = [
     "This is built to be faster and lighter than other embedding libraries e.g. Transformers, Sentence-Transformers, etc.",
     "fastembed is supported by and maintained by Qdrant.",
 ]

{fastembed_gpu-0.4.2 → fastembed_gpu-0.5.0}/fastembed/common/model_management.py

@@ -3,27 +3,31 @@ import time
 import shutil
 import tarfile
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Optional
 
 import requests
 from huggingface_hub import snapshot_download
-from huggingface_hub.utils import RepositoryNotFoundError
+from huggingface_hub.utils import (
+    RepositoryNotFoundError,
+    disable_progress_bars,
+    enable_progress_bars,
+)
 from loguru import logger
 from tqdm import tqdm
 
 
 class ModelManagement:
     @classmethod
-    def list_supported_models(cls) -> List[Dict[str, Any]]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """Lists the supported models.
 
         Returns:
-            List[Dict[str, Any]]: A list of dictionaries containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
         """
         raise NotImplementedError()
 
     @classmethod
-    def _get_model_description(cls, model_name: str) -> Dict[str, Any]:
+    def _get_model_description(cls, model_name: str) -> dict[str, Any]:
         """
         Gets the model description from the model_name.
 
@@ -34,7 +38,7 @@ class ModelManagement:
             ValueError: If the model_name is not supported.
 
         Returns:
-            Dict[str, Any]: The model description.
+            dict[str, Any]: The model description.
         """
         for model in cls.list_supported_models():
             if model_name.lower() == model["model"].lower():
@@ -93,8 +97,8 @@ class ModelManagement:
     def download_files_from_huggingface(
         cls,
         hf_source_repo: str,
-        cache_dir: Optional[str] = None,
-        extra_patterns: Optional[List[str]] = None,
+        cache_dir: str,
+        extra_patterns: Optional[list[str]] = None,
         local_files_only: bool = False,
         **kwargs,
     ) -> str:
@@ -103,7 +107,7 @@ class ModelManagement:
         Args:
             hf_source_repo (str): Name of the model on HuggingFace Hub, e.g. "qdrant/all-MiniLM-L6-v2-onnx".
             cache_dir (Optional[str]): The path to the cache directory.
-            extra_patterns (Optional[List[str]]): extra patterns to allow in the snapshot download, typically
+            extra_patterns (Optional[list[str]]): extra patterns to allow in the snapshot download, typically
                 includes the required model files.
             local_files_only (bool, optional): Whether to only use local files. Defaults to False.
         Returns:
@@ -119,6 +123,12 @@ class ModelManagement:
         if extra_patterns is not None:
             allow_patterns.extend(extra_patterns)
 
+        snapshot_dir = Path(cache_dir) / f"models--{hf_source_repo.replace('/', '--')}"
+        is_cached = snapshot_dir.exists()
+
+        if is_cached:
+            disable_progress_bars()
+
         return snapshot_download(
             repo_id=hf_source_repo,
             allow_patterns=allow_patterns,
@@ -211,13 +221,13 @@ class ModelManagement:
 
     @classmethod
     def download_model(
-        cls, model: Dict[str, Any], cache_dir: Path, retries: int = 3, **kwargs
+        cls, model: dict[str, Any], cache_dir: Path, retries: int = 3, **kwargs
     ) -> Path:
         """
         Downloads a model from HuggingFace Hub or Google Cloud Storage.
 
         Args:
-            model (Dict[str, Any]): The model description.
+            model (dict[str, Any]): The model description.
                 Example:
                 ```
                 {
@@ -265,6 +275,8 @@ class ModelManagement:
                             f"Could not download model from HuggingFace: {e} "
                             "Falling back to other sources."
                         )
+                finally:
+                    enable_progress_bars()
             if url_source or local_files_only:
                 try:
                     return cls.retrieve_model_gcs(

{fastembed_gpu-0.4.2 → fastembed_gpu-0.5.0}/fastembed/common/onnx_model.py

@@ -1,7 +1,7 @@
 import warnings
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Dict, Generic, Iterable, Optional, Sequence, Tuple, Type, TypeVar
+from typing import Any, Generic, Iterable, Optional, Sequence, Type, TypeVar
 
 import numpy as np
 import onnxruntime as ort
@@ -33,8 +33,8 @@ class OnnxModel(Generic[T]):
         self.tokenizer = None
 
     def _preprocess_onnx_input(
-        self, onnx_input: Dict[str, np.ndarray], **kwargs
-    ) -> Dict[str, np.ndarray]:
+        self, onnx_input: dict[str, np.ndarray], **kwargs
+    ) -> dict[str, np.ndarray]:
         """
         Preprocess the onnx input.
         """
@@ -52,6 +52,13 @@ class OnnxModel(Generic[T]):
         model_path = model_dir / model_file
         # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
 
+        if cuda and providers is not None:
+            warnings.warn(
+                f"`cuda` and `providers` are mutually exclusive parameters, cuda: {cuda}, providers: {providers}",
+                category=UserWarning,
+                stacklevel=6,
+            )
+
         if providers is not None:
             onnx_providers = list(providers)
         elif cuda:
@@ -121,5 +128,5 @@ class EmbeddingWorker(Worker):
     def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker":
         return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
 
-    def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         raise NotImplementedError("Subclasses must implement this method")

{fastembed_gpu-0.4.2 → fastembed_gpu-0.5.0}/fastembed/common/preprocessor_utils.py

@@ -1,6 +1,6 @@
 import json
 from pathlib import Path
-from typing import Tuple
+
 from tokenizers import AddedToken, Tokenizer
 
 from fastembed.image.transform.operators import Compose
@@ -17,7 +17,7 @@ def load_special_tokens(model_dir: Path) -> dict:
     return tokens_map
 
 
-def load_tokenizer(model_dir: Path) -> Tuple[Tokenizer, dict]:
+def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict]:
     config_path = model_dir / "config.json"
     if not config_path.exists():
         raise ValueError(f"Could not find config.json in {model_dir}")

{fastembed_gpu-0.4.2 → fastembed_gpu-0.5.0}/fastembed/common/types.py

@@ -1,7 +1,7 @@
 import os
 import sys
 from PIL import Image
-from typing import Any, Dict, Iterable, Tuple, Union
+from typing import Any, Iterable, Union
 
 if sys.version_info >= (3, 10):
     from typing import TypeAlias
@@ -13,4 +13,4 @@ PathInput: TypeAlias = Union[str, os.PathLike]
 PilInput: TypeAlias = Union[Image.Image, Iterable[Image.Image]]
 ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput], PilInput]
 
-OnnxProvider: TypeAlias = Union[str, Tuple[str, Dict[Any, Any]]]
+OnnxProvider: TypeAlias = Union[str, tuple[str, dict[Any, Any]]]

{fastembed_gpu-0.4.2 → fastembed_gpu-0.5.0}/fastembed/common/utils.py

@@ -1,13 +1,13 @@
 import os
+import sys
+import re
 import tempfile
-from itertools import islice
+import unicodedata
 from pathlib import Path
+from itertools import islice
 from typing import Generator, Iterable, Optional, Union
-import unicodedata
-import sys
+
 import numpy as np
-import re
-from typing import Set
 
 
 def normalize(input_array, p=2, dim=1, eps=1e-12) -> np.ndarray:
@@ -45,7 +45,7 @@ def define_cache_dir(cache_dir: Optional[str] = None) -> Path:
     return cache_path
 
 
-def get_all_punctuation() -> Set[str]:
+def get_all_punctuation() -> set[str]:
     return set(
         chr(i) for i in range(sys.maxunicode) if unicodedata.category(chr(i)).startswith("P")
     )

{fastembed_gpu-0.4.2 → fastembed_gpu-0.5.0}/fastembed/image/image_embedding.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Iterable, List, Optional, Sequence, Type
+from typing import Any, Iterable, Optional, Sequence, Type
 
 import numpy as np
 
@@ -8,15 +8,15 @@ from fastembed.image.onnx_embedding import OnnxImageEmbedding
 
 
 class ImageEmbedding(ImageEmbeddingBase):
-    EMBEDDINGS_REGISTRY: List[Type[ImageEmbeddingBase]] = [OnnxImageEmbedding]
+    EMBEDDINGS_REGISTRY: list[Type[ImageEmbeddingBase]] = [OnnxImageEmbedding]
 
     @classmethod
-    def list_supported_models(cls) -> List[Dict[str, Any]]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """
         Lists the supported models.
 
         Returns:
-            List[Dict[str, Any]]: A list of dictionaries containing the model information.
+            list[dict[str, Any]]: A list of dictionaries containing the model information.
 
             Example:
                 ```
@@ -47,7 +47,7 @@ class ImageEmbedding(ImageEmbeddingBase):
         threads: Optional[int] = None,
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
-        device_ids: Optional[List[int]] = None,
+        device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
         **kwargs,
     ):

{fastembed_gpu-0.4.2 → fastembed_gpu-0.5.0}/fastembed/image/onnx_embedding.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Iterable, List, Optional, Sequence, Type
+from typing import Any, Iterable, Optional, Sequence, Type
 
 import numpy as np
 
@@ -53,6 +53,17 @@ supported_onnx_models = [
         },
         "model_file": "model.onnx",
     },
+    {
+        "model": "jinaai/jina-clip-v1",
+        "dim": 768,
+        "description": "Image embeddings, Multimodal (text&image), 2024 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.34,
+        "sources": {
+            "hf": "jinaai/jina-clip-v1",
+        },
+        "model_file": "onnx/vision_model.onnx",
+    },
 ]
 
 
@@ -64,7 +75,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
         threads: Optional[int] = None,
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
-        device_ids: Optional[List[int]] = None,
+        device_ids: Optional[list[int]] = None,
         lazy_load: bool = False,
         device_id: Optional[int] = None,
         **kwargs,
@@ -80,7 +91,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
                 Mutually exclusive with the `cuda` and `device_ids` arguments. Defaults to None.
             cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
                 Defaults to False.
-            device_ids (Optional[List[int]], optional): The list of device ids to use for data parallel processing in
+            device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
                 workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
             lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
                 Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
@@ -129,12 +140,12 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
         )
 
     @classmethod
-    def list_supported_models(cls) -> List[Dict[str, Any]]:
+    def list_supported_models(cls) -> list[dict[str, Any]]:
         """
         Lists the supported models.
 
         Returns:
-            List[Dict[str, Any]]: A list of dictionaries containing the model information.
+            list[Dict[str, Any]]: A list of dictionaries containing the model information.
         """
         return supported_onnx_models
 
@@ -178,8 +189,8 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
         return OnnxImageEmbeddingWorker
 
     def _preprocess_onnx_input(
-        self, onnx_input: Dict[str, np.ndarray], **kwargs
-    ) -> Dict[str, np.ndarray]:
+        self, onnx_input: dict[str, np.ndarray], **kwargs
+    ) -> dict[str, np.ndarray]:
         """
         Preprocess the onnx input.
         """

{fastembed_gpu-0.4.2 → fastembed_gpu-0.5.0}/fastembed/image/onnx_image_model.py

@@ -2,7 +2,7 @@ import contextlib
 import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type
+from typing import Any, Iterable, Optional, Sequence, Type
 
 import numpy as np
 from PIL import Image
@@ -29,8 +29,8 @@ class OnnxImageModel(OnnxModel[T]):
         self.processor = None
 
     def _preprocess_onnx_input(
-        self, onnx_input: Dict[str, np.ndarray], **kwargs
-    ) -> Dict[str, np.ndarray]:
+        self, onnx_input: dict[str, np.ndarray], **kwargs
+    ) -> dict[str, np.ndarray]:
         """
         Preprocess the onnx input.
         """
@@ -58,10 +58,10 @@ class OnnxImageModel(OnnxModel[T]):
     def load_onnx_model(self) -> None:
         raise NotImplementedError("Subclasses must implement this method")
 
-    def _build_onnx_input(self, encoded: np.ndarray) -> Dict[str, np.ndarray]:
+    def _build_onnx_input(self, encoded: np.ndarray) -> dict[str, np.ndarray]:
         return {node.name: encoded for node in self.model.get_inputs()}
 
-    def onnx_embed(self, images: List[ImageInput], **kwargs) -> OnnxOutputContext:
+    def onnx_embed(self, images: list[ImageInput], **kwargs) -> OnnxOutputContext:
         with contextlib.ExitStack():
             image_files = [
                 Image.open(image) if not isinstance(image, Image.Image) else image
@@ -83,7 +83,7 @@ class OnnxImageModel(OnnxModel[T]):
         parallel: Optional[int] = None,
         providers: Optional[Sequence[OnnxProvider]] = None,
         cuda: bool = False,
-        device_ids: Optional[List[int]] = None,
+        device_ids: Optional[list[int]] = None,
         **kwargs,
     ) -> Iterable[T]:
         is_small = False
@@ -125,7 +125,7 @@ class OnnxImageModel(OnnxModel[T]):
 
 
 class ImageEmbeddingWorker(EmbeddingWorker):
-    def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
             embeddings = self.model.onnx_embed(batch)
             yield idx, embeddings

{fastembed_gpu-0.4.2 → fastembed_gpu-0.5.0}/fastembed/image/transform/functional.py

@@ -1,4 +1,4 @@
-from typing import Sized, Tuple, Union
+from typing import Sized, Union
 
 import numpy as np
 from PIL import Image
@@ -14,7 +14,7 @@ def convert_to_rgb(image: Image.Image) -> Image.Image:
 
 def center_crop(
     image: Union[Image.Image, np.ndarray],
-    size: Tuple[int, int],
+    size: tuple[int, int],
 ) -> np.ndarray:
     if isinstance(image, np.ndarray):
         _, orig_height, orig_width = image.shape
@@ -62,8 +62,8 @@ def center_crop(
 
 def normalize(
     image: np.ndarray,
-    mean=Union[float, np.ndarray],
-    std=Union[float, np.ndarray],
+    mean: Union[float, np.ndarray],
+    std: Union[float, np.ndarray],
 ) -> np.ndarray:
     if not isinstance(image, np.ndarray):
         raise ValueError("image must be a numpy array")
@@ -96,10 +96,10 @@ def normalize(
 
 
 def resize(
-    image: Image,
-    size: Union[int, Tuple[int, int]],
-    resample: Image.Resampling = Image.Resampling.BILINEAR,
-) -> Image:
+    image: Image.Image,
+    size: Union[int, tuple[int, int]],
+    resample: Union[int, Image.Resampling] = Image.Resampling.BILINEAR,
+) -> Image.Image:
     if isinstance(size, tuple):
         return image.resize(size, resample)
 
@@ -122,3 +122,29 @@ def pil2ndarray(image: Union[Image.Image, np.ndarray]):
     if isinstance(image, Image.Image):
         return np.asarray(image).transpose((2, 0, 1))
     return image
+
+
+def pad2square(
+    image: Image.Image,
+    size: int,
+    fill_color: Union[str, int, tuple[int, ...]] = 0,
+) -> Image.Image:
+    height, width = image.height, image.width
+
+    left, right = 0, width
+    top, bottom = 0, height
+
+    crop_required = False
+    if width > size:
+        left = (width - size) // 2
+        right = left + size
+        crop_required = True
+
+    if height > size:
+        top = (height - size) // 2
+        bottom = top + size
+        crop_required = True
+
+    new_image = Image.new(mode="RGB", size=(size, size), color=fill_color)
+    new_image.paste(image.crop((left, top, right, bottom)) if crop_required else image)
+    return new_image