Method retrieve

index/dense/libs.py:203–227

(self, qid2emb)

        
    
    def retrieve(self, qid2emb):
        res = defaultdict(dict)
        for qid in tqdm(qid2emb):
            query_emb = qid2emb[qid]
            query_emb = query_emb.view(1, -1)
            row, col = torch.nonzero(query_emb, as_tuple=True)
            values = query_emb[tensor_to_list(row), tensor_to_list(col)]
            threshold = 0
            filtered_indexes, scores = self.numba_score_float(self.numba_index_doc_ids,
                                                                self.numba_index_doc_values,
                                                                col.cpu().numpy(),
                                                                values.cpu().numpy().astype(np.float32),
                                                                threshold=threshold,
                                                                size_collection=self.sparse_index.nb_docs())
            # threshold set to 0 by default, could be better
            filtered_indexes, scores = self.select_topk(filtered_indexes, scores, k=self.top_k)
            for id_, sc in zip(filtered_indexes, scores):
                res[str(qid)][str(self.doc_ids[id_])] = float(sc)
            

        with open(os.path.join(self.retrieval_output_path, "run.json"), "w") as f:
            json.dump(res, f)
        print("Write the retrieval result to {} successfully.".format(self.retrieval_output_path))

        return res
        

PyTorch_over_1_6 = float((torch.__version__.split('.')[1])) >= 6 and float((torch.__version__.split('.')[0])) >= 1