Repository: deepset-ai/haystack

Issue:
SentenceTransformersRanker's `predict_batch` returns wrong number of documents
**Describe the bug**

When using the `predict_batch` method of the SentenceTransformersRanker, the wrong length of documents is returned when a large(r) `top_k` is used, eg. 50 or 100.

**Expected behavior**

`predict_batch` should return the `top_k` number of documents, assuming `top_k` is smaller than the document store size.

**To Reproduce**

```python3
from haystack import Document
from haystack.nodes import SentenceTransformersRanker

docs = [Document(content=f"test {number}") for number in range(100)]
ranker = SentenceTransformersRanker(
    model_name_or_path="cross-encoder/ms-marco-MiniLM-L-4-v2"
)

print(len(ranker.predict("where is test 3?", docs, top_k=50)))
# >>> 50

print(len(ranker.predict_batch(["where is test 3?"], docs, batch_size=32, top_k=50)))
# >>> 4
```

Upon closer inspection, I think I found the cause. This is because the `similarity_scores` transient variable is used instead of the `preds` list for processing results post-batching.

https://github.com/deepset-ai/haystack/blob/722dead1b26e19501f1494b0b1372b2d76ac3b51/haystack/nodes/ranker/sentence_transformers.py#L219-L248

Commit https://github.com/atreyasha/haystack/commit/9c76a9ac67e53147feea2658b31068279c1c2c6c in my fork of `haystack` should fix the issue.

**FAQ Check**
- [x] Have you had a look at [our new FAQ page](https://docs.haystack.deepset.ai/docs/faq)?

**System:**
 - OS: Linux x86_64
 - Haystack version (commit or version number): `farm-haystack[faiss-gpu]==1.3.2`


Discussion:
@atreyasha can you open a PR please, you seem to be on the right track! 

Resolving patch:
diff --git a/haystack/nodes/ranker/sentence_transformers.py b/haystack/nodes/ranker/sentence_transformers.py
--- a/haystack/nodes/ranker/sentence_transformers.py
+++ b/haystack/nodes/ranker/sentence_transformers.py
@@ -225,7 +225,7 @@ def predict_batch(
         logits_dim = similarity_scores.shape[1]  # [batch_size, logits_dim]
         if single_list_of_docs:
             sorted_scores_and_documents = sorted(
-                zip(similarity_scores, documents),
+                zip(preds, documents),
                 key=lambda similarity_document_tuple:
                 # assume the last element in logits represents the `has_answer` label
                 similarity_document_tuple[0][-1] if logits_dim >= 2 else similarity_document_tuple[0],
@@ -244,7 +244,7 @@ def predict_batch(
             right_idx = 0
             for number in number_of_docs:
                 right_idx = left_idx + number
-                grouped_predictions.append(similarity_scores[left_idx:right_idx])
+                grouped_predictions.append(preds[left_idx:right_idx])
                 left_idx = right_idx
 
             result = []


Test patch:
diff --git a/test/nodes/test_ranker.py b/test/nodes/test_ranker.py
--- a/test/nodes/test_ranker.py
+++ b/test/nodes/test_ranker.py
@@ -225,3 +225,11 @@ def test_ranker_returns_raw_score_for_two_logits(ranker_two_logits):
     score = results[0].score
     precomputed_score = -3.61354
     assert math.isclose(precomputed_score, score, rel_tol=0.001)
+
+
+def test_predict_batch_returns_correct_number_of_docs(ranker):
+    docs = [Document(content=f"test {number}") for number in range(5)]
+
+    assert len(ranker.predict("where is test 3?", docs, top_k=4)) == 4
+
+    assert len(ranker.predict_batch(["where is test 3?"], docs, batch_size=2, top_k=4)) == 4