fastembed-bio 0.1.0__py3-none-any.whl

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -0,0 +1,204 @@
+import os
+from multiprocessing import get_all_start_methods
+from pathlib import Path
+from typing import Any, Iterable, Sequence, Type
+
+import numpy as np
+from tokenizers import Encoding
+
+from fastembed.common.onnx_model import (
+    EmbeddingWorker,
+    OnnxModel,
+    OnnxOutputContext,
+    OnnxProvider,
+)
+from fastembed.common.types import NumpyArray, Device
+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: list[str] | None = 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: int | None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_id: int | None = None,
+        extra_session_options: dict[str, Any] | None = None,
+    ) -> None:
+        super()._load_onnx_model(
+            model_dir=model_dir,
+            model_file=model_file,
+            threads=threads,
+            providers=providers,
+            cuda=cuda,
+            device_id=device_id,
+            extra_session_options=extra_session_options,
+        )
+        self.tokenizer, _ = load_tokenizer(model_dir=model_dir)
+        assert self.tokenizer is not None
+
+    def tokenize(self, pairs: list[tuple[str, str]], **_: Any) -> list[Encoding]:
+        return self.tokenizer.encode_batch(pairs)  # type: ignore[union-attr]
+
+    def _build_onnx_input(self, tokenized_input: list[Encoding]) -> dict[str, NumpyArray]:
+        input_names: set[str] = {node.name for node in self.model.get_inputs()}  # type: ignore[union-attr]
+        inputs: dict[str, NumpyArray] = {
+            "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)  # type: ignore[union-attr]
+        relevant_output = outputs[0]
+        scores: NumpyArray = 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: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = None,
+        local_files_only: bool = False,
+        specific_model_path: str | None = None,
+        extra_session_options: dict[str, Any] | None = 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,
+                "local_files_only": local_files_only,
+                "specific_model_path": specific_model_path,
+                **kwargs,
+            }
+
+            if extra_session_options is not None:
+                params.update(extra_session_options)
+
+            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)  # type: ignore
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[float]:
+        """Post-process the ONNX model output to convert it into a usable format.
+
+        Args:
+            output (OnnxOutputContext): The raw output from the ONNX model.
+            **kwargs: Additional keyword arguments that may be needed by specific implementations.
+
+        Returns:
+            Iterable[float]: Post-processed output as an iterable of float values.
+        """
+        raise NotImplementedError("Subclasses must implement this method")
+
+    def _preprocess_onnx_input(
+        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+    ) -> dict[str, NumpyArray]:
+        """
+        Preprocess the onnx input.
+        """
+        return onnx_input
+
+    def _token_count(
+        self, pairs: Iterable[tuple[str, str]], batch_size: int = 1024, **_: Any
+    ) -> int:
+        if not hasattr(self, "model") or self.model is None:
+            self.load_onnx_model()  # loads the tokenizer as well
+
+        token_num = 0
+        assert self.tokenizer is not None
+        for batch in iter_batch(pairs, batch_size):
+            for tokens in self.tokenizer.encode_batch(batch):
+                token_num += sum(tokens.attention_mask)
+
+        return token_num
+
+
+class TextRerankerWorker(EmbeddingWorker[float]):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ):
+        self.model: OnnxCrossEncoderModel
+        super().__init__(model_name, cache_dir, **kwargs)
+
+    def init_embedding(
+        self,
+        model_name: str,
+        cache_dir: str,
+        **kwargs: Any,
+    ) -> OnnxCrossEncoderModel:
+        raise NotImplementedError()
+
+    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,178 @@
+from typing import Any, Iterable, Sequence, Type
+from dataclasses import asdict
+
+from fastembed.common import OnnxProvider
+from fastembed.common.types import Device
+from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
+from fastembed.rerank.cross_encoder.custom_text_cross_encoder import CustomTextCrossEncoder
+
+from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase
+from fastembed.common.model_description import (
+    ModelSource,
+    BaseModelDescription,
+)
+
+
+class TextCrossEncoder(TextCrossEncoderBase):
+    CROSS_ENCODER_REGISTRY: list[Type[TextCrossEncoderBase]] = [
+        OnnxTextCrossEncoder,
+        CustomTextCrossEncoder,
+    ]
+
+    @classmethod
+    def list_supported_models(cls) -> list[dict[str, Any]]:
+        """Lists the supported models.
+
+        Returns:
+            list[BaseModelDescription]: 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",
+                    }
+                ]
+                ```
+        """
+        return [asdict(model) for model in cls._list_supported_models()]
+
+    @classmethod
+    def _list_supported_models(cls) -> list[BaseModelDescription]:
+        result: list[BaseModelDescription] = []
+        for encoder in cls.CROSS_ENCODER_REGISTRY:
+            result.extend(encoder._list_supported_models())
+        return result
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        providers: Sequence[OnnxProvider] | None = None,
+        cuda: bool | Device = Device.AUTO,
+        device_ids: list[int] | None = 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: int | None = 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
+        )
+
+    @classmethod
+    def add_custom_model(
+        cls,
+        model: str,
+        sources: ModelSource,
+        model_file: str = "onnx/model.onnx",
+        description: str = "",
+        license: str = "",
+        size_in_gb: float = 0.0,
+        additional_files: list[str] | None = None,
+    ) -> None:
+        registered_models = cls._list_supported_models()
+        for registered_model in registered_models:
+            if model == registered_model.model:
+                raise ValueError(
+                    f"Model {model} is already registered in CrossEncoderModel, if you still want to add this model, "
+                    f"please use another model name"
+                )
+
+        CustomTextCrossEncoder.add_model(
+            BaseModelDescription(
+                model=model,
+                sources=sources,
+                model_file=model_file,
+                description=description,
+                license=license,
+                size_in_GB=size_in_gb,
+                additional_files=additional_files or [],
+            )
+        )
+
+    def token_count(
+        self, pairs: Iterable[tuple[str, str]], batch_size: int = 1024, **kwargs: Any
+    ) -> int:
+        """Returns the number of tokens in the pairs.
+
+        Args:
+            pairs: Iterable of tuples, where each tuple contains a query and a document to be tokenized
+            batch_size: Batch size for tokenizing
+
+        Returns:
+            token count: overall number of tokens in the pairs
+        """
+        return self.model.token_count(pairs, batch_size=batch_size, **kwargs)

fastembed/rerank/cross_encoder/text_cross_encoder_base.py

@@ -0,0 +1,63 @@
+from typing import Any, Iterable
+
+from fastembed.common.model_description import BaseModelDescription
+from fastembed.common.model_management import ModelManagement
+
+
+class TextCrossEncoderBase(ModelManagement[BaseModelDescription]):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        **kwargs: Any,
+    ):
+        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: Any,
+    ) -> 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: int | None = 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")
+
+    def token_count(self, pairs: Iterable[tuple[str, str]], **kwargs: Any) -> int:
+        """Returns the number of tokens in the pairs."""
+        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"]