Spaces:
Sleeping
Sleeping
| """Metric tests against hand-computed values. | |
| Eval code that is silently wrong is worse than no eval, so every metric is | |
| checked against a case whose answer can be worked out on paper. | |
| """ | |
| from __future__ import annotations | |
| import numpy as np | |
| import pytest | |
| from mss.evaluate import ( | |
| average_precision, | |
| evaluate_retrieval, | |
| markdown_table, | |
| ndcg_at_k, | |
| random_baseline, | |
| rank_matrix, | |
| recall_at_k, | |
| reciprocal_rank, | |
| ) | |
| def test_recall_counts_relevant_items_found(): | |
| ranking = np.array([3, 1, 4, 0, 2]) | |
| assert recall_at_k(ranking, {3}, 1) == 1.0 | |
| assert recall_at_k(ranking, {4}, 1) == 0.0 | |
| assert recall_at_k(ranking, {4}, 3) == 1.0 | |
| # two relevant, one inside the cut -> half of them recalled | |
| assert recall_at_k(ranking, {3, 2}, 3) == 0.5 | |
| def test_reciprocal_rank_uses_the_first_hit(): | |
| ranking = np.array([3, 1, 4, 0, 2]) | |
| assert reciprocal_rank(ranking, {3}) == 1.0 | |
| assert reciprocal_rank(ranking, {1}) == pytest.approx(1 / 2) | |
| assert reciprocal_rank(ranking, {2}) == pytest.approx(1 / 5) | |
| # first of several relevant items is what counts | |
| assert reciprocal_rank(ranking, {4, 2}) == pytest.approx(1 / 3) | |
| assert reciprocal_rank(ranking, set()) == 0.0 | |
| def test_average_precision_hand_computed(): | |
| ranking = np.array([0, 1, 2, 3, 4]) | |
| # relevant at positions 1 and 3 (1-based): precision 1/1 and 2/3 | |
| assert average_precision(ranking, {0, 2}) == pytest.approx((1.0 + 2 / 3) / 2) | |
| # a single relevant item makes AP equal to the reciprocal rank | |
| assert average_precision(ranking, {2}) == pytest.approx(1 / 3) | |
| def test_average_precision_penalises_unretrieved_relevants(): | |
| # only the top 3 can ever be hit here; the 4th relevant item is missing | |
| ranking = np.array([0, 1, 2]) | |
| ap = average_precision(ranking, {0, 1, 2, 99}) | |
| assert ap == pytest.approx((1.0 + 1.0 + 1.0) / 4) | |
| def test_ndcg_is_one_for_a_perfect_ranking(): | |
| ranking = np.array([0, 1, 2, 3]) | |
| assert ndcg_at_k(ranking, {0, 1}, 4) == pytest.approx(1.0) | |
| def test_ndcg_hand_computed_for_an_imperfect_ranking(): | |
| ranking = np.array([0, 1, 2, 3]) | |
| # one relevant item at position 2 -> DCG = 1/log2(3), IDCG = 1/log2(2) = 1 | |
| assert ndcg_at_k(ranking, {1}, 4) == pytest.approx(1 / np.log2(3)) | |
| def test_ndcg_ignores_relevants_beyond_the_cut(): | |
| ranking = np.array([9, 8, 7, 0]) | |
| assert ndcg_at_k(ranking, {0}, 3) == 0.0 | |
| def test_rank_matrix_orders_by_cosine(): | |
| items = np.eye(4, dtype=np.float32) | |
| queries = np.array([[0, 0, 1, 0]], dtype=np.float32) | |
| assert rank_matrix(queries, items)[0][0] == 2 | |
| def test_rank_matrix_is_deterministic_under_ties(): | |
| items = np.ones((5, 3), dtype=np.float32) | |
| queries = np.ones((1, 3), dtype=np.float32) | |
| a = rank_matrix(queries, items) | |
| b = rank_matrix(queries, items) | |
| np.testing.assert_array_equal(a, b) | |
| np.testing.assert_array_equal(a[0], np.arange(5)) | |
| def test_perfect_retrieval_scores_one(): | |
| vecs = np.eye(6, dtype=np.float32) | |
| res = evaluate_retrieval(vecs, vecs, [{i} for i in range(6)]) | |
| assert res.metrics["R@1"] == pytest.approx(1.0) | |
| assert res.metrics["MRR"] == pytest.approx(1.0) | |
| assert res.metrics["mAP"] == pytest.approx(1.0) | |
| assert res.metrics["median_rank"] == 1.0 | |
| def test_metrics_are_bounded_and_ordered(): | |
| rng = np.random.default_rng(0) | |
| q = rng.standard_normal((30, 32)).astype(np.float32) | |
| it = rng.standard_normal((80, 32)).astype(np.float32) | |
| q /= np.linalg.norm(q, axis=1, keepdims=True) | |
| it /= np.linalg.norm(it, axis=1, keepdims=True) | |
| res = evaluate_retrieval(q, it, [{i} for i in range(30)]) | |
| for key, val in res.metrics.items(): | |
| if key != "median_rank": | |
| assert 0.0 <= val <= 1.0, key | |
| # recall cannot shrink as the cut widens | |
| assert res.metrics["R@1"] <= res.metrics["R@5"] <= res.metrics["R@10"] | |
| def test_random_baseline_lands_near_theory(): | |
| n_items = 1000 | |
| res = random_baseline(n_items, [{i} for i in range(400)], seed=1) | |
| # with one relevant item among n, expected R@k is k/n | |
| assert res.metrics["R@10"] == pytest.approx(10 / n_items, abs=0.01) | |
| assert res.metrics["median_rank"] == n_items / 2 | |
| def test_markdown_table_has_a_row_per_result(): | |
| rng = np.random.default_rng(2) | |
| v = rng.standard_normal((5, 8)).astype(np.float32) | |
| v /= np.linalg.norm(v, axis=1, keepdims=True) | |
| r1 = evaluate_retrieval(v, v, [{i} for i in range(5)], name="a") | |
| r2 = random_baseline(5, [{i} for i in range(5)]) | |
| table = markdown_table([r1, r2]) | |
| assert table.count("\n") == 3 # header + separator + 2 rows | |
| assert "| a |" in table | |