Repository: deepset-ai/haystack
Pull request: fix: Fix `TableTextRetriever` for input consisting of tables only
### Related Issues
- fixes #4013 

### Proposed Changes:
* The `language_model3` within `TriAdaptiveModel` expects kwargs which don't contain the "passage_" prefix. I changed these to conform with the function kwargs.
* The `TableTextRetriever` is composed of an underlying `BertModel` which in turn contains a `BertEmbedding` layer which expects a 2D batch shape of [batch_size, max_seq_len] but the `TableTextRetriever` performed inference on a 3D tensor with shape [batch_size, number_of_hard_negative_passages, max_seq_len]. I changed this to conform with the 2D expectation by adding this change to `retriever. _get_predictions()` 

### How did you test it?
I tested it locally using the code snippet provided in the issue.

### 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:
`TypeError` when updating embeddings with `TableTextRetriever`
**Describe the bug**
When calling update_embeddings using the TableTextRetriever, we run into a `TypeError`.

**Error message**
```
Updating Embedding:   0%|          | 0/1 [00:00<?, ? docs/s]
Traceback (most recent call last):
  File "/Users/bogdan/Repositories/haystack/bogdan_scripts/discord_tabletextret.py", line 13, in <module>
    doc_store.update_embeddings(retriever)
  File "/Users/bogdan/Repositories/haystack/haystack/document_stores/memory.py", line 531, in update_embeddings
    embeddings = retriever.embed_documents(document_batch)
  File "/Users/bogdan/Repositories/haystack/haystack/nodes/retriever/dense.py", line 1214, in embed_documents
    embeddings = self._get_predictions(model_input)["passages"]
  File "/Users/bogdan/Repositories/haystack/haystack/nodes/retriever/dense.py", line 1132, in _get_predictions
    query_embeddings, passage_embeddings = self.model.forward(**batch)[0]
  File "/Users/bogdan/Repositories/haystack/haystack/modeling/model/triadaptive_model.py", line 252, in forward
    pooled_output = self.forward_lm(**kwargs)
  File "/Users/bogdan/Repositories/haystack/haystack/modeling/model/triadaptive_model.py", line 312, in forward_lm
    pooled_output2, _ = self.language_model3(
  File "/opt/homebrew/Caskroom/miniforge/base/envs/haystack/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
TypeError: forward() got an unexpected keyword argument 'passage_input_ids'
```

**To Reproduce**
```python
from haystack.document_stores import InMemoryDocumentStore
from haystack.nodes import TableTextRetriever
from haystack import Document
import pandas as pd

table = pd.DataFrame(columns=["id", "text"], data=[["1", "This is a test"], ["2", "This is another test"]])
doc = Document(content=table, content_type="table")

doc_store = InMemoryDocumentStore()
retriever = TableTextRetriever(document_store=doc_store)

doc_store.write_documents([doc])
doc_store.update_embeddings(retriever)
```

Hints from issue discussion:
Hey, I'm a research engineer currently working on a mix of ml research and software engineering with a focus on language modelling. I'd like to get involved and take this issue if it's still free?
Hey @jebbbbb, we would be more than happy to get your contribution here! Feel free to reach out whenever you come across a problem :)
Hey @bogdankostic, I've changed my username to this one now. Yes sure, thanks!
Hi @bogdankostic, I've opened a PR #4048 to address this issue

Commit message: test: add test for ttr + dataframe case

Patch to review:
diff --git a/haystack/modeling/model/language_model.py b/haystack/modeling/model/language_model.py
index a2290d4665..82f3186fe2 100644
--- a/haystack/modeling/model/language_model.py
+++ b/haystack/modeling/model/language_model.py
@@ -705,9 +705,9 @@ def forward(
         :param input_ids: The IDs of each token in the input sequence. It's a tensor of shape [batch_size, number_of_hard_negative, max_seq_len].
         :param segment_ids: The ID of the segment. For example, in next sentence prediction, the tokens in the
            first sentence are marked with 0 and the tokens in the second sentence are marked with 1.
-           It is a tensor of shape [batch_size, number_of_hard_negative_passages, max_seq_len].
+           It is a tensor of shape [batch_size, max_seq_len].
         :param attention_mask: A mask that assigns 1 to valid input tokens and 0 to padding tokens
-           of shape [batch_size,  number_of_hard_negative_passages, max_seq_len].
+           of shape [batch_size, max_seq_len].
         :param output_hidden_states: whether to add the hidden states along with the pooled output
         :param output_attentions: unused
         :return: Embeddings for each token in the input sequence.
diff --git a/haystack/modeling/model/triadaptive_model.py b/haystack/modeling/model/triadaptive_model.py
index 47244316eb..139290e22a 100644
--- a/haystack/modeling/model/triadaptive_model.py
+++ b/haystack/modeling/model/triadaptive_model.py
@@ -310,9 +310,9 @@ def forward_lm(self, **kwargs):
             # Current batch consists of only tables
             if all(table_mask):
                 pooled_output2, _ = self.language_model3(
-                    passage_input_ids=kwargs["passage_input_ids"],
-                    passage_segment_ids=kwargs["table_segment_ids"],
-                    passage_attention_mask=kwargs["passage_attention_mask"],
+                    input_ids=kwargs["passage_input_ids"],
+                    segment_ids=kwargs["table_segment_ids"],
+                    attention_mask=kwargs["passage_attention_mask"],
                     output_hidden_states=False,
                     output_attentions=False,
                 )
diff --git a/haystack/nodes/retriever/dense.py b/haystack/nodes/retriever/dense.py
index 3761151280..db5d77f1fe 100644
--- a/haystack/nodes/retriever/dense.py
+++ b/haystack/nodes/retriever/dense.py
@@ -1125,7 +1125,11 @@ def _get_predictions(self, dicts: List[Dict[str, Any]]) -> Dict[str, np.ndarray]
             disable=disable_tqdm,
         ) as progress_bar:
             for batch in data_loader:
