fastembed-gpu 0.5.1__py3-none-any.whl

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -0,0 +1,150 @@
+import os
+from multiprocessing import get_all_start_methods
+from pathlib import Path
+from typing import Any, Iterable, Optional, Sequence, Type
+
+import numpy as np
+from tokenizers import Encoding
+
+from fastembed.common.onnx_model import (
+    EmbeddingWorker,
+    OnnxModel,
+    OnnxOutputContext,
+    OnnxProvider,
+)
+from fastembed.common.preprocessor_utils import load_tokenizer
+from fastembed.common.utils import iter_batch
+from fastembed.parallel_processor import ParallelWorkerPool
+
+
+class OnnxCrossEncoderModel(OnnxModel[float]):
+    ONNX_OUTPUT_NAMES: Optional[list[str]] = None
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["TextRerankerWorker"]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    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, _ = load_tokenizer(model_dir=model_dir)
+
+    def tokenize(self, pairs: list[tuple[str, str]], **_: Any) -> list[Encoding]:
+        return self.tokenizer.encode_batch(pairs)
+
+    def _build_onnx_input(self, tokenized_input):
+        input_names = {node.name for node in self.model.get_inputs()}
+        inputs = {
+            "input_ids": np.array([enc.ids for enc in tokenized_input], dtype=np.int64),
+        }
+        if "token_type_ids" in input_names:
+            inputs["token_type_ids"] = np.array(
+                [enc.type_ids for enc in tokenized_input], dtype=np.int64
+            )
+        if "attention_mask" in input_names:
+            inputs["attention_mask"] = np.array(
+                [enc.attention_mask for enc in tokenized_input], dtype=np.int64
+            )
+        return inputs
+
+    def onnx_embed(self, query: str, documents: list[str], **kwargs: Any) -> OnnxOutputContext:
+        pairs = [(query, doc) for doc in documents]
+        return self.onnx_embed_pairs(pairs, **kwargs)
+
+    def onnx_embed_pairs(self, pairs: list[tuple[str, str]], **kwargs: Any) -> OnnxOutputContext:
+        tokenized_input = self.tokenize(pairs, **kwargs)
+        inputs = self._build_onnx_input(tokenized_input)
+        onnx_input = self._preprocess_onnx_input(inputs, **kwargs)
+        outputs = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)
+        relevant_output = outputs[0]
+        scores = relevant_output[:, 0]
+        return OnnxOutputContext(model_output=scores)
+
+    def _rerank_documents(
+        self, query: str, documents: Iterable[str], batch_size: int, **kwargs: Any
+    ) -> Iterable[float]:
+        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(query, batch, **kwargs))
+
+    def _rerank_pairs(
+        self,
+        model_name: str,
+        cache_dir: str,
+        pairs: Iterable[tuple[str, str]],
+        batch_size: int,
+        parallel: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[list[int]] = None,
+        **kwargs: Any,
+    ) -> Iterable[float]:
+        is_small = False
+
+        if isinstance(pairs, tuple):
+            pairs = [pairs]
+            is_small = True
+
+        if isinstance(pairs, list):
+            if len(pairs) < 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(pairs, batch_size):
+                yield from self._post_process_onnx_output(self.onnx_embed_pairs(batch, **kwargs))
+        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(pairs, batch_size), **params):
+                yield from self._post_process_onnx_output(batch)
+
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[float]:
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def _preprocess_onnx_input(
+        self, onnx_input: dict[str, np.ndarray], **kwargs: Any
+    ) -> dict[str, np.ndarray]:
+        """
+        Preprocess the onnx input.
+        """
+        return onnx_input
+
+
+class TextRerankerWorker(EmbeddingWorker):
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
+        for idx, batch in items:
+            onnx_output = self.model.onnx_embed_pairs(batch)
+            yield idx, onnx_output

fastembed/rerank/cross_encoder/text_cross_encoder.py

