Add retriever benchmarking

haystack/database/faiss.py

@@ -109,11 +109,17 @@ def update_embeddings(self, retriever: BaseRetriever, index: Optional[str] = Non
         :param index: Index name to update
         :return: None
         """
-        # Some FAISS indexes(like the default HNSWx) do not support removing vectors, so a new index is created.
-        faiss_index = self._create_new_index(vector_size=self.vector_size)
-        index = index or self.index
 
+        index = index or self.index
         documents = self.get_all_documents(index=index)
+
+        if len(documents) == 0:
+            logger.warning("Calling DocumentStore.update_embeddings() on an empty index")
+            self.faiss_index = None
+            return
+
+        # Some FAISS indexes(like the default HNSWx) do not support removing vectors, so a new index is created.
+        faiss_index = self._create_new_index(vector_size=self.vector_size)
         logger.info(f"Updating embeddings for {len(documents)} docs ...")
         embeddings = retriever.embed_passages(documents)  # type: ignore
         assert len(documents) == len(embeddings)
@@ -127,14 +133,14 @@ def update_embeddings(self, retriever: BaseRetriever, index: Optional[str] = Non
             hnsw_vectors = self._get_hnsw_vectors(embeddings=embeddings, phi=phi)
             faiss_index.add(hnsw_vectors)
 
-        doc_meta_to_update = []
-        for vector_id, doc in enumerate(documents[i : i + self.index_buffer_size]):
-            meta = doc.meta or {}
-            meta["vector_id"] = vector_id
-            doc_meta_to_update.append((doc.id, meta))
+            doc_meta_to_update = []
+            for vector_id, doc in enumerate(documents[i : i + self.index_buffer_size]):
+                meta = doc.meta or {}
+                meta["vector_id"] = vector_id
+                doc_meta_to_update.append((doc.id, meta))
 
-        for doc_id, meta in doc_meta_to_update:
-            super(FAISSDocumentStore, self).update_document_meta(id=doc_id, meta=meta)
+            for doc_id, meta in doc_meta_to_update:
+                super(FAISSDocumentStore, self).update_document_meta(id=doc_id, meta=meta)
 
         self.faiss_index = faiss_index
 

haystack/retriever/base.py

@@ -1,6 +1,8 @@
 from abc import ABC, abstractmethod
 from typing import List
 import logging
+from time import perf_counter
+from functools import wraps
 
 from haystack.database.base import Document
 from haystack.database.base import BaseDocumentStore
@@ -15,6 +17,18 @@ class BaseRetriever(ABC):
     def retrieve(self, query: str, filters: dict = None, top_k: int = 10, index: str = None) -> List[Document]:
         pass
 
+    def timing(self, fn):
+        @wraps(fn)
+        def wrapper(*args, **kwargs):
+            if "retrieve_time" not in self.__dict__:
+                self.retrieve_time = 0
+            tic = perf_counter()
+            ret = fn(*args, **kwargs)
+            toc = perf_counter()
+            self.retrieve_time += toc - tic
+            return ret
+        return wrapper
+
     def eval(
         self,
         label_index: str = "label",
@@ -45,6 +59,8 @@ def eval(
         # Extract all questions for evaluation
         filters = {"origin": [label_origin]}
 
+        timed_retrieve = self.timing(self.retrieve)
+
         labels = self.document_store.get_all_labels_aggregated(index=label_index, filters=filters)
 
         correct_retrievals = 0
@@ -62,7 +78,7 @@ def eval(
         # Option 1: Open-domain evaluation by checking if the answer string is in the retrieved docs
         if open_domain:
             for question, gold_answers in question_label_dict.items():
-                retrieved_docs = self.retrieve(question, top_k=top_k, index=doc_index)
+                retrieved_docs = timed_retrieve(question, top_k=top_k, index=doc_index)
                 # check if correct doc in retrieved docs
                 for doc_idx, doc in enumerate(retrieved_docs):
                     for gold_answer in gold_answers:
@@ -73,7 +89,7 @@ def eval(
         # Option 2: Strict evaluation by document ids that are listed in the labels
         else:
             for question, gold_ids in question_label_dict.items():
-                retrieved_docs = self.retrieve(question, top_k=top_k, index=doc_index)
+                retrieved_docs = timed_retrieve(question, top_k=top_k, index=doc_index)
                 # check if correct doc in retrieved docs
                 for doc_idx, doc in enumerate(retrieved_docs):
                     for gold_id in gold_ids:
@@ -89,4 +105,4 @@ def eval(
         logger.info((f"For {correct_retrievals} out of {number_of_questions} questions ({recall:.2%}), the answer was in"
                      f" the top-{top_k} candidate passages selected by the retriever."))
 
-        return {"recall": recall, "map": mean_avg_precision}
+        return {"recall": recall, "map": mean_avg_precision, "retrieve_time": self.retrieve_time}

test/benchmarks/run.py

@@ -0,0 +1,88 @@
+from haystack.indexing.utils import fetch_archive_from_http
+import os
+from haystack.database.sql import SQLDocumentStore
+from haystack.database.memory import InMemoryDocumentStore
+from haystack.database.elasticsearch import Elasticsearch, ElasticsearchDocumentStore
+from haystack.database.faiss import FAISSDocumentStore
+from haystack.retriever.sparse import ElasticsearchRetriever, TfidfRetriever
+from haystack.retriever.dense import DensePassageRetriever
+from time import perf_counter
+
+from pathlib import Path
+
+
+retriever_doc_stores = [("elastic", "elasticsearch"),
+                        ("dpr", "faiss")]
+reader_models = [""]
+reader_type = ["farm", "transformers"]
+
+data_dir = Path("../../data/nq")
+filename = "nq_dev_subset_v3.json"
+s3_url = "https://s3.eu-central-1.amazonaws.com/deepset.ai-farm-qa/datasets/nq_dev_subset_v3.json.zip"
+doc_index = "eval_document"
+label_index = "label"
+
+def prepare_data(data_dir):
+    fetch_archive_from_http(url=s3_url, output_dir=data_dir)
+
+def get_document_store(document_store_type):
+    """ TODO This method is taken from test/conftest.py but maybe should be within Haystack.
+    Perhaps a class method of DocStore that just takes string for type of DocStore"""
+    if document_store_type == "sql":
+        if os.path.exists("haystack_test.db"):
+            os.remove("haystack_test.db")
+        document_store = SQLDocumentStore(url="sqlite:///haystack_test.db")
+    elif document_store_type == "memory":
+        document_store = InMemoryDocumentStore()
+    elif document_store_type == "elasticsearch":
+        # make sure we start from a fresh index
+        client = Elasticsearch()
+        client.indices.delete(index='haystack_test*', ignore=[404])
+        document_store = ElasticsearchDocumentStore(index="haystack_test")
+    elif document_store_type == "faiss":
+        if os.path.exists("haystack_test_faiss.db"):
+            os.remove("haystack_test_faiss.db")
+        document_store = FAISSDocumentStore(sql_url="sqlite:///haystack_test_faiss.db")
+    else:
+        raise Exception(f"No document store fixture for '{document_store_type}'")
+    return document_store
+
+def get_retriever(retriever_name, doc_store):
+    if retriever_name == "elastic":
+        return ElasticsearchRetriever(doc_store)
+    if retriever_name == "tfidf":
+        return TfidfRetriever(doc_store)
+    if retriever_name == "dpr":
+        return DensePassageRetriever(document_store=doc_store,
+                                      query_embedding_model="facebook/dpr-question_encoder-single-nq-base",
+                                      passage_embedding_model="facebook/dpr-ctx_encoder-single-nq-base",
+                                      use_gpu=True)
+
+
+def benchmark_indexing(doc_store, data_dir, filename, retriever):
+    doc_store.delete_all_documents(index=doc_index)
+    doc_store.delete_all_documents(index=label_index)
+    tic = perf_counter()
+    doc_store.add_eval_data(data_dir / filename)
+    try:
+        doc_store.update_embeddings(retriever, index=doc_index)
+    except AttributeError:
+        pass
+    toc = perf_counter()
+    time = toc - tic
+    return doc_store, time
+
+
+def main():
+    # prepare_data(data_dir)
+    for retriever_name, doc_store_name in retriever_doc_stores:
+        doc_store = get_document_store(doc_store_name)
+        retriever = get_retriever(retriever_name, doc_store)
+        doc_store, indexing_time = benchmark_indexing(doc_store, data_dir, filename, retriever)
+        results = retriever.eval()
+        results["indexing_time"] = indexing_time
+        print(results)
+
+if __name__ == "__main__":
+    main()
+