|
| 1 | +import unittest |
| 2 | + |
| 3 | +import numpy as np |
| 4 | + |
| 5 | +import hnswlib |
| 6 | + |
| 7 | + |
| 8 | +class RandomSelfTestCase(unittest.TestCase): |
| 9 | + def testMetadata(self): |
| 10 | + |
| 11 | + dim = 16 |
| 12 | + num_elements = 10000 |
| 13 | + |
| 14 | + # Generating sample data |
| 15 | + data = np.float32(np.random.random((num_elements, dim))) |
| 16 | + |
| 17 | + # Declaring index |
| 18 | + p = hnswlib.Index(space='l2', dim=dim) # possible options are l2, cosine or ip |
| 19 | + |
| 20 | + # Initing index |
| 21 | + # max_elements - the maximum number of elements, should be known beforehand |
| 22 | + # (probably will be made optional in the future) |
| 23 | + # |
| 24 | + # ef_construction - controls index search speed/build speed tradeoff |
| 25 | + # M - is tightly connected with internal dimensionality of the data |
| 26 | + # stronlgy affects the memory consumption |
| 27 | + |
| 28 | + p.init_index(max_elements=num_elements, ef_construction=100, M=16) |
| 29 | + |
| 30 | + # Controlling the recall by setting ef: |
| 31 | + # higher ef leads to better accuracy, but slower search |
| 32 | + p.set_ef(100) |
| 33 | + |
| 34 | + p.set_num_threads(4) # by default using all available cores |
| 35 | + |
| 36 | + print("Adding all elements (%d)" % (len(data))) |
| 37 | + p.add_items(data) |
| 38 | + |
| 39 | + # test methods |
| 40 | + self.assertEqual(p.get_max_elements(), num_elements) |
| 41 | + self.assertEqual(p.get_current_count(), num_elements) |
| 42 | + |
| 43 | + # test properties |
| 44 | + self.assertEqual(p.space, 'l2') |
| 45 | + self.assertEqual(p.dim, dim) |
| 46 | + self.assertEqual(p.M, 16) |
| 47 | + self.assertEqual(p.ef_construction, 100) |
| 48 | + self.assertEqual(p.max_elements, num_elements) |
| 49 | + self.assertEqual(p.element_count, num_elements) |
0 commit comments