mteb

1
from __future__ import annotations
2

3
from typing import List
4

5
import numpy as np
6

7
from mteb.evaluation.evaluators import ClusteringEvaluator
8

9

10
class TestClusteringEvaluator:
11
    def test_clustering_v_measure(self):
12
        class Model:
13
            def encode(self, sentences: List[str], batch_size=32) -> np.ndarray:
14
                return np.eye(len(sentences))
15

16
        model = Model()
17
        sentences = ["dog walked home", "cat walked home", "robot walked to the park"]
18
        clusterer = ClusteringEvaluator(sentences=sentences, labels=[1, 2, 3])
19
        result = clusterer(model)
20

21
        assert result == {"v_measure": 1.0}
22