fastembed-gpu 0.5.1__py3-none-any.whl

fastembed/text/onnx_embedding.py

@@ -0,0 +1,333 @@
+from typing import Any, Iterable, Optional, Sequence, Type, Union
+
+import numpy as np
+
+from fastembed.common import OnnxProvider
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.utils import define_cache_dir, normalize
+from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
+from fastembed.text.text_embedding_base import TextEmbeddingBase
+
+supported_onnx_models = [
+    {
+        "model": "BAAI/bge-base-en",
+        "dim": 768,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2023 year.",
+        "license": "mit",
+        "size_in_GB": 0.42,
+        "sources": {
+            "hf": "Qdrant/fast-bge-base-en",
+            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz",
+        },
+        "model_file": "model_optimized.onnx",
+    },
+    {
+        "model": "BAAI/bge-base-en-v1.5",
+        "dim": 768,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: not so necessary, 2023 year.",
+        "license": "mit",
+        "size_in_GB": 0.21,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
+            "hf": "qdrant/bge-base-en-v1.5-onnx-q",
+        },
+        "model_file": "model_optimized.onnx",
+    },
+    {
+        "model": "BAAI/bge-large-en-v1.5",
+        "dim": 1024,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: not so necessary, 2023 year.",
+        "license": "mit",
+        "size_in_GB": 1.20,
+        "sources": {
+            "hf": "qdrant/bge-large-en-v1.5-onnx",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "BAAI/bge-small-en",
+        "dim": 384,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2023 year.",
+        "license": "mit",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "Qdrant/bge-small-en",
+            "url": "https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
+        },
+        "model_file": "model_optimized.onnx",
+    },
+    {
+        "model": "BAAI/bge-small-en-v1.5",
+        "dim": 384,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: not so necessary, 2023 year.",
+        "license": "mit",
+        "size_in_GB": 0.067,
+        "sources": {
+            "hf": "qdrant/bge-small-en-v1.5-onnx-q",
+        },
+        "model_file": "model_optimized.onnx",
+    },
+    {
+        "model": "BAAI/bge-small-zh-v1.5",
+        "dim": 512,
+        "description": "Text embeddings, Unimodal (text), Chinese, 512 input tokens truncation, Prefixes for queries/documents: not so necessary, 2023 year.",
+        "license": "mit",
+        "size_in_GB": 0.09,
+        "sources": {
+            "hf": "Qdrant/bge-small-zh-v1.5",
+            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
+        },
+        "model_file": "model_optimized.onnx",
+    },
+    {
+        "model": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
+        "dim": 384,
+        "description": "Text embeddings, Unimodal (text), Multilingual (~50 languages), 512 input tokens truncation, Prefixes for queries/documents: not necessary, 2019 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.22,
+        "sources": {
+            "hf": "qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q",
+        },
+        "model_file": "model_optimized.onnx",
+    },
+    {
+        "model": "thenlper/gte-large",
+        "dim": 1024,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: not necessary, 2023 year.",
+        "license": "mit",
+        "size_in_GB": 1.20,
+        "sources": {
+            "hf": "qdrant/gte-large-onnx",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "mixedbread-ai/mxbai-embed-large-v1",
+        "dim": 1024,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.64,
+        "sources": {
+            "hf": "mixedbread-ai/mxbai-embed-large-v1",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-xs",
+        "dim": 384,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.09,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-xs",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-s",
+        "dim": 384,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-s",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-m",
+        "dim": 768,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.43,
+        "sources": {
+            "hf": "Snowflake/snowflake-arctic-embed-m",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-m-long",
+        "dim": 768,
+        "description": "Text embeddings, Unimodal (text), English, 2048 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.54,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-m-long",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "snowflake/snowflake-arctic-embed-l",
+        "dim": 1024,
+        "description": "Text embeddings, Unimodal (text), English, 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 1.02,
+        "sources": {
+            "hf": "snowflake/snowflake-arctic-embed-l",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "jinaai/jina-clip-v1",
+        "dim": 768,
+        "description": "Text embeddings, Multimodal (text&image), English, Prefixes for queries/documents: not necessary, 2024 year",
+        "license": "apache-2.0",
+        "size_in_GB": 0.55,
+        "sources": {
+            "hf": "jinaai/jina-clip-v1",
+        },
+        "model_file": "onnx/text_model.onnx",
+    },
+]
+
+
+class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
+    """Implementation of the Flag Embedding model."""
+
+    @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 __init__(
+        self,
+        model_name: str = "BAAI/bge-small-en-v1.5",
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[list[int]] = None,
+        lazy_load: bool = False,
+        device_id: Optional[int] = 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.
+            providers (Optional[Sequence[OnnxProvider]], optional): The list of onnxruntime providers to use.
+                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
+                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.
+            device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
+
+        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)
+        self.providers = providers
+        self.lazy_load = lazy_load
+
+        # List of device ids, that can be used for data parallel processing in workers
+        self.device_ids = device_ids
+        self.cuda = cuda
+
+        # This device_id will be used if we need to load model in current process
+        if device_id is not None:
+            self.device_id = device_id
+        elif self.device_ids is not None:
+            self.device_id = self.device_ids[0]
+        else:
+            self.device_id = None
+
+        self.model_description = self._get_model_description(model_name)
+        self.cache_dir = define_cache_dir(cache_dir)
+        self._model_dir = self.download_model(
+            self.model_description, self.cache_dir, local_files_only=self._local_files_only
+        )
+
+        if not self.lazy_load:
+            self.load_onnx_model()
+
+    def embed(
+        self,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        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:
+            documents: Iterator of documents or single document 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_documents(
+            model_name=self.model_name,
+            cache_dir=str(self.cache_dir),
+            documents=documents,
+            batch_size=batch_size,
+            parallel=parallel,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_ids=self.device_ids,
+            **kwargs,
+        )
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
+        return OnnxTextEmbeddingWorker
+
+    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]:
+        embeddings = output.model_output
+        if embeddings.ndim == 3:  # (batch_size, seq_len, embedding_dim)
+            processed_embeddings = embeddings[:, 0]
+        elif embeddings.ndim == 2:  # (batch_size, embedding_dim)
+            processed_embeddings = embeddings
+        else:
+            raise ValueError(f"Unsupported embedding shape: {embeddings.shape}")
+        return normalize(processed_embeddings).astype(np.float32)
+
+    def load_onnx_model(self) -> None:
+        self._load_onnx_model(
+            model_dir=self._model_dir,
+            model_file=self.model_description["model_file"],
+            threads=self.threads,
+            providers=self.providers,
+            cuda=self.cuda,
+            device_id=self.device_id,
+        )
+
+
+class OnnxTextEmbeddingWorker(TextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs,
+    ) -> OnnxTextEmbedding:
+        return OnnxTextEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/text/onnx_text_model.py

@@ -0,0 +1,145 @@
+import os
+from multiprocessing import get_all_start_methods
+from pathlib import Path
+from typing import Any, Iterable, Optional, Sequence, Type, Union
+
+import numpy as np
+from tokenizers import Encoding
+
+from fastembed.common import OnnxProvider
+from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
+from fastembed.common.preprocessor_utils import load_tokenizer
+from fastembed.common.utils import iter_batch
+from fastembed.parallel_processor import ParallelWorkerPool
+
+
+class OnnxTextModel(OnnxModel[T]):
+    ONNX_OUTPUT_NAMES: Optional[list[str]] = None
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
+        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:
+        super().__init__()
+        self.tokenizer = None
+        self.special_token_to_id = {}
+
+    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,
+        cuda: bool = False,
+        device_id: Optional[int] = None,
+    ) -> None:
+        super()._load_onnx_model(
+            model_dir=model_dir,
+            model_file=model_file,
+            threads=threads,
+            providers=providers,
+            cuda=cuda,
+            device_id=device_id,
+        )
+        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
+
+    def load_onnx_model(self) -> None:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def tokenize(self, documents: list[str], **kwargs) -> list[Encoding]:
+        return self.tokenizer.encode_batch(documents)
+
+    def onnx_embed(
+        self,
+        documents: list[str],
+        **kwargs,
+    ) -> OnnxOutputContext:
+        encoded = self.tokenize(documents, **kwargs)
+        input_ids = np.array([e.ids for e in encoded])
+        attention_mask = np.array([e.attention_mask for e in encoded])
+        input_names = {node.name for node in self.model.get_inputs()}
+        onnx_input = {
+            "input_ids": np.array(input_ids, dtype=np.int64),
+        }
+        if "attention_mask" in input_names:
+            onnx_input["attention_mask"] = np.array(attention_mask, dtype=np.int64)
+        if "token_type_ids" in input_names:
+            onnx_input["token_type_ids"] = np.array(
+                [np.zeros(len(e), dtype=np.int64) for e in input_ids], dtype=np.int64
+            )
+        onnx_input = self._preprocess_onnx_input(onnx_input, **kwargs)
+
+        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)
+        return OnnxOutputContext(
+            model_output=model_output[0],
+            attention_mask=onnx_input.get("attention_mask", attention_mask),
+            input_ids=onnx_input.get("input_ids", input_ids),
+        )
+
+    def _embed_documents(
+        self,
+        model_name: str,
+        cache_dir: str,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[list[int]] = None,
+        **kwargs,
+    ) -> Iterable[T]:
+        is_small = False
+
+        if isinstance(documents, str):
+            documents = [documents]
+            is_small = True
+
+        if isinstance(documents, list):
+            if len(documents) < batch_size:
+                is_small = True
+
+        if parallel is None or is_small:
+            if not hasattr(self, "model") or self.model is None:
+                self.load_onnx_model()
+            for batch in iter_batch(documents, batch_size):
+                yield from self._post_process_onnx_output(self.onnx_embed(batch))
+        else:
+            if parallel == 0:
+                parallel = os.cpu_count()
+
+            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+            params = {
+                "model_name": model_name,
+                "cache_dir": cache_dir,
+                "providers": providers,
+                **kwargs,
+            }
+
+            pool = ParallelWorkerPool(
+                num_workers=parallel or 1,
+                worker=self._get_worker_class(),
+                cuda=cuda,
+                device_ids=device_ids,
+                start_method=start_method,
+            )
+            for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
+                yield from self._post_process_onnx_output(batch)
+
+
+class TextEmbeddingWorker(EmbeddingWorker):
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
+        for idx, batch in items:
+            onnx_output = self.model.onnx_embed(batch)
+            yield idx, onnx_output

fastembed/text/pooled_embedding.py

@@ -0,0 +1,92 @@
+from typing import Any, Iterable, Type
+
+import numpy as np
+
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
+
+supported_pooled_models = [
+    {
+        "model": "nomic-ai/nomic-embed-text-v1.5",
+        "dim": 768,
+        "description": "Text embeddings, Multimodal (text, image), English, 8192 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.52,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1.5",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+    {
+        "model": "nomic-ai/nomic-embed-text-v1.5-Q",
+        "dim": 768,
+        "description": "Text embeddings, Multimodal (text, image), English, 8192 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.13,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1.5",
+        },
+        "model_file": "onnx/model_quantized.onnx",
+    },
+    {
+        "model": "nomic-ai/nomic-embed-text-v1",
+        "dim": 768,
+        "description": "Text embeddings, Multimodal (text, image), English, 8192 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.52,
+        "sources": {
+            "hf": "nomic-ai/nomic-embed-text-v1",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+]
+
+
+class PooledEmbedding(OnnxTextEmbedding):
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
+        return PooledEmbeddingWorker
+
+    @classmethod
+    def mean_pooling(cls, model_output: np.ndarray, attention_mask: np.ndarray) -> np.ndarray:
+        token_embeddings = model_output
+        input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
+        input_mask_expanded = np.tile(input_mask_expanded, (1, 1, token_embeddings.shape[-1]))
+        input_mask_expanded = input_mask_expanded.astype(float)
+        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
+        sum_mask = np.sum(input_mask_expanded, axis=1)
+        pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)
+        return pooled_embeddings
+
+    @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_pooled_models
+
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
+        if output.attention_mask is None:
+            raise ValueError("attention_mask must be provided for document post-processing")
+
+        embeddings = output.model_output
+        attn_mask = output.attention_mask
+        return self.mean_pooling(embeddings, attn_mask).astype(np.float32)
+
+
+class PooledEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs,
+    ) -> OnnxTextEmbedding:
+        return PooledEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )