[ENH] -- CrossEncoder.rank

sentence_transformers/cross_encoder/CrossEncoder.py

@@ -532,7 +532,7 @@ def rank(
         """
         query_doc_pairs = [[query, doc] for doc in documents]
         scores = self.predict(
-            query_doc_pairs,
+            sentences=query_doc_pairs,
             batch_size=batch_size,
             show_progress_bar=show_progress_bar,
             num_workers=num_workers,
@@ -543,11 +543,10 @@ def rank(
         )
 
         results = []
-        for i in range(len(scores)):
+        for i, score in enumerate(scores):
+            results.append({"corpus_id": i, "score": score})
             if return_documents:
-                results.append({"corpus_id": i, "score": scores[i], "text": documents[i]})
-            else:
-                results.append({"corpus_id": i, "score": scores[i]})
+                results[-1].update({"text": documents[i]})
 
         results = sorted(results, key=lambda x: x["score"], reverse=True)
         return results[:top_k]

tests/test_cross_encoder.py

@@ -13,6 +13,7 @@
 import numpy as np
 import pytest
 import torch
+from pytest import FixtureRequest
 from torch.utils.data import DataLoader
 
 from sentence_transformers import CrossEncoder, util
@@ -133,7 +134,12 @@ def test_load_with_revision() -> None:
     assert not torch.equal(main_prob, older_model.predict(test_sentences, convert_to_tensor=True))
 
 
-def test_rank() -> None:
+@pytest.mark.parametrize(
+    argnames="return_documents",
+    argvalues=[True, False],
+    ids=["return-docs", "no-return-docs"],
+)
+def test_rank(return_documents: bool, request: FixtureRequest) -> None:
     model = CrossEncoder("cross-encoder/stsb-distilroberta-base")
     # We want to compute the similarity between the query sentence
     query = "A man is eating pasta."
@@ -153,7 +159,10 @@ def test_rank() -> None:
     expected_ranking = [0, 1, 3, 6, 2, 5, 7, 4, 8]
 
     # 1. We rank all sentences in the corpus for the query
-    ranks = model.rank(query, corpus)
+    ranks = model.rank(query=query, documents=corpus, return_documents=return_documents)
+    if request.node.callspec.id == "return-docs":
+        assert {*corpus} == {rank.get("text") for rank in ranks}
+
     pred_ranking = [rank["corpus_id"] for rank in ranks]
     assert pred_ranking == expected_ranking