diff --git a/fastembed/rerank/cross_encoder/onnx_text_model.py b/fastembed/rerank/cross_encoder/onnx_text_model.py index 55f3ea85..1bfef206 100644 --- a/fastembed/rerank/cross_encoder/onnx_text_model.py +++ b/fastembed/rerank/cross_encoder/onnx_text_model.py @@ -103,11 +103,16 @@ def _rerank_pairs( ) -> Iterable[float]: is_small = False - if isinstance(pairs, tuple): - pairs = [pairs] + if ( + isinstance(pairs, tuple) + and len(pairs) == 2 + and isinstance(pairs[0], str) + and isinstance(pairs[1], str) + ): + pairs = [(pairs[0], pairs[1])] is_small = True - if isinstance(pairs, list): + if isinstance(pairs, (list, tuple)): if len(pairs) < batch_size: is_small = True diff --git a/tests/test_text_cross_encoder.py b/tests/test_text_cross_encoder.py index 4d0d5b7d..5da7df48 100644 --- a/tests/test_text_cross_encoder.py +++ b/tests/test_text_cross_encoder.py @@ -93,6 +93,32 @@ def test_batch_rerank(model_cache, model_name: str) -> None: ), f"Model: {model_name}, Scores: {scores}, Expected: {canonical_scores}" +@pytest.mark.parametrize("parallel", [None, 2]) +@pytest.mark.parametrize("pair_count", [0, 1, 2, 3]) +def test_rerank_tuple_of_pairs(model_cache, pair_count: int, parallel: int | None) -> None: + model_name = "Xenova/ms-marco-MiniLM-L-6-v2" + query = "What is the capital of France?" + documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."] * 2 + pairs = tuple((query, document) for document in documents[:pair_count]) + + with model_cache(model_name) as model: + scores = np.array(list(model.rerank_pairs(pairs, batch_size=2, parallel=parallel))) + + expected = np.tile(CANONICAL_SCORE_VALUES[model_name], 2)[:pair_count] + np.testing.assert_allclose(scores, expected, atol=1e-3) + + +@pytest.mark.parametrize("parallel", [None, 2]) +def test_rerank_single_pair(model_cache, parallel: int | None) -> None: + model_name = "Xenova/ms-marco-MiniLM-L-6-v2" + pair = ("What is the capital of France?", "Paris is the capital of France.") + + with model_cache(model_name) as model: + scores = np.array(list(model.rerank_pairs(pair, batch_size=1, parallel=parallel))) + + np.testing.assert_allclose(scores, CANONICAL_SCORE_VALUES[model_name][:1], atol=1e-3) + + @pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"]) def test_lazy_load(model_name: str) -> None: is_ci = os.getenv("CI")