fix bug

configs/datasets.yaml

@@ -14,10 +14,10 @@ RETRIEVE:
     use_cluster_to_filter: False # 过滤器中使用聚类算法
     cite_type: "all_cite_id_list"
     limit_num: 100  # 限制entity对应的paper数量
-    sn_num_for_entity: 5 # SN搜索的文章数量，扩充entity
+    sn_num_for_entity: 3 # SN搜索的文章数量，扩充entity
     kg_jump_num: 1   # 跳数
     kg_cover_num: 3  # entity重合数量
-    sum_paper_num: 30  # 最多检索到的paper数量
+    sum_paper_num: 50  # 最多检索到的paper数量
     sn_retrieve_paper_num: 55 # 通过SN检索到的文章
     cocite_top_k: 1
     need_normalize: True

src/retriever.py

@@ -26,7 +26,7 @@ def main(ctx):
 @click.option(
     "-c",
     "--config-path",
-    default="../configs/datasets.yaml",
+    default="./configs/datasets.yaml",
     type=click.File(),
     required=True,
     help="Dataset configuration file in YAML",

src/utils/paper_client.py

@@ -130,7 +130,6 @@ class PaperClient:
                         related_entities.add(entity)
 
             return list(related_entities)
-
         related_entities = bfs_query(entity_name, n, k)
         if entity_name in related_entities:
             related_entities.remove(entity_name)
@@ -541,6 +540,7 @@ class PaperClient:
         data = {"nodes": [], "relationships": []}
         query = """
             MATCH (e:Entity)-[r:RELATED_TO]->(p:Paper)
+            WHERE p.venue_name='iclr' and p.year='2024'
             RETURN p, e, r
         """
         results = graph.run(query)
@@ -572,6 +572,7 @@ class PaperClient:
             WHERE p.venue_name='acl' and p.year='2024'
             RETURN p
         """
+        """
         results = graph.run(query)
         for record in tqdm(results):
             paper_node = record["p"]
@@ -581,6 +582,7 @@ class PaperClient:
                 "label": "Paper",
                 "properties": dict(paper_node)
             })
+        """
         # 去除重复节点
         # data["nodes"] = [dict(t) for t in {tuple(d.items()) for d in data["nodes"]}]
         unique_nodes = []

src/utils/paper_retriever.py

@@ -124,7 +124,7 @@ class Retriever(object):
         )
         sum_paper_num = 0
         for key, value in entity_paper_num_dict.items():
-            if sum_paper_num <= 100:
+            if sum_paper_num <= self.config.RETRIEVE.sum_paper_num:
                 sum_paper_num += value
                 new_entities.append(key)
             elif (
@@ -188,35 +188,27 @@ class Retriever(object):
         return similarity
 
     def cal_related_score(
-        self, context, related_paper_id_list, entities=None, type_name="motivation"
+        self, context, related_paper_id_list, entities=None, 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)
-            logger.debug("get entity from context: {}".format(entities))
         origin_vector = self.embedding_model.encode(
             context, convert_to_tensor=True, device=self.device
         ).unsqueeze(0)
-        related_contexts = [
+        context_embeddings = [
             self.paper_client.get_paper_attribute(paper_id, type_name)
             for paper_id in related_paper_id_list
         ]
-        if len(related_contexts) > 0:
-            context_embeddings = self.embedding_model.encode(
-                related_contexts,
-                batch_size=512,
-                convert_to_tensor=True,
-                device=self.device,
-            )
+        if len(context_embeddings) > 0:
+            context_embeddings = torch.tensor(context_embeddings).to(self.device)
             score_1 = torch.nn.functional.cosine_similarity(
                 origin_vector, context_embeddings
             )
             score_1 = score_1.cpu().numpy()
             if self.config.RETRIEVE.need_normalize:
                 score_1 = score_1 / np.max(score_1)
-        # score_2 not enable
-        # if self.config.RETRIEVE.beta != 0:
         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(