Repository: deepset-ai/haystack

Issue:
Feature Request: Add index parameter to TFiDF retriever
**Problem**
When using the inmemory docstore on a non standard index (e.g. for evaluation) we cannot use the TFiDF Retriever, because you cannot set an index there.

**Solution**
Lets add the index option to the TFiDF retriever please.

**Background**
I would like the inmemory store + TFiDF to be used for fast haystack examples (without "complicated" ES or FAISS setup)



Discussion:
Related? #1637 
Nice one! Yes it is. I stumbled upon this exact same error by adding the data in an InMemoryStore at a custom index in combination with TFiDFRetriever.
Though in #1637 I do not see a custom index being used...
What's the status on this? I'm actually facing the same problem 🥲 
And also, what would be a workaround?
While thinking about #3447, I decided to revamp this issue.

## Document store and retrievers
From what I understood, in Haystack:
- document stores: should store documents and their representations (such as embeddings)
- retrievers: should be used to query the document stores, in order to provide the most relevant documents
(currently, in some cases, they also build document representations that are then saved to a document store. See #2403)

This holds true for most of the cases, with some exceptions:
- graphs: in the current implementations, they combine the features of a document store and a retriever
- TF-IDF retriever: it uses a document store as a source, but it builds an internal representation for documents (based on a `DataFrame` and a `TfidfVectorizer`)

## How to make the TF-IDF retriever to support different indexes (accept the `index` parameter)?
I see two alternative options:
- find a way to store the representation in the document store
clean from a logical point of view, but hard to implement since this retriever support several different document stores
- modify the current implementation of TF-IDF to create a representation for every index
simpler but dirtier

## Support for BM25Retriever in `InMemoryDocumentStore` (#3447)
It seems related to the same topics.
However, unlike TF-IDF, in Haystack BM25 is a pure retriever (it completely relies on document stores to save document representations).
Based on the proposals made to use external libraries (rank-bm25, gensim) it would seem that even in that case, we want to create a specific implementation of BM25 that builds its own internal representation of the documents...

Sorry if I misunderstood something...

@ZanSara @Timoeller WDYT?

Hey @anakin87! Great analysis, as usual! This is going to be useful even for me to explain the situation to others :rocket:

## Mandatory premise

So, first things first: I believe the current abstraction of Retrievers is fundamentally wrong. As you noticed, some Retriever rely fully on the docstore for the retrieval steps, others have their own internal representation stored (usually) in memory. The whole thing at some point will need to be re-evaluated and clarified (a topic already raised, as you noticed, in https://github.com/deepset-ai/haystack/issues/2403) by assigning to document store, retrievers and embedders consistent and distinct responsibilities.

However, we need to get stuff done with the current architecture for now :sweat_smile: 

## How to make the TF-IDF retriever to support different indexes (accept the index parameter)?

>    - find a way to store the representation in the document store
>      clean from a logical point of view, but hard to implement since this retriever support several different document stores

Technically all docstores should support indices already, so this should not pose too many challenges. However, I have not verified this in practice.

>    - modify the current implementation of TF-IDF to create a representation for every index
>      simpler but dirtier

In case the first solution proves too tough, this is a viable alternative. it might make TFIDFRetriever slow as a snail on large collections, but honestly it's already slow in such conditions, so I don't see it as a big issue :sweat_smile: 

## Support for BM25Retriever in InMemoryDocumentStore (https://github.com/deepset-ai/haystack/issues/3447)

> It seems related to the same topics.
> However, unlike TF-IDF, in Haystack BM25 is a pure retriever (it completely relies on document stores to save document representations).
> Based on the proposals made to use external libraries (rank-bm25, gensim) it would seem that even in that case, we want to create a specific implementation of BM25 that builds its own internal representation of the documents...

For as odd as it sounds, I think the main challenge of adding BM25 support to `InMemoryDocumentStore` is to figure out where to put the BM25 retrieval step. Currently I tend to believe it should be added as a method of `InMemoryDocumentStore`, which can then store the representation alongside the documents. `BM25Retriever` might remain almost unaffected. Please tag me early on in the PR if you decide to go for this, so we can review together how the implementation will look like.

I hope this is helpful, let me know if I forgot to address something!

Resolving patch:
diff --git a/haystack/nodes/retriever/sparse.py b/haystack/nodes/retriever/sparse.py
--- a/haystack/nodes/retriever/sparse.py
+++ b/haystack/nodes/retriever/sparse.py
@@ -1,5 +1,5 @@
 # mypy: disable-error-code=override
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union, Any
 
 import logging
 from collections import OrderedDict, namedtuple
@@ -27,7 +27,7 @@ def __init__(
         scale_score: bool = True,
     ):
         """
-        :param document_store: an instance of one of the following DocumentStores to retrieve from: ElasticsearchDocumentStore, OpenSearchDocumentStore and OpenDistroElasticsearchDocumentStore.
+        :param document_store: An instance of one of the following DocumentStores to retrieve from: InMemoryDocumentStore, ElasticsearchDocumentStore, OpenSearchDocumentStore, and OpenDistroElasticsearchDocumentStore.
             If None, a document store must be passed to the retrieve method for this Retriever to work.
         :param all_terms_must_match: Whether all terms of the query must match the document.
                                      If true all query terms must be present in a document in order to be retrieved (i.e the AND operator is being used implicitly between query terms: "cozy fish restaurant" -> "cozy AND fish AND restaurant").
@@ -460,15 +460,18 @@ def __init__(self, document_store: Optional[BaseDocumentStore] = None, top_k: in
         self.document_store = document_store
         self.top_k = top_k
         self.auto_fit = auto_fit
-        self.document_count = 0
+        self.dataframes: Dict[str, pd.DataFrame] = {}
+        self.tfidf_matrices: Dict[str, Any] = {}
+        self.document_counts: Dict[str, int] = {}
         if document_store and document_store.get_document_count():
             self.fit(document_store=document_store)
 
-    def _get_all_paragraphs(self, document_store: BaseDocumentStore) -> List[Paragraph]:
+    def _get_all_paragraphs(self, document_store: BaseDocumentStore, index: Optional[str]) -> List[Paragraph]:
         """
         Split the list of documents in paragraphs
         """
-        documents = document_store.get_all_documents()
+        index = index or document_store.index
+        documents = document_store.get_all_documents(index=index)
 
         paragraphs = []
         p_id = 0
@@ -483,11 +486,11 @@ def _get_all_paragraphs(self, document_store: BaseDocumentStore) -> List[Paragra
         logger.info("Found %s candidate paragraphs from %s docs in DB", len(paragraphs), len(documents))
         return paragraphs
 
-    def _calc_scores(self, queries: Union[str, List[str]]) -> List[Dict[int, float]]:
+    def _calc_scores(self, queries: List[str], index: str) -> List[Dict[int, float]]:
         if isinstance(queries, str):
             queries = [queries]
         question_vector = self.vectorizer.transform(queries)
-        doc_scores_per_query = self.tfidf_matrix.dot(question_vector.T).T.toarray()
+        doc_scores_per_query = self.tfidf_matrices[index].dot(question_vector.T).T.toarray()
         doc_scores_per_query = [
             [(doc_idx, doc_score) for doc_idx, doc_score in enumerate(doc_scores)]
             for doc_scores in doc_scores_per_query
@@ -522,41 +525,44 @@ def retrieve(
 
         :param document_store: the docstore to use for retrieval. If `None`, the one given in the `__init__` is used instead.
         """
+        if filters:
+            raise NotImplementedError("TfidfRetriever doesn't support filters.")
+        if scale_score:
+            raise NotImplementedError("TfidfRetriever doesn't support scaling score to the unit interval.")
+
+        document_store = document_store or self.document_store
         if document_store is None:
-            document_store = self.document_store
-            if document_store is None:
-                raise ValueError(
-                    "This Retriever was not initialized with a Document Store. Provide one to the retrieve() method."
-                )
-        else:
-            self.fit(document_store=document_store)
+            raise ValueError(
+                "Pass a DocumentStore in the retrieve() method. The Retriever needs a Document Store instance to work."
+            )
+        index = index or document_store.index
+        if index is None:
+            raise ValueError(
+                "Both the `index` parameter passed to the `retrieve` method and the default `index` of the Document store are null. Pass a non-null `index` value."
+            )
 
         if self.auto_fit:
-            if document_store.get_document_count(headers=headers) != self.document_count:
-                # run fit() to update self.df, self.tfidf_matrix and self.document_count
+            if (
+                index not in self.document_counts
+                or document_store.get_document_count(headers=headers, index=index) != self.document_counts[index]
+            ):
+                # run fit() to update self.dataframes, self.tfidf_matrices and self.document_counts
                 logger.warning(
                     "Indexed documents have been updated and fit() method needs to be run before retrieval. Running it now."
                 )
-                self.fit(document_store=document_store)
-        if self.df is None:
+                self.fit(document_store=document_store, index=index)
+        if self.dataframes[index] is None:
             raise DocumentStoreError(
-                "Retrieval requires dataframe df and tf-idf matrix but fit() did not calculate them probably due to an empty document store."
+                "Retrieval requires dataframe and tf-idf matrix but fit() did not calculate them probably due to an empty document store."
             )
 
-        if filters:
-            raise NotImplementedError("Filters are not implemented in TfidfRetriever.")
-        if index:
-            raise NotImplementedError("Switching index is not supported in TfidfRetriever.")
-        if scale_score:
-            raise NotImplementedError("Scaling score to the unit interval is not supported in TfidfRetriever.")
-
         if top_k is None:
             top_k = self.top_k
         # get scores
-        indices_and_scores = self._calc_scores(query)
+        indices_and_scores = self._calc_scores(queries=[query], index=index)
 
         # rank paragraphs
-        df_sliced = self.df.loc[indices_and_scores[0].keys()]
+        df_sliced = self.dataframes[index].loc[indices_and_scores[0].keys()]
         df_sliced = df_sliced[:top_k]
 
         logger.debug(
@@ -580,7 +586,7 @@ def retrieve(
 
     def retrieve_batch(
         self,
-        queries: Union[str, List[str]],
+        queries: List[str],
         filters: Optional[Union[FilterType, List[Optional[FilterType]]]] = None,
         top_k: Optional[int] = None,
         index: Optional[str] = None,
@@ -606,41 +612,45 @@ def retrieve_batch(
                             Otherwise raw similarity scores (e.g. cosine or dot_product) will be used.
         :param document_store: the docstore to use for retrieval. If `None`, the one given in the `__init__` is used instead.
         """
+        if filters:
+            raise NotImplementedError("TfidfRetriever doesn't support filters.")
+        if scale_score:
+            raise NotImplementedError("TfidfRetriever doesn't support scaling score to the unit interval.")
+
+        document_store = document_store or self.document_store
         if document_store is None:
-            document_store = self.document_store
-            if document_store is None:
-                raise ValueError(
-                    "This Retriever was not initialized with a Document Store. Provide one to the retrieve() method."
-                )
-        else:
-            self.fit(document_store=document_store)
+            raise ValueError(
+                "Pass a DocumentStore in the retrieve_batch() method. The Retriever needs a Document Store instance to work."
+            )
+
+        index = index or document_store.index
+        if index is None:
+            raise ValueError(
+                "Both the `index` parameter passed to the `retrieve_batch` method and the default `index` of the Document store are null. Pass a non-null `index` value."
+            )
 
         if self.auto_fit:
-            if document_store.get_document_count(headers=headers) != self.document_count:
-                # run fit() to update self.df, self.tfidf_matrix and self.document_count
+            if (
+                index not in self.document_counts
+                or document_store.get_document_count(headers=headers, index=index) != self.document_counts[index]
+            ):
+                # run fit() to update self.dataframes, self.tfidf_matrices and self.document_counts
                 logger.warning(
                     "Indexed documents have been updated and fit() method needs to be run before retrieval. Running it now."
                 )
-                self.fit(document_store=document_store)
-        if self.df is None:
+                self.fit(document_store=document_store, index=index)
+        if self.dataframes[index] is None:
             raise DocumentStoreError(
-                "Retrieval requires dataframe df and tf-idf matrix but fit() did not calculate them probably due to an empty document store."
+                "Retrieval requires dataframe and tf-idf matrix but fit() did not calculate them probably because of an empty document store."
             )
 
-        if filters:
-            raise NotImplementedError("Filters are not implemented in TfidfRetriever.")
-        if index:
-            raise NotImplementedError("Switching index is not supported in TfidfRetriever.")
-        if scale_score:
-            raise NotImplementedError("Scaling score to the unit interval is not supported in TfidfRetriever.")
-
         if top_k is None:
             top_k = self.top_k
 
-        indices_and_scores = self._calc_scores(queries)
+        indices_and_scores = self._calc_scores(queries=queries, index=index)
         all_documents = []
         for query_result in indices_and_scores:
-            df_sliced = self.df.loc[query_result.keys()]
+            df_sliced = self.dataframes[index].loc[query_result.keys()]
             df_sliced = df_sliced[:top_k]
             logger.debug(
                 "Identified %s candidates via retriever:\n%s",
@@ -661,7 +671,7 @@ def retrieve_batch(
 
         return all_documents
 
-    def fit(self, document_store: BaseDocumentStore):
+    def fit(self, document_store: BaseDocumentStore, index: Optional[str] = None):
         """
         Performing training on this class according to the TF-IDF algorithm.
         """
@@ -669,11 +679,22 @@ def fit(self, document_store: BaseDocumentStore):
             raise ValueError(
                 "This Retriever was not initialized with a Document Store. Provide one to the fit() method."
             )
-        paragraphs = self._get_all_paragraphs(document_store=document_store)
+
+        index = index or document_store.index
+        if index is None:
+            raise ValueError(
+                "Both the `index` parameter passed to the `fit` method and the default `index` of the Document store are null. Pass a non-null `index` value."
+            )
+
+        paragraphs = self._get_all_paragraphs(document_store=document_store, index=index)
         if not paragraphs or len(paragraphs) == 0:
             raise DocumentStoreError("Fit method called with empty document store")
 
-        self.df = pd.DataFrame.from_dict(paragraphs)
-        self.df["content"] = self.df["content"].apply(lambda x: " ".join(x))  # pylint: disable=unnecessary-lambda
-        self.tfidf_matrix = self.vectorizer.fit_transform(self.df["content"])
-        self.document_count = document_store.get_document_count()
+        df = pd.DataFrame.from_dict(paragraphs)
+        df["content"] = df["content"].apply(" ".join)
+        self.dataframes[index] = df
+
+        tfidf_matrix = self.vectorizer.fit_transform(df["content"])
+        self.tfidf_matrices[index] = tfidf_matrix
+
+        self.document_counts[index] = document_store.get_document_count(index=index)


Test patch:
diff --git a/test/nodes/test_retriever.py b/test/nodes/test_retriever.py
--- a/test/nodes/test_retriever.py
+++ b/test/nodes/test_retriever.py
@@ -47,22 +47,18 @@
         ("embedding", "milvus"),
         ("bm25", "elasticsearch"),
         ("bm25", "memory"),
-        ("es_filter_only", "elasticsearch"),
         ("tfidf", "memory"),
     ],
     indirect=True,
 )
 def test_retrieval_without_filters(retriever_with_docs: BaseRetriever, document_store_with_docs: BaseDocumentStore):
-    if not isinstance(retriever_with_docs, (BM25Retriever, FilterRetriever, TfidfRetriever)):
+    if not isinstance(retriever_with_docs, (BM25Retriever, TfidfRetriever)):
         document_store_with_docs.update_embeddings(retriever_with_docs)
 
-    # NOTE: FilterRetriever simply returns all documents matching a filter,
-    # so without filters applied it does nothing
-    if not isinstance(retriever_with_docs, FilterRetriever):
-        res = retriever_with_docs.retrieve(query="Who lives in Berlin?")
-        assert res[0].content == "My name is Carla and I live in Berlin"
-        assert len(res) == 5
-        assert res[0].meta["name"] == "filename1"
+    res = retriever_with_docs.retrieve(query="Who lives in Berlin?")
+    assert res[0].content == "My name is Carla and I live in Berlin"
+    assert len(res) == 5
+    assert res[0].meta["name"] == "filename1"
 
 
 @pytest.mark.parametrize(
@@ -104,6 +100,21 @@ def test_retrieval_with_filters(retriever_with_docs: BaseRetriever, document_sto
     assert len(result) == 0
 
 
+def test_tfidf_retriever_multiple_indexes():
+    docs_index_0 = [Document(content="test_1"), Document(content="test_2"), Document(content="test_3")]
+    docs_index_1 = [Document(content="test_4"), Document(content="test_5")]
+    ds = InMemoryDocumentStore(index="index_0")
+    tfidf_retriever = TfidfRetriever(document_store=ds)
+
+    ds.write_documents(docs_index_0)
+    tfidf_retriever.fit(ds, index="index_0")
+    ds.write_documents(docs_index_1, index="index_1")
+    tfidf_retriever.fit(ds, index="index_1")
+
+    assert tfidf_retriever.document_counts["index_0"] == ds.get_document_count(index="index_0")
+    assert tfidf_retriever.document_counts["index_1"] == ds.get_document_count(index="index_1")
+
+
 class MockBaseRetriever(MockRetriever):
     def __init__(self, document_store: BaseDocumentStore, mock_document: Document):
         self.document_store = document_store