Skip to content

fix: score a tuple of pairs in rerank_pairs instead of treating it as one pair #797

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
michaelkillgta wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: score a tuple of pairs in rerank_pairs instead of treating it as one pair #797

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 8 additions & 3 deletions fastembed/rerank/cross_encoder/onnx_text_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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:
Expand Down
21 changes: 21 additions & 0 deletions tests/test_text_cross_encoder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down