update config

configs/datasets.yaml

@@ -9,6 +9,9 @@ ARTICLE:
     summarizing_prompt: ./assets/prompt/summarizing.xml
 
 RETRIEVE:
+    retriever_name: "SNKG"
+    use_cocite: True
+    use_cluster_to_filter: True # 过滤器中使用聚类算法
     cite_type: "all_cite_id_list"
     limit_num: 100  # 限制entity对应的paper数量
     sn_num_for_entity: 5 # SN搜索的文章数量，扩充entity
@@ -19,8 +22,6 @@ RETRIEVE:
     sum_paper_num: 100  # 最多检索到的paper数量
     sn_retrieve_paper_num: 55 # 通过SN检索到的文章
     cocite_top_k: 1
-    use_cocite: True
-    use_cluster_to_filter: True # 过滤器中使用聚类算法
     need_normalize: True
     alpha: 1
     beta: 0

src/generator.py

@@ -342,7 +342,7 @@ def backtracking(config_path, ids_path, retriever_name, brainstorm_mode, use_cue
     logger.info("\nretrieve name : {}".format(retriever_name))
     logger.info("Loaded configuration:\n{}".format(OmegaConf.to_yaml(config)))
     api_helper = APIHelper(config)
-    paper_client = PaperClient(config)
+    paper_client = PaperClient()
     eval_data = []
     processed_ids = set()
     cur_num = 0

src/pages/button_interface.py

@@ -8,17 +8,15 @@ class Backend(object):
     def __init__(self) -> None:
         CONFIG_PATH = "./configs/datasets.yaml"
         EXAMPLE_PATH = "./assets/data/example.json"
-        RETRIEVER_NAME = "SNKG"
         USE_INSPIRATION = True
         BRAINSTORM_MODE = "mode_c"
 
         self.config = ConfigReader.load(CONFIG_PATH)
+        RETRIEVER_NAME = self.config.RETRIEVE.retriever_name
         self.api_helper = APIHelper(self.config)
         self.retriever_factory = RetrieverFactory.get_retriever_factory().create_retriever(
             RETRIEVER_NAME,
-            self.config,
-            use_cocite=self.config.RETRIEVE.use_cocite,
-            use_cluster_to_filter=self.config.RETRIEVE.use_cluster_to_filter,
+            self.config
         )
         self.idea_generator = IdeaGenerator(self.config, None)
         self.use_inspiration = USE_INSPIRATION

src/retriever.py

@@ -38,31 +38,16 @@ def main(ctx):
     required=True,
     help="Dataset configuration file in YAML",
 )
-@click.option(
-    "-r",
-    "--retriever-name",
-    default="SNKG",
-    type=str,
-    required=True,
-    help="Retrieve method",
-)
-@click.option(
-    "--co-cite",
-    is_flag=True,
-    help="Whether to use co-citation, defaults to False",
-)
-@click.option(
-    "--cluster-to-filter",
-    is_flag=True,
-    help="Whether to use cluster-to-filter, defaults to False",
-)
 def retrieve(
-    config_path, ids_path, retriever_name, co_cite, cluster_to_filter, **kwargs
+    config_path, ids_path, **kwargs
 ):
     check_env()
     check_embedding()
     config = ConfigReader.load(config_path, **kwargs)
     log_dir = config.DEFAULT.log_dir
+    retriever_name = config.RETRIEVE.retriever_name
+    cluster_to_filter = config.RETRIEVE.use_cluster_to_filter
+    co_cite = config.RETRIEVE.use_cocite
     if not os.path.exists(log_dir):
         os.makedirs(log_dir)
         print(f"Created log directory: {log_dir}")
@@ -73,10 +58,10 @@ def retrieve(
         ),
     )
     logger.add(log_file, level=config.DEFAULT.log_level)
-    logger.info("\nretrieve name : {}".format(retriever_name))
+    logger.info("=== Retriever name : {} ===".format(retriever_name))
     logger.info("Loaded configuration:\n{}".format(OmegaConf.to_yaml(config)))
     api_helper = APIHelper(config)
-    paper_client = PaperClient(config)
+    paper_client = PaperClient()
     precision = 0
     filtered_precision = 0
     recall = 0
