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

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

@@ -1,12 +1,12 @@
 Metadata-Version: 2.1
 Name: fastembed-gpu
-Version: 0.2.7
+Version: 0.3.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: NirantK
-Author-email: nirant.bits@gmail.com
+Author: Qdrant Team
+Author-email: info@qdrant.tech
 Requires-Python: >=3.8.0,<3.13
 Classifier: License :: Other/Proprietary License
 Classifier: Programming Language :: Python :: 3
@@ -15,13 +15,17 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
-Requires-Dist: huggingface-hub (>=0.20,<0.21)
+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: onnxruntime-gpu (>=1.17.0,<2.0.0)
+Requires-Dist: pillow (>=10.3.0,<11.0.0)
 Requires-Dist: requests (>=2.31,<3.0)
-Requires-Dist: tokenizers (>=0.15,<0.16)
+Requires-Dist: snowballstemmer (>=2.2.0,<3.0.0)
+Requires-Dist: tokenizers (>=0.15,<1.0)
 Requires-Dist: tqdm (>=4.66,<5.0)
 Project-URL: Repository, https://github.com/qdrant/fastembed
 Description-Content-Type: text/markdown
@@ -78,17 +82,22 @@ len(embeddings_list[0]) # Vector of 384 dimensions
 
 ### ⚡️ FastEmbed on a GPU
 
-FastEmbed supports running on GPU devices. It requires installation of the `fastembed-gpu` package.
-Make sure not to have the `fastembed` package installed, as it might interfere with the `fastembed-gpu` package.
+FastEmbed supports running on GPU devices.
+It requires installation of the `fastembed-gpu` package.
 
 ```bash
 pip install fastembed-gpu
-``` 
+```
+
+Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for the detailed instructions and CUDA 12.x support.
 
 ```python
 from fastembed import TextEmbedding
 
-embedding_model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5", providers=["CUDAExecutionProvider"])
+embedding_model = TextEmbedding(
+    model_name="BAAI/bge-small-en-v1.5", 
+    providers=["CUDAExecutionProvider"]
+)
 print("The model BAAI/bge-small-en-v1.5 is ready to use on a GPU.")
 
 ```

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

@@ -50,17 +50,22 @@ len(embeddings_list[0]) # Vector of 384 dimensions
 
 ### ⚡️ FastEmbed on a GPU
 
-FastEmbed supports running on GPU devices. It requires installation of the `fastembed-gpu` package.
-Make sure not to have the `fastembed` package installed, as it might interfere with the `fastembed-gpu` package.
+FastEmbed supports running on GPU devices.
+It requires installation of the `fastembed-gpu` package.
 
 ```bash
 pip install fastembed-gpu
-``` 
+```
+
+Check our [example](https://qdrant.github.io/fastembed/examples/FastEmbed_GPU/) for the detailed instructions and CUDA 12.x support.
 
 ```python
 from fastembed import TextEmbedding
 
-embedding_model = TextEmbedding(model_name="BAAI/bge-small-en-v1.5", providers=["CUDAExecutionProvider"])
+embedding_model = TextEmbedding(
+    model_name="BAAI/bge-small-en-v1.5", 
+    providers=["CUDAExecutionProvider"]
+)
 print("The model BAAI/bge-small-en-v1.5 is ready to use on a GPU.")
 
 ```

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

@@ -1,7 +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
 
 try:
     version = importlib.metadata.version("fastembed")
@@ -9,4 +11,10 @@ except importlib.metadata.PackageNotFoundError as _:
     version = importlib.metadata.version("fastembed-gpu")
 
 __version__ = version
-__all__ = ["TextEmbedding", "SparseTextEmbedding", "SparseEmbedding"]
+__all__ = [
+    "TextEmbedding",
+    "SparseTextEmbedding",
+    "SparseEmbedding",
+    "ImageEmbedding",
+    "LateInteractionTextEmbedding",
+]

fastembed_gpu-0.3.0/fastembed/common/__init__.py

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

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

@@ -108,6 +108,7 @@ class ModelManagement:
             "tokenizer.json",
             "tokenizer_config.json",
             "special_tokens_map.json",
+            "preprocessor_config.json",
         ]
         if extra_patterns is not None:
             allow_patterns.extend(extra_patterns)

fastembed_gpu-0.3.0/fastembed/common/onnx_model.py

