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

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)