-                batch = {key: batch[key].to(self.devices[0]) for key in batch}
+                # LanguageModel expects [batch, seq] not [batch, num_hard, seq]
+                for key in batch:
+                    batch[key].to(self.devices[0])
+                    if batch[key].ndim == 3:
+                        batch[key] = batch[key].squeeze(1)
 
                 # get logits
                 with torch.inference_mode():
diff --git a/test/document_stores/test_document_store.py b/test/document_stores/test_document_store.py
index f5d8c1a251..5b9bf370f5 100644
--- a/test/document_stores/test_document_store.py
+++ b/test/document_stores/test_document_store.py
@@ -219,6 +219,17 @@ def test_update_embeddings_table_text_retriever(document_store, retriever):
     )
 
 
+@pytest.mark.parametrize("retriever", ["table_text_retriever"], indirect=True)
+@pytest.mark.embedding_dim(512)
+def test_update_embeddings_table_text_retriever_with_dataframe(retriever):
+    table = pd.DataFrame(columns=["id", "text"], data=[["1", "This is a test"], ["2", "This is another test"]])
+    doc = Document(content=table, content_type="table")
+    doc_store = InMemoryDocumentStore()
+    retriever = retriever(document_store=doc_store)
+    doc_store.write_documents([doc])
+    doc_store.update_embeddings(retriever)
+
+
 @pytest.mark.parametrize("document_store_type", ["elasticsearch", "memory"])
 def test_custom_embedding_field(document_store_type, tmp_path):
     document_store = get_document_store(


Code review comments:
diff_hunk: @@ -310,9 +310,9 @@ def forward_lm(self, **kwargs):
             # Current batch consists of only tables
             if all(table_mask):
                 pooled_output2, _ = self.language_model3(
-                    passage_input_ids=kwargs["passage_input_ids"],
-                    passage_segment_ids=kwargs["table_segment_ids"],
-                    passage_attention_mask=kwargs["passage_attention_mask"],
+                    input_ids=kwargs["passage_input_ids"],
+                    segment_ids=kwargs["table_segment_ids"],
+                    attention_mask=kwargs["passage_attention_mask"],
                     output_hidden_states=False,
                     output_attentions=False,
                 )
path: haystack/modeling/model/triadaptive_model.py
text: @user1:
Let's make the input two-dimensional here, similar to the case where we only have text Documents.
```suggestion
            if all(table_mask):
                # Make input two-dimensional
                max_seq_len = kwargs["table_input_ids"].shape[-1]
                input_ids = kwargs["table_input_ids"].view(-1, max_seq_len)
                attention_mask = kwargs["table_attention_mask"].view(-1, max_seq_len)
                segment_ids = kwargs["table_segment_ids"].view(-1, max_seq_len)

                pooled_output2, _ = self.language_model3(
                    input_ids=input_ids,
                    segment_ids=segment_ids,
                    attention_mask=attention_mask,
                    output_hidden_states=False,
                    output_attentions=False,
                )
```

diff_hunk: @@ -1125,7 +1125,11 @@ def _get_predictions(self, dicts: List[Dict[str, Any]]) -> Dict[str, np.ndarray]
             disable=disable_tqdm,
         ) as progress_bar:
             for batch in data_loader:
-                batch = {key: batch[key].to(self.devices[0]) for key in batch}
+                # LanguageModel expects [batch, seq] not [batch, num_hard, seq]
+                for key in batch:
+                    batch[key].to(self.devices[0])
+                    if batch[key].ndim == 3:
+                        batch[key] = batch[key].squeeze(1)
path: haystack/nodes/retriever/dense.py
text: @user1:
I wouldn't correct the tensor shape in the `TableTextRetriever` but in the `TriAdaptiveModel`. For consistency with our existing code base, I would do this the same way as we're doing it in [this line](https://github.com/deepset-ai/haystack/blob/b9ab7b3ca2d548b3669038d837b83e6558fa8c6d/haystack/modeling/model/triadaptive_model.py#L374) and the four following ones for text only inputs. (See the suggested change in `triadaptive_model.py`.)

diff_hunk: @@ -219,6 +219,17 @@ def test_update_embeddings_table_text_retriever(document_store, retriever):
     )
 
 
+@pytest.mark.parametrize("retriever", ["table_text_retriever"], indirect=True)
+@pytest.mark.embedding_dim(512)
+def test_update_embeddings_table_text_retriever_with_dataframe(retriever):
+    table = pd.DataFrame(columns=["id", "text"], data=[["1", "This is a test"], ["2", "This is another test"]])
+    doc = Document(content=table, content_type="table")
+    doc_store = InMemoryDocumentStore()
+    retriever = retriever(document_store=doc_store)
+    doc_store.write_documents([doc])
+    doc_store.update_embeddings(retriever)
+
+
path: test/document_stores/test_document_store.py
text: @user2:
Let's rename this test to `test_table_text_retriever_embedding_only_table`, move it to `test/nodes/test_retriever.py` and add the `@user1.mark.integration` decorator.

Also, let's add a similar test `test_table_text_retriever_embedding_only_text` where we test the `TableTextRetriever` with only Documents of type text.