| 128 | self.arrow_writer = pa.ipc.new_stream(self.arrow_file, schema) |
| 129 | |
| 130 | def get_knns(self, queries, ignore_first=False): |
| 131 | if not self.knn_gpu: |
| 132 | queries = queries.cpu() |
| 133 | dists, knns = self.index.search(queries.to(torch.float32), self.k) |
| 134 | dists, knns = dists.to(self.device), knns.to(self.device) |
| 135 | |
| 136 | # If we need to ignore the first nearest neighbor |
| 137 | # There seems to be a bug of faiss 1.11.0 cuvs, that the searched result isn't sorted by distance, please use faiss 1.12.0 w/o cuvs instead |
| 138 | if ignore_first: |
| 139 | return dists[:,1:], knns[:,1:] |
| 140 | else: |
| 141 | return dists, knns |
| 142 | |
| 143 | def knns_to_probs(self, knns, neg_dists): |
| 144 | """Compute kNN probability distribution following the reference implementation""" |