add kwargs to index search (#5628)

* add kwargs to search

* format

Author: SaulLu

src/datasets/search.py

@@ -57,14 +57,14 @@ class BatchedNearestExamplesResults(NamedTuple):
 class BaseIndex:
     """Base class for indexing"""
 
-    def search(self, query, k: int = 10) -> SearchResults:
+    def search(self, query, k: int = 10, **kwargs) -> SearchResults:
         """
         To implement.
         This method has to return the scores and the indices of the retrieved examples given a certain query.
         """
         raise NotImplementedError
 
-    def search_batch(self, queries, k: int = 10) -> BatchedSearchResults:
+    def search_batch(self, queries, k: int = 10, **kwargs) -> BatchedSearchResults:
         """Find the nearest examples indices to the query.
 
         Args:
@@ -176,7 +176,7 @@ def passage_generator():
             )
         logger.info(f"Indexed {successes:d} documents")
 
-    def search(self, query: str, k=10) -> SearchResults:
+    def search(self, query: str, k=10, **kwargs) -> SearchResults:
         """Find the nearest examples indices to the query.
 
         Args:
@@ -190,16 +190,17 @@ def search(self, query: str, k=10) -> SearchResults:
         response = self.es_client.search(
             index=self.es_index_name,
             body={"query": {"multi_match": {"query": query, "fields": ["text"], "type": "cross_fields"}}, "size": k},
+            **kwargs,
         )
         hits = response["hits"]["hits"]
         return SearchResults([hit["_score"] for hit in hits], [int(hit["_id"]) for hit in hits])
 
-    def search_batch(self, queries, k: int = 10, max_workers=10) -> BatchedSearchResults:
+    def search_batch(self, queries, k: int = 10, max_workers=10, **kwargs) -> BatchedSearchResults:
         import concurrent.futures
 
         total_scores, total_indices = [None] * len(queries), [None] * len(queries)
         with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
-            future_to_index = {executor.submit(self.search, query, k): i for i, query in enumerate(queries)}
+            future_to_index = {executor.submit(self.search, query, k, **kwargs): i for i, query in enumerate(queries)}
             for future in concurrent.futures.as_completed(future_to_index):
                 index = future_to_index[future]
                 results: SearchResults = future.result()
@@ -337,7 +338,7 @@ def _faiss_index_to_device(index: "faiss.Index", device: Optional[Union[int, Lis
 
         return index
 
-    def search(self, query: np.array, k=10) -> SearchResults:
+    def search(self, query: np.array, k=10, **kwargs) -> SearchResults:
         """Find the nearest examples indices to the query.
 
         Args:
@@ -354,10 +355,10 @@ def search(self, query: np.array, k=10) -> SearchResults:
         queries = query.reshape(1, -1)
         if not queries.flags.c_contiguous:
             queries = np.asarray(queries, order="C")
-        scores, indices = self.faiss_index.search(queries, k)
+        scores, indices = self.faiss_index.search(queries, k, **kwargs)
         return SearchResults(scores[0], indices[0].astype(int))
 
-    def search_batch(self, queries: np.array, k=10) -> BatchedSearchResults:
+    def search_batch(self, queries: np.array, k=10, **kwargs) -> BatchedSearchResults:
         """Find the nearest examples indices to the queries.
 
         Args:
@@ -372,7 +373,7 @@ def search_batch(self, queries: np.array, k=10) -> BatchedSearchResults:
             raise ValueError("Shape of query must be 2D")
         if not queries.flags.c_contiguous:
             queries = np.asarray(queries, order="C")
-        scores, indices = self.faiss_index.search(queries, k)
+        scores, indices = self.faiss_index.search(queries, k, **kwargs)
         return BatchedSearchResults(scores, indices.astype(int))
 
     def save(self, file: Union[str, PurePath]):
@@ -667,7 +668,7 @@ def drop_index(self, index_name: str):
         """
         del self._indexes[index_name]
 
-    def search(self, index_name: str, query: Union[str, np.array], k: int = 10) -> SearchResults:
+    def search(self, index_name: str, query: Union[str, np.array], k: int = 10, **kwargs) -> SearchResults:
         """Find the nearest examples indices in the dataset to the query.
 
         Args:
@@ -683,9 +684,11 @@ def search(self, index_name: str, query: Union[str, np.array], k: int = 10) -> S
             - indices (`List[List[int]]`): The indices of the retrieved examples.
         """
         self._check_index_is_initialized(index_name)
-        return self._indexes[index_name].search(query, k)
+        return self._indexes[index_name].search(query, k, **kwargs)
 
-    def search_batch(self, index_name: str, queries: Union[List[str], np.array], k: int = 10) -> BatchedSearchResults:
+    def search_batch(
+        self, index_name: str, queries: Union[List[str], np.array], k: int = 10, **kwargs
+    ) -> BatchedSearchResults:
         """Find the nearest examples indices in the dataset to the query.
 
         Args:
@@ -701,10 +704,10 @@ def search_batch(self, index_name: str, queries: Union[List[str], np.array], k:
             - total_indices (`List[List[int]]`): The indices of the retrieved examples per query.
         """
         self._check_index_is_initialized(index_name)
-        return self._indexes[index_name].search_batch(queries, k)
+        return self._indexes[index_name].search_batch(queries, k, **kwargs)
 
     def get_nearest_examples(
-        self, index_name: str, query: Union[str, np.array], k: int = 10
+        self, index_name: str, query: Union[str, np.array], k: int = 10, **kwargs
     ) -> NearestExamplesResults:
         """Find the nearest examples in the dataset to the query.
 
@@ -721,12 +724,12 @@ def get_nearest_examples(
             - examples (`dict`): The retrieved examples.
         """
         self._check_index_is_initialized(index_name)
-        scores, indices = self.search(index_name, query, k)
+        scores, indices = self.search(index_name, query, k, **kwargs)
         top_indices = [i for i in indices if i >= 0]
         return NearestExamplesResults(scores[: len(top_indices)], self[top_indices])
 
     def get_nearest_examples_batch(
-        self, index_name: str, queries: Union[List[str], np.array], k: int = 10
+        self, index_name: str, queries: Union[List[str], np.array], k: int = 10, **kwargs
     ) -> BatchedNearestExamplesResults:
         """Find the nearest examples in the dataset to the query.
 
@@ -743,7 +746,7 @@ def get_nearest_examples_batch(
             - total_examples (`List[dict]`): The retrieved examples per query.
         """
         self._check_index_is_initialized(index_name)
-        total_scores, total_indices = self.search_batch(index_name, queries, k)
+        total_scores, total_indices = self.search_batch(index_name, queries, k, **kwargs)
         total_scores = [
             scores_i[: len([i for i in indices_i if i >= 0])]
             for scores_i, indices_i in zip(total_scores, total_indices)

tests/test_search.py

@@ -188,6 +188,13 @@ def test_elasticsearch(self):
             self.assertEqual(scores[0], 1)
             self.assertEqual(indices[0], 0)
 
+            # single query with timeout
+            query = "foo"
+            mocked_search.return_value = {"hits": {"hits": [{"_score": 1, "_id": 0}]}}
+            scores, indices = index.search(query, request_timeout=30)
+            self.assertEqual(scores[0], 1)
+            self.assertEqual(indices[0], 0)
+
             # batched queries
             queries = ["foo", "bar", "foobar"]
             mocked_search.return_value = {"hits": {"hits": [{"_score": 1, "_id": 1}]}}
@@ -196,3 +203,12 @@ def test_elasticsearch(self):
             best_indices = [indices[0] for indices in total_indices]
             self.assertGreater(np.min(best_scores), 0)
             self.assertListEqual([1, 1, 1], best_indices)
+
+            # batched queries with timeout
+            queries = ["foo", "bar", "foobar"]
+            mocked_search.return_value = {"hits": {"hits": [{"_score": 1, "_id": 1}]}}
+            total_scores, total_indices = index.search_batch(queries, request_timeout=30)
+            best_scores = [scores[0] for scores in total_scores]
+            best_indices = [indices[0] for indices in total_indices]
+            self.assertGreater(np.min(best_scores), 0)
+            self.assertListEqual([1, 1, 1], best_indices)