speed up

src/utils/paper_retriever.py

@@ -188,15 +188,11 @@ class Retriever(object):
         return similarity
 
     def cal_related_score(
-        self, context, related_paper_id_list, entities=None, type_name="embedding"
+        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)))
-        if entities is None:
-            entities = self.api_helper.generate_entity_list(context)
-        origin_vector = self.embedding_model.encode(
-            context, convert_to_tensor=True, device=self.device
-        ).unsqueeze(0)
+        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
@@ -275,11 +271,10 @@ class Retriever(object):
                             break
             return paper_id_list
 
-    def cosine_similarity_search(self, context, k=1, type_name="embedding"):
+    def cosine_similarity_search(self, embedding, k=1, type_name="embedding"):
         """
         return related paper: list
         """
-        embedding = self.embedding_model.encode(context)
         result = self.paper_client.cosine_similarity_search(
             embedding, k, type_name=type_name
         )
@@ -506,8 +501,9 @@ class SNRetriever(Retriever):
 
     def retrieve_paper(self, bg):
         entities = []
+        embedding = self.embedding_model.encode(bg, device=self.device)
         sn_paper_id_list = self.cosine_similarity_search(
-            context=bg,
+            embedding=embedding,
             k=self.config.RETRIEVE.sn_retrieve_paper_num,
         )
         related_paper = set()
@@ -524,6 +520,7 @@ class SNRetriever(Retriever):
         related_paper = list(related_paper)
         logger.debug(f"paper num before filter: {len(related_paper)}")
         result = {
+            "embedding": embedding,
             "paper": related_paper,
             "entities": entities,
             "cocite_paper": list(cocite_id_set),
@@ -548,7 +545,7 @@ class SNRetriever(Retriever):
         related_paper_id_list = retrieve_result["paper"]
         retrieve_paper_num = len(related_paper_id_list)
         _, _, score_all_dict = self.cal_related_score(
-            bg, related_paper_id_list=related_paper_id_list, entities=entities
+            retrieve_result["embedding"], related_paper_id_list=related_paper_id_list
         )
         top_k_matrix = {}
         recall = 0
@@ -626,8 +623,9 @@ class KGRetriever(Retriever):
         retrieve_result = self.retrieve_paper(entities)
         related_paper_id_list = retrieve_result["paper"]
         retrieve_paper_num = len(related_paper_id_list)
+        embedding = self.embedding_model.encode(bg, device=self.device)
         _, _, score_all_dict = self.cal_related_score(
-            bg, related_paper_id_list=related_paper_id_list, entities=entities
+            embedding, related_paper_id_list=related_paper_id_list
         )
         top_k_matrix = {}
         recall = 0
@@ -668,8 +666,9 @@ class SNKGRetriever(Retriever):
 
     def retrieve_paper(self, bg, entities):
         sn_entities = []
+        embedding = self.embedding_model.encode(bg, device=self.device)
         sn_paper_id_list = self.cosine_similarity_search(
-            context=bg, k=self.config.RETRIEVE.sn_num_for_entity
+            embedding, k=self.config.RETRIEVE.sn_num_for_entity
         )
         related_paper = set()
         related_paper.update(sn_paper_id_list)
@@ -689,6 +688,7 @@ class SNKGRetriever(Retriever):
             related_paper = related_paper.union(cocite_id_set)
         related_paper = list(related_paper)
         result = {
+            "embedding": embedding,
             "paper": related_paper,
             "entities": entities,
             "cocite_paper": list(cocite_id_set),
@@ -717,7 +717,7 @@ class SNKGRetriever(Retriever):
         retrieve_paper_num = len(related_paper_id_list)
         logger.info("=== Begin cal related paper score ===")
         _, _, score_all_dict = self.cal_related_score(
-            bg, related_paper_id_list=related_paper_id_list, entities=entities
+            retrieve_result["embedding"], related_paper_id_list=related_paper_id_list
         )
         logger.info("=== End cal related paper score ===")
         top_k_matrix = {}