Skip to content

Commit 9576b4e

Browse files
fix(eval): load MSMARCO remote qrels when dataset_dir is not set
When no dataset_dir is given, MSMARCOEvalDataLoader._load_remote_qrels reads qrels without saving them, but the check on qrels_save_path was inverted. For the dl19/dl20 splits (and document dev) it skipped the downloaded TREC file and iterated the undefined HF `qrels`, raising UnboundLocalError; for passage dev it tried to open(None). Flip the condition to match the save_dir branch and add a regression test that mocks the download.
1 parent fd1a2bd commit 9576b4e

3 files changed

Lines changed: 25 additions & 1 deletion

File tree

‎FlagEmbedding/evaluation/msmarco/data_loader.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -179,7 +179,7 @@ def _load_remote_qrels(
179179
logging.info(f"{self.eval_name} {dataset_name} qrels saved to {save_path}")
180180
else:
181181
qrels_dict = {}
182-
if qrels_save_path is None:
182+
if qrels_save_path is not None:
183183
with open(qrels_save_path, "r", encoding="utf-8") as f:
184184
for line in tqdm(f.readlines(), desc="Loading qrels"):
185185
qid, _, docid, rel = line.strip().split()

‎tests/README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ This directory contains tests for the FlagEmbedding library, including compatibi
1010
- `test_tokenizer_pad_compat.py`: Tests normalization of sorted per-example tokenized records before `tokenizer.pad`, including the local BGE reranker tokenizer when available.
1111
- `test_infer_embedder_basic.py`: Tests basic functionality of BGE embedder models with a small public checkpoint.
1212
- `test_infer_reranker_basic.py`: Tests basic functionality of reranker models.
13+
- `test_msmarco_data_loader.py`: Tests loading MSMARCO qrels from remote without a local `dataset_dir`, using a mocked download.
1314

1415
## Running Tests
1516

‎tests/test_msmarco_data_loader.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
"""
2+
Test loading MSMARCO qrels from remote without a local ``dataset_dir``.
3+
4+
The download is mocked with a tiny TREC-style qrels file, so no network access is needed.
5+
"""
6+
7+
import pytest
8+
9+
from FlagEmbedding.evaluation.msmarco import MSMARCOEvalDataLoader
10+
11+
12+
@pytest.mark.parametrize("dataset_name", ["passage", "document"])
13+
@pytest.mark.parametrize("split", ["dl19", "dl20"])
14+
def test_load_remote_qrels_without_save_dir(tmp_path, monkeypatch, dataset_name, split):
15+
qrels_file = tmp_path / "qrels.txt"
16+
qrels_file.write_text("1 Q0 doc-1 2\n1 Q0 doc-2 0\n2 Q0 doc-3 1\n", encoding="utf-8")
17+
18+
data_loader = MSMARCOEvalDataLoader(eval_name="msmarco", cache_dir=str(tmp_path))
19+
monkeypatch.setattr(data_loader, "_download_file", lambda *args, **kwargs: str(qrels_file))
20+
21+
qrels = data_loader.load_qrels(dataset_name=dataset_name, split=split)
22+
23+
assert dict(qrels) == {"1": {"doc-1": 2, "doc-2": 0}, "2": {"doc-3": 1}}

0 commit comments

Comments
 (0)