@@ -0,0 +1,120 @@
+from typing import Any, Iterable, Optional, Sequence, Type
+
+from fastembed.common import OnnxProvider
+from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
+from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase
+
+
+class TextCrossEncoder(TextCrossEncoderBase):
+    CROSS_ENCODER_REGISTRY: list[Type[TextCrossEncoderBase]] = [
+        OnnxTextCrossEncoder,
+    ]
+
+    @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": "Xenova/ms-marco-MiniLM-L-6-v2",
+                        "size_in_GB": 0.08,
+                        "sources": {
+                            "hf": "Xenova/ms-marco-MiniLM-L-6-v2",
+                        },
+                        "model_file": "onnx/model.onnx",
+                        "description": "MiniLM-L-6-v2 model optimized for re-ranking tasks.",
+                        "license": "apache-2.0",
+                    }
+                ]
+                ```
+        """
+        result = []
+        for encoder in cls.CROSS_ENCODER_REGISTRY:
+            result.extend(encoder.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,
+        cuda: bool = False,
+        device_ids: Optional[list[int]] = None,
+        lazy_load: bool = False,
+        **kwargs: Any,
+    ):
+        super().__init__(model_name, cache_dir, threads, **kwargs)
+
+        for CROSS_ENCODER_TYPE in self.CROSS_ENCODER_REGISTRY:
+            supported_models = CROSS_ENCODER_TYPE.list_supported_models()
+            if any(model_name.lower() == model["model"].lower() for model in supported_models):
+                self.model = CROSS_ENCODER_TYPE(
+                    model_name=model_name,
+                    cache_dir=cache_dir,
+                    threads=threads,
+                    providers=providers,
+                    cuda=cuda,
+                    device_ids=device_ids,
+                    lazy_load=lazy_load,
+                    **kwargs,
+                )
+                return
+
+        raise ValueError(
+            f"Model {model_name} is not supported in TextCrossEncoder."
+            "Please check the supported models using `TextCrossEncoder.list_supported_models()`"
+        )
+
+    def rerank(
+        self, query: str, documents: Iterable[str], batch_size: int = 64, **kwargs: Any
+    ) -> Iterable[float]:
+        """Rerank a list of documents based on a query.
+
+        Args:
+            query: Query to rerank the documents against
+            documents: Iterator of documents to rerank
+            batch_size: Batch size for reranking
+
+        Returns:
+            Iterable of scores for each document
+        """
+        yield from self.model.rerank(query, documents, batch_size=batch_size, **kwargs)
+
+    def rerank_pairs(
+        self,
+        pairs: Iterable[tuple[str, str]],
+        batch_size: int = 64,
+        parallel: Optional[int] = None,
+        **kwargs: Any,
+    ) -> Iterable[float]:
+        """
+        Rerank a list of query-document pairs.
+
+        Args:
+            pairs (Iterable[tuple[str, str]]): An iterable of tuples, where each tuple contains a query and a document
+                to be scored together.
+            batch_size (int, optional): The number of query-document pairs to process in a single batch. Defaults to 64.
+            parallel (Optional[int], optional): The number of parallel processes to use for reranking.
+                If None, parallelization is disabled. Defaults to None.
+            **kwargs (Any): Additional arguments to pass to the underlying reranking model.
+
+        Returns:
+            Iterable[float]: An iterable of scores corresponding to each query-document pair in the input.
+            Higher scores indicate a stronger match between the query and the document.
+
+        Example:
+            >>> encoder = TextCrossEncoder("Xenova/ms-marco-MiniLM-L-6-v2")
+            >>> pairs = [("What is AI?", "Artificial intelligence is ..."), ("What is ML?", "Machine learning is ...")]
+            >>> scores = list(encoder.rerank_pairs(pairs))
+            >>> print(list(map(lambda x: round(x, 2), scores)))
+            [-1.24, -10.6]
+        """
+        yield from self.model.rerank_pairs(
+            pairs, batch_size=batch_size, parallel=parallel, **kwargs
+        )

fastembed/rerank/cross_encoder/text_cross_encoder_base.py

@@ -0,0 +1,58 @@
+from typing import Any, Iterable, Optional
+
+from fastembed.common.model_management import ModelManagement
+
+
+class TextCrossEncoderBase(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 rerank(
+        self,
+        query: str,
+        documents: Iterable[str],
+        batch_size: int = 64,
+        **kwargs,
+    ) -> Iterable[float]:
+        """Rerank a list of documents given a query.
+
+        Args:
+            query (str): The query to rerank the documents.
+            documents (Iterable[str]): The list of texts to rerank.
+            batch_size (int): The batch size to use for reranking.
+            **kwargs: Additional keyword argument to pass to the rerank method.
+
+        Yields:
+            Iterable[float]: The scores of the reranked the documents.
+        """
+        raise NotImplementedError("This method should be overridden by subclasses")
+
+    def rerank_pairs(
+        self,
+        pairs: Iterable[tuple[str, str]],
+        batch_size: int = 64,
+        parallel: Optional[int] = None,
+        **kwargs: Any,
+    ) -> Iterable[float]:
+        """Rerank query-document pairs.
+        Args:
+            pairs (Iterable[tuple[str, str]]): Query-document pairs to rerank
+            batch_size (int): The batch size to use for reranking.
+            parallel: 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 rerank method.
+        Yields:
+            Iterable[float]: Scores for each individual pair
+        """
+        raise NotImplementedError("This method should be overridden by subclasses")

fastembed/sparse/__init__.py

@@ -0,0 +1,4 @@
+from fastembed.sparse.sparse_embedding_base import SparseEmbedding
+from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
+
+__all__ = ["SparseEmbedding", "SparseTextEmbedding"]

fastembed/sparse/bm25.py

@@ -0,0 +1,347 @@
+import os
+from collections import defaultdict
+from multiprocessing import get_all_start_methods
+from pathlib import Path
+from typing import Any, Iterable, Optional, Type, Union
+
+import mmh3
+import numpy as np
+from py_rust_stemmers import SnowballStemmer
+from fastembed.common.utils import (
+    define_cache_dir,
+    iter_batch,
+    get_all_punctuation,
+    remove_non_alphanumeric,
+)
+from fastembed.parallel_processor import ParallelWorkerPool, Worker
+from fastembed.sparse.sparse_embedding_base import (
+    SparseEmbedding,
+    SparseTextEmbeddingBase,
+)
+from fastembed.sparse.utils.tokenizer import SimpleTokenizer
+
+supported_languages = [
+    "arabic",
+    "azerbaijani",
+    "basque",
+    "bengali",
+    "catalan",
+    "chinese",
+    "danish",
+    "dutch",
+    "english",
+    "finnish",
+    "french",
+    "german",
+    "greek",
+    "hebrew",
+    "hinglish",
+    "hungarian",
+    "indonesian",
+    "italian",
+    "kazakh",
+    "nepali",
+    "norwegian",
+    "portuguese",
+    "romanian",
+    "russian",
+    "slovene",
+    "spanish",
+    "swedish",
+    "tajik",
+    "turkish",
+]
+
+supported_bm25_models = [
+    {
+        "model": "Qdrant/bm25",
+        "description": "BM25 as sparse embeddings meant to be used with Qdrant",
+        "license": "apache-2.0",
+        "size_in_GB": 0.01,
+        "sources": {
+            "hf": "Qdrant/bm25",
+        },
+        "model_file": "mock.file",  # bm25 does not require a model, so we just use a mock
+        "additional_files": [f"{lang}.txt" for lang in supported_languages],
+        "requires_idf": True,
+    },
+]
+
+
+class Bm25(SparseTextEmbeddingBase):
+    """Implements traditional BM25 in a form of sparse embeddings.
+    Uses a count of tokens in the document to evaluate the importance of the token.
+
+    WARNING: This model is expected to be used with `modifier="idf"` in the sparse vector index of Qdrant.
+
+    BM25 formula:
+
+    score(q, d) = SUM[ IDF(q_i) * (f(q_i, d) * (k + 1)) / (f(q_i, d) + k * (1 - b + b * (|d| / avg_len))) ],
+
+    where IDF is the inverse document frequency, computed on Qdrant's side
+    f(q_i, d) is the term frequency of the token q_i in the document d
+    k, b, avg_len are hyperparameters, described below.
+
+    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.
+        k (float, optional): The k parameter in the BM25 formula. Defines the saturation of the term frequency.
+            I.e. defines how fast the moment when additional terms stop to increase the score. Defaults to 1.2.
+        b (float, optional): The b parameter in the BM25 formula. Defines the importance of the document length.
+            Defaults to 0.75.
+        avg_len (float, optional): The average length of the documents in the corpus. Defaults to 256.0.
+        language (str): Specifies the language for the stemmer.
+        disable_stemmer (bool): Disable the stemmer.
+    Raises:
+        ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        k: float = 1.2,
+        b: float = 0.75,
+        avg_len: float = 256.0,
+        language: str = "english",
+        token_max_length: int = 40,
+        disable_stemmer: bool = False,
+        **kwargs,
+    ):
+        super().__init__(model_name, cache_dir, **kwargs)
+
+        if language not in supported_languages:
+            raise ValueError(f"{language} language is not supported")
+        else:
+            self.language = language
+
+        self.k = k
+        self.b = b
+        self.avg_len = avg_len
+
+        model_description = self._get_model_description(model_name)
+        self.cache_dir = define_cache_dir(cache_dir)
+
+        self._model_dir = self.download_model(
+            model_description, self.cache_dir, local_files_only=self._local_files_only
+        )
+
+        self.token_max_length = token_max_length
+        self.punctuation = set(get_all_punctuation())
+        self.disable_stemmer = disable_stemmer
+
+        if disable_stemmer:
+            self.stopwords = set()
+            self.stemmer = None
+        else:
+            self.stopwords = set(self._load_stopwords(self._model_dir, self.language))
+            self.stemmer = SnowballStemmer(language)
+
+        self.tokenizer = SimpleTokenizer
+
+    @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_bm25_models
+
+    @classmethod
+    def _load_stopwords(cls, model_dir: Path, language: str) -> list[str]:
+        stopwords_path = model_dir / f"{language}.txt"
+        if not stopwords_path.exists():
+            return []
+
+        with open(stopwords_path, "r") as f:
+            return f.read().splitlines()
+
+    def _embed_documents(
+        self,
+        model_name: str,
+        cache_dir: str,
+        documents: Union[str, Iterable[str]],
+        batch_size: int = 256,
+        parallel: Optional[int] = None,
+    ) -> Iterable[SparseEmbedding]:
+        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:
+            for batch in iter_batch(documents, batch_size):
+                yield from self.raw_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,
+                "k": self.k,
+                "b": self.b,
+                "avg_len": self.avg_len,
+                "language": self.language,
+                "token_max_length": self.token_max_length,
+                "disable_stemmer": self.disable_stemmer,
+            }
+            pool = ParallelWorkerPool(
+                num_workers=parallel or 1,
+                worker=self._get_worker_class(),
+                start_method=start_method,
+            )
+            for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
+                for record in batch:
+                    yield record
+
+    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,
+        )
+
+    def _stem(self, tokens: list[str]) -> list[str]:
+        stemmed_tokens = []
+        for token in tokens:
+            lower_token = token.lower()
+
+            if token in self.punctuation:
+                continue
+
+            if lower_token in self.stopwords:
+                continue
+
+            if len(token) > self.token_max_length:
+                continue
+
+            stemmed_token = self.stemmer.stem_word(lower_token) if self.stemmer else lower_token
+
+            if stemmed_token:
+                stemmed_tokens.append(stemmed_token)
+        return stemmed_tokens
+
+    def raw_embed(
+        self,
+        documents: list[str],
+    ) -> list[SparseEmbedding]:
+        embeddings = []
+        for document in documents:
+            document = remove_non_alphanumeric(document)
+            tokens = self.tokenizer.tokenize(document)
+            stemmed_tokens = self._stem(tokens)
+            token_id2value = self._term_frequency(stemmed_tokens)
+            embeddings.append(SparseEmbedding.from_dict(token_id2value))
+        return embeddings
+
+    def _term_frequency(self, tokens: list[str]) -> dict[int, float]:
+        """Calculate the term frequency part of the BM25 formula.
+
+        (
+            f(q_i, d) * (k + 1)
+        ) / (
+            f(q_i, d) + k * (1 - b + b * (|d| / avg_len))
+        )
+
+        Args:
+            tokens (list[str]): The list of tokens in the document.
+
+        Returns:
+            dict[int, float]: The token_id to term frequency mapping.
+        """
+        tf_map = {}
+        counter = defaultdict(int)
+        for stemmed_token in tokens:
+            counter[stemmed_token] += 1
+
+        doc_len = len(tokens)
+        for stemmed_token in counter:
+            token_id = self.compute_token_id(stemmed_token)
+            num_occurrences = counter[stemmed_token]
+            tf_map[token_id] = num_occurrences * (self.k + 1)
+            tf_map[token_id] /= num_occurrences + self.k * (
+                1 - self.b + self.b * doc_len / self.avg_len
+            )
+        return tf_map
+
+    @classmethod
+    def compute_token_id(cls, token: str) -> int:
+        return abs(mmh3.hash(token))
+
+    def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]:
+        """To emulate BM25 behaviour, we don't need to use weights in the query, and
+        it's enough to just hash the tokens and assign a weight of 1.0 to them.
+        """
+        if isinstance(query, str):
+            query = [query]
+
+        for text in query:
+            text = remove_non_alphanumeric(text)
+            tokens = self.tokenizer.tokenize(text)
+            stemmed_tokens = self._stem(tokens)
+            token_ids = np.array(
+                list(set(self.compute_token_id(token) for token in stemmed_tokens)),
+                dtype=np.int32,
+            )
+            values = np.ones_like(token_ids)
+            yield SparseEmbedding(indices=token_ids, values=values)
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["Bm25Worker"]:
+        return Bm25Worker
+
+
+class Bm25Worker(Worker):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs,
+    ):
+        self.model = self.init_embedding(model_name, cache_dir, **kwargs)
+
+    @classmethod
+    def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "Bm25Worker":
+        return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
+
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
+        for idx, batch in items:
+            onnx_output = self.model.raw_embed(batch)
+            yield idx, onnx_output
+
+    @staticmethod
+    def init_embedding(model_name: str, cache_dir: str, **kwargs) -> Bm25:
+        return Bm25(model_name=model_name, cache_dir=cache_dir, **kwargs)