fastembed-gpu 0.5.1__py3-none-any.whl

fastembed/sparse/splade_pp.py

@@ -0,0 +1,180 @@
+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
+from fastembed.sparse.sparse_embedding_base import (
+    SparseEmbedding,
+    SparseTextEmbeddingBase,
+)
+from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker
+
+supported_splade_models = [
+    {
+        "model": "prithivida/Splade_PP_en_v1",
+        "vocab_size": 30522,
+        "description": "Independent Implementation of SPLADE++ Model for English.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.532,
+        "sources": {
+            "hf": "Qdrant/SPLADE_PP_en_v1",
+        },
+        "model_file": "model.onnx",
+    },
+    {
+        "model": "prithvida/Splade_PP_en_v1",
+        "vocab_size": 30522,
+        "description": "Independent Implementation of SPLADE++ Model for English.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.532,
+        "sources": {
+            "hf": "Qdrant/SPLADE_PP_en_v1",
+        },
+        "model_file": "model.onnx",
+    },
+]
+
+
+class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[SparseEmbedding]:
+        if output.attention_mask is None:
+            raise ValueError("attention_mask must be provided for document post-processing")
+
+        relu_log = np.log(1 + np.maximum(output.model_output, 0))
+
+        weighted_log = relu_log * np.expand_dims(output.attention_mask, axis=-1)
+
+        scores = np.max(weighted_log, axis=1)
+
+        # Score matrix of shape (batch_size, vocab_size)
+        # Most of the values are 0, only a few are non-zero
+        for row_scores in scores:
+            indices = row_scores.nonzero()[0]
+            scores = row_scores[indices]
+            yield SparseEmbedding(values=scores, indices=indices)
+
+    @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_splade_models
+
+    def __init__(
+        self,
+        model_name: str,
+        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 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,
+        )
+
+    def embed(
+        self,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+        **kwargs,
+    ) -> Iterable[SparseEmbedding]:
+        """
+        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 SpladePPEmbeddingWorker
+
+
+class SpladePPEmbeddingWorker(TextEmbeddingWorker):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> SpladePP:
+        return SpladePP(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/sparse/utils/tokenizer.py

@@ -0,0 +1,120 @@
+# This code is a modified copy of the `NLTKWordTokenizer` class from `NLTK` library.
+
+import re
+
+
+class SimpleTokenizer:
+    @staticmethod
+    def tokenize(text: str) -> list[str]:
+        text = re.sub(r"[^\w]", " ", text.lower())
+        text = re.sub(r"\s+", " ", text)
+
+        return text.strip().split()
+
+
+class WordTokenizer:
+    """The tokenizer is "destructive" such that the regexes applied will munge the
+    input string to a state beyond re-construction.
+    """
+
+    # Starting quotes.
+    STARTING_QUOTES = [
+        (re.compile("([«“‘„]|[`]+)", re.U), r" \1 "),
+        (re.compile(r"^\""), r"``"),
+        (re.compile(r"(``)"), r" \1 "),
+        (re.compile(r"([ \(\[{<])(\"|\'{2})"), r"\1 `` "),
+        (re.compile(r"(?i)(\')(?!re|ve|ll|m|t|s|d|n)(\w)\b", re.U), r"\1 \2"),
+    ]
+
+    # Ending quotes.
+    ENDING_QUOTES = [
+        (re.compile("([»”’])", re.U), r" \1 "),
+        (re.compile(r"''"), " '' "),
+        (re.compile(r'"'), " '' "),
+        (re.compile(r"([^' ])('[sS]|'[mM]|'[dD]|') "), r"\1 \2 "),
+        (re.compile(r"([^' ])('ll|'LL|'re|'RE|'ve|'VE|n't|N'T) "), r"\1 \2 "),
+    ]
+
+    # Punctuation.
+    PUNCTUATION = [
+        (re.compile(r'([^\.])(\.)([\]\)}>"\'' "»”’ " r"]*)\s*$", re.U), r"\1 \2 \3 "),
+        (re.compile(r"([:,])([^\d])"), r" \1 \2"),
+        (re.compile(r"([:,])$"), r" \1 "),
+        (
+            re.compile(r"\.{2,}", re.U),
+            r" \g<0> ",
+        ),
+        (re.compile(r"[;@#$%&]"), r" \g<0> "),
+        (
+            re.compile(r'([^\.])(\.)([\]\)}>"\']*)\s*$'),
+            r"\1 \2\3 ",
+        ),  # Handles the final period.
+        (re.compile(r"[?!]"), r" \g<0> "),
+        (re.compile(r"([^'])' "), r"\1 ' "),
+        (
+            re.compile(r"[*]", re.U),
+            r" \g<0> ",
+        ),
+    ]
+
+    # Pads parentheses
+    PARENS_BRACKETS = (re.compile(r"[\]\[\(\)\{\}\<\>]"), r" \g<0> ")
+    DOUBLE_DASHES = (re.compile(r"--"), r" -- ")
+
+    # List of contractions adapted from Robert MacIntyre's tokenizer.
+    CONTRACTIONS2 = [
+        re.compile(pattern)
+        for pattern in (
+            r"(?i)\b(can)(?#X)(not)\b",
+            r"(?i)\b(d)(?#X)('ye)\b",
+            r"(?i)\b(gim)(?#X)(me)\b",
+            r"(?i)\b(gon)(?#X)(na)\b",
+            r"(?i)\b(got)(?#X)(ta)\b",
+            r"(?i)\b(lem)(?#X)(me)\b",
+            r"(?i)\b(more)(?#X)('n)\b",
+            r"(?i)\b(wan)(?#X)(na)(?=\s)",
+        )
+    ]
+    CONTRACTIONS3 = [
+        re.compile(pattern) for pattern in (r"(?i) ('t)(?#X)(is)\b", r"(?i) ('t)(?#X)(was)\b")
+    ]
+
+    @classmethod
+    def tokenize(cls, text: str) -> list[str]:
+        """Return a tokenized copy of `text`.
+
+        >>> s = '''Good muffins cost $3.88 (roughly 3,36 euros)\nin New York.'''
+        >>> WordTokenizer().tokenize(s)
+        ['Good', 'muffins', 'cost', '$', '3.88', '(', 'roughly', '3,36', 'euros', ')', 'in', 'New', 'York', '.']
+
+        Args:
+            text: The text to be tokenized.
+
+        Returns:
+            A list of tokens.
+        """
+        for regexp, substitution in cls.STARTING_QUOTES:
+            text = regexp.sub(substitution, text)
+
+        for regexp, substitution in cls.PUNCTUATION:
+            text = regexp.sub(substitution, text)
+
+        # Handles parentheses.
+        regexp, substitution = cls.PARENS_BRACKETS
+        text = regexp.sub(substitution, text)
+
+        # Handles double dash.
+        regexp, substitution = cls.DOUBLE_DASHES
+        text = regexp.sub(substitution, text)
+
+        # add extra space to make things easier
+        text = " " + text + " "
+
+        for regexp, substitution in cls.ENDING_QUOTES:
+            text = regexp.sub(substitution, text)
+
+        for regexp in cls.CONTRACTIONS2:
+            text = regexp.sub(r" \1 \2 ", text)
+        for regexp in cls.CONTRACTIONS3:
+            text = regexp.sub(r" \1 \2 ", text)
+        return text.split()

fastembed/text/__init__.py

@@ -0,0 +1,3 @@
+from fastembed.text.text_embedding import TextEmbedding
+
+__all__ = ["TextEmbedding"]

fastembed/text/clip_embedding.py

@@ -0,0 +1,54 @@
+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_clip_models = [
+    {
+        "model": "Qdrant/clip-ViT-B-32-text",
+        "dim": 512,
+        "description": "Text embeddings, Multimodal (text&image), English, 77 input tokens truncation, Prefixes for queries/documents: not necessary, 2021 year",
+        "license": "mit",
+        "size_in_GB": 0.25,
+        "sources": {
+            "hf": "Qdrant/clip-ViT-B-32-text",
+        },
+        "model_file": "model.onnx",
+    },
+]
+
+
+class CLIPOnnxEmbedding(OnnxTextEmbedding):
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
+        return CLIPEmbeddingWorker
+
+    @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_clip_models
+
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
+        return output.model_output
+
+
+class CLIPEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs,
+    ) -> OnnxTextEmbedding:
+        return CLIPOnnxEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/text/e5_onnx_embedding.py

@@ -0,0 +1,72 @@
+from typing import Any, Type
+
+import numpy as np
+
+from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
+
+supported_multilingual_e5_models = [
+    {
+        "model": "intfloat/multilingual-e5-large",
+        "dim": 1024,
+        "description": "Text embeddings, Unimodal (text), Multilingual (~100 languages), 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
+        "license": "mit",
+        "size_in_GB": 2.24,
+        "sources": {
+            "url": "https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
+            "hf": "qdrant/multilingual-e5-large-onnx",
+        },
+        "model_file": "model.onnx",
+        "additional_files": ["model.onnx_data"],
+    },
+    {
+        "model": "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
+        "dim": 768,
+        "description": "Text embeddings, Unimodal (text), Multilingual (~50 languages), 384 input tokens truncation, Prefixes for queries/documents: not necessary, 2021 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 1.00,
+        "sources": {
+            "hf": "xenova/paraphrase-multilingual-mpnet-base-v2",
+        },
+        "model_file": "onnx/model.onnx",
+    },
+]
+
+
+class E5OnnxEmbedding(OnnxTextEmbedding):
+    @classmethod
+    def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
+        return E5OnnxEmbeddingWorker
+
+    @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_multilingual_e5_models
+
+    def _preprocess_onnx_input(
+        self, onnx_input: dict[str, np.ndarray], **kwargs
+    ) -> dict[str, np.ndarray]:
+        """
+        Preprocess the onnx input.
+        """
+        onnx_input.pop("token_type_ids", None)
+        return onnx_input
+
+
+class E5OnnxEmbeddingWorker(OnnxTextEmbeddingWorker):
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs,
+    ) -> E5OnnxEmbedding:
+        return E5OnnxEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )