mteb

1
from __future__ import annotations
2

3
import pytest
4

5
from mteb.evaluation.evaluators import PairClassificationEvaluator
6

7
TOL = 0.0001
8

9

10
class TestPairClassificationEvaluator:
11
    def test_accuracy(self):
12
        scores = [6.12, 5.39, 5.28, 5.94, 6.34, 6.47, 7.88, 6.62, 8.04, 5.9]
13
        labels = [0, 0, 0, 0, 1, 0, 0, 0, 1, 0]
14
        high_score_more_similar = True
15
        acc, acc_threshold = PairClassificationEvaluator.find_best_acc_and_threshold(
16
            scores, labels, high_score_more_similar
17
        )
18
        assert acc == pytest.approx(0.9, TOL)
19
        assert acc_threshold == pytest.approx(7.95999, TOL)
20

21
    def test_f1(self):
22
        scores = [6.12, 5.39, 5.28, 5.94, 6.34, 6.47, 7.88, 6.62, 8.04, 5.9]
23
        labels = [0, 0, 0, 0, 1, 0, 0, 0, 1, 0]
24
        high_score_more_similar = True
25

26
        f1, precision, recall, f1_threshold = (
27
            PairClassificationEvaluator.find_best_f1_and_threshold(
28
                scores, labels, high_score_more_similar
29
            )
30
        )
31
        assert f1 == pytest.approx(0.66666, TOL)
32
        assert precision == pytest.approx(1.0, TOL)
33
        assert recall == pytest.approx(0.5, TOL)
34
        assert f1_threshold == pytest.approx(7.95999, TOL)
35

36
    def test_ap(self):
37
        scores = [6.12, 5.39, 5.28, 5.94, 6.34, 6.47, 7.88, 6.62, 8.04, 5.9]
38
        labels = [0, 0, 0, 0, 1, 0, 0, 0, 1, 0]
39
        high_score_more_similar = True
40
        ap = PairClassificationEvaluator.ap_score(
41
            scores, labels, high_score_more_similar
42
        )
43
        assert ap == pytest.approx(0.7, TOL)
44