@@ -0,0 +1,113 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, Generic, Iterable, Optional, Tuple, Type, TypeVar, Sequence
+import warnings
+
+import numpy as np
+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")
+
+
+@dataclass
+class OnnxOutputContext:
+    model_output: np.ndarray
+    attention_mask: Optional[np.ndarray] = None
+    input_ids: Optional[np.ndarray] = None
+
+
+class OnnxModel(Generic[T]):
+    @classmethod
+    def _get_worker_class(cls) -> Type["EmbeddingWorker"]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def __init__(self) -> None:
+        self.model = None
+        self.tokenizer = None
+
+    def _preprocess_onnx_input(
+        self, onnx_input: Dict[str, np.ndarray], **kwargs
+    ) -> Dict[str, np.ndarray]:
+        """
+        Preprocess the onnx input.
+        """
+        return onnx_input
+
+    def load_onnx_model(
+        self,
+        model_dir: Path,
+        model_file: str,
+        threads: Optional[int],
+        providers: Optional[Sequence[OnnxProvider]] = None,
+    ) -> None:
+        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)
+        available_providers = ort.get_available_providers()
+        requested_provider_names = []
+        for provider in onnx_providers:
+            # check providers available
+            provider_name = provider if isinstance(provider, str) else provider[0]
+            requested_provider_names.append(provider_name)
+            if provider_name not in available_providers:
+                raise ValueError(
+                    f"Provider {provider_name} is not available. Available providers: {available_providers}"
+                )
+
+        so = ort.SessionOptions()
+        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+        if threads is not None:
+            so.intra_op_num_threads = threads
+            so.inter_op_num_threads = threads
+
+        self.model = ort.InferenceSession(
+            str(model_path), providers=onnx_providers, sess_options=so
+        )
+        if "CUDAExecutionProvider" in requested_provider_names:
+            current_providers = self.model.get_providers()
+            if "CUDAExecutionProvider" not in current_providers:
+                warnings.warn(
+                    f"Attempt to set CUDAExecutionProvider failed. Current providers: {current_providers}."
+                    "If you are using CUDA 12.x, install onnxruntime-gpu via "
+                    "`pip install onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/`",
+                    RuntimeWarning,
+                )
+
+    def onnx_embed(self, *args, **kwargs) -> OnnxOutputContext:
+        raise NotImplementedError("Subclasses must implement this method")
+
+
+class EmbeddingWorker(Worker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+    ) -> OnnxModel:
+        raise NotImplementedError()
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str,
+    ):
+        self.model = self.init_embedding(model_name, cache_dir)
+
+    @classmethod
+    def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker":
+        return cls(
+            model_name=model_name,
+            cache_dir=cache_dir,
+        )
+
+    def process(self, items: Iterable[Tuple[int, Any]]) -> Iterable[Tuple[int, Any]]:
+        raise NotImplementedError("Subclasses must implement this method")

fastembed_gpu-0.2.7/fastembed/common/models.py → fastembed_gpu-0.3.0/fastembed/common/preprocessor_utils.py

@@ -1,11 +1,24 @@
 import json
 from pathlib import Path
+from typing import Tuple
 
-import numpy as np
 from tokenizers import Tokenizer, AddedToken
 
+from fastembed.image.transform.operators import Compose
 
-def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tokenizer:
+
+def load_special_tokens(model_dir: Path) -> dict:
+    tokens_map_path = model_dir / "special_tokens_map.json"
+    if not tokens_map_path.exists():
+        raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
+
+    with open(str(tokens_map_path)) as tokens_map_file:
+        tokens_map = json.load(tokens_map_file)
+
+    return tokens_map
+
+
+def load_tokenizer(model_dir: Path, max_length: int = 512) -> 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}")
@@ -18,18 +31,13 @@ def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tokenizer:
     if not tokenizer_config_path.exists():
         raise ValueError(f"Could not find tokenizer_config.json in {model_dir}")
 
-    tokens_map_path = model_dir / "special_tokens_map.json"
-    if not tokens_map_path.exists():
-        raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")
-
     with open(str(config_path)) as config_file:
         config = json.load(config_file)
 
     with open(str(tokenizer_config_path)) as tokenizer_config_file:
         tokenizer_config = json.load(tokenizer_config_file)
 
