milvus-io /
milvus
| @@ 786-812 (lines=27) @@ | ||
| 783 | # ResultHeap |
|
| 784 | ########################################### |
|
| 785 | ||
| 786 | class ResultHeap: |
|
| 787 | """Accumulate query results from a sliced dataset. The final result will |
|
| 788 | be in self.D, self.I.""" |
|
| 789 | ||
| 790 | def __init__(self, nq, k): |
|
| 791 | " nq: number of query vectors, k: number of results per query " |
|
| 792 | self.I = np.zeros((nq, k), dtype='int64') |
|
| 793 | self.D = np.zeros((nq, k), dtype='float32') |
|
| 794 | self.nq, self.k = nq, k |
|
| 795 | heaps = float_maxheap_array_t() |
|
| 796 | heaps.k = k |
|
| 797 | heaps.nh = nq |
|
| 798 | heaps.val = swig_ptr(self.D) |
|
| 799 | heaps.ids = swig_ptr(self.I) |
|
| 800 | heaps.heapify() |
|
| 801 | self.heaps = heaps |
|
| 802 | ||
| 803 | def add_result(self, D, I): |
|
| 804 | """D, I do not need to be in a particular order (heap or sorted)""" |
|
| 805 | assert D.shape == (self.nq, self.k) |
|
| 806 | assert I.shape == (self.nq, self.k) |
|
| 807 | self.heaps.addn_with_ids( |
|
| 808 | self.k, faiss.swig_ptr(D), |
|
| 809 | faiss.swig_ptr(I), self.k) |
|
| 810 | ||
| 811 | def finalize(self): |
|
| 812 | self.heaps.reorder() |
|
| 813 | ||
| @@ 84-109 (lines=26) @@ | ||
| 81 | print('set nprobe=', nprobe) |
|
| 82 | faiss.ParameterSpace().set_index_parameter(index_1, 'nprobe', nprobe) |
|
| 83 | ||
| 84 | class ResultHeap: |
|
| 85 | """ Combine query results from a sliced dataset """ |
|
| 86 | ||
| 87 | def __init__(self, nq, k): |
|
| 88 | " nq: number of query vectors, k: number of results per query " |
|
| 89 | self.I = np.zeros((nq, k), dtype='int64') |
|
| 90 | self.D = np.zeros((nq, k), dtype='float32') |
|
| 91 | self.nq, self.k = nq, k |
|
| 92 | heaps = faiss.float_maxheap_array_t() |
|
| 93 | heaps.k = k |
|
| 94 | heaps.nh = nq |
|
| 95 | heaps.val = faiss.swig_ptr(self.D) |
|
| 96 | heaps.ids = faiss.swig_ptr(self.I) |
|
| 97 | heaps.heapify() |
|
| 98 | self.heaps = heaps |
|
| 99 | ||
| 100 | def add_batch_result(self, D, I, i0): |
|
| 101 | assert D.shape == (self.nq, self.k) |
|
| 102 | assert I.shape == (self.nq, self.k) |
|
| 103 | I += i0 |
|
| 104 | self.heaps.addn_with_ids( |
|
| 105 | self.k, faiss.swig_ptr(D), |
|
| 106 | faiss.swig_ptr(I), self.k) |
|
| 107 | ||
| 108 | def finalize(self): |
|
| 109 | self.heaps.reorder() |
|
| 110 | ||
| 111 | stats = faiss.cvar.indexIVF_stats |
|
| 112 | stats.reset() |
|
| @@ 76-101 (lines=26) @@ | ||
| 73 | return np.ascontiguousarray(x, dtype='float32') |
|
| 74 | ||
| 75 | ||
| 76 | class ResultHeap: |
|
| 77 | """ Combine query results from a sliced dataset """ |
|
| 78 | ||
| 79 | def __init__(self, nq, k): |
|
| 80 | " nq: number of query vectors, k: number of results per query " |
|
| 81 | self.I = np.zeros((nq, k), dtype='int64') |
|
| 82 | self.D = np.zeros((nq, k), dtype='float32') |
|
| 83 | self.nq, self.k = nq, k |
|
| 84 | heaps = faiss.float_maxheap_array_t() |
|
| 85 | heaps.k = k |
|
| 86 | heaps.nh = nq |
|
| 87 | heaps.val = faiss.swig_ptr(self.D) |
|
| 88 | heaps.ids = faiss.swig_ptr(self.I) |
|
| 89 | heaps.heapify() |
|
| 90 | self.heaps = heaps |
|
| 91 | ||
| 92 | def add_batch_result(self, D, I, i0): |
|
| 93 | assert D.shape == (self.nq, self.k) |
|
| 94 | assert I.shape == (self.nq, self.k) |
|
| 95 | I += i0 |
|
| 96 | self.heaps.addn_with_ids( |
|
| 97 | self.k, faiss.swig_ptr(D), |
|
| 98 | faiss.swig_ptr(I), self.k) |
|
| 99 | ||
| 100 | def finalize(self): |
|
| 101 | self.heaps.reorder() |
|
| 102 | ||
| 103 | ||
| 104 | def compute_GT_sliced(xb, xq, k): |
|
| @@ 75-100 (lines=26) @@ | ||
| 72 | return np.ascontiguousarray(x, dtype='float32') |
|
| 73 | ||
| 74 | ||
| 75 | class ResultHeap: |
|
| 76 | """ Combine query results from a sliced dataset """ |
|
| 77 | ||
| 78 | def __init__(self, nq, k): |
|
| 79 | " nq: number of query vectors, k: number of results per query " |
|
| 80 | self.I = np.zeros((nq, k), dtype='int64') |
|
| 81 | self.D = np.zeros((nq, k), dtype='float32') |
|
| 82 | self.nq, self.k = nq, k |
|
| 83 | heaps = faiss.float_maxheap_array_t() |
|
| 84 | heaps.k = k |
|
| 85 | heaps.nh = nq |
|
| 86 | heaps.val = faiss.swig_ptr(self.D) |
|
| 87 | heaps.ids = faiss.swig_ptr(self.I) |
|
| 88 | heaps.heapify() |
|
| 89 | self.heaps = heaps |
|
| 90 | ||
| 91 | def add_batch_result(self, D, I, i0): |
|
| 92 | assert D.shape == (self.nq, self.k) |
|
| 93 | assert I.shape == (self.nq, self.k) |
|
| 94 | I += i0 |
|
| 95 | self.heaps.addn_with_ids( |
|
| 96 | self.k, faiss.swig_ptr(D), |
|
| 97 | faiss.swig_ptr(I), self.k) |
|
| 98 | ||
| 99 | def finalize(self): |
|
| 100 | self.heaps.reorder() |
|
| 101 | ||
| 102 | ||
| 103 | ||
| @@ 72-97 (lines=26) @@ | ||
| 69 | ||
| 70 | ||
| 71 | ||
| 72 | class ResultHeap: |
|
| 73 | """ Combine query results from a sliced dataset (for k-nn search) """ |
|
| 74 | ||
| 75 | def __init__(self, nq, k): |
|
| 76 | " nq: number of query vectors, k: number of results per query " |
|
| 77 | self.I = np.zeros((nq, k), dtype='int64') |
|
| 78 | self.D = np.zeros((nq, k), dtype='float32') |
|
| 79 | self.nq, self.k = nq, k |
|
| 80 | heaps = faiss.float_maxheap_array_t() |
|
| 81 | heaps.k = k |
|
| 82 | heaps.nh = nq |
|
| 83 | heaps.val = faiss.swig_ptr(self.D) |
|
| 84 | heaps.ids = faiss.swig_ptr(self.I) |
|
| 85 | heaps.heapify() |
|
| 86 | self.heaps = heaps |
|
| 87 | ||
| 88 | def add_batch_result(self, D, I, i0): |
|
| 89 | assert D.shape == (self.nq, self.k) |
|
| 90 | assert I.shape == (self.nq, self.k) |
|
| 91 | I += i0 |
|
| 92 | self.heaps.addn_with_ids( |
|
| 93 | self.k, faiss.swig_ptr(D), |
|
| 94 | faiss.swig_ptr(I), self.k) |
|
| 95 | ||
| 96 | def finalize(self): |
|
| 97 | self.heaps.reorder() |
|
| 98 | ||
| 99 | def distribute_weights(weights, nbin): |
|
| 100 | """ assign a set of weights to a smaller set of bins to balance them """ |
|
