Skip to content

fix: support HF datasets in dense exact search#219

Open
GioiaZheng wants to merge 1 commit into
beir-cellar:mainfrom
GioiaZheng:fix/hf-loader-dense-search
Open

fix: support HF datasets in dense exact search#219
GioiaZheng wants to merge 1 commit into
beir-cellar:mainfrom
GioiaZheng:fix/hf-loader-dense-search

Conversation

@GioiaZheng

Copy link
Copy Markdown

Why

HFDataLoader returns Hugging Face Dataset objects for corpus and queries, while DenseRetrievalExactSearch expected dict-like inputs and called queries.keys(). This makes the documented HF loader dense retrieval path fail before retrieval starts.

Closes #208.

What changed

  • Add input preparation helpers for dict inputs and dataset-like inputs with id/text columns.
  • Reuse the same preparation path in search() and encode().
  • Preserve chunked corpus encoding by sorting dataset row indices and materializing only the current batch passed to encode_corpus().
  • Add focused tests for HF Dataset inputs in dense exact search and encoding.

Validation

  • .\.venv\Scripts\python.exe -m pytest tests\retrieval\search\dense\test_exact_search_hf_dataset.py
  • .\.venv\Scripts\python.exe -m ruff check beir\retrieval\search\dense\exact_search.py tests\retrieval\search\dense\test_exact_search_hf_dataset.py
  • git diff --check HEAD~1..HEAD

Risk / follow-up

Low scope: this targets non-streaming HF Dataset objects, which is the path used by the current HF loader dense retrieval example. Streaming IterableDataset support is left separate because the current loader path still relies on length and indexed access in multiple places.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

HFDataLoader returns HuggingFace Dataset, but EvaluateRetrieval expects dict (AttributeError: 'Dataset' object has no attribute 'keys')

1 participant