Repository: deepset-ai/haystack
Pull request: feat: add support for `BM25Retriever` in `InMemoryDocumentStore`
### Related Issues
- fixes #3447

Only a first draft...

### Checklist
- [X] I have read the [contributors guidelines](https://github.com/deepset-ai/haystack/blob/main/CONTRIBUTING.md) and the [code of conduct](https://github.com/deepset-ai/haystack/blob/main/code_of_conduct.txt)
- [X] I have updated the related issue with new insights and changes
- [ ] I added tests that demonstrate the correct behavior of the change
- [X] I've used the [conventional commit convention](https://www.conventionalcommits.org/en/v1.0.0/) for my PR title
- [x] I documented my code
- [X] I ran [pre-commit hooks](https://github.com/deepset-ai/haystack/blob/main/CONTRIBUTING.md#installation) and fixed any issue


Problem statement:
Add support for BM25Retriever in InMemoryDocumentStore
**Is your feature request related to a problem? Please describe.**
Many of our tutorials are using the ElasticsearchDocumentStore. While this is a good choice for a production system, it can be quite cumbersome to set up and run during your "first minutes with haystack". It would be awesome to use the `InMemoryDocumentStore` instead for the first steps of users (Tutorial 1, Quick Start ...).  The only thing that I see holding us back here: `InMemoryDocumentStore` doesn't support our `BM25Retriever` which is a fast and easy retriever to get started with (and therefore a good choice in tutorials). Switching to another retriever might complicate tutorials, slow them down or reduce the quality of answers users get.

**Describe the solution you'd like**
Supporting the usage of the `BM25Retriever` in combination with the `InMemoryDocumentStore`

**Describe alternatives you've considered**
Using TFIDFRetriever but I am concerned about the quality of results and leading our users into a wrong direction here.

**Priority**
I don't think this feature is urgent but it might be a helpful step when we want to improve the early user experience

Hints from issue discussion:
Some additional context (same idea from @bglearning): https://github.com/deepset-ai/haystack-tutorials/pull/44#issuecomment-1282170131
BM25 was recently added to gensim: https://github.com/RaRe-Technologies/gensim/pull/3304, we might use this.
@ZanSara: please advise me on how I can contribute to the haystack, in general, or to this issue, in particular. Thanks :)
Hello @vtharmalingam! Start from here: https://github.com/deepset-ai/haystack/blob/main/CONTRIBUTING.md
Hi @vtharmalingam great to hear that you would like to contribute on this issue! Feel free to open a draft pull request early on so that we can help with feedback. Please let me know if you need any other help.
I am also interested in contributing to this feature!
Anyway, @vtharmalingam takes precedence.

After the discussions linked above, I want to bring **some ideas**.

- I tend to think that the retriever should only query the document store. On the other hand, the ds is responsible for storing documents and their representations (such as BM25). _Unlike the TF-IDF retriever implementation_
- To do that, `InMemoryDocumentStore` should be a subclass of `KeywordDocumentStore` (instead of `BaseDocumentStore`) and implement the methods `query` and `query_batch`.
- I like the library [rank_bm25](https://github.com/dorianbrown/rank_bm25): it is simple and lightweight. It must be said that it is less performant than Gensim at retrieval time, but maybe we could accept that, as this document store is not meant for production use cases.
- About the practical implementation, I imagine that for every index, we could have an instance of the BM25 class: 
`self.indexes[index].bm25 = BM25Okapi(tokenized_corpus)`
- A doubt: should the BM25 representation be generated by default when working with the `InMemoryDocumentStore`?
It may not be essential if we are using the IMDS for dense retrieval or for TF-IDF retrieval; in those cases, BM25 computation would make the document store unnecessarily slow.

WDYT?
Hey @anakin87! Sounds good :slightly_frowning_face: 

> - I tend to think that the retriever should only query the document store. On the other hand, the ds is responsible for storing documents and their representations (such as BM25). _Unlike the TF-IDF retriever implementation_
> - To do that, `InMemoryDocumentStore` should be a subclass of `KeywordDocumentStore` (instead of `BaseDocumentStore`) and implement the methods `query` and `query_batch`.

Fully agree!

> - I like the library [rank_bm25](https://github.com/dorianbrown/rank_bm25): it is simple and lightweight. It must be said that it is less performant than Gensim at retrieval time, but maybe we could accept that, as this document store is not meant for production use cases.

I also agree on this. It bring in no dependencies, which is a relief :smile: For now at least I imagine it would be a nice compromise. We should make sure that for workloads like the one in the tutorials it doesn't become too slow, however. Let's see after how many documents the 1s threshold for retrieval is reached.

> - About the practical implementation, I imagine that for every index, we could have an instance of the BM25 class:
`self.indexes[index].bm25 = BM25Okapi(tokenized_corpus)`

I trust you on this for now, I haven't read the `rank_bm25` docs yet. in principle sounds good to me :+1: 

> - A doubt: should the BM25 representation be generated by default when working with the `InMemoryDocumentStore`?
It may not be essential if we are using the IMDS for dense retrieval or for TF-IDF retrieval; in those cases, BM25 computation would make the document store unnecessarily slow.

This issue really has no good answer in the current architecture, but so is the amazing `retriever.document_store.update_embeddings(retriever=retriever)` incantation, right? :smiling_face_with_tear: My opinion is to **not** make it automatic. But we can do better than the `update_embedding` stuff above: let's make a `__init__` param for this!

Something like `docstore = InMemoryDocumentStore(bm25_ranks=True)`. Let's also leave the `update_bm25_ranks()` method public, so if one wants to manually update the ranks, it can still be done. But I'd like it to be an init param so one can set it and forget.

Then we can do the same to the other docstores for `update_embeddings`... but in another PR I guess :smile: 
I'm starting to work on this... 🛠️

Commit message: better docstrings; revert not running tests

Patch to review:
diff --git a/haystack/document_stores/memory.py b/haystack/document_stores/memory.py
index dd61bd9286..091c13fc94 100644
--- a/haystack/document_stores/memory.py
+++ b/haystack/document_stores/memory.py
@@ -1,17 +1,24 @@
 from typing import Any, Dict, List, Optional, Union, Generator
 
+try:
+    from typing import Literal
+except ImportError:
+    from typing_extensions import Literal  # type: ignore
+
 import time
 import logging
 from copy import deepcopy
 from collections import defaultdict
+import re
 
 import numpy as np
 import torch
 from tqdm import tqdm
+import rank_bm25
 
 from haystack.schema import Document, Label
-from haystack.errors import DuplicateDocumentError
-from haystack.document_stores import BaseDocumentStore
+from haystack.errors import DuplicateDocumentError, DocumentStoreError
+from haystack.document_stores import KeywordDocumentStore
 from haystack.document_stores.base import get_batches_from_generator
 from haystack.modeling.utils import initialize_device_settings
 from haystack.document_stores.filter_utils import LogicalFilterClause
@@ -20,7 +27,8 @@
 logger = logging.getLogger(__name__)
 
 
-class InMemoryDocumentStore(BaseDocumentStore):
+class InMemoryDocumentStore(KeywordDocumentStore):
+    # pylint: disable=R0904
     """
     In-memory document store
     """
@@ -38,6 +46,10 @@ def __init__(
         use_gpu: bool = True,
         scoring_batch_size: int = 500000,
         devices: Optional[List[Union[str, torch.device]]] = None,
+        use_bm25: bool = False,
+        bm25_tokenization_regex: str = r"(?u)\b\w\w+\b",
+        bm25_algorithm: Literal["BM25Okapi", "BM25L", "BM25Plus"] = "BM25Okapi",
+        bm25_parameters: dict = {},
     ):
         """
         :param index: The documents are scoped to an index attribute that can be used when writing, querying,
@@ -67,6 +79,14 @@ def __init__(
                         A list containing torch device objects and/or strings is supported (For example
                         [torch.device('cuda:0'), "mps", "cuda:1"]). When specifying `use_gpu=False` the devices
                         parameter is not used and a single cpu device is used for inference.
+        :param use_bm25: Whether to build a sparse representation of documents based on BM25.
+                         `use_bm25=True` is required to connect `BM25Retriever` to this Document Store.
+        :param bm25_tokenization_regex: The regular expression to use for tokenization of the text.
+        :param bm25_algorithm: The specific BM25 implementation to adopt.
+                               Parameter options : ( 'BM25Okapi', 'BM25L', 'BM25Plus')
+        :param bm25_parameters: Parameters for BM25 implementation in a dictionary format.
+                                For example: {'k1':1.5, 'b':0.75, 'epsilon':0.25}
+                                You can learn more about these parameters by visiting https://github.com/dorianbrown/rank_bm25
         """
         super().__init__()
 
@@ -81,6 +101,11 @@ def __init__(
         self.duplicate_documents = duplicate_documents
         self.use_gpu = use_gpu
         self.scoring_batch_size = scoring_batch_size
+        self.use_bm25 = use_bm25
+        self.bm25_tokenization_regex = bm25_tokenization_regex
+        self.bm25_algorithm = bm25_algorithm
+        self.bm25_parameters = bm25_parameters
+        self.bm25: Dict[str, Any] = {}
 
         self.devices, _ = initialize_device_settings(devices=devices, use_cuda=self.use_gpu, multi_gpu=False)
         if len(self.devices) > 1:
@@ -134,6 +159,7 @@ def write_documents(
             Document.from_dict(d, field_map=field_map) if isinstance(d, dict) else d for d in documents
         ]
         documents_objects = self._drop_duplicate_documents(documents=documents_objects)
+        modified_documents = 0
         for document in documents_objects:
             if document.id in self.indexes[index]:
                 if duplicate_documents == "fail":
@@ -146,6 +172,25 @@ def write_documents(
                     )
                     continue
             self.indexes[index][document.id] = document
+            modified_documents += 1
+
+        if self.use_bm25 is True and modified_documents > 0:
+            self.update_bm25(index=index)
+
+    def update_bm25(self, index: Optional[str] = None):
+        """
+        Updates the BM25 sparse representation in the the document store.
+
+        :param index: Index name for which the BM25 representation is to be updated. If set to None, the default self.index is used.
+        """
+        index = index or self.index
+        tokenizer = re.compile(self.bm25_tokenization_regex).findall
+
+        logger.info("Updating BM25 representation...")
+        all_documents = self.get_all_documents(index=index)
+        tokenized_corpus = [tokenizer(doc.content.lower()) for doc in all_documents]
+        bm25_class = getattr(rank_bm25, self.bm25_algorithm)
+        self.bm25[index] = bm25_class(tokenized_corpus, **self.bm25_parameters)
 
     def _create_document_field_map(self):
         return {self.embedding_field: "embedding"}
@@ -759,12 +804,16 @@ def delete_documents(
         index = index or self.index
         if not filters and not ids:
             self.indexes[index] = {}
+            if index in self.bm25:
+                self.bm25[index] = {}
             return
         docs_to_delete = self.get_all_documents(index=index, filters=filters)
         if ids:
             docs_to_delete = [doc for doc in docs_to_delete if doc.id in ids]
         for doc in docs_to_delete:
             del self.indexes[index][doc.id]
+        if self.use_bm25 is True and len(docs_to_delete) > 0:
+            self.update_bm25(index=index)
 
     def delete_index(self, index: str):
         """
@@ -777,6 +826,9 @@ def delete_index(self, index: str):
             del self.indexes[index]
             logger.info("Index '%s' deleted.", index)
 
+        if index in self.bm25:
+            del self.bm25[index]
+
     def delete_labels(
         self,
         index: Optional[str] = None,
@@ -828,3 +880,112 @@ def delete_labels(
             labels_to_delete = [label for label in labels_to_delete if label.id in ids]
         for label in labels_to_delete:
             del self.indexes[index][label.id]
+
+    def query(
+        self,
+        query: Optional[str],
+        filters: Optional[Dict[str, Union[Dict, List, str, int, float, bool]]] = None,
+        top_k: int = 10,
+        custom_query: Optional[str] = None,
+        index: Optional[str] = None,
+        headers: Optional[Dict[str, str]] = None,
+        all_terms_must_match: bool = False,
+        scale_score: bool = False,
+    ) -> List[Document]:
+        """
+        Scan through documents in DocumentStore and return a small number documents
+        that are most relevant to the query as defined by the BM25 algorithm.
+        :param query: The query.
+        :param top_k: How many documents to return per query.
+        :param index: The name of the index in the DocumentStore from which to retrieve documents.
+        """
+
+        if headers:
+            logger.warning("InMemoryDocumentStore does not support headers. This parameter is ignored.")
+        if custom_query:
+            logger.warning("InMemoryDocumentStore does not support custom_query. This parameter is ignored.")
+        if all_terms_must_match is True:
+            logger.warning("InMemoryDocumentStore does not support all_terms_must_match. This parameter is ignored.")
+        if filters:
+            logger.warning(
+                "InMemoryDocumentStore does not support filters for BM25 retrieval. This parameter is ignored."
+            )
+        if scale_score is True:
+            logger.warning(
+                "InMemoryDocumentStore does not support scale_score for BM25 retrieval. This parameter is ignored."
+            )
+
+        index = index or self.index
+        if index not in self.bm25:
+            raise DocumentStoreError(
+                f"No BM25 representation found for the index: {index}. The Document store should be initialized with use_bm25=True"
+            )
+
+        if query is None:
+            query = ""
+
+        tokenizer = re.compile(self.bm25_tokenization_regex).findall
+        tokenized_query = tokenizer(query.lower())
+        docs_scores = self.bm25[index].get_scores(tokenized_query)
+        top_docs_positions = np.argsort(docs_scores)[::-1][:top_k]
+
+        docs_list = list(self.indexes[index].values())
+        top_docs = []
+        for i in top_docs_positions:
+            doc = docs_list[i]
+            doc.score = docs_scores[i]
+            top_docs.append(doc)
+
+        return top_docs
+
+    def query_batch(
+        self,
+        queries: List[str],
+        filters: Optional[
+            Union[
+                Dict[str, Union[Dict, List, str, int, float, bool]],
+                List[Dict[str, Union[Dict, List, str, int, float, bool]]],
+            ]
+        ] = None,
+        top_k: int = 10,
+        custom_query: Optional[str] = None,
+        index: Optional[str] = None,
+        headers: Optional[Dict[str, str]] = None,
+        all_terms_must_match: bool = False,
+        scale_score: bool = False,
+    ) -> List[List[Document]]:
+        """
+        Scan through documents in DocumentStore and return a small number documents
+        that are most relevant to the provided queries as defined by keyword matching algorithms like BM25.
+        This method lets you find relevant documents for list of query strings (output: List of Lists of Documents).
+        :param query: The query.
+        :param top_k: How many documents to return per query.
+        :param index: The name of the index in the DocumentStore from which to retrieve documents.
+        """
+
+        if headers:
+            logger.warning("InMemoryDocumentStore does not support headers. This parameter is ignored.")
+        if custom_query:
+            logger.warning("InMemoryDocumentStore does not support custom_query. This parameter is ignored.")
+        if all_terms_must_match is True:
+            logger.warning("InMemoryDocumentStore does not support all_terms_must_match. This parameter is ignored.")
+        if filters:
+            logger.warning(
+                "InMemoryDocumentStore does not support filters for BM25 retrieval. This parameter is ignored."
+            )
+        if scale_score is True:
+            logger.warning(
+                "InMemoryDocumentStore does not support scale_score for BM25 retrieval. This parameter is ignored."
+            )
+
+        index = index or self.index
+        if index not in self.bm25:
+            raise DocumentStoreError(
+                f"No BM25 representation found for the index: {index}. The Document store should be initialized with use_bm25=True"
+            )
+
+        result_documents = []
+        for query in queries:
+            result_documents.append(self.query(query=query, top_k=top_k, index=index))
+
+        return result_documents
diff --git a/pyproject.toml b/pyproject.toml
index 5121505a45..8ad68f4cee 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -54,6 +54,7 @@ dependencies = [
   "transformers==4.21.2",
   "nltk",
   "pandas",
+  "rank_bm25",
 
   # Utils
   "dill",  # pickle extension for (de-)serialization


Code review comments:
diff_hunk: @@ -146,6 +172,25 @@ def write_documents(
                     )
                     continue
             self.indexes[index][document.id] = document
+            modified_documents += 1
+
+        if self.use_bm25 is True and modified_documents > 0:
+            self.update_bm25(index=index)
+
+    def update_bm25(self, index: Optional[str] = None):
+        """
+        Updates the BM25 sparse representation in the the document store.
+
+        :param index: Index name for which the BM25 representation is to be updated. If set to None, the default self.index is used.
+        """
+        index = index or self.index
+        tokenizer = re.compile(self.bm25_tokenization_regex).findall
path: haystack/document_stores/memory.py
text: @user3:
[I'm sure this is a bit of an early optimization, but here it goes :smile:]

I've seen you are using this regex in at least two places. We could speed up this a tiny bit by saving the compiled regex. At first I though we should just store the compile regex in the init method, but then, if someone tries to set it directly with `docstore.bm25_tokenization_regex = "a_new_regex"`, stuff will break.

So how about using a property?

```python
@user1
def bm25_tokenization_regex(self):
    return self._tokenizer

@user2.setter
def bm25_tokenization_regex(self, regex_string: str):
    self._tokenizer = re.compile(regex_string).findall
```
(untested)

 By adding these two method somewhere in the class, you can simply remove this line and use `self.bm25_tokenization_regex` directly instead of `tokenizer` in the line below.




@author:
I like the idea!

diff_hunk: @@ -146,6 +172,25 @@ def write_documents(
                     )
                     continue
             self.indexes[index][document.id] = document
+            modified_documents += 1
+
+        if self.use_bm25 is True and modified_documents > 0:
+            self.update_bm25(index=index)
+
+    def update_bm25(self, index: Optional[str] = None):
+        """
+        Updates the BM25 sparse representation in the the document store.
+
+        :param index: Index name for which the BM25 representation is to be updated. If set to None, the default self.index is used.
+        """
+        index = index or self.index
+        tokenizer = re.compile(self.bm25_tokenization_regex).findall
+
+        logger.info("Updating BM25 representation...")
+        all_documents = self.get_all_documents(index=index)
+        tokenized_corpus = [tokenizer(doc.content.lower()) for doc in all_documents]
+        bm25_class = getattr(rank_bm25, self.bm25_algorithm)
+        self.bm25[index] = bm25_class(tokenized_corpus, **self.bm25_parameters)
path: haystack/document_stores/memory.py
text: @user1:
If we change line 106 as suggested, this would become
```suggestion
        self.bm25[index] = self.bm25_algorithm(tokenized_corpus, **self.bm25_parameters)
```

diff_hunk: @@ -828,3 +880,112 @@ def delete_labels(
             labels_to_delete = [label for label in labels_to_delete if label.id in ids]
         for label in labels_to_delete:
             del self.indexes[index][label.id]
+
+    def query(
+        self,
+        query: Optional[str],
+        filters: Optional[Dict[str, Union[Dict, List, str, int, float, bool]]] = None,
+        top_k: int = 10,
+        custom_query: Optional[str] = None,
+        index: Optional[str] = None,
+        headers: Optional[Dict[str, str]] = None,
+        all_terms_must_match: bool = False,
+        scale_score: bool = False,
+    ) -> List[Document]:
+        """
+        Scan through documents in DocumentStore and return a small number documents
+        that are most relevant to the query as defined by the BM25 algorithm.
+        :param query: The query.
+        :param top_k: How many documents to return per query.
+        :param index: The name of the index in the DocumentStore from which to retrieve documents.
+        """
+
+        if headers:
+            logger.warning("InMemoryDocumentStore does not support headers. This parameter is ignored.")
+        if custom_query:
+            logger.warning("InMemoryDocumentStore does not support custom_query. This parameter is ignored.")
+        if all_terms_must_match is True:
+            logger.warning("InMemoryDocumentStore does not support all_terms_must_match. This parameter is ignored.")
+        if filters:
+            logger.warning(
+                "InMemoryDocumentStore does not support filters for BM25 retrieval. This parameter is ignored."
+            )
+        if scale_score is True:
+            logger.warning(
+                "InMemoryDocumentStore does not support scale_score for BM25 retrieval. This parameter is ignored."
+            )
+
+        index = index or self.index
+        if index not in self.bm25:
+            raise DocumentStoreError(
+                f"No BM25 representation found for the index: {index}. The Document store should be initialized with use_bm25=True"
+            )
+
+        if query is None:
+            query = ""
path: haystack/document_stores/memory.py
text: @user1:
Just a nitpick: `if query is None`, is any document going to be retrieved? We could cut it short:
```suggestion
        if query is None:
            return []
```

diff_hunk: @@ -146,6 +172,25 @@ def write_documents(
                     )
                     continue
             self.indexes[index][document.id] = document
+            modified_documents += 1
+
+        if self.use_bm25 is True and modified_documents > 0:
+            self.update_bm25(index=index)
+
+    def update_bm25(self, index: Optional[str] = None):
+        """
+        Updates the BM25 sparse representation in the the document store.
+
+        :param index: Index name for which the BM25 representation is to be updated. If set to None, the default self.index is used.
+        """
+        index = index or self.index
+        tokenizer = re.compile(self.bm25_tokenization_regex).findall
+
+        logger.info("Updating BM25 representation...")
+        all_documents = self.get_all_documents(index=index)
+        tokenized_corpus = [tokenizer(doc.content.lower()) for doc in all_documents]
path: haystack/document_stores/memory.py
text: @user1:
If you want a progress bar to show, we can use `tqdm` here! Something like (untested):

```python
tokenized_corpus = [
    tokenizer(doc.content.lower()) 
    for doc in tqdm(self.get_all_documents(index=index), unit=" docs", desc="Updating BM25 representation...")
]
```