AnswerDotAI / AnswerDotAI/RAGatouille

Batch search in-memory without index results in TypeError if you pass document_metadatas

Open
#79 2 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
4k
Forks
276
PR merge metrics
No merged PRs in 30d

Description

My code:

```python
def rag(args: SearchArgs) -> List[List[SearchResult]]:
RAG = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0", verbose=0)
metadata = [doc["metadata"] for doc in args["docs"]]
RAG.encode(
[doc["content"] for doc in args["docs"]],
document_metadatas=metadata,
)
k = args.get("k", 5)
search_results = RAG.search_encoded_docs(query=args["query"], k=k)
# ....
return search_results

res = rag(
{
"query": ["Hello World", "This is a test"],
"docs": [
{"content": "doc1", "metadata": {"id": 0}},
{"content": "doc2", "metadata": {"id": 1}},
],
}
)
```

Error:

```
Documents encoded!
Traceback (most recent call last):
File "/Users/james/Projects/TS/open-recommender/packages/cli/src/rag/rag.py", line 996, in
res = rag(
^^^^
File "/Users/james/Projects/TS/open-recommender/packages/cli/src/rag/rag.py", line 37, in rag
search_results = RAG.search_encoded_docs(query=args["query"], k=k)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/james/Projects/TS/open-recommender/env/lib/python3.11/site-packages/ragatouille/RAGPretrainedModel.py", line 377, in search_encoded_docs
return self.model.search_encoded_docs(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/james/Projects/TS/open-recommender/env/lib/python3.11/site-packages/ragatouille/models/colbert.py", line 722, in search_encoded_docs
result["result_index"]
~~~~~~^^^^^^^^^^^^^^^^
TypeError: list indices must be integers or slices, not str
```

One thing I found confusing, and which seems to be causing this error too, is that `search_encoded_docs` returns either `List[SearchResult]` or a nested `List[List[SearchResult]]` depending on the number of queries. Maybe it would be better to just always return `List[List[SearchResult]]` otherwise we have to implement conditional logic based on the number of queries we pass, which we might not know ahead of time?

The error does not happen if I only pass one query, or if I don't pass `document_metadatas`.

Workaround:

```python
def rag(args: SearchArgs) -> List[List[SearchResult]]:
RAG = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0", verbose=0)
metadatas = [doc["metadata"] for doc in args["docs"]]
RAG.encode(
[doc["content"] for doc in args["docs"]],
# document_metadatas=metadatas,
)
k = args.get("k", 5)
search_results = RAG.search_encoded_docs(query=args["query"], k=k)
if type(search_results[0]) is not list:
for result in search_results:
result["metadata"] = metadatas[result["result_index"]]
return [search_results]
else:
for batch in search_results:
for result in batch:
result["metadata"] = metadatas[result["result_index"]]
return search_results
```

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by reproducing the two-query example with document_metadatas, then inspect search_encoded_docs in ragatouille/RAGPretrainedModel.py and ragatouille/models/colbert.py, where the traceback shows the failure. Confirm the behavior for both single and batched queries; done means batched searches with metadata no longer raise the TypeError and the returned result shape is consistent or clearly documented.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
search
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.