Add Support for Cherry-Picking Specific Documents via Document IDs

[leagrieder]

🔍 PR: Add Support for Cherry-Picking Specific Documents via Document IDs

Summary

This PR adds the ability to strictly filter the hybrid retrieval to user-specified document IDs. Rather than searching over the entire collection and later patching in documents, the hybrid search itself applies a filter expression (id in [...]) to ensure that only the desired doc IDs are considered.

By merging the doc ID restriction directly into both dense and sparse searches, we retain the standard hybrid ranking logic but limit it to the documents the user requested. This is particularly beneficial for scenarios where a user (e.g. MOOVE) wants deterministic or curated doc sets, while still leveraging ML-based ranking (e.g. MOORE) to order those docs by relevance.

✅ Main Changes

1. Support for Document IDs in the Retrieval Flow

   • retriever.py: The retrieve() method now has an optional document_ids argument. If provided, it attaches a filter expression (e.g., id in ["doc1", "doc2", ...]) to the dense and sparse search requests.
   • read_queries() (in run_rag.py): Now returns a list of dictionaries, each of which can include both the user query text and an optional list of document IDs.
   • run_retriever.py: The retrieve() function parses --document-ids (if any) into a Python list before calling the updated retriever, which then restricts search accordingly.

2. Removal of Old “Patch-In” Logic

   • The earlier approach of retrieving everything and then fetch-missing-docs is gone. Now, if doc IDs are given, we filter within the search step.
   • The helper get_documents_by_ids() remains in the codebase if we need direct doc lookups in the future.

3. CLI Behavior

   • The --document-ids/-d flag can optionally be provided as a comma-separated list (e.g., doc1,doc2,doc3) in the retrieve part of the CLI.

4. Bug Fixes & Minor Improvements

   • Ensured read_queries() can return dictionaries with document_ids as well as input text.
   • Updated the logic in _get_relevant_documents() to pass document_ids along to retrieve() in order to keep it consistent with the new approach.
   • Logging improvements for better transparency (e.g. showing the expr filter in logs, etc.).

🧪 Usage Example

mmore retrieve \
  --config-file index_config.yaml \
  --input-file queries.jsonl \
  --output-file results.json \
  --document-ids doc1,doc2,doc3

• The system will read queries from queries.jsonl.
• If --document-ids is supplied, it will filter to doc1, doc2, doc3 only.
• Hybrid search ranks these filtered docs by relevance and saves the final results to results.json.

[fabnemEPFL]

Why not just use retriever.retrieve(…) with the empty doc_ids_list? It should work too

[leagrieder]

Good Catch!!
And actually, it was my mistake to bypass the invoke function and call retrieve directly when doc ids were specified because doing so skipped the callback manager ( for logging, error handling, performance tracking, and consistent and robust retrieval process).

Although everything seemed to work at first glance, I later realized that by not using invoke like it was initially done in the initial codebase I was missing out on these important features. So I have updated the code now such that it always calls invoke but with optional document_ids in kwargs