(self, index_dir_path, retrieval_output_path, dim_voc, top_k)
| 186 | return filtered_indexes, -scores[filtered_indexes] |
| 187 | |
| 188 | def __init__(self, index_dir_path, retrieval_output_path, dim_voc, top_k): |
| 189 | self.sparse_index = IndexDictOfArray(index_dir_path, dim_voc=dim_voc) |
| 190 | self.doc_ids = pickle.load(open(os.path.join(index_dir_path, "doc_ids.pkl"), "rb")) |
| 191 | self.top_k = top_k |
| 192 | self.retrieval_output_path = retrieval_output_path |
| 193 | |
| 194 | # convert to numba |
| 195 | self.numba_index_doc_ids = numba.typed.Dict() |
| 196 | self.numba_index_doc_values = numba.typed.Dict() |
| 197 | for key, value in self.sparse_index.index_doc_id.items(): |
| 198 | self.numba_index_doc_ids[key] = value |
| 199 | for key, value in self.sparse_index.index_doc_value.items(): |
| 200 | self.numba_index_doc_values[key] = value |
| 201 | |
| 202 | |
| 203 | def retrieve(self, qid2emb): |
nothing calls this directly
no test coverage detected