upstream-archive/fastembed
1
0
Fork 0
mirror of https://github.com/qdrant/fastembed.git synced 2026-07-23 11:20:51 -05:00
Code Issues Packages Projects Releases Wiki Activity

Cross encoders parallelism (#419)

Browse Source 
* 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 <dimitriy_rudenko@mail.ru>

* Test fix

* Update fastembed/rerank/cross_encoder/text_cross_encoder.py

Co-authored-by: George <george.panchuk@qdrant.tech>

* Update fastembed/rerank/cross_encoder/text_cross_encoder.py

Co-authored-by: George <george.panchuk@qdrant.tech>

* Update fastembed/rerank/cross_encoder/text_cross_encoder.py

Co-authored-by: George <george.panchuk@qdrant.tech>

* Update fastembed/rerank/cross_encoder/text_cross_encoder_base.py

Co-authored-by: George <george.panchuk@qdrant.tech>

* 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 <george.panchuk@qdrant.tech>
This commit is contained in:
Dmitrii Ogn
2024-12-16 21:59:34 +03:00
committed by GitHub
parent 3b5e4c8722
commit c8b1a18cfc
5 changed files with 273 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

59
fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py
View File

@@ -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,
)

112
fastembed/rerank/cross_encoder/onnx_text_model.py
View File

@@ -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

41
fastembed/rerank/cross_encoder/text_cross_encoder.py
View File

@@ -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
)

27
fastembed/rerank/cross_encoder/text_cross_encoder_base.py
View File

@@ -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")

89
tests/test_text_cross_encoder.py
View File

@@ -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)