diff --git a/fastembed/rerank/cross_encoder/onnx_text_model.py b/fastembed/rerank/cross_encoder/onnx_text_model.py index 6768d002..da854772 100644 --- a/fastembed/rerank/cross_encoder/onnx_text_model.py +++ b/fastembed/rerank/cross_encoder/onnx_text_model.py @@ -1,7 +1,7 @@ import os from multiprocessing import get_all_start_methods from pathlib import Path -from typing import Any, Iterable, Sequence, Type +from typing import Any, Iterable, Sequence, Type, cast import numpy as np from tokenizers import Encoding @@ -118,8 +118,13 @@ def _rerank_pairs( is_small = False if isinstance(pairs, tuple): - pairs = [pairs] - is_small = True + # a single (query, document) pair is a tuple of two strings, + # while a tuple of pairs is just another sequence of pairs + if len(pairs) == 2 and all(isinstance(item, str) for item in pairs): + pairs = [cast(tuple[str, str], pairs)] + is_small = True + else: + pairs = list(pairs) if isinstance(pairs, list): if len(pairs) < batch_size: diff --git a/tests/test_text_cross_encoder.py b/tests/test_text_cross_encoder.py index 85988ca8..bb3b98a9 100644 --- a/tests/test_text_cross_encoder.py +++ b/tests/test_text_cross_encoder.py @@ -135,6 +135,27 @@ def test_rerank_pairs_parallel(model_cache, model_name: str) -> None: ), f"Model: {model_name}, Scores (Parallel): {scores_parallel}, Expected: {canonical_scores}" +@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"]) +def test_rerank_pairs_tuple_input(model_cache, model_name: str) -> None: + with model_cache(model_name) as model: + query = "What is the capital of France?" + documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."] + pairs = [(query, doc) for doc in documents] + expected = np.array(list(model.rerank_pairs(pairs))) + + # a tuple of pairs is an iterable of pairs, not a single pair + scores_tuple = np.array(list(model.rerank_pairs(tuple(pairs)))) + assert np.allclose(scores_tuple, expected, atol=1e-5) + + three_pairs = tuple(pairs + [pairs[0]]) + scores_three = np.array(list(model.rerank_pairs(three_pairs))) + assert np.allclose(scores_three, np.append(expected, expected[0]), atol=1e-5) + + # a lone (query, document) pair is still accepted + scores_single = np.array(list(model.rerank_pairs(pairs[0]))) + assert np.allclose(scores_single, expected[:1], atol=1e-5) + + @pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"]) def test_token_count(model_cache, model_name: str) -> None: with model_cache(model_name) as model: