From 9576b4ee3ef66d89f4fbc7780728a9f02a74092c Mon Sep 17 00:00:00 2001 From: Mohammad Hijjawi Date: Mon, 28 Sep 2026 17:03:36 +0100 Subject: [PATCH] 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. --- .../evaluation/msmarco/data_loader.py | 2 +- tests/README.md | 1 + tests/test_msmarco_data_loader.py | 23 +++++++++++++++++++ 3 files changed, 25 insertions(+), 1 deletion(-) create mode 100644 tests/test_msmarco_data_loader.py diff --git a/FlagEmbedding/evaluation/msmarco/data_loader.py b/FlagEmbedding/evaluation/msmarco/data_loader.py index 8f90191c7..688fe448d 100644 --- a/FlagEmbedding/evaluation/msmarco/data_loader.py +++ b/FlagEmbedding/evaluation/msmarco/data_loader.py @@ -179,7 +179,7 @@ def _load_remote_qrels( logging.info(f"{self.eval_name} {dataset_name} qrels saved to {save_path}") else: qrels_dict = {} - if qrels_save_path is None: + if qrels_save_path is not None: with open(qrels_save_path, "r", encoding="utf-8") as f: for line in tqdm(f.readlines(), desc="Loading qrels"): qid, _, docid, rel = line.strip().split() diff --git a/tests/README.md b/tests/README.md index 9ab0f3638..d9f448d54 100644 --- a/tests/README.md +++ b/tests/README.md @@ -10,6 +10,7 @@ This directory contains tests for the FlagEmbedding library, including compatibi - `test_tokenizer_pad_compat.py`: Tests normalization of sorted per-example tokenized records before `tokenizer.pad`, including the local BGE reranker tokenizer when available. - `test_infer_embedder_basic.py`: Tests basic functionality of BGE embedder models with a small public checkpoint. - `test_infer_reranker_basic.py`: Tests basic functionality of reranker models. +- `test_msmarco_data_loader.py`: Tests loading MSMARCO qrels from remote without a local `dataset_dir`, using a mocked download. ## Running Tests diff --git a/tests/test_msmarco_data_loader.py b/tests/test_msmarco_data_loader.py new file mode 100644 index 000000000..218ec3817 --- /dev/null +++ b/tests/test_msmarco_data_loader.py @@ -0,0 +1,23 @@ +""" +Test loading MSMARCO qrels from remote without a local ``dataset_dir``. + +The download is mocked with a tiny TREC-style qrels file, so no network access is needed. +""" + +import pytest + +from FlagEmbedding.evaluation.msmarco import MSMARCOEvalDataLoader + + +@pytest.mark.parametrize("dataset_name", ["passage", "document"]) +@pytest.mark.parametrize("split", ["dl19", "dl20"]) +def test_load_remote_qrels_without_save_dir(tmp_path, monkeypatch, dataset_name, split): + qrels_file = tmp_path / "qrels.txt" + qrels_file.write_text("1 Q0 doc-1 2\n1 Q0 doc-2 0\n2 Q0 doc-3 1\n", encoding="utf-8") + + data_loader = MSMARCOEvalDataLoader(eval_name="msmarco", cache_dir=str(tmp_path)) + monkeypatch.setattr(data_loader, "_download_file", lambda *args, **kwargs: str(qrels_file)) + + qrels = data_loader.load_qrels(dataset_name=dataset_name, split=split) + + assert dict(qrels) == {"1": {"doc-1": 2, "doc-2": 0}, "2": {"doc-3": 1}}