-    with open(str(tokens_map_path)) as tokens_map_file:
-        tokens_map = json.load(tokens_map_file)
+    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))
@@ -43,12 +51,24 @@ def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tokenizer:
         elif isinstance(token, dict):
             tokenizer.add_special_tokens([AddedToken(**token)])
 
-    return tokenizer
+    special_token_to_id = {}
+
+    for token in tokens_map.values():
+        if isinstance(token, str):
+            special_token_to_id[token] = tokenizer.token_to_id(token)
+        elif isinstance(token, dict):
+            token_str = token.get("content", "")
+            special_token_to_id[token_str] = tokenizer.token_to_id(token_str)
+
+    return tokenizer, special_token_to_id
+
 
+def load_preprocessor(model_dir: Path) -> Compose:
+    preprocessor_config_path = model_dir / "preprocessor_config.json"
+    if not preprocessor_config_path.exists():
+        raise ValueError(f"Could not find preprocessor_config.json in {model_dir}")
 
-def normalize(input_array, p=2, dim=1, eps=1e-12) -> np.ndarray:
-    # Calculate the Lp norm along the specified dimension
-    norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
-    norm = np.maximum(norm, eps)  # Avoid division by zero
-    normalized_array = input_array / norm
-    return normalized_array
+    with open(str(preprocessor_config_path)) as preprocessor_config_file:
+        preprocessor_config = json.load(preprocessor_config_file)
+        transforms = Compose.from_config(preprocessor_config)
+    return transforms

fastembed_gpu-0.3.0/fastembed/common/types.py

@@ -0,0 +1,14 @@
+import os
+import sys
+from typing import Union, Iterable, Tuple, Dict, Any
+
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias
+
+
+PathInput: TypeAlias = Union[str, os.PathLike]
+ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput]]
+
+OnnxProvider: TypeAlias = Union[str, Tuple[str, Dict[Any, Any]]]

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

@@ -4,6 +4,16 @@ from itertools import islice
 from pathlib import Path
 from typing import Union, Iterable, Generator, Optional
 
