Repository: deepset-ai/haystack
Pull request: refactor: `InMemoryDocumentStore` - manage documents without embedding & fix mypy errors
### Related Issues
- fixes #4080

### Proposed Changes:
- implement @bogdankostic's proposal: 
> ignoring the Documents without embeddings during retrieval. (We should probably print a warning message for this case.) 
- remove mypy ignore comments and fix the errors
- fix `get_scores_numpy `and `get_scores_torch` methods (there were some errors also pointed out by mypy)

### How did you test it?
- CI
- added new test for documents without embeddings

### 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
- [X] I added tests that demonstrate the correct behavior of the change
- [X] I've used one of the [conventional commit types](https://www.conventionalcommits.org/en/v1.0.0/) for my PR title: `fix:`, `feat:`, `build:`, `chore:`, `ci:`, `docs:`, `style:`, `refactor:`, `perf:`, `test:`.
- [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:
Improve error message for `InMemoryDocumentStore` for Documents without embeddings
**Describe the bug**
If we use the `InMemoryDocumentStore` in combination with a dense retriever and try to do retrieval when at least one of the indexed Documents doesn't contain an embedding, the error message isn't very useful to the user.

**Error message**
```
Traceback (most recent call last):
  File "/Users/bogdan/Repositories/haystack/haystack/pipelines/base.py", line 530, in run
    node_output, stream_id = self._run_node(node_id, node_input)
  File "/Users/bogdan/Repositories/haystack/haystack/pipelines/base.py", line 457, in _run_node
    return self.graph.nodes[node_id]["component"]._dispatch_run(**node_input)
  File "/Users/bogdan/Repositories/haystack/haystack/nodes/base.py", line 201, in _dispatch_run
    return self._dispatch_run_general(self.run, **kwargs)
  File "/Users/bogdan/Repositories/haystack/haystack/nodes/base.py", line 245, in _dispatch_run_general
    output, stream = run_method(**run_inputs, **run_params)
  File "/Users/bogdan/Repositories/haystack/haystack/nodes/retriever/base.py", line 290, in run
    output, stream = run_query_timed(
  File "/Users/bogdan/Repositories/haystack/haystack/nodes/retriever/base.py", line 111, in wrapper
    ret = fn(*args, **kwargs)
  File "/Users/bogdan/Repositories/haystack/haystack/nodes/retriever/base.py", line 344, in run_query
    documents = self.retrieve(
  File "/Users/bogdan/Repositories/haystack/haystack/nodes/retriever/dense.py", line 1647, in retrieve
    documents = document_store.query_by_embedding(
  File "/Users/bogdan/Repositories/haystack/haystack/document_stores/memory.py", line 452, in query_by_embedding
    scores = self.get_scores(query_emb, document_to_search)
  File "/Users/bogdan/Repositories/haystack/haystack/document_stores/memory.py", line 355, in get_scores
    scores = self.get_scores_numpy(query_emb, document_to_search)
  File "/Users/bogdan/Repositories/haystack/haystack/document_stores/memory.py", line 333, in get_scores_numpy
    doc_embeds = doc_embeds.unsqueeze(dim=0)  # type: ignore [attr-defined]
AttributeError: 'numpy.ndarray' object has no attribute 'unsqueeze'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/bogdan/Repositories/haystack/bogdan_scripts/discord_error.py", line 21, in <module>
    result = query_pipeline.run(query=query)  # this line is where the error pops up in collab
  File "/Users/bogdan/Repositories/haystack/haystack/pipelines/standard_pipelines.py", line 364, in run
    output = self.pipeline.run(query=query, params=params, debug=debug)
  File "/Users/bogdan/Repositories/haystack/haystack/pipelines/base.py", line 535, in run
    raise Exception(
Exception: Exception while running node 'Retriever': 'numpy.ndarray' object has no attribute 'unsqueeze'
```

**Expected behavior**
I see two possible ways of how we can improve this:
- Raising an error stating that at least one of the Documents doesn't contain an embedding and that the user needs to call `update_embeddings` in order to be able to do dense retrieval
- Ignoring the Documents without embeddings during retrieval. (We should probably print a warning message for this case.)

**To Reproduce**
```python
from haystack.document_stores import InMemoryDocumentStore
from haystack.nodes import EmbeddingRetriever
from haystack.pipelines import DocumentSearchPipeline
from haystack import Document

document_store = InMemoryDocumentStore(embedding_dim=1024)

# OpenAI EmbeddingRetriever
retriever = EmbeddingRetriever(
    document_store=document_store,
    batch_size=8,
    embedding_model="ada",
    api_key="<api_key>",
    max_seq_len=1024,
)
document_store.write_documents([Document(content="test Document")])

query = "<my_query>"
query_pipeline = DocumentSearchPipeline(retriever=retriever)

result = query_pipeline.run(query=query)  # this line is where the error pops up in collab
```

Hints from issue discussion:
Hey @bogdankostic!
Which approach do you prefer?

- raising an error
- skip the document and print a warning
I'd prefer the second approach.
@bogdankostic I started to study the code a bit and I found that `numpy.ndarray` does not have the `unsqueeze` method.
See also mypy warnings in https://github.com/deepset-ai/haystack/issues/2351.

My impression is that the following lines https://github.com/deepset-ai/haystack/blob/adf4a3ea2f88e302b687499acf9e9bedb6ad0935/haystack/document_stores/memory.py#L335-L336
are executed only when there is only one document with no embeddings. So we rarely see this error.

In the `get_scores_numpy` and `get_scores_torch` methods,
there are several silenced mypy warnings.
So my impression is that these methods need good refactoring.

I can try to do that, although it won't be child's play... :smile: 
I will be submitting a PR in a few hours...

Commit message: try to replace error with warning

Patch to review:
diff --git a/haystack/document_stores/memory.py b/haystack/document_stores/memory.py
index bf3029c745..87e4082fb4 100644
--- a/haystack/document_stores/memory.py
+++ b/haystack/document_stores/memory.py
@@ -281,47 +281,57 @@ def get_document_by_id(
         else:
             return None
 
-    def get_documents_by_id(self, ids: List[str], index: Optional[str] = None) -> List[Document]:  # type: ignore
+    def get_documents_by_id(
+        self,
+        ids: List[str],
+        index: Optional[str] = None,
+        batch_size: Optional[int] = None,
+        headers: Optional[Dict[str, str]] = None,
+    ) -> List[Document]:
         """
         Fetch documents by specifying a list of text id strings.
         """
+        if headers:
+            raise NotImplementedError("InMemoryDocumentStore does not support headers.")
+        if batch_size:
+            logger.warning(
+                "InMemoryDocumentStore does not support batching in `get_documents_by_id` method. This parameter is ignored."
+            )
         index = index or self.index
         documents = [self.indexes[index][id] for id in ids]
         return documents
 
-    def get_scores_torch(self, query_emb: np.ndarray, document_to_search: List[Document]) -> List[float]:
+    def get_scores_torch(self, query_emb: np.ndarray, documents_to_search: List[Document]) -> List[float]:
         """
         Calculate similarity scores between query embedding and a list of documents using torch.
 
         :param query_emb: Embedding of the query (e.g. gathered from DPR)
-        :param document_to_search: List of documents to compare `query_emb` against.
+        :param documents_to_search: List of documents to compare `query_emb` against.
         """
-        query_emb = torch.tensor(query_emb, dtype=torch.float).to(self.main_device)  # type: ignore [assignment]
-        if len(query_emb.shape) == 1:
-            query_emb = query_emb.unsqueeze(dim=0)  # type: ignore [attr-defined]
+        query_emb_tensor = torch.tensor(query_emb, dtype=torch.float).to(self.main_device)
+        if len(query_emb_tensor.shape) == 1:
+            query_emb_tensor = query_emb_tensor.unsqueeze(dim=0)
 
-        doc_embeds = np.array([doc.embedding for doc in document_to_search])
-        doc_embeds = torch.as_tensor(doc_embeds, dtype=torch.float)  # type: ignore [assignment]
-        if len(doc_embeds.shape) == 1 and doc_embeds.shape[0] == 1:
-            doc_embeds = doc_embeds.unsqueeze(dim=0)  # type: ignore [attr-defined]
-        elif len(doc_embeds.shape) == 1 and doc_embeds.shape[0] == 0:
+        doc_embeds = np.array([doc.embedding for doc in documents_to_search])
+        doc_embeds_tensor = torch.as_tensor(doc_embeds, dtype=torch.float)
+        if len(doc_embeds_tensor.shape) == 1 and doc_embeds_tensor.shape[0] == 0:
             return []
 
         if self.similarity == "cosine":
             # cosine similarity is just a normed dot product
-            query_emb_norm = torch.norm(query_emb, dim=1)
-            query_emb = torch.div(query_emb, query_emb_norm)  # type: ignore [assignment,arg-type]
+            query_emb_norm = torch.norm(query_emb_tensor, dim=1)
+            query_emb_tensor = torch.div(query_emb_tensor, query_emb_norm)
 
-            doc_embeds_norms = torch.norm(doc_embeds, dim=1)
-            doc_embeds = torch.div(doc_embeds.T, doc_embeds_norms).T  # type: ignore [assignment,arg-type]
+            doc_embeds_norms = torch.norm(doc_embeds_tensor, dim=1)
+            doc_embeds_tensor = torch.div(doc_embeds_tensor.T, doc_embeds_norms).T
 
         curr_pos = 0
-        scores = []  # type: ignore [var-annotated]
-        while curr_pos < len(doc_embeds):
-            doc_embeds_slice = doc_embeds[curr_pos : curr_pos + self.scoring_batch_size]
-            doc_embeds_slice = doc_embeds_slice.to(self.main_device)  # type: ignore [attr-defined]
+        scores: List[float] = []
+        while curr_pos < len(doc_embeds_tensor):
+            doc_embeds_slice = doc_embeds_tensor[curr_pos : curr_pos + self.scoring_batch_size]
+            doc_embeds_slice = doc_embeds_slice.to(self.main_device)
             with torch.inference_mode():
-                slice_scores = torch.matmul(doc_embeds_slice, query_emb.T).cpu()  # type: ignore [arg-type,arg-type]
+                slice_scores = torch.matmul(doc_embeds_slice, query_emb_tensor.T).cpu()
                 slice_scores = slice_scores.squeeze(dim=1)
                 slice_scores = slice_scores.numpy().tolist()
 
@@ -330,20 +340,18 @@ def get_scores_torch(self, query_emb: np.ndarray, document_to_search: List[Docum
 
         return scores
 
-    def get_scores_numpy(self, query_emb: np.ndarray, document_to_search: List[Document]) -> List[float]:
+    def get_scores_numpy(self, query_emb: np.ndarray, documents_to_search: List[Document]) -> List[float]:
         """
         Calculate similarity scores between query embedding and a list of documents using numpy.
 
         :param query_emb: Embedding of the query (e.g. gathered from DPR)
-        :param document_to_search: List of documents to compare `query_emb` against.
+        :param documents_to_search: List of documents to compare `query_emb` against.
         """
         if len(query_emb.shape) == 1:
             query_emb = np.expand_dims(query_emb, 0)
 
-        doc_embeds = np.array([doc.embedding for doc in document_to_search])
-        if len(doc_embeds.shape) == 1 and doc_embeds.shape[0] == 1:
-            doc_embeds = doc_embeds.unsqueeze(dim=0)  # type: ignore [attr-defined]
-        elif len(doc_embeds.shape) == 1 and doc_embeds.shape[0] == 0:
+        doc_embeds = np.array([doc.embedding for doc in documents_to_search])
+        if len(doc_embeds.shape) == 1 and doc_embeds.shape[0] == 0:
             return []
 
         if self.similarity == "cosine":
@@ -360,11 +368,11 @@ def get_scores_numpy(self, query_emb: np.ndarray, document_to_search: List[Docum
 
         return scores
 
-    def get_scores(self, query_emb: np.ndarray, document_to_search: List[Document]) -> List[float]:
+    def get_scores(self, query_emb: np.ndarray, documents_to_search: List[Document]) -> List[float]:
         if self.main_device.type == "cuda":
-            scores = self.get_scores_torch(query_emb, document_to_search)
+            scores = self.get_scores_torch(query_emb, documents_to_search)
         else:
-            scores = self.get_scores_numpy(query_emb, document_to_search)
+            scores = self.get_scores_numpy(query_emb, documents_to_search)
 
         return scores
 
@@ -460,11 +468,17 @@ def query_by_embedding(
         if query_emb is None:
             return []
 
-        document_to_search = self.get_all_documents(index=index, filters=filters, return_embedding=True)
-        scores = self.get_scores(query_emb, document_to_search)
+        documents = self.get_all_documents(index=index, filters=filters, return_embedding=True)
+        documents_with_embeddings = [doc for doc in documents if doc.embedding is not None]
+        if len(documents) != len(documents_with_embeddings):
+            logger.warning(
+                "Skipping some of your documents that don't have embeddings. "
+                "To generate embeddings, run the document store's update_embeddings() method."
+            )
+        scores = self.get_scores(query_emb, documents_with_embeddings)
 
         candidate_docs = []
-        for doc, score in zip(document_to_search, scores):
+        for doc, score in zip(documents_with_embeddings, scores):
             curr_meta = deepcopy(doc.meta)
             new_document = Document(
                 id=doc.id, content=doc.content, content_type=doc.content_type, meta=curr_meta, embedding=doc.embedding
diff --git a/test/document_stores/test_memory.py b/test/document_stores/test_memory.py
index 92d3196f4d..d04bb1fd35 100644
--- a/test/document_stores/test_memory.py
+++ b/test/document_stores/test_memory.py
@@ -3,6 +3,7 @@
 import pandas as pd
 import pytest
 from rank_bm25 import BM25
+import numpy as np
 
 from haystack.document_stores.memory import InMemoryDocumentStore
 from haystack.schema import Document
@@ -112,3 +113,15 @@ def test_memory_query_by_embedding_batch(self, ds, documents):
         for docs, query_emb in zip(docs_batch, query_embs):
             assert len(docs) == 5
             assert (docs[0].embedding == query_emb).all()
+
+    @pytest.mark.integration
+    def test_memory_query_by_embedding_docs_wo_embeddings(self, ds, caplog):
+        # write document but don't update embeddings
+        ds.write_documents([Document(content="test Document")])
+
+        query_embedding = np.random.rand(768).astype(np.float32)
+
+        with caplog.at_level(logging.WARNING):
+            docs = ds.query_by_embedding(query_emb=query_embedding, top_k=1)
+            assert "Skipping some of your documents that don't have embeddings" in caplog.text
+        assert len(docs) == 0


Code review comments:
diff_hunk: @@ -281,47 +281,57 @@ def get_document_by_id(
         else:
             return None
 
-    def get_documents_by_id(self, ids: List[str], index: Optional[str] = None) -> List[Document]:  # type: ignore
+    def get_documents_by_id(
+        self,
+        ids: List[str],
+        index: Optional[str] = None,
+        batch_size: Optional[int] = None,
+        headers: Optional[Dict[str, str]] = None,
+    ) -> List[Document]:
         """
         Fetch documents by specifying a list of text id strings.
         """
+        if headers:
+            raise NotImplementedError("InMemoryDocumentStore does not support headers.")
+        if batch_size:
+            logger.warning(
+                "InMemoryDocumentStore does not support batching in `get_documents_by_id` method. This parameter is ignored."
+            )
         index = index or self.index
         documents = [self.indexes[index][id] for id in ids]
         return documents
 
-    def get_scores_torch(self, query_emb: np.ndarray, document_to_search: List[Document]) -> List[float]:
+    def get_scores_torch(self, query_emb: np.ndarray, documents_to_search: List[Document]) -> List[float]:
         """
         Calculate similarity scores between query embedding and a list of documents using torch.
 
         :param query_emb: Embedding of the query (e.g. gathered from DPR)
-        :param document_to_search: List of documents to compare `query_emb` against.
+        :param documents_to_search: List of documents to compare `query_emb` against.
         """
-        query_emb = torch.tensor(query_emb, dtype=torch.float).to(self.main_device)  # type: ignore [assignment]
-        if len(query_emb.shape) == 1:
-            query_emb = query_emb.unsqueeze(dim=0)  # type: ignore [attr-defined]
+        query_emb_tensor = torch.tensor(query_emb, dtype=torch.float).to(self.main_device)
+        if len(query_emb_tensor.shape) == 1:
+            query_emb_tensor = query_emb_tensor.unsqueeze(dim=0)
 
-        doc_embeds = np.array([doc.embedding for doc in document_to_search])
-        doc_embeds = torch.as_tensor(doc_embeds, dtype=torch.float)  # type: ignore [assignment]
-        if len(doc_embeds.shape) == 1 and doc_embeds.shape[0] == 1:
-            doc_embeds = doc_embeds.unsqueeze(dim=0)  # type: ignore [attr-defined]
-        elif len(doc_embeds.shape) == 1 and doc_embeds.shape[0] == 0:
+        doc_embeds = np.array([doc.embedding for doc in documents_to_search])
+        doc_embeds_tensor = torch.as_tensor(doc_embeds, dtype=torch.float)
+        if len(doc_embeds_tensor.shape) == 1 and doc_embeds_tensor.shape[0] == 0:
path: haystack/document_stores/memory.py
text: @user1:
`doc_embeds_tensor.shape[0]` is equivalent to `len(doc_embeds_tensor)` here, correct ([numpy docs](https://numpy.org/devdocs/reference/generated/numpy.shape.html))?
What if `len(doc_embeds_tensor.shape) == 1 and doc_embeds_tensor.shape[0] == 1`? In the past, we then called `unsqueeze(dim=0)`. Do I understand correctly that we don't need that step and can continue with the execution?
Maybe we can add a comment to the `if len(doc_embeds_tensor.shape) == 1 and doc_embeds_tensor.shape[0] == 0:` then. Explaining that no scores are returned if none of the provided documents has an embedding? That would make the code easier to read. What do you think?

@author:
Hello @user2-risch!
Thanks for the review and sorry for not being understandable. 🙂 

If you look at https://github.com/deepset-ai/haystack/issues/4080#issuecomment-1423152835, what emerges is that the original code was a bit buggy. I did a too quick fix, which was not very clear and maybe not so correct.

- Now I tried to write more expressive code: for example, I used `ndim` instead of `len(array.shape)`.
- I also reintegrated that `unsqueeze`/`expand_dims` part in a more reasonable way, I hope.  In any case, I'm having a hard time imagining when that part will be executed, now that we filter documents without embeddings.
- I also added the suggested comment.

Now my code needs another careful review, also because I'm not a tensor wizard 🧙‍♂️.