hnsw-python/experiments/connectivity_experiment.py at master · generall/hnsw-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
import pickle

import numpy as np


from cat_hnsw.benchmark.runner import BaseExperiment
from cat_hnsw.hnsw import HNSW
from cat_hnsw.hnsw_cat import HNSWCat


class ConnectivityExperiment(BaseExperiment):

    def __init__(self, experiment_name):
        super().__init__(experiment_name, num_elements=10000)

    def generate_index_class(self, param):
        return HNSW('cosine', m=param, ef=self.ef)

    def load_index(self):
        with open(self.index_path, 'rb') as fr:
            hnsw_n: HNSW = pickle.load(fr)

        hnsw_n2 = HNSWCat('cosine').init_from_existing(hnsw_n)
        return hnsw_n2

    def get_mask(self, index, experiment_param, variable_param):
        all_mask = np.arange(0, index.data.shape[0]) % 100 > variable_param
        return all_mask


if __name__ == "__main__":
    experiment = ConnectivityExperiment(
        "connectivity_m0",
    )

    experiment.run_accuracy_test(
        'm0_8',
        experiment_param=8,
        variable_params=list(range(50, 99)),
        attempts_per_value=100
    )

    experiment.run_accuracy_test(
        'm0_16',
        experiment_param=16,
        variable_params=list(range(50, 99)),
        attempts_per_value=100
    )

    experiment.run_accuracy_test(
        'm0_24',
        experiment_param=24,
        variable_params=list(range(50, 99)),
        attempts_per_value=100
    )

    experiment.run_accuracy_test(
        'm0_32',
        experiment_param=32,
        variable_params=list(range(50, 99)),
        attempts_per_value=100
    )