+import numpy as np
+
+
+def normalize(input_array, p=2, dim=1, eps=1e-12) -> np.ndarray:
+    # Calculate the Lp norm along the specified dimension
+    norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
+    norm = np.maximum(norm, eps)  # Avoid division by zero
+    normalized_array = input_array / norm
+    return normalized_array
+
 
 def iter_batch(iterable: Union[Iterable, Generator], size: int) -> Iterable:
     """

fastembed_gpu-0.3.0/fastembed/image/__init__.py

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

fastembed_gpu-0.3.0/fastembed/image/image_embedding.py

@@ -0,0 +1,87 @@
+from typing import Any, Dict, Iterable, List, Optional, Type, Sequence
+
+import numpy as np
+
+from fastembed.common import ImageInput, OnnxProvider
+from fastembed.image.image_embedding_base import ImageEmbeddingBase
+from fastembed.image.onnx_embedding import OnnxImageEmbedding
+
+
+class ImageEmbedding(ImageEmbeddingBase):
+    EMBEDDINGS_REGISTRY: List[Type[ImageEmbeddingBase]] = [OnnxImageEmbedding]
+
+    @classmethod
+    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.
+
+            Example:
+                ```
+                [
+                    {
+                        "model": "Qdrant/clip-ViT-B-32-vision",
+                        "dim": 512,
+                        "description": "CLIP vision encoder based on ViT-B/32",
+                        "size_in_GB": 0.33,
+                        "sources": {
+                            "hf": "Qdrant/clip-ViT-B-32-vision",
+                        },
+                        "model_file": "model.onnx",
+                    }
+                ]
+                ```
+        """
+        result = []
+        for embedding in cls.EMBEDDINGS_REGISTRY:
+            result.extend(embedding.list_supported_models())
+        return result
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        **kwargs,
+    ):
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+
+        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):
+                self.model = EMBEDDING_MODEL_TYPE(
+                    model_name, cache_dir, threads, providers=providers, **kwargs
+                )
+                return
+
+        raise ValueError(
+            f"Model {model_name} is not supported in TextEmbedding."
+            "Please check the supported models using `TextEmbedding.list_supported_models()`"
+        )
+
+    def embed(
+        self,
+        images: ImageInput,
+        batch_size: int = 16,
+        parallel: Optional[int] = None,
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
+        """
+        Encode a list of documents into list of embeddings.
+        We use mean pooling with attention so that the model can handle variable-length inputs.
+
+        Args:
+            images: Iterator of image paths or single image path to embed
+            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
+            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.
+
+        Returns:
+            List of embeddings, one per document
+        """
+        yield from self.model.embed(images, batch_size, parallel, **kwargs)

fastembed_gpu-0.3.0/fastembed/image/image_embedding_base.py

@@ -0,0 +1,39 @@
+from typing import Iterable, Optional
+
+import numpy as np
+
+from fastembed.common.model_management import ModelManagement
+from fastembed.common.types import ImageInput
+
+
+class ImageEmbeddingBase(ModelManagement):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        **kwargs,
+    ):
+        self.model_name = model_name
+        self.cache_dir = cache_dir
+        self.threads = threads
+        self._local_files_only = kwargs.pop("local_files_only", False)
+
+    def embed(
+        self,
+        images: ImageInput,
+        batch_size: int = 16,
+        parallel: Optional[int] = None,
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
+        """
+        Embeds a list of images into a list of embeddings.
+
+        Args:
+            images - The list of image paths to preprocess and embed.
+            **kwargs: Additional keyword argument to pass to the embed method.
+
+        Yields:
+            Iterable[np.ndarray]: The embeddings.
+        """
+        raise NotImplementedError()

fastembed_gpu-0.3.0/fastembed/image/onnx_embedding.py

@@ -0,0 +1,134 @@
+from typing import Dict, Optional, Iterable, Type, List, Any, Sequence
+
+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.image.image_embedding_base import ImageEmbeddingBase
+from fastembed.image.onnx_image_model import OnnxImageModel, ImageEmbeddingWorker
+
+supported_onnx_models = [
+    {
+        "model": "Qdrant/clip-ViT-B-32-vision",
+        "dim": 512,
+        "description": "CLIP vision encoder based on ViT-B/32",
+        "size_in_GB": 0.34,
+        "sources": {
+            "hf": "Qdrant/clip-ViT-B-32-vision",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "Qdrant/resnet50-onnx",
+        "dim": 2048,
+        "description": "ResNet-50 from `Deep Residual Learning for Image Recognition <https://arxiv.org/abs/1512.03385>`__.",
+        "size_in_GB": 0.1,
+        "sources": {
+            "hf": "Qdrant/resnet50-onnx",
+        },
+        "model_file": "model.onnx",
+    },
+]
+
+
+class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            model_name (str): The name of the model to use.
+            cache_dir (str, optional): The path to the cache directory.
+                                       Can be set using the `FASTEMBED_CACHE_PATH` env variable.
+                                       Defaults to `fastembed_cache` in the system's temp directory.
+            threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
+
+        Raises:
+            ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+        """
+
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+
+        model_description = self._get_model_description(model_name)
+        cache_dir = define_cache_dir(cache_dir)
+        model_dir = self.download_model(
+            model_description, cache_dir, local_files_only=self._local_files_only
+        )
+
+        self.load_onnx_model(
+            model_dir=model_dir,
+            model_file=model_description["model_file"],
+            threads=threads,
+            providers=providers,
+        )
+
+    @classmethod
+    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.
+        """
+        return supported_onnx_models
+
+    def embed(
+        self,
+        images: ImageInput,
+        batch_size: int = 16,
+        parallel: Optional[int] = None,
+        **kwargs,
+    ) -> Iterable[np.ndarray]:
+        """
+        Encode a list of images into list of embeddings.
+        We use mean pooling with attention so that the model can handle variable-length inputs.
+
+        Args:
+            images: Iterator of image paths or single image path to embed
+            batch_size: Batch size for encoding -- higher values will use more memory, but be faster
+            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.
+
+        Returns:
+            List of embeddings, one per document
+        """
+        yield from self._embed_images(
+            model_name=self.model_name,
+            cache_dir=str(self.cache_dir),
+            images=images,
+            batch_size=batch_size,
+            parallel=parallel,
+        )
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker"]:
+        return OnnxImageEmbeddingWorker
+
+    def _preprocess_onnx_input(
+        self, onnx_input: Dict[str, np.ndarray], **kwargs
+    ) -> Dict[str, np.ndarray]:
+        """
+        Preprocess the onnx input.
+        """
+
+        return onnx_input
+
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
+        return normalize(output.model_output).astype(np.float32)
+
+
+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)