feat: add query_by_embedding_batch #3546
Conversation
| body.append(headers) | ||
| body.append(cur_query_body) | ||
|
|
||
| logger.debug("Retriever query: %s", body) |
There was a problem hiding this comment.
What issue are we expecting to debug with this log? Do we need to add more details to debug this issue or this is enough?
There was a problem hiding this comment.
I just copied it from query_batch I guess. But OpenSearch's and elastic search's client already produce these logs. So you can achieve the same by just running something like
logging.getLogger("opensearch").setLevel(logging.DEBUG)I'll remove it.
| scale_score=scale_score, | ||
| ) | ||
| documents = document_store.query_by_embedding_batch( | ||
| query_embs=query_embeddings, # type: ignore |
There was a problem hiding this comment.
Would this stop MultiModal retriever from using all the document stores?
There was a problem hiding this comment.
My understanding is that any document store still works because the default implementation is just to run query_by_embedding and only OpenSearchDocumentStore and ElasticsearchDocumentStore have a special implementation. Our tests use the MultiModalRetriever only with InMemoryDocumentStore if I am not mistaken. So we should double check that manually before merging. @tstadel You could try to run the tests involving MultiModalRetriever in test_retriever.py with ElasticsearchDocumentStore, maybe?
There was a problem hiding this comment.
@julian-risch @mayankjobanputra Ok I ran all multimodal tests of test_retriever.py locally using OpenSearchDocumentStore and added another one that uses retrieve_batch. All tests were fine. I added the retrieve_batch test using InMemoryDocumentStore as all the others.
I know that #type: ignore is not really nice here, so I resolved the query_embs one by allowing for multi-dimensional np.ndarrays.
For the other one I'll open another PR to resolve this. It would be just too big for this one I guess.
|
Nice improvement idea! Changes look good to me. We should do manual checks with the MultiModalRetriever as Mayank pointed out. Other than that I am curious whether you have an idea about the expected speed increase? We don't have an easy-to-run benchmark script at the moment but it might become a topic for the next quarter... |
Ok, I'll try to do that this week. Regarding speed increase, it depends on how you choose the cluster setup. But we noticed a 49% latency drop in total on a quite powerfuI cluster when running multi-embedding queries (i.e. we run multiple queries (up to 50) in parallel and combine the results using reciprocal rank fusion) on a dataset with ~30 mio. documents*. Similar results should be possible for evaluating dense retrieval pipelines using
|
Related Issues
Proposed Changes:
query_by_embeddinglikequery_batchfor DocumentStoresquery_by_embeddingOpenSearchDocumentStoreandElasticsearchDocumentStoreimplement it to usemsearchlogger.debug("Retriever query: %s", body)fromSearchEngineDocumentStoreas the same info is already being logged by theopensearchandelasticsearchclientsHow did you test it?
run_batchof pipeline and retrieversNotes for the reviewer
Checklist