@@ -90,9 +75,7 @@ def retrieve(
     # Init Retriever
     rt = RetrieverFactory.get_retriever_factory().create_retriever(
         retriever_name,
-        config,
-        use_cocite=co_cite,
-        use_cluster_to_filter=cluster_to_filter,
+        config
     )
     for line in ids_path:
         paper = json.loads(line)
@@ -108,7 +91,7 @@ def retrieve(
             entities = paper["entities"]
         else:
             entities = api_helper.generate_entity_list(bg)
-        logger.info("origin entities from background: {}".format(entities))
+        logger.info("\norigin entities from background: {}".format(entities))
         cite_type = config.RETRIEVE.cite_type
         if cite_type in paper and len(paper[cite_type]) >= 5:
             target_paper_id_list = paper[cite_type]

src/utils/hash.py

@@ -62,10 +62,12 @@ class EmbeddingModel:
     def __new__(cls, config):
         if cls._instance is None:
             cls._instance = super(EmbeddingModel, cls).__new__(cls)
+            device = "cuda" if torch.cuda.is_available() else "cpu"
             cls._instance.embedding_model = SentenceTransformer(
                 model_name_or_path=get_dir(config.DEFAULT.embedding),
-                device="cuda" if torch.cuda.is_available() else "cpu",
+                device=device,
             )
+            print(f"==== using device {device} ====")
         return cls._instance
 
 def get_embedding_model(config):

src/utils/paper_retriever.py

@@ -84,10 +84,10 @@ class Retriever(object):
     __metaclass__ = ABCMeta
     retriever_name = "BASE"
 
-    def __init__(self, config, use_cocite=False, use_cluster_to_filter=False):
+    def __init__(self, config):
         self.config = config
-        self.use_cocite = use_cocite
-        self.use_cluster_to_filter = use_cluster_to_filter
+        self.use_cocite = config.RETRIEVE.use_cocite
+        self.use_cluster_to_filter = config.RETRIEVE.use_cluster_to_filter
         self.paper_client = PaperClient()
         self.cocite = CoCite()
         self.api_helper = APIHelper(config=config)
@@ -389,7 +389,9 @@ class Retriever(object):
         logger.debug(f"target label counts : {target_label_counts}")
         target_label_list = list(target_label_counts.keys())
         max_k = max(self.config.RETRIEVE.top_k_list)
+        logger.info("=== Begin filter related paper ===")
         max_k_paper_id_list = self.filter_related_paper(score_all_dict, top_k=max_k)
+        logger.info("=== End filter related paper ===")
         for k in self.config.RETRIEVE.top_k_list:
             # 前top k 的文章
             top_k = min(k, len(max_k_paper_id_list))
@@ -507,8 +509,8 @@ class autoregister:
 
 @autoregister("SN")
 class SNRetriever(Retriever):
-    def __init__(self, config, use_cocite=False, use_cluster_to_filter=False):
-        super().__init__(config, use_cocite, use_cluster_to_filter)
+    def __init__(self, config):
+        super().__init__(config)
 
     def retrieve_paper(self, bg):
         entities = []
@@ -590,8 +592,8 @@ class SNRetriever(Retriever):
 
 @autoregister("KG")
 class KGRetriever(Retriever):
-    def __init__(self, config, use_cocite=False, use_cluster_to_filter=False):
-        super().__init__(config, use_cocite, use_cluster_to_filter)
+    def __init__(self, config):
+        super().__init__(config)
 
     def retrieve_paper(self, entities):
         new_entities = self.retrieve_entities_by_enties(entities)
@@ -669,8 +671,8 @@ class KGRetriever(Retriever):
 
 @autoregister("SNKG")
 class SNKGRetriever(Retriever):
-    def __init__(self, config, use_cocite=False, use_cluster_to_filter=False):
-        super().__init__(config, use_cocite, use_cluster_to_filter)
+    def __init__(self, config):
+        super().__init__(config)
 
     def retrieve_paper(self, bg, entities):
         sn_entities = []
@@ -721,9 +723,11 @@ class SNKGRetriever(Retriever):
         retrieve_result = self.retrieve_paper(bg, entities)
         related_paper_id_list = retrieve_result["paper"]
         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
         )
+        logger.info("=== End cal related paper score ===")
         top_k_matrix = {}
         recall = 0
         precision = 0
@@ -738,7 +742,9 @@ class SNKGRetriever(Retriever):
             logger.debug("before filter:")
             logger.debug(f"Recall: {recall:.3f}")
             logger.debug(f"Precision: {precision:.3f}")
+        logger.info("=== Begin filter related paper score ===")
         related_paper = self.filter_related_paper(score_all_dict, top_k)
+        logger.info("=== End filter related paper score ===")
         related_paper = self.update_related_paper(related_paper)
         result = {
             "recall": recall,