mteb

1
from __future__ import annotations
2

3
import pytest
4

5
from mteb.evaluation.evaluators import RerankingEvaluator
6

7
TOL = 0.0001
8

9

10
class TestRerankingEvaluator:
11
    def setup_method(self):
12
        """Setup any state tied to the execution of the given method in a class.
13

14
        setup_method is invoked for every test method of a class.
15
        """
16
        self.evaluator = RerankingEvaluator([])
17

18
    def test_mrr_at_k(self):
19
        is_relevant = [1, 1, 1, 0, 0, 0, 0, 0, 0]
20
        pred_ranking = [5, 2, 6, 1, 3, 4, 7, 8, 9]
21

22
        assert self.evaluator.mrr_at_k_score(
23
            is_relevant, pred_ranking, 10
24
        ) == pytest.approx(0.5, TOL)
25
        assert self.evaluator.mrr_at_k_score(
26
            is_relevant, pred_ranking, 3
27
        ) == pytest.approx(0.5, TOL)
28
        assert self.evaluator.mrr_at_k_score(
29
            is_relevant, pred_ranking, 1
30
        ) == pytest.approx(0, TOL)
31

32
    def test_map(self):
33
        is_relevant = [1, 1, 1, 0, 0]
34
        pred_scores = [0.75, 0.93, 0.85, 0.76, 0.75]
35

36
        assert self.evaluator.ap_score(is_relevant, pred_scores) == pytest.approx(
37
            0.86666, TOL
38
        )
39