diff --git a/FlagEmbedding/abc/evaluation/utils.py b/FlagEmbedding/abc/evaluation/utils.py index 9ed3545dc..6583c22d8 100644 --- a/FlagEmbedding/abc/evaluation/utils.py +++ b/FlagEmbedding/abc/evaluation/utils.py @@ -31,8 +31,12 @@ def evaluate_mrr( k_max, top_hits = max(k_values), {} for query_id, doc_scores in results.items(): + # A score-only sort keeps dict order, so a tied relevant document can + # rank behind an earlier id. Break ties by document id, higher first. top_hits[query_id] = sorted( - doc_scores.items(), key=lambda item: item[1], reverse=True + doc_scores.items(), + key=lambda item: (item[1], item[0]), + reverse=True, )[0:k_max] for query_id in top_hits: @@ -77,7 +81,11 @@ def evaluate_recall_cap( logging.info("\n") for query_id, doc_scores in results.items(): - top_hits = sorted(doc_scores.items(), key=lambda item: item[1], reverse=True)[0:k_max] + top_hits = sorted( + doc_scores.items(), + key=lambda item: (item[1], item[0]), + reverse=True, + )[0:k_max] query_relevant_docs = [doc_id for doc_id in qrels[query_id] if qrels[query_id][doc_id] > 0] for k in k_values: retrieved_docs = [row[0] for row in top_hits[0:k] if qrels[query_id].get(row[0], 0) > 0] diff --git a/tests/test_evaluate_mrr.py b/tests/test_evaluate_mrr.py new file mode 100644 index 000000000..3bd662dd9 --- /dev/null +++ b/tests/test_evaluate_mrr.py @@ -0,0 +1,20 @@ +from FlagEmbedding.abc.evaluation.utils import evaluate_mrr, evaluate_recall_cap + + +def test_a_tied_relevant_doc_is_ranked_by_id(): + qrels = {"q": {"b": 1, "a": 0}} + results = {"q": {"a": 1.0, "b": 1.0}} + + mrr = evaluate_mrr(qrels, results, [1]) + assert mrr["MRR@1"] == 1.0 + + capped = evaluate_recall_cap(qrels, results, [1]) + assert capped["R_cap@1"] == 1.0 + + +def test_a_higher_score_still_ranks_first(): + qrels = {"q": {"b": 1}} + results = {"q": {"a": 0.2, "b": 0.9}} + + mrr = evaluate_mrr(qrels, results, [1]) + assert mrr["MRR@1"] == 1.0