speed up calculate score

src/utils/paper_client.py

@@ -79,6 +79,21 @@ class PaperClient:
             logger.error(f"paper id {paper_id} get {attribute_name} failed.")
             return None
 
+    def get_papers_attribute(self, paper_id_list, attribute_name):
+        query = """
+            UNWIND $paper_ids AS paper_id
+            MATCH (p:Paper {hash_id: paper_id})
+            RETURN p.hash_id AS hash_id, p[$attribute_name] AS attributeValue
+        """
+        with self.driver.session() as session:
+            result = session.execute_read(
+                lambda tx: tx.run(
+                    query, paper_ids=paper_id_list, attribute_name=attribute_name
+                ).data()
+            )
+        paper_attributes = [record["attributeValue"] for record in result]
+        return paper_attributes
+
     def get_paper_by_attribute(self, attribute_name, anttribute_value):
         query = f"""
             MATCH (p:Paper {{{attribute_name}: '{anttribute_value}'}})

src/utils/paper_retriever.py

@@ -184,12 +184,11 @@ class Retriever(object):
         self, embedding, related_paper_id_list, type_name="embedding"
     ):
         score_1 = np.zeros((len(related_paper_id_list)))
-        score_2 = np.zeros((len(related_paper_id_list)))
+        # score_2 = np.zeros((len(related_paper_id_list)))
         origin_vector = torch.tensor(embedding).to(self.device).unsqueeze(0)
-        context_embeddings = [
-            self.paper_client.get_paper_attribute(paper_id, type_name)
-            for paper_id in related_paper_id_list
-        ]
+        context_embeddings = self.paper_client.get_papers_attribute(
+            related_paper_id_list, type_name
+        )
         if len(context_embeddings) > 0:
             context_embeddings = torch.tensor(context_embeddings).to(self.device)
             score_1 = torch.nn.functional.cosine_similarity(
@@ -198,8 +197,9 @@ class Retriever(object):
             score_1 = score_1.cpu().numpy()
             if self.config.RETRIEVE.need_normalize:
                 score_1 = score_1 / np.max(score_1)
-        score_sn_dict = dict(zip(related_paper_id_list, score_1))
-        score_en_dict = dict(zip(related_paper_id_list, score_2))
+        score_all_dict = dict(zip(related_paper_id_list, score_1))
+        # score_en_dict = dict(zip(related_paper_id_list, score_2))
+        """
         score_all_dict = dict(
             zip(
                 related_paper_id_list,
@@ -207,7 +207,8 @@ class Retriever(object):
                 + score_2 * self.config.RETRIEVE.beta,
             )
         )
-        return score_sn_dict, score_en_dict, score_all_dict
+        """
+        return {}, {}, score_all_dict
 
     def filter_related_paper(self, score_dict, top_k):
         if len(score_dict) <= top_k: