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}}