From c8b1a18cfc3bdaa99c19e305e4ee813885c00583 Mon Sep 17 00:00:00 2001 From: Dmitrii Ogn Date: Mon, 16 Dec 2024 21:59:34 +0300 Subject: [PATCH] Cross encoders parallelism (#419) * Merge master * rerank_pairs interface + parallelism support * remove test notebook * Removed unused code * New tests for cross encoders and new interface * Importing Self fix. We will need it for mypy support in newer versions * Removed Self typing * Removed non-needed changes from text * Isort + black * wip: start reviewing (#420) Co-authored-by: Dmitrii Ogn * Test fix * Update fastembed/rerank/cross_encoder/text_cross_encoder.py Co-authored-by: George * Update fastembed/rerank/cross_encoder/text_cross_encoder.py Co-authored-by: George * Update fastembed/rerank/cross_encoder/text_cross_encoder.py Co-authored-by: George * Update fastembed/rerank/cross_encoder/text_cross_encoder_base.py Co-authored-by: George * Test for parallel processing + bugfix of PosixPath passing * Removed non-needed import and added docstring * Typing fix + argument passing * Test parametrization Moved to selected models set to test * Run base test on all models * Typing fix + improvement of input_names check * nit: fix post process, update docstring, update tokenize, remove redundant imports --------- Co-authored-by: George --- .../cross_encoder/onnx_text_cross_encoder.py | 59 ++++++++- .../rerank/cross_encoder/onnx_text_model.py | 112 +++++++++++++++--- .../cross_encoder/text_cross_encoder.py | 41 ++++++- .../cross_encoder/text_cross_encoder_base.py | 27 ++++- tests/test_text_cross_encoder.py | 89 ++++++++++---- 5 files changed, 273 insertions(+), 55 deletions(-) diff --git a/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py b/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py index 3a4087d..3b71671 100644 --- a/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py +++ b/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py @@ -1,11 +1,15 @@ -from typing import Iterable, Any, Sequence, Optional +from typing import Any, Iterable, Optional, Sequence, Type from loguru import logger from fastembed.common import OnnxProvider -from fastembed.rerank.cross_encoder.onnx_text_model import OnnxCrossEncoderModel -from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase +from fastembed.common.onnx_model import OnnxOutputContext from fastembed.common.utils import define_cache_dir +from fastembed.rerank.cross_encoder.onnx_text_model import ( + OnnxCrossEncoderModel, + TextRerankerWorker, +) +from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase supported_onnx_models = [ { @@ -91,7 +95,7 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel): device_ids: Optional[list[int]] = None, lazy_load: bool = False, device_id: Optional[int] = None, - **kwargs, + **kwargs: Any, ): """ Args: @@ -138,7 +142,9 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel): 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 + self.model_description, + self.cache_dir, + local_files_only=self._local_files_only, ) if not self.lazy_load: @@ -159,7 +165,7 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel): query: str, documents: Iterable[str], batch_size: int = 64, - **kwargs, + **kwargs: Any, ) -> Iterable[float]: """Reranks documents based on their relevance to a given query. @@ -175,3 +181,44 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel): yield from self._rerank_documents( query=query, documents=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]: + yield from self._rerank_pairs( + model_name=self.model_name, + cache_dir=str(self.cache_dir), + pairs=pairs, + 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[TextRerankerWorker]: + return TextCrossEncoderWorker + + def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[float]: + return (float(elem) for elem in output.model_output) + + +class TextCrossEncoderWorker(TextRerankerWorker): + def init_embedding( + self, + model_name: str, + cache_dir: str, + **kwargs, + ) -> OnnxTextCrossEncoder: + return OnnxTextCrossEncoder( + model_name=model_name, + cache_dir=cache_dir, + threads=1, + **kwargs, + ) diff --git a/fastembed/rerank/cross_encoder/onnx_text_model.py b/fastembed/rerank/cross_encoder/onnx_text_model.py index 86499a0..f3a02b6 100644 --- a/fastembed/rerank/cross_encoder/onnx_text_model.py +++ b/fastembed/rerank/cross_encoder/onnx_text_model.py @@ -1,17 +1,29 @@ -from typing import Sequence, Optional, Iterable +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 OnnxModel, OnnxProvider, OnnxOutputContext +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): +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, @@ -31,40 +43,108 @@ class OnnxCrossEncoderModel(OnnxModel): ) self.tokenizer, _ = load_tokenizer(model_dir=model_dir) - def tokenize(self, query: str, documents: list[str], **kwargs) -> list[Encoding]: - return self.tokenizer.encode_batch([(query, doc) for doc in documents]) - - def onnx_embed(self, query: str, documents: list[str], **kwargs) -> OnnxOutputContext: - tokenized_input = self.tokenize(query, documents, **kwargs) + 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), - "attention_mask": np.array( - [enc.attention_mask for enc in tokenized_input], dtype=np.int64 - ), } - input_names = {node.name for node in self.model.get_inputs()} 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) - return OnnxOutputContext(model_output=outputs[0][:, 0].tolist()) + 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 + 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.onnx_embed(query, batch, **kwargs).model_output + 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 + 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 diff --git a/fastembed/rerank/cross_encoder/text_cross_encoder.py b/fastembed/rerank/cross_encoder/text_cross_encoder.py index 481423c..e6d0b2e 100644 --- a/fastembed/rerank/cross_encoder/text_cross_encoder.py +++ b/fastembed/rerank/cross_encoder/text_cross_encoder.py @@ -1,8 +1,8 @@ from typing import Any, Iterable, Optional, Sequence, Type -from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase -from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder 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): @@ -47,7 +47,7 @@ class TextCrossEncoder(TextCrossEncoderBase): cuda: bool = False, device_ids: Optional[list[int]] = None, lazy_load: bool = False, - **kwargs, + **kwargs: Any, ): super().__init__(model_name, cache_dir, threads, **kwargs) @@ -72,7 +72,7 @@ class TextCrossEncoder(TextCrossEncoderBase): ) def rerank( - self, query: str, documents: Iterable[str], batch_size: int = 64, **kwargs + self, query: str, documents: Iterable[str], batch_size: int = 64, **kwargs: Any ) -> Iterable[float]: """Rerank a list of documents based on a query. @@ -85,3 +85,36 @@ class TextCrossEncoder(TextCrossEncoderBase): 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 + ) diff --git a/fastembed/rerank/cross_encoder/text_cross_encoder_base.py b/fastembed/rerank/cross_encoder/text_cross_encoder_base.py index 7a494ce..79d0210 100644 --- a/fastembed/rerank/cross_encoder/text_cross_encoder_base.py +++ b/fastembed/rerank/cross_encoder/text_cross_encoder_base.py @@ -1,4 +1,4 @@ -from typing import Iterable, Optional +from typing import Any, Iterable, Optional from fastembed.common.model_management import ModelManagement @@ -23,7 +23,7 @@ class TextCrossEncoderBase(ModelManagement): batch_size: int = 64, **kwargs, ) -> Iterable[float]: - """Reranks a list of documents given a query. + """Rerank a list of documents given a query. Args: query (str): The query to rerank the documents. @@ -32,6 +32,27 @@ class TextCrossEncoderBase(ModelManagement): **kwargs: Additional keyword argument to pass to the rerank method. Yields: - Iterable[float]: The scores of reranked the documents. + 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") diff --git a/tests/test_text_cross_encoder.py b/tests/test_text_cross_encoder.py index c449f11..15ab373 100644 --- a/tests/test_text_cross_encoder.py +++ b/tests/test_text_cross_encoder.py @@ -15,36 +15,43 @@ CANONICAL_SCORE_VALUES = { "jinaai/jina-reranker-v2-base-multilingual": np.array([1.6533, -1.6455]), } - -def test_rerank(): - is_ci = os.getenv("CI") - - for model_desc in TextCrossEncoder.list_supported_models(): - if not is_ci and model_desc["size_in_GB"] > 1: - continue - - model_name = model_desc["model"] - model = TextCrossEncoder(model_name=model_name) - - query = "What is the capital of France?" - documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."] - scores = np.array(list(model.rerank(query, documents))) - - canonical_scores = CANONICAL_SCORE_VALUES[model_name] - assert np.allclose( - scores, canonical_scores, atol=1e-3 - ), f"Model: {model_name}, Scores: {scores}, Expected: {canonical_scores}" - if is_ci: - delete_model_cache(model.model._model_dir) +SELECTED_MODELS = { + "Xenova": "Xenova/ms-marco-MiniLM-L-6-v2", + "BAAI": "BAAI/bge-reranker-base", + "jinaai": "jinaai/jina-reranker-v1-tiny-en", +} @pytest.mark.parametrize( "model_name", - [ - model_desc["model"] - for model_desc in TextCrossEncoder.list_supported_models() - if model_desc["size_in_GB"] < 1 and model_desc["model"] in CANONICAL_SCORE_VALUES.keys() - ], + [model_name for model_name in CANONICAL_SCORE_VALUES], +) +def test_rerank(model_name): + is_ci = os.getenv("CI") + + model = TextCrossEncoder(model_name=model_name) + + query = "What is the capital of France?" + documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."] + scores = np.array(list(model.rerank(query, documents))) + + pairs = [(query, doc) for doc in documents] + scores2 = np.array(list(model.rerank_pairs(pairs))) + assert np.allclose( + scores, scores2, atol=1e-5 + ), f"Model: {model_name}, Scores: {scores}, Scores2: {scores2}" + + canonical_scores = CANONICAL_SCORE_VALUES[model_name] + assert np.allclose( + scores, canonical_scores, atol=1e-3 + ), f"Model: {model_name}, Scores: {scores}, Expected: {canonical_scores}" + if is_ci: + delete_model_cache(model.model._model_dir) + + +@pytest.mark.parametrize( + "model_name", + [model_name for model_name in SELECTED_MODELS.values()], ) def test_batch_rerank(model_name): is_ci = os.getenv("CI") @@ -55,6 +62,12 @@ def test_batch_rerank(model_name): documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."] * 50 scores = np.array(list(model.rerank(query, documents, batch_size=10))) + pairs = [(query, doc) for doc in documents] + scores2 = np.array(list(model.rerank_pairs(pairs))) + assert np.allclose( + scores, scores2, atol=1e-5 + ), f"Model: {model_name}, Scores: {scores}, Scores2: {scores2}" + canonical_scores = np.tile(CANONICAL_SCORE_VALUES[model_name], 50) assert scores.shape == canonical_scores.shape, f"Unexpected shape for model {model_name}" @@ -80,3 +93,27 @@ def test_lazy_load(model_name): if is_ci: delete_model_cache(model.model._model_dir) + + +@pytest.mark.parametrize( + "model_name", + [model_name for model_name in SELECTED_MODELS.values()], +) +def test_rerank_pairs_parallel(model_name): + is_ci = os.getenv("CI") + + model = TextCrossEncoder(model_name=model_name) + query = "What is the capital of France?" + documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."] * 10 + pairs = [(query, doc) for doc in documents] + scores_parallel = np.array(list(model.rerank_pairs(pairs, parallel=2, batch_size=10))) + scores_sequential = np.array(list(model.rerank_pairs(pairs, batch_size=10))) + assert np.allclose( + scores_parallel, scores_sequential, atol=1e-5 + ), f"Model: {model_name}, Scores (Parallel): {scores_parallel}, Scores (Sequential): {scores_sequential}" + canonical_scores = CANONICAL_SCORE_VALUES[model_name] + assert np.allclose( + scores_parallel[: len(canonical_scores)], canonical_scores, atol=1e-3 + ), f"Model: {model_name}, Scores (Parallel): {scores_parallel}, Expected: {canonical_scores}" + if is_ci: + delete_model_cache(model.model._model_dir)