fix: handle single-query exact search

beir/retrieval/search/dense/exact_search.py

@@ -99,7 +99,7 @@ def search(
             # Get top-k values
             cos_scores_top_k_values, cos_scores_top_k_idx = torch.topk(
                 cos_scores,
-                min(top_k + 1, len(cos_scores[1])),
+                min(top_k + 1, cos_scores.shape[1]),
                 dim=1,
                 largest=True,
                 sorted=return_sorted,

tests/retrieval/search/dense/test_exact_search.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import torch
+
+from beir.retrieval.search.dense import DenseRetrievalExactSearch
+
+
+class DummyModel:
+    def encode_queries(self, queries, **kwargs):
+        assert queries == ["alpha"]
+        return torch.tensor([[1.0, 0.0]])
+
+    def encode_corpus(self, corpus, **kwargs):
+        assert [doc["text"] for doc in corpus] == ["alpha", "beta"]
+        return torch.tensor([[1.0, 0.0], [0.0, 1.0]])
+
+
+def test_exact_search_handles_single_query():
+    retriever = DenseRetrievalExactSearch(DummyModel(), show_progress_bar=False)
+
+    results = retriever.search(
+        corpus={
+            "d1": {"title": "", "text": "alpha"},
+            "d2": {"title": "", "text": "beta"},
+        },
+        queries={"q1": "alpha"},
+        top_k=1,
+        score_function="dot",
+    )
+
+    assert results == {"q1": {"d1": 1.0}}