reduce neo4j query time in retrieve

app.py

@@ -1,25 +1,32 @@
 import sys
+
 sys.path.append("./src")
 import streamlit as st
-from app_pages import button_interface, step_by_step_generation, one_click_generation, homepage
+from app_pages import (
+    button_interface,
+    step_by_step_generation,
+    one_click_generation,
+    homepage,
+)
 from app_pages.locale import _
-from utils.hash import check_env, check_embedding
 
 if __name__ == "__main__":
-    check_env()
-    check_embedding() 
     backend = button_interface.Backend()
     # backend = None
     st.set_page_config(layout="wide")
-    if "language" not in st.session_state:
-        st.session_state["language"] = "zh"
+
+    # st.logo("./assets/pic/logo.jpg", size="large")
     def fn1():
         one_click_generation.one_click_generation(backend)
+
     def fn2():
         step_by_step_generation.step_by_step_generation(backend)
-    pg = st.navigation([
-        st.Page(homepage.home_page, title=_("🏠️ Homepage")),
-        st.Page(fn1, title=_("💧 One-click Generation")),
-        st.Page(fn2, title=_("💦 Step-by-step Generation")),
-    ])
-    pg.run()
+
+    pg = st.navigation(
+        [
+            st.Page(homepage.home_page, title=_("🏠️ Homepage")),
+            st.Page(fn1, title=_("💧 One-click Generation")),
+            st.Page(fn2, title=_("💦 Step-by-step Generation")),
+        ]
+    )
+    pg.run()

src/app_pages/button_interface.py

@@ -2,8 +2,10 @@ import json
 from utils.paper_retriever import RetrieverFactory
 from utils.llms_api import APIHelper
 from utils.header import ConfigReader
+from utils.hash import check_env, check_embedding
 from generator import IdeaGenerator
 
+
 class Backend(object):
     def __init__(self) -> None:
         CONFIG_PATH = "./configs/datasets.yaml"
@@ -12,11 +14,14 @@ class Backend(object):
         BRAINSTORM_MODE = "mode_c"
 
         self.config = ConfigReader.load(CONFIG_PATH)
+        check_env()
+        check_embedding(self.config.DEFAULT.embedding)
         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
+        self.retriever_factory = (
+            RetrieverFactory.get_retriever_factory().create_retriever(
+                RETRIEVER_NAME, self.config
+            )
         )
         self.idea_generator = IdeaGenerator(self.config, None)
         self.use_inspiration = USE_INSPIRATION
@@ -33,14 +38,14 @@ class Backend(object):
             return []
 
     def background2brainstorm_callback(self, background, json_strs=None):
-        if json_strs is not None: # only for DEBUG_MODE
+        if json_strs is not None:  # only for DEBUG_MODE
             json_contents = json.loads(json_strs)
             return json_contents["brainstorm"]
         else:
             return self.api_helper.generate_brainstorm(background)
 
     def brainstorm2entities_callback(self, background, brainstorm, json_strs=None):
-        if json_strs is not None: # only for DEBUG_MODE
+        if json_strs is not None:  # only for DEBUG_MODE
             json_contents = json.loads(json_strs)
             entities_bg = json_contents["entities_bg"]
             entities_bs = json_contents["entities_bs"]
@@ -71,13 +76,17 @@ class Backend(object):
             for i, p in enumerate(result["related_paper"]):
                 res.append(str(p))
         else:
-            result = self.retriever_factory.retrieve(background, entities, need_evaluate=False, target_paper_id_list=[])
+            result = self.retriever_factory.retrieve(
+                background, entities, need_evaluate=False, target_paper_id_list=[]
+            )
             res = []
             for i, p in enumerate(result["related_paper"]):
                 res.append(f'{p["title"]}. {p["venue_name"].upper()} {p["year"]}.')
         return res, result["related_paper"]
 
-    def literature2initial_ideas_callback(self, background, brainstorms, retrieved_literature, json_strs=None):
+    def literature2initial_ideas_callback(
+        self, background, brainstorms, retrieved_literature, json_strs=None
+    ):
         if json_strs is not None:
             json_contents = json.loads(json_strs)
             return json_contents["median"]["initial_idea"]
@@ -86,15 +95,16 @@ class Backend(object):
             self.idea_generator.brainstorm = brainstorms
             if self.use_inspiration:
                 message_input, idea_modified, median = (
-                self.idea_generator.generate_by_inspiration(
-                    background, "new_idea", self.brainstorm_mode, False)
+                    self.idea_generator.generate_by_inspiration(
+                        background, "new_idea", self.brainstorm_mode, False
+                    )
                 )
             else:
                 message_input, idea_modified, median = self.idea_generator.generate(
                     background, "new_idea", self.brainstorm_mode, False
                 )
             return median["initial_idea"], idea_modified
-        
+
     def initial2final_callback(self, initial_ideas, final_ideas, json_strs=None):
         if json_strs is not None:
             json_contents = json.loads(json_strs)
@@ -107,6 +117,7 @@ class Backend(object):
             return self.examples[i].get("background", "Background not found.")
         else:
             return "Example not found. Please select a valid index."
+
     #     return ("The application scope of large-scale language models such as GPT-4 and LLaMA "
     # "has rapidly expanded, demonstrating powerful capabilities in natural language processing "
     # "and multimodal tasks. However, as the size and complexity of the models increase, understanding "

src/generator.py

@@ -10,6 +10,7 @@ import warnings
 import time
 import os
 from utils.hash import check_env, check_embedding
+
 warnings.filterwarnings("ignore")
 
 
@@ -24,9 +25,14 @@ def extract_problem(problem, background):
         research_problem = background
     return research_problem
 
+
 class IdeaGenerator:
     def __init__(
-        self, config, paper_list: list[dict] = [], cue_words: list = None, brainstorm: str = None
+        self,
+        config,
+        paper_list: list[dict] = [],
+        cue_words: list = None,
+        brainstorm: str = None,
     ) -> None:
         self.api_helper = APIHelper(config)
         self.paper_list = paper_list
@@ -58,7 +64,9 @@ class IdeaGenerator:
         idea = self.api_helper.generate_idea_with_cue_words(
             problem, self.paper_list, self.cue_words
         )
-        idea_filtered = self.api_helper.integrate_idea(background, self.brainstorm, idea)
+        idea_filtered = self.api_helper.integrate_idea(
+            background, self.brainstorm, idea
+        )
         return message_input, problem, idea, idea_filtered
 
     def generate_without_cue_words_bs(self, background: str):
@@ -66,7 +74,9 @@ class IdeaGenerator:
             background, self.paper_list
         )
         idea = self.api_helper.generate_idea(problem, self.paper_list)
-        idea_filtered = self.api_helper.integrate_idea(background, self.brainstorm, idea)
+        idea_filtered = self.api_helper.integrate_idea(
+            background, self.brainstorm, idea
+        )
         return message_input, problem, idea, idea_filtered
 
     def generate_with_cue_words_ins(self, background: str):
@@ -93,16 +103,12 @@ class IdeaGenerator:
         research_problem = extract_problem(problem, background)
         inspirations = []
         for paper in self.paper_list:
-            inspiration = self.api_helper.generate_inspiration(
-                research_problem, paper
-            )
+            inspiration = self.api_helper.generate_inspiration(research_problem, paper)
             inspirations.append(inspiration)
-        idea = self.api_helper.generate_idea_by_inspiration(
-            problem, inspirations
-        )
+        idea = self.api_helper.generate_idea_by_inspiration(problem, inspirations)
         idea_filtered = self.api_helper.filter_idea(idea, background)
         return message_input, problem, inspirations, idea, idea_filtered
-    
+
     def generate_with_cue_words_ins_bs(self, background: str):
         problem, message_input = self.api_helper.generate_problem_with_cue_words(
             background, self.paper_list, self.cue_words
@@ -117,7 +123,9 @@ class IdeaGenerator:
         idea = self.api_helper.generate_idea_by_inspiration_with_cue_words(
             problem, inspirations, self.cue_words
         )
-        idea_filtered = self.api_helper.integrate_idea(background, self.brainstorm, idea)
+        idea_filtered = self.api_helper.integrate_idea(
+            background, self.brainstorm, idea
+        )
         return message_input, problem, inspirations, idea, idea_filtered
 
     def generate_without_cue_words_ins_bs(self, background: str):
@@ -127,14 +135,12 @@ class IdeaGenerator:
         research_problem = extract_problem(problem, background)
         inspirations = []
         for paper in self.paper_list:
-            inspiration = self.api_helper.generate_inspiration(
-                research_problem, paper
-            )
+            inspiration = self.api_helper.generate_inspiration(research_problem, paper)
             inspirations.append(inspiration)
-        idea = self.api_helper.generate_idea_by_inspiration(
-            problem, inspirations
+        idea = self.api_helper.generate_idea_by_inspiration(problem, inspirations)
+        idea_filtered = self.api_helper.integrate_idea(
+            background, self.brainstorm, idea
         )
-        idea_filtered = self.api_helper.integrate_idea(background, self.brainstorm, idea)
         return message_input, problem, inspirations, idea, idea_filtered
 
     def generate(
@@ -151,44 +157,34 @@ class IdeaGenerator:
             mode_name = "Generate new idea"
         if bs_mode == "mode_a":
             if use_cue_words:
-                logger.info("{} using brainstorm_mode_a with cue words.".format(mode_name))
-                (
-                    message_input, 
-                    problem, 
-                    idea, 
-                    idea_filtered
-                ) = (
+                logger.info(
+                    "{} using brainstorm_mode_a with cue words.".format(mode_name)
+                )
+                (message_input, problem, idea, idea_filtered) = (
                     self.generate_with_cue_words(background)
                 )
             else:
-                logger.info("{} using brainstorm_mode_a without cue words.".format(mode_name))
-                (
-                    message_input, 
-                    problem, 
-                    idea, 
-                    idea_filtered
-                ) = (
+                logger.info(
+                    "{} using brainstorm_mode_a without cue words.".format(mode_name)
+                )
+                (message_input, problem, idea, idea_filtered) = (
                     self.generate_without_cue_words(background)
                 )
         elif bs_mode == "mode_b" or bs_mode == "mode_c":
             if use_cue_words:
-                logger.info("{} using brainstorm_{} with cue words.".format(mode_name, bs_mode))
-                (
-                    message_input, 
-                    problem, 
-                    idea, 
-                    idea_filtered
-                ) = (
+                logger.info(
+                    "{} using brainstorm_{} with cue words.".format(mode_name, bs_mode)
+                )
+                (message_input, problem, idea, idea_filtered) = (
                     self.generate_with_cue_words_bs(background)
                 )
             else:
-                logger.info("{} using brainstorm_{} without cue words.".format(mode_name, bs_mode))
-                (
-                    message_input, 
-                    problem, 
-                    idea, 
-                    idea_filtered
-                ) = (
+                logger.info(
+                    "{} using brainstorm_{} without cue words.".format(
+                        mode_name, bs_mode
+                    )
+                )
+                (message_input, problem, idea, idea_filtered) = (
                     self.generate_without_cue_words_bs(background)
                 )
 
@@ -214,48 +210,34 @@ class IdeaGenerator:
             mode_name = "Generate new idea"
         if bs_mode == "mode_a":
             if use_cue_words:
-                logger.info("{} using brainstorm_mode_a with cue words.".format(mode_name))
-                (
-                    message_input, 
-                    problem,
-                    inspirations, 
-                    idea, 
-                    idea_filtered
-                ) = (
+                logger.info(
+                    "{} using brainstorm_mode_a with cue words.".format(mode_name)
+                )
+                (message_input, problem, inspirations, idea, idea_filtered) = (
                     self.generate_with_cue_words_ins(background)
                 )
             else:
-                logger.info("{} using brainstorm_mode_a without cue words.".format(mode_name))
-                (
-                    message_input, 
-                    problem, 
-                    inspirations,
-                    idea, 
-                    idea_filtered
-                ) = (
+                logger.info(
+                    "{} using brainstorm_mode_a without cue words.".format(mode_name)
+                )
+                (message_input, problem, inspirations, idea, idea_filtered) = (
                     self.generate_without_cue_words_ins(background)
                 )
         elif bs_mode == "mode_b" or bs_mode == "mode_c":
             if use_cue_words:
-                logger.info("{} using brainstorm_{} with cue words.".format(mode_name, bs_mode))
-                (
-                    message_input, 
-                    problem, 
-                    inspirations,
-                    idea, 
-                    idea_filtered
-                ) = (
+                logger.info(
+                    "{} using brainstorm_{} with cue words.".format(mode_name, bs_mode)
+                )
+                (message_input, problem, inspirations, idea, idea_filtered) = (
                     self.generate_with_cue_words_ins_bs(background)
                 )
             else:
-                logger.info("{} using brainstorm_{} without cue words.".format(mode_name, bs_mode))
-                (
-                    message_input, 
-                    problem, 
-                    inspirations,
-                    idea, 
-                    idea_filtered
-                ) = (
+                logger.info(
+                    "{} using brainstorm_{} without cue words.".format(
+                        mode_name, bs_mode
+                    )
+                )
+                (message_input, problem, inspirations, idea, idea_filtered) = (
                     self.generate_without_cue_words_ins_bs(background)
                 )
 
@@ -330,9 +312,18 @@ def main(ctx):
     required=False,
     help="The number of papers you want to process",
 )
-def backtracking(config_path, ids_path, retriever_name, brainstorm_mode, use_cue_words, use_inspiration, num, **kwargs):
+def backtracking(
+    config_path,
+    ids_path,
+    retriever_name,
+    brainstorm_mode,
+    use_cue_words,
+    use_inspiration,
+    num,
+    **kwargs,
+):
     check_env()
-    check_embedding() 
+    check_embedding()
     # Configuration
     config = ConfigReader.load(config_path, **kwargs)
     logger.add(
@@ -349,7 +340,10 @@ def backtracking(config_path, ids_path, retriever_name, brainstorm_mode, use_cue
     batch_size = 2
     output_dir = "./assets/output_idea/"
     os.makedirs(output_dir, exist_ok=True)
-    output_file = os.path.join(output_dir, f"output_backtracking_{brainstorm_mode}_cue_{use_cue_words}_ins_{use_inspiration}.json")
+    output_file = os.path.join(
+        output_dir,
+        f"output_backtracking_{brainstorm_mode}_cue_{use_cue_words}_ins_{use_inspiration}.json",
+    )
     if os.path.exists(output_file):
         with open(output_file, "r", encoding="utf-8") as f:
             try:
@@ -388,7 +382,7 @@ def backtracking(config_path, ids_path, retriever_name, brainstorm_mode, use_cue
         if brainstorm_mode == "mode_c":
             entities_bs = api_helper.generate_entity_list(brainstorm, 10)
             logger.debug("Original entities from brainstorm: {}".format(entities_bs))
-            entities_all = list(set(entities)|set(entities_bs))
+            entities_all = list(set(entities) | set(entities_bs))
         else:
             entities_bs = None
             entities_all = entities
@@ -404,8 +398,7 @@ def backtracking(config_path, ids_path, retriever_name, brainstorm_mode, use_cue
             continue
         # 3. 检索相关论文
         rt = RetrieverFactory.get_retriever_factory().create_retriever(
-            retriever_name, 
-            config
+            retriever_name, config
         )
         result = rt.retrieve(
             bg, entities_all, need_evaluate=False, target_paper_id_list=[]
@@ -438,7 +431,7 @@ def backtracking(config_path, ids_path, retriever_name, brainstorm_mode, use_cue
                 "hash_id": paper["hash_id"],
                 "background": bg,
                 "entities_bg": entities,
-                "brainstorm" : brainstorm,
+                "brainstorm": brainstorm,
                 "entities_bs": entities_bs,
                 "entities_rt": entities_rt,
                 "related_paper": [p["hash_id"] for p in related_paper],
@@ -467,6 +460,7 @@ def backtracking(config_path, ids_path, retriever_name, brainstorm_mode, use_cue
     ) as f:
         json.dump(eval_data, f, ensure_ascii=False, indent=4)
 
+
 @main.command()
 @click.option(
     "-c",
@@ -512,9 +506,16 @@ def backtracking(config_path, ids_path, retriever_name, brainstorm_mode, use_cue
     required=False,
     help="The number of data you want to process",
 )
-def new_idea(config_path, ids_path, retriever_name, brainstorm_mode, use_inspiration, num, **kwargs):
+def new_idea(
+    config_path,
+    ids_path,
+    retriever_name,
+    brainstorm_mode,
+    use_inspiration,
+    num,
+    **kwargs,
+):
     check_env()
-    check_embedding()
     logger.add(
         "log/generate_{}_{}.log".format(time.time(), retriever_name), level="DEBUG"
     )  # 添加文件输出
@@ -522,6 +523,7 @@ def new_idea(config_path, ids_path, retriever_name, brainstorm_mode, use_inspira
     # Configuration
     config = ConfigReader.load(config_path, **kwargs)
     api_helper = APIHelper(config)
+    check_embedding(config.DEFAULT.embedding)
     eval_data = []
     cur_num = 0
     data_num = 0
@@ -529,7 +531,9 @@ def new_idea(config_path, ids_path, retriever_name, brainstorm_mode, use_inspira
     bg_ids = set()
     output_dir = "./assets/output_idea/"
     os.makedirs(output_dir, exist_ok=True)
-    output_file = os.path.join(output_dir, f"output_new_idea_{brainstorm_mode}_ins_{use_inspiration}.json")
+    output_file = os.path.join(
+        output_dir, f"output_new_idea_{brainstorm_mode}_ins_{use_inspiration}.json"
+    )
     if os.path.exists(output_file):
         with open(output_file, "r", encoding="utf-8") as f:
             try:
@@ -538,7 +542,7 @@ def new_idea(config_path, ids_path, retriever_name, brainstorm_mode, use_inspira
                 cur_num = len(eval_data)
             except json.JSONDecodeError:
                 eval_data = []
-    print(f"{cur_num} datas have been processed.")
+    logger.debug(f"{cur_num} datas have been processed.")
     for line in ids_path:
         # 解析每行的JSON数据
         data = json.loads(line)
@@ -568,16 +572,17 @@ def new_idea(config_path, ids_path, retriever_name, brainstorm_mode, use_inspira
         if brainstorm_mode == "mode_c":
             entities_bs = api_helper.generate_entity_list(brainstorm, 10)
             logger.debug("Original entities from brainstorm: {}".format(entities_bs))
-            entities_all = list(set(entities)|set(entities_bs))
+            entities_all = list(set(entities) | set(entities_bs))
         else:
             entities_bs = None
             entities_all = entities
         # 2. 检索相关论文
         rt = RetrieverFactory.get_retriever_factory().create_retriever(
-            retriever_name,
-            config
+            retriever_name, config
+        )
+        result = rt.retrieve(
+            bg, entities_all, need_evaluate=False, target_paper_id_list=[]
         )
-        result = rt.retrieve(bg, entities_all, need_evaluate=False, target_paper_id_list=[])
         related_paper = result["related_paper"]
         logger.info("Find {} related papers...".format(len(related_paper)))
         entities_rt = result["entities"]
@@ -597,7 +602,7 @@ def new_idea(config_path, ids_path, retriever_name, brainstorm_mode, use_inspira
             {
                 "background": bg,
                 "entities_bg": entities,
-                "brainstorm" : brainstorm,
+                "brainstorm": brainstorm,
                 "entities_bs": entities_bs,
                 "entities_rt": entities_rt,
                 "related_paper": [p["hash_id"] for p in related_paper],
@@ -621,5 +626,6 @@ def new_idea(config_path, ids_path, retriever_name, brainstorm_mode, use_inspira
     with open(output_file, "w", encoding="utf-8") as f:
         json.dump(eval_data, f, ensure_ascii=False, indent=4)
 
+
 if __name__ == "__main__":
     main()

src/paper_manager.py

@@ -389,10 +389,8 @@ class PaperManager:
                     )
 
         if need_summary:
-            print(paper.keys())
             if not self.check_parse(paper):
                 logger.error(f"paper {paper['hash_id']} need parse first...")
-
             result = self.api_helper(
                 paper["title"], paper["abstract"], paper["introduction"]
             )
@@ -628,9 +626,11 @@ class PaperManager:
 
     def insert_embedding(self, hash_id=None):
         self.paper_client.add_paper_abstract_embedding(self.embedding_model, hash_id)
-        # self.client.add_paper_bg_embedding(self.embedding_model, hash_id)
-        # self.client.add_paper_contribution_embedding(self.embedding_model, hash_id)
-        # self.client.add_paper_summary_embedding(self.embedding_model, hash_id)
+        # self.paper_client.add_paper_bg_embedding(self.embedding_model, hash_id)
+        # self.paper_client.add_paper_contribution_embedding(
+        #     self.embedding_model, hash_id
+        # )
+        # self.paper_client.add_paper_summary_embedding(self.embedding_model, hash_id)
 
     def cosine_similarity_search(self, data_type, context, k=1):
         """
@@ -837,8 +837,9 @@ def local(config_path, year, venue_name, output, **kwargs):
         os.makedirs(os.path.dirname(output_path))
     config = ConfigReader.load(config_path, output_path=output_path, **kwargs)
     pm = PaperManager(config, venue_name, year)
+    print("###")
     pm.update_paper_from_json_to_json(
-        need_download=True, need_parse=True, need_summary=True, need_ground_truth=True
+        need_download=True, need_parse=True, need_summary=True
     )
 
 

src/retriever.py

@@ -41,9 +41,9 @@ def main(ctx):
 def retrieve(
     config_path, ids_path, **kwargs
 ):
-    check_env()
-    check_embedding()
     config = ConfigReader.load(config_path, **kwargs)
+    check_embedding(config.DEFAULT.embedding)
+    check_env()
     log_dir = config.DEFAULT.log_dir
     retriever_name = config.RETRIEVE.retriever_name
     cluster_to_filter = config.RETRIEVE.use_cluster_to_filter

src/utils/api/__init__.py

@@ -22,8 +22,10 @@ Creation Date : 2024-10-29
 
 Author : Frank Kang([email protected])
 """
+
 from .base_helper import HelperCompany
 from .openai_helper import OpenAIHelper  # noqa: F401, ensure autoregister
 from .zhipuai_helper import ZhipuAIHelper  # noqa: F401, ensure autoregister
+from .local_helper import LocalHelper  # noqa: F401, ensure autoregister
 
 __all__ = ["HelperCompany"]

src/utils/api/base_helper.py

@@ -17,6 +17,9 @@ from abc import ABCMeta
 from typing_extensions import Literal, override
 from ..base_company import BaseCompany
 from typing import Union
+import requests
+import json
+from requests.exceptions import RequestException
 
 
 class NotGiven:
@@ -109,6 +112,31 @@ class BaseHelper:
         self.base_url = base_url
         self.client = None
 
+    def apply_for_service(self, data_param, max_attempts=4):
+        attempt = 0
+        while attempt < max_attempts:
+            try:
+                # print(f"尝试 #{attempt + 1}")
+                r = requests.post(
+                    self.base_url + "/llm/generate",
+                    headers={"Content-Type": "application/json"},
+                    data=json.dumps(data_param),
+                )
+                # 检查请求是否成功
+                if r.status_code == 200:
+                    # print("服务请求成功。")
+                    response = r.json()["data"]["output"]
+                    return response  # 或者根据需要返回其他内容
+                else:
+                    print("服务请求失败，响应状态码：", response.status_code)
+            except RequestException as e:
+                print("请求发生错误：", e)
+
+            attempt += 1
+            if attempt == max_attempts:
+                print("达到最大尝试次数，服务请求失败。")
+                return None  # 或者根据你的情况抛出异常
+
     def create(
         self,
         *args,
@@ -124,7 +152,7 @@ class BaseHelper:
         extra_headers: None | NotGiven = None,
         extra_body: None | NotGiven = None,
         timeout: float | None | NotGiven = None,
-        **kwargs
+        **kwargs,
     ):
         """
         Creates a model response for the given chat conversation.
@@ -187,20 +215,44 @@ class BaseHelper:
 
             timeout: Override the client-level default timeout for this request, in seconds
         """
-        return self.client.chat.completions.create(
-            *args,
-            model=self.model,
-            messages=messages,
-            stream=stream,
-            temperature=temperature,
-            top_p=top_p,
-            max_tokens=max_tokens,
-            seed=seed,
-            stop=stop,
-            tools=tools,
-            tool_choice=tool_choice,
-            extra_headers=extra_headers,
-            extra_body=extra_body,
-            timeout=timeout,
-            **kwargs
-        )
+        if self.model != "local":
+            return (
+                self.client.chat.completions.create(
+                    *args,
+                    model=self.model,
+                    messages=messages,
+                    stream=stream,
+                    temperature=temperature,
+                    top_p=top_p,
+                    max_tokens=max_tokens,
+                    seed=seed,
+                    stop=stop,
+                    tools=tools,
+                    tool_choice=tool_choice,
+                    extra_headers=extra_headers,
+                    extra_body=extra_body,
+                    timeout=timeout,
+                    **kwargs,
+                )
+                .choices[0]
+                .message.content
+            )
+        else:
+            default_system = "You are a helpful assistant."
+            input_content = ""
+            for message in messages:
+                if message["role"] == "system":
+                    default_system = message["content"]
+                else:
+                    input_content += message["content"]
+            data_param = {}
+            data_param["input"] = input_content
+            data_param["serviceParams"] = {"stream": False, "system": default_system}
+            data_param["ModelParams"] = {
+                "temperature": 0.8,
+                "presence_penalty": 2.0,
+                "frequency_penalty": 0.0,
+                "top_p": 0.8,
+            }
+            response = self.apply_for_service(data_param)
+            return response

src/utils/api/local_helper.py

@@ -0,0 +1,39 @@
+r"""_summary_
+-*- coding: utf-8 -*-
+
+Module : data.utils.api.zhipuai_helper
+
+File Name : zhipuai_helper.py
+
+Description : Helper class for ZhipuAI interface, generally not used directly.
+    For example:
+    ```
+    from data.utils.api import HelperCompany
+    helper = HelperCompany.get()['ZhipuAI']
+    ...
+    ```
+   
+Creation Date : 2024-11-28
+
+Author : lihuigu([email protected])
+"""
+
+from .base_helper import register_helper, BaseHelper
+
+
+@register_helper("Local")
+class LocalHelper(BaseHelper):
+    """_summary_
+
+    Helper class for ZhipuAI interface, generally not used directly.
+
+    For example:
+    ```
+    from data.utils.api import HelperCompany
+    helper = HelperCompany.get()['Local']
+    ...
+    ```
+    """
+
+    def __init__(self, api_key, model, base_url=None, timeout=None):
+        super().__init__(api_key, model, base_url)

src/utils/hash.py

@@ -12,18 +12,35 @@ ENV_CHECKED = False
 EMBEDDING_CHECKED = False
 
 
-def check_embedding():
+def check_embedding(repo_id):
+    print("=== check embedding model ===")
     global EMBEDDING_CHECKED
     if not EMBEDDING_CHECKED:
         # Define the repository and files to download
-        repo_id = "sentence-transformers/all-MiniLM-L6-v2"  # "BAAI/bge-small-en-v1.5"
         local_dir = f"./assets/model/{repo_id}"
-        files_to_download = [
-            "config.json",
-            "pytorch_model.bin",
-            "tokenizer_config.json",
-            "vocab.txt",
-        ]
+        if repo_id in [
+            "sentence-transformers/all-MiniLM-L6-v2",
+            "BAAI/bge-small-en-v1.5",
+            "BAAAI/llm_embedder",
+        ]:
+            # repo_id = "sentence-transformers/all-MiniLM-L6-v2"
+            # repo_id = "BAAI/bge-small-en-v1.5"
+            files_to_download = [
+                "config.json",
+                "pytorch_model.bin",
+                "tokenizer_config.json",
+                "vocab.txt",
+            ]
+        elif repo_id in ["Alibaba-NLP/gte-base-en-v1.5"]:
+            files_to_download = [
+                "config.json",
+                "model.safetensors",
+                "modules.json",
+                "tokenizer.json",
+                "sentence_bert_config.json",
+                "tokenizer_config.json",
+                "vocab.txt",
+            ]
         # Download each file and save it to the /model/bge directory
         for file_name in files_to_download:
             if not os.path.exists(os.path.join(local_dir, file_name)):
@@ -47,12 +64,15 @@ def check_env():
             "NEO4J_PASSWD",
             "MODEL_NAME",
             "MODEL_TYPE",
-            "MODEL_API_KEY",
             "BASE_URL",
         ]
         for env_name in env_name_list:
             if env_name not in os.environ or os.environ[env_name] == "":
                 raise ValueError(f"{env_name} is not set...")
+        if os.environ["MODEL_TYPE"] != "Local":
+            env_name = "MODEL_API_KEY"
+            if env_name not in os.environ or os.environ[env_name] == "":
+                raise ValueError(f"{env_name} is not set...")
         ENV_CHECKED = True
 
 
@@ -61,16 +81,21 @@ class EmbeddingModel:
 
     def __new__(cls, config):
         if cls._instance is None:
+            local_dir = f"./assets/model/{config.DEFAULT.embedding}"
             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),
+                model_name_or_path=get_dir(local_dir),
                 device=device,
+                trust_remote_code=True,
             )
             print(f"==== using device {device} ====")
         return cls._instance
 
+
 def get_embedding_model(config):
+    print("=== get embedding model ===")
+    check_embedding(config.DEFAULT.embedding)
     return EmbeddingModel(config).embedding_model
 
 

src/utils/llms_api.py

@@ -49,7 +49,10 @@ class APIHelper(object):
     def get_helper(self):
         MODEL_TYPE = os.environ["MODEL_TYPE"]
         MODEL_NAME = os.environ["MODEL_NAME"]
-        MODEL_API_KEY = os.environ["MODEL_API_KEY"]
+        if MODEL_NAME != "local":
+            MODEL_API_KEY = os.environ["MODEL_API_KEY"]
+        else:
+            MODEL_API_KEY = ""
         BASE_URL = os.environ["BASE_URL"]
         return HelperCompany.get()[MODEL_TYPE](
             MODEL_API_KEY, MODEL_NAME, BASE_URL, timeout=None
@@ -64,6 +67,8 @@ class APIHelper(object):
             "glm4-air",
             "qwen-max",
             "qwen-plus",
+            "gpt-4o-mini",
+            "local",
         ]:
             raise ValueError(f"Check model name...")
 
@@ -78,13 +83,13 @@ class APIHelper(object):
             response1 = self.generator.create(
                 messages=message,
             )
-            summary = clean_text(response1.choices[0].message.content)
+            summary = clean_text(response1)
             message.append({"role": "assistant", "content": summary})
             message.append(self.prompt.queries[1][0]())
             response2 = self.generator.create(
                 messages=message,
             )
-            detail = response2.choices[0].message.content
+            detail = response2
             motivation = clean_text(detail.split(TAG_moti)[1].split(TAG_contr)[0])
             contribution = clean_text(detail.split(TAG_contr)[1])
             result = {
@@ -116,7 +121,7 @@ class APIHelper(object):
             response = self.generator.create(
                 messages=message,
             )
-            entities = response.choices[0].message.content
+            entities = response
             entity_list = entities.strip().split(", ")
             clean_entity_list = []
             for entity in entity_list:
@@ -151,7 +156,7 @@ class APIHelper(object):
             response_brainstorming = self.generator.create(
                 messages=message,
             )
-            brainstorming_ideas = response_brainstorming.choices[0].message.content
+            brainstorming_ideas = response_brainstorming
 
         except Exception:
             traceback.print_exc()
@@ -178,7 +183,7 @@ class APIHelper(object):
             response = self.generator.create(
                 messages=message,
             )
-            problem = response.choices[0].message.content
+            problem = response
         except Exception:
             traceback.print_exc()
             return None
@@ -207,7 +212,7 @@ class APIHelper(object):
             response = self.generator.create(
                 messages=message,
             )
-            problem = response.choices[0].message.content
+            problem = response
         except Exception:
             traceback.print_exc()
             return None
@@ -228,7 +233,7 @@ class APIHelper(object):
             response = self.generator.create(
                 messages=message,
             )
-            inspiration = response.choices[0].message.content
+            inspiration = response
         except Exception:
             traceback.print_exc()
             return None
@@ -254,7 +259,7 @@ class APIHelper(object):
             response = self.generator.create(
                 messages=message,
             )
-            inspiration = response.choices[0].message.content
+            inspiration = response
         except Exception:
             traceback.print_exc()
             return None
@@ -282,7 +287,7 @@ class APIHelper(object):
             response = self.generator.create(
                 messages=message,
             )
-            idea = response.choices[0].message.content
+            idea = response
         except Exception:
             traceback.print_exc()
             return None
@@ -314,7 +319,7 @@ class APIHelper(object):
             response = self.generator.create(
                 messages=message,
             )
-            idea = response.choices[0].message.content
+            idea = response
         except Exception:
             traceback.print_exc()
             return None
@@ -340,7 +345,7 @@ class APIHelper(object):
             response = self.generator.create(
                 messages=message,
             )
-            idea = response.choices[0].message.content
+            idea = response
         except Exception:
             traceback.print_exc()
             return None
@@ -372,7 +377,7 @@ class APIHelper(object):
             response = self.generator.create(
                 messages=message,
             )
-            idea = response.choices[0].message.content
+            idea = response
         except Exception:
             traceback.print_exc()
             return None
@@ -391,7 +396,7 @@ class APIHelper(object):
             response = self.generator.create(
                 messages=message,
             )
-            idea = response.choices[0].message.content
+            idea = response
         except Exception:
             traceback.print_exc()
             return None
@@ -413,7 +418,7 @@ class APIHelper(object):
             response = self.generator.create(
                 messages=message,
             )
-            idea_filtered = response.choices[0].message.content
+            idea_filtered = response
         except Exception:
             traceback.print_exc()
             return None
@@ -435,7 +440,7 @@ class APIHelper(object):
             response = self.generator.create(
                 messages=message,
             )
-            idea_modified = response.choices[0].message.content
+            idea_modified = response
         except Exception:
             traceback.print_exc()
             return None
@@ -454,7 +459,7 @@ class APIHelper(object):
             response = self.generator.create(
                 messages=message,
             )
-            ground_truth = response.choices[0].message.content
+            ground_truth = response
         except Exception:
             traceback.print_exc()
         return ground_truth
@@ -469,7 +474,7 @@ class APIHelper(object):
             response = self.generator.create(
                 messages=message,
             )
-            idea_norm = response.choices[0].message.content
+            idea_norm = response
         except Exception:
             traceback.print_exc()
             return None
@@ -492,7 +497,7 @@ class APIHelper(object):
                 messages=message,
                 max_tokens=10,
             )
-            index = response.choices[0].message.content
+            index = response
         except Exception:
             traceback.print_exc()
             return None
@@ -509,7 +514,7 @@ class APIHelper(object):
                 messages=message,
                 max_tokens=10,
             )
-            score = response.choices[0].message.content
+            score = response
         except Exception:
             traceback.print_exc()
             return None
@@ -548,7 +553,7 @@ class APIHelper(object):
                 stop=None,
                 seed=0,
             )
-            content = response.choices[0].message.content
+            content = response
             new_msg_history = new_msg_history + [
                 {"role": "assistant", "content": content}
             ]
@@ -577,7 +582,7 @@ class APIHelper(object):
             response = self.generator.create(
                 messages=message,
             )
-            result = response.choices[0].message.content
+            result = response
         except Exception:
             traceback.print_exc()
             return None
@@ -601,7 +606,7 @@ class APIHelper(object):
             response = self.generator.create(
                 messages=message,
             )
-            result = response.choices[0].message.content
+            result = response
         except Exception:
             traceback.print_exc()
             return None
@@ -625,7 +630,7 @@ class APIHelper(object):
             response = self.generator.create(
                 messages=message,
             )
-            result = response.choices[0].message.content
+            result = response
         except Exception:
             traceback.print_exc()
             return None
@@ -649,7 +654,7 @@ class APIHelper(object):
             response = self.generator.create(
                 messages=message,
             )
-            result = response.choices[0].message.content
+            result = response
         except Exception:
             traceback.print_exc()
             return None
@@ -673,7 +678,7 @@ class APIHelper(object):
             response = self.generator.create(
                 messages=message,
             )
-            result = response.choices[0].message.content
+            result = response
         except Exception:
             traceback.print_exc()
             return None

src/utils/paper_client.py

@@ -8,6 +8,7 @@ from collections import defaultdict, deque
 from py2neo import Graph, Node, Relationship
 from loguru import logger
 
+
 class PaperClient:
     _instance = None
     _initialized = False
@@ -43,10 +44,28 @@ class PaperClient:
         with self.driver.session() as session:
             result = session.execute_read(lambda tx: tx.run(query).data())
         if result:
-            paper_from_client = result[0]['p']
+            paper_from_client = result[0]["p"]
             if paper_from_client is not None:
                 paper.update(paper_from_client)
-    
+
+    def update_papers_from_client(self, paper_id_list):
+        query = """
+            UNWIND $papers AS paper
+            MATCH (p:Paper {hash_id: paper.hash_id})
+            RETURN p as result
+        """
+        paper_data = [
+            {
+                "hash_id": hash_id,
+            }
+            for hash_id in paper_id_list
+        ]
+        with self.driver.session() as session:
+            result = session.execute_read(
+                lambda tx: tx.run(query, papers=paper_data).data()
+            )
+        return [r["result"] for r in result]
+
     def get_paper_attribute(self, paper_id, attribute_name):
         query = f"""
             MATCH (p:Paper {{hash_id: {paper_id}}})
@@ -55,11 +74,11 @@ class PaperClient:
         with self.driver.session() as session:
             result = session.execute_read(lambda tx: tx.run(query).data())
         if result:
-            return result[0]['attributeValue']
+            return result[0]["attributeValue"]
         else:
             logger.error(f"paper id {paper_id} get {attribute_name} failed.")
             return None
-    
+
     def get_paper_by_attribute(self, attribute_name, anttribute_value):
         query = f"""
             MATCH (p:Paper {{{attribute_name}: '{anttribute_value}'}})
@@ -68,7 +87,7 @@ class PaperClient:
         with self.driver.session() as session:
             result = session.execute_read(lambda tx: tx.run(query).data())
         if result:
-            return result[0]['p']
+            return result[0]["p"]
         else:
             return None
 
@@ -81,71 +100,50 @@ class PaperClient:
             RETURN p.hash_id as hash_id
             """
         with self.driver.session() as session:
-            result = session.execute_read(lambda tx: tx.run(query, entity=entity).data())
+            result = session.execute_read(
+                lambda tx: tx.run(query, entity=entity).data()
+            )
         if result:
-            return [record['hash_id'] for record in result]
+            return [record["hash_id"] for record in result]
         else:
             return []
-    
-    def find_related_entities_by_entity(self, entity_name, n=1, k=3, relation_name="related"):
-        # relation_name = "related"
-        def bfs_query(entity_name, n, k):
-            queue = deque([(entity_name, 0)])  
-            visited = set([entity_name])  
-            related_entities = set()  
-
-            while queue:
-                batch_queue = [queue.popleft() for _ in range(len(queue))]  
-                batch_entities = [item[0] for item in batch_queue]
-                batch_depths = [item[1] for item in batch_queue]
-
-                if all(depth >= n for depth in batch_depths):
-                    continue
-                if relation_name == "related":
-                    query = """
-                        UNWIND $batch_entities AS entity_name
-                        MATCH (e1:Entity {name: entity_name})-[:RELATED_TO]->(p:Paper)<-[:RELATED_TO]-(e2:Entity)
-                        WHERE e1 <> e2
-                        WITH e1, e2, COUNT(p) AS common_papers, entity_name
-                        WHERE common_papers > $k
-                        RETURN e2.name AS entities, entity_name AS source_entity, common_papers
-                    """
-                elif relation_name == "connect":
-                    query = """
-                        UNWIND $batch_entities AS entity_name
-                        MATCH (e1:Entity {name: entity_name})-[r:CONNECT]-(e2:Entity)
-                        WHERE e1 <> e2 and r.strength >= $k
-                        WITH e1, e2, entity_name
-                        RETURN e2.name AS entities, entity_name AS source_entity
-                    """
-                with self.driver.session() as session:
-                    result = session.execute_read(lambda tx: tx.run(query, batch_entities=batch_entities, k=k).data())
-
-                for record in result:
-                    entity = record['entities']
-                    source_entity = record['source_entity']
-                    if entity not in visited:
-                        visited.add(entity)
-                        queue.append((entity, batch_depths[batch_entities.index(source_entity)] + 1))
-                        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)
-        return related_entities
-
-    def find_entities_by_paper(self, hash_id: int):
+
+    def find_related_entities_by_entity_list(
+        self, entity_names, n=1, k=3, relation_name="related"
+    ):
+        related_entities = set()
         query = """
-            MATCH (e:Entity)-[:RELATED_TO]->(p:Paper {hash_id: $hash_id})
-            RETURN e.name AS entity_name
+            UNWIND $batch_entities AS entity_name
+            MATCH (e1:Entity {name: entity_name})-[:RELATED_TO]->(p:Paper)<-[:RELATED_TO]-(e2:Entity)
+            WHERE e1 <> e2
+            WITH e1, e2, COUNT(p) AS common_papers, entity_name
+            WHERE common_papers > $k
+            RETURN e2.name AS entities, entity_name AS source_entity, common_papers
         """
         with self.driver.session() as session:
-            result = session.execute_read(lambda tx: tx.run(query, hash_id=hash_id).data())
-        if result:
-            return [record['entity_name'] for record in result]
-        else:
-            return []
+            result = session.execute_read(
+                lambda tx: tx.run(query, batch_entities=entity_names, k=k).data()
+            )
+        for record in result:
+            entity = record["entities"]
+            related_entities.add(entity)
+        return list(related_entities)
+
+    def find_entities_by_paper_list(self, hash_ids: list):
+        query = """
+            UNWIND $hash_ids AS hash_id
+            MATCH (e:Entity)-[:RELATED_TO]->(p:Paper {hash_id: hash_id})
+            RETURN hash_id, e.name AS entity_name
+        """
+        with self.driver.session() as session:
+            result = session.execute_read(
+                lambda tx: tx.run(query, hash_ids=hash_ids).data()
+            )
+        # 按照每个 hash_id 分组实体
+        entity_list = []
+        for record in result:
+            entity_list.append(record["entity_name"])
+        return entity_list
 
     def find_paper_by_entity(self, entity_name):
         query = """
@@ -153,18 +151,19 @@ class PaperClient:
             RETURN p.hash_id AS hash_id
         """
         with self.driver.session() as session:
-            result = session.execute_read(lambda tx: tx.run(query, entity_name=entity_name).data())
+            result = session.execute_read(
+                lambda tx: tx.run(query, entity_name=entity_name).data()
+            )
         if result:
-            return [record['hash_id'] for record in result]
+            return [record["hash_id"] for record in result]
         else:
             return []
-        
+
     # TODO: @云翔
     # 增加通过entity返回包含entity语句的功能
     def find_sentence_by_entity(self, entity_name):
         # Return: list(str)
         return []
-    
 
     def find_sentences_by_entity(self, entity_name):
         query = """
@@ -178,14 +177,25 @@ class PaperClient:
             p.hash_id AS hash_id
         """
         sentences = []
-        
+
         with self.driver.session() as session:
-            result = session.execute_read(lambda tx: tx.run(query, entity_name=entity_name).data())
+            result = session.execute_read(
+                lambda tx: tx.run(query, entity_name=entity_name).data()
+            )
         for record in result:
-            for key in ['abstract', 'introduction', 'methodology']:
+            for key in ["abstract", "introduction", "methodology"]:
                 if record[key]:
-                    filtered_sentences = [sentence.strip() + '.' for sentence in record[key].split('.') if entity_name in sentence]
-                    sentences.extend([f"{record['hash_id']}: {sentence}" for sentence in filtered_sentences])
+                    filtered_sentences = [
+                        sentence.strip() + "."
+                        for sentence in record[key].split(".")
+                        if entity_name in sentence
+                    ]
+                    sentences.extend(
+                        [
+                            f"{record['hash_id']}: {sentence}"
+                            for sentence in filtered_sentences
+                        ]
+                    )
 
         return sentences
 
@@ -194,9 +204,11 @@ class PaperClient:
             MATCH (n:Paper) where n.year=$year and n.venue_name=$venue_name return n
         """
         with self.driver.session() as session:
-            result = session.execute_read(lambda tx: tx.run(query, year=year, venue_name=venue_name).data())
+            result = session.execute_read(
+                lambda tx: tx.run(query, year=year, venue_name=venue_name).data()
+            )
         if result:
-            return [record['n'] for record in result]
+            return [record["n"] for record in result]
         else:
             return []
 
@@ -230,7 +242,26 @@ class PaperClient:
             RETURN p
             """
         with self.driver.session() as session:
-            result = session.execute_write(lambda tx: tx.run(query, hash_id=paper["hash_id"], venue_name=paper["venue_name"], year=paper["year"], title=paper["title"], pdf_url=paper["pdf_url"], abstract=paper["abstract"], introduction=paper["introduction"], reference=paper["reference"], summary=paper["summary"], motivation=paper["motivation"], contribution=paper["contribution"], methodology=paper["methodology"], ground_truth=paper["ground_truth"], reference_filter=paper["reference_filter"], conclusions=paper["conclusions"]).data())
+            result = session.execute_write(
+                lambda tx: tx.run(
+                    query,
+                    hash_id=paper["hash_id"],
+                    venue_name=paper["venue_name"],
+                    year=paper["year"],
+                    title=paper["title"],
+                    pdf_url=paper["pdf_url"],
+                    abstract=paper["abstract"],
+                    introduction=paper["introduction"],
+                    reference=paper["reference"],
+                    summary=paper["summary"],
+                    motivation=paper["motivation"],
+                    contribution=paper["contribution"],
+                    methodology=paper["methodology"],
+                    ground_truth=paper["ground_truth"],
+                    reference_filter=paper["reference_filter"],
+                    conclusions=paper["conclusions"],
+                ).data()
+            )
 
     def check_entity_node_count(self, hash_id: int):
         query_check_count = """
@@ -239,7 +270,9 @@ class PaperClient:
         """
         with self.driver.session() as session:
             # Check the number of related entities
-            result = session.execute_read(lambda tx: tx.run(query_check_count, hash_id=hash_id).data())
+            result = session.execute_read(
+                lambda tx: tx.run(query_check_count, hash_id=hash_id).data()
+            )
             if result[0]["entity_count"] > 3:
                 return False
         return True
@@ -254,16 +287,30 @@ class PaperClient:
         """
         with self.driver.session() as session:
             for entity_name in entities:
-                result = session.execute_write(lambda tx: tx.run(query, entity_name=entity_name, hash_id=hash_id).data())
-        
+                result = session.execute_write(
+                    lambda tx: tx.run(
+                        query, entity_name=entity_name, hash_id=hash_id
+                    ).data()
+                )
+
     def add_paper_citation(self, paper: dict):
         query = """
             MERGE (p:Paper {hash_id: $hash_id}) ON MATCH SET p.cite_id_list = $cite_id_list, p.entities = $entities, p.all_cite_id_list = $all_cite_id_list
             """
         with self.driver.session() as session:
-            result = session.execute_write(lambda tx: tx.run(query, hash_id=paper["hash_id"], cite_id_list=paper["cite_id_list"], entities=paper["entities"], all_cite_id_list=paper["all_cite_id_list"]).data())
-
-    def add_paper_abstract_embedding(self, embedding_model, hash_id=None):
+            result = session.execute_write(
+                lambda tx: tx.run(
+                    query,
+                    hash_id=paper["hash_id"],
+                    cite_id_list=paper["cite_id_list"],
+                    entities=paper["entities"],
+                    all_cite_id_list=paper["all_cite_id_list"],
+                ).data()
+            )
+
+    def add_paper_abstract_embedding(
+        self, embedding_model, hash_id=None, batch_size=512
+    ):
         if hash_id is not None:
             query = """
             MATCH (p:Paper {hash_id: $hash_id})
@@ -271,119 +318,302 @@ class PaperClient:
             RETURN p.abstract AS context, p.hash_id AS hash_id, p.title AS title
             """
             with self.driver.session() as session:
-                results = session.execute_write(lambda tx: tx.run(query, hash_id=hash_id).data())
-        else:
+                results = session.execute_write(
+                    lambda tx: tx.run(query, hash_id=hash_id).data()
+                )
+            contexts = [result["title"] + result["context"] for result in results]
+            paper_ids = [result["hash_id"] for result in results]
+            context_embeddings = embedding_model.encode(
+                contexts, convert_to_tensor=True, device=self.device
+            )
             query = """
+                MERGE (p:Paper {hash_id: $hash_id})
+                ON CREATE SET p.abstract_embedding = $embedding
+                ON MATCH SET p.abstract_embedding = $embedding
+            """
+            for idx, hash_id in tqdm(enumerate(paper_ids)):
+                embedding = (
+                    context_embeddings[idx].detach().cpu().numpy().flatten().tolist()
+                )
+                with self.driver.session() as session:
+                    results = session.execute_write(
+                        lambda tx: tx.run(
+                            query, hash_id=hash_id, embedding=embedding
+                        ).data()
+                    )
+            return
+        offset = 0
+        while True:
+            query = f"""
             MATCH (p:Paper)
             WHERE p.abstract IS NOT NULL
             RETURN p.abstract AS context, p.hash_id AS hash_id, p.title AS title
+            SKIP $offset LIMIT $batch_size
             """
             with self.driver.session() as session:
-                results = session.execute_write(lambda tx: tx.run(query).data())
-        contexts = [result["title"] + result["context"] for result in results]
-        paper_ids = [result["hash_id"] for result in results]
-        context_embeddings = embedding_model.encode(contexts, batch_size=512, convert_to_tensor=True, device=self.device)
-        query = """
-            MERGE (p:Paper {hash_id: $hash_id})
-            ON CREATE SET p.abstract_embedding = $embedding
-            ON MATCH SET p.abstract_embedding = $embedding
-        """
-        for idx, hash_id in tqdm(enumerate(paper_ids)):
-            embedding = context_embeddings[idx].detach().cpu().numpy().flatten().tolist()
+                results = session.execute_write(
+                    lambda tx: tx.run(
+                        query, offset=offset, batch_size=batch_size
+                    ).data()
+                )
+            if not results:
+                break
+            contexts = [result["title"] + result["context"] for result in results]
+            paper_ids = [result["hash_id"] for result in results]
+            context_embeddings = embedding_model.encode(
+                contexts,
+                batch_size=batch_size,
+                convert_to_tensor=True,
+                device=self.device,
+            )
+            write_query = """
+            UNWIND $data AS row
+            MERGE (p:Paper {hash_id: row.hash_id})
+            ON CREATE SET p.abstract_embedding = row.embedding
+            ON MATCH SET p.abstract_embedding = row.embedding
+            """
+            data_to_write = []
+            for idx, hash_id in enumerate(paper_ids):
+                embedding = (
+                    context_embeddings[idx].detach().cpu().numpy().flatten().tolist()
+                )
+                data_to_write.append({"hash_id": hash_id, "embedding": embedding})
             with self.driver.session() as session:
-                results = session.execute_write(lambda tx: tx.run(query, hash_id=hash_id, embedding=embedding).data())
+                session.execute_write(
+                    lambda tx: tx.run(write_query, data=data_to_write)
+                )
+            offset += batch_size
+            logger.info(f"== Processed batch starting at offset {offset} ==")
 
-    def add_paper_bg_embedding(self, embedding_model, hash_id=None):
+    def add_paper_bg_embedding(self, embedding_model, hash_id=None, batch_size=512):
         if hash_id is not None:
             query = """
             MATCH (p:Paper {hash_id: $hash_id})
             WHERE p.motivation IS NOT NULL
-            RETURN p.motivation AS context, p.hash_id AS hash_id
+            RETURN p.motivation AS context, p.hash_id AS hash_id, p.title AS title
             """
             with self.driver.session() as session:
-                results = session.execute_write(lambda tx: tx.run(query, hash_id=hash_id).data())
-        else:
+                results = session.execute_write(
+                    lambda tx: tx.run(query, hash_id=hash_id).data()
+                )
+            contexts = [result["context"] for result in results]
+            paper_ids = [result["hash_id"] for result in results]
+            context_embeddings = embedding_model.encode(
+                contexts, convert_to_tensor=True, device=self.device
+            )
             query = """
+                MERGE (p:Paper {hash_id: $hash_id})
+                ON CREATE SET p.motivation_embedding = $embedding
+                ON MATCH SET p.motivation_embedding = $embedding
+            """
+            for idx, hash_id in tqdm(enumerate(paper_ids)):
+                embedding = (
+                    context_embeddings[idx].detach().cpu().numpy().flatten().tolist()
+                )
+                with self.driver.session() as session:
+                    results = session.execute_write(
+                        lambda tx: tx.run(
+                            query, hash_id=hash_id, embedding=embedding
+                        ).data()
+                    )
+            return
+        offset = 0
+        while True:
+            query = f"""
             MATCH (p:Paper)
             WHERE p.motivation IS NOT NULL
-            RETURN p.motivation AS context, p.hash_id AS hash_id
+            RETURN p.motivation AS context, p.hash_id AS hash_id, p.title AS title
+            SKIP $offset LIMIT $batch_size
             """
             with self.driver.session() as session:
-                results = session.execute_write(lambda tx: tx.run(query).data())
-        contexts = [result["context"] for result in results]
-        paper_ids = [result["hash_id"] for result in results]
-        context_embeddings = embedding_model.encode(contexts, batch_size=256, convert_to_tensor=True, device=self.device)
-        query = """
-            MERGE (p:Paper {hash_id: $hash_id})
-            ON CREATE SET p.embedding = $embedding
-            ON MATCH SET p.embedding = $embedding
-        """
-        for idx, hash_id in tqdm(enumerate(paper_ids)):
-            embedding = context_embeddings[idx].detach().cpu().numpy().flatten().tolist()
+                results = session.execute_write(
+                    lambda tx: tx.run(
+                        query, offset=offset, batch_size=batch_size
+                    ).data()
+                )
+            if not results:
+                break
+            contexts = [result["title"] + result["context"] for result in results]
+            paper_ids = [result["hash_id"] for result in results]
+            context_embeddings = embedding_model.encode(
+                contexts,
+                batch_size=batch_size,
+                convert_to_tensor=True,
+                device=self.device,
+            )
+            write_query = """
+            UNWIND $data AS row
+            MERGE (p:Paper {hash_id: row.hash_id})
+            ON CREATE SET p.motivation_embedding = row.embedding
+            ON MATCH SET p.motivation_embedding = row.embedding
+            """
+            data_to_write = []
+            for idx, hash_id in enumerate(paper_ids):
+                embedding = (
+                    context_embeddings[idx].detach().cpu().numpy().flatten().tolist()
+                )
+                data_to_write.append({"hash_id": hash_id, "embedding": embedding})
             with self.driver.session() as session:
-                results = session.execute_write(lambda tx: tx.run(query, hash_id=hash_id, embedding=embedding).data())
-
-    def add_paper_contribution_embedding(self, embedding_model, hash_id=None):
+                session.execute_write(
+                    lambda tx: tx.run(write_query, data=data_to_write)
+                )
+            offset += batch_size
+            logger.info(f"== Processed batch starting at offset {offset} ==")
+
+    def add_paper_contribution_embedding(
+        self, embedding_model, hash_id=None, batch_size=512
+    ):
         if hash_id is not None:
             query = """
             MATCH (p:Paper {hash_id: $hash_id})
             WHERE p.contribution IS NOT NULL
-            RETURN p.contribution AS context, p.hash_id AS hash_id
+            RETURN p.contribution AS context, p.hash_id AS hash_id, p.title AS title
             """
             with self.driver.session() as session:
-                results = session.execute_write(lambda tx: tx.run(query, hash_id=hash_id).data())
-        else:
+                results = session.execute_write(
+                    lambda tx: tx.run(query, hash_id=hash_id).data()
+                )
+            contexts = [result["context"] for result in results]
+            paper_ids = [result["hash_id"] for result in results]
+            context_embeddings = embedding_model.encode(
+                contexts, convert_to_tensor=True, device=self.device
+            )
             query = """
+                MERGE (p:Paper {hash_id: $hash_id})
+                ON CREATE SET p.contribution_embedding = $embedding
+                ON MATCH SET p.contribution_embedding = $embedding
+            """
+            for idx, hash_id in tqdm(enumerate(paper_ids)):
+                embedding = (
+                    context_embeddings[idx].detach().cpu().numpy().flatten().tolist()
+                )
+                with self.driver.session() as session:
+                    results = session.execute_write(
+                        lambda tx: tx.run(
+                            query, hash_id=hash_id, embedding=embedding
+                        ).data()
+                    )
+            return
+        offset = 0
+        while True:
+            query = f"""
             MATCH (p:Paper)
             WHERE p.contribution IS NOT NULL
-            RETURN p.contribution AS context, p.hash_id AS hash_id
+            RETURN p.contribution AS context, p.hash_id AS hash_id, p.title AS title
+            SKIP $offset LIMIT $batch_size
             """
             with self.driver.session() as session:
-                results = session.execute_write(lambda tx: tx.run(query).data())
-        contexts = [result["context"] for result in results]
-        paper_ids = [result["hash_id"] for result in results]
-        context_embeddings = embedding_model.encode(contexts, batch_size=256, convert_to_tensor=True, device=self.device)
-        query = """
-            MERGE (p:Paper {hash_id: $hash_id})
-            ON CREATE SET p.contribution_embedding = $embedding
-            ON MATCH SET p.contribution_embedding = $embedding
-        """
-        for idx, hash_id in tqdm(enumerate(paper_ids)):
-            embedding = context_embeddings[idx].detach().cpu().numpy().flatten().tolist()
+                results = session.execute_write(
+                    lambda tx: tx.run(
+                        query, offset=offset, batch_size=batch_size
+                    ).data()
+                )
+            if not results:
+                break
+            contexts = [result["context"] for result in results]
+            paper_ids = [result["hash_id"] for result in results]
+            context_embeddings = embedding_model.encode(
+                contexts,
+                batch_size=batch_size,
+                convert_to_tensor=True,
+                device=self.device,
+            )
+            write_query = """
+            UNWIND $data AS row
+            MERGE (p:Paper {hash_id: row.hash_id})
+            ON CREATE SET p.contribution_embedding = row.embedding
+            ON MATCH SET p.contribution_embedding = row.embedding
+            """
+            data_to_write = []
+            for idx, hash_id in enumerate(paper_ids):
+                embedding = (
+                    context_embeddings[idx].detach().cpu().numpy().flatten().tolist()
+                )
+                data_to_write.append({"hash_id": hash_id, "embedding": embedding})
             with self.driver.session() as session:
-                results = session.execute_write(lambda tx: tx.run(query, hash_id=hash_id, embedding=embedding).data())
-    
-
-    def add_paper_summary_embedding(self, embedding_model, hash_id=None):
+                session.execute_write(
+                    lambda tx: tx.run(write_query, data=data_to_write)
+                )
+            offset += batch_size
+            logger.info(f"== Processed batch starting at offset {offset} ==")
+
+    def add_paper_summary_embedding(
+        self, embedding_model, hash_id=None, batch_size=512
+    ):
         if hash_id is not None:
             query = """
             MATCH (p:Paper {hash_id: $hash_id})
             WHERE p.summary IS NOT NULL
-            RETURN p.summary AS context, p.hash_id AS hash_id
+            RETURN p.summary AS context, p.hash_id AS hash_id, p.title AS title
             """
             with self.driver.session() as session:
-                results = session.execute_write(lambda tx: tx.run(query, hash_id=hash_id).data())
-        else:
+                results = session.execute_write(
+                    lambda tx: tx.run(query, hash_id=hash_id).data()
+                )
+            contexts = [result["context"] for result in results]
+            paper_ids = [result["hash_id"] for result in results]
+            context_embeddings = embedding_model.encode(
+                contexts, convert_to_tensor=True, device=self.device
+            )
             query = """
+                MERGE (p:Paper {hash_id: $hash_id})
+                ON CREATE SET p.summary_embedding = $embedding
+                ON MATCH SET p.summary_embedding = $embedding
+            """
+            for idx, hash_id in tqdm(enumerate(paper_ids)):
+                embedding = (
+                    context_embeddings[idx].detach().cpu().numpy().flatten().tolist()
+                )
+                with self.driver.session() as session:
+                    results = session.execute_write(
+                        lambda tx: tx.run(
+                            query, hash_id=hash_id, embedding=embedding
+                        ).data()
+                    )
+            return
+        offset = 0
+        while True:
+            query = f"""
             MATCH (p:Paper)
             WHERE p.summary IS NOT NULL
-            RETURN p.summary AS context, p.hash_id AS hash_id
+            RETURN p.summary AS context, p.hash_id AS hash_id, p.title AS title
+            SKIP $offset LIMIT $batch_size
             """
             with self.driver.session() as session:
-                results = session.execute_write(lambda tx: tx.run(query).data())
-        contexts = [result["context"] for result in results]
-        paper_ids = [result["hash_id"] for result in results]
-        context_embeddings = embedding_model.encode(contexts, batch_size=256, convert_to_tensor=True, device=self.device)
-        query = """
-            MERGE (p:Paper {hash_id: $hash_id})
-            ON CREATE SET p.summary_embedding = $embedding
-            ON MATCH SET p.summary_embedding = $embedding
-        """
-        for idx, hash_id in tqdm(enumerate(paper_ids)):
-            embedding = context_embeddings[idx].detach().cpu().numpy().flatten().tolist()
+                results = session.execute_write(
+                    lambda tx: tx.run(
+                        query, offset=offset, batch_size=batch_size
+                    ).data()
+                )
+            if not results:
+                break
+            contexts = [result["context"] for result in results]
+            paper_ids = [result["hash_id"] for result in results]
+            context_embeddings = embedding_model.encode(
+                contexts,
+                batch_size=batch_size,
+                convert_to_tensor=True,
+                device=self.device,
+            )
+            write_query = """
+            UNWIND $data AS row
+            MERGE (p:Paper {hash_id: row.hash_id})
+            ON CREATE SET p.summary_embedding = row.embedding
+            ON MATCH SET p.summary_embedding = row.embedding
+            """
+            data_to_write = []
+            for idx, hash_id in enumerate(paper_ids):
+                embedding = (
+                    context_embeddings[idx].detach().cpu().numpy().flatten().tolist()
+                )
+                data_to_write.append({"hash_id": hash_id, "embedding": embedding})
             with self.driver.session() as session:
-                results = session.execute_write(lambda tx: tx.run(query, hash_id=hash_id, embedding=embedding).data())
-        
+                session.execute_write(
+                    lambda tx: tx.run(write_query, data=data_to_write)
+                )
+            offset += batch_size
+            logger.info(f"== Processed batch starting at offset {offset} ==")
+
     def cosine_similarity_search(self, embedding, k=1, type_name="embedding"):
         query = f"""
             MATCH (paper:Paper)
@@ -394,8 +624,10 @@ class PaperClient:
             ORDER BY score DESC LIMIT {k}
             """
         with self.driver.session() as session:
-            results = session.execute_read(lambda tx: tx.run(query, embedding=embedding).data())
-        related_paper = [] 
+            results = session.execute_read(
+                lambda tx: tx.run(query, embedding=embedding).data()
+            )
+        related_paper = []
         for result in results:
             related_paper.append(result["paper"]["hash_id"])
         return related_paper
@@ -417,7 +649,7 @@ class PaperClient:
             """
         with self.driver.session() as session:
             session.execute_write(lambda tx: tx.run(query).data())
-    
+
     def filter_paper_id_list(self, paper_id_list, year="2024"):
         if not paper_id_list:
             return []
@@ -429,12 +661,14 @@ class PaperClient:
             RETURN p.hash_id AS hash_id
         """
         with self.driver.session() as session:
-            result = session.execute_read(lambda tx: tx.run(query, paper_id_list=paper_id_list, year=year).data())
-        
-        existing_paper_ids = [record['hash_id'] for record in result]
+            result = session.execute_read(
+                lambda tx: tx.run(query, paper_id_list=paper_id_list, year=year).data()
+            )
+
+        existing_paper_ids = [record["hash_id"] for record in result]
         existing_paper_ids = list(set(existing_paper_ids))
         return existing_paper_ids
-    
+
     def check_index_exists(self):
         query = "SHOW INDEXES"
         with self.driver.session() as session:
@@ -451,7 +685,7 @@ class PaperClient:
         """
         with self.driver.session() as session:
             session.execute_write(lambda tx: tx.run(query).data())
-    
+
     def get_entity_related_paper_num(self, entity_name):
         query = """
             MATCH (e:Entity {name: $entity_name})-[:RELATED_TO]->(p:Paper)
@@ -459,10 +693,30 @@ class PaperClient:
             RETURN PaperCount
         """
         with self.driver.session() as session:
-            result = session.execute_read(lambda tx: tx.run(query, entity_name=entity_name).data())
-        paper_num = result[0]['PaperCount']
+            result = session.execute_read(
+                lambda tx: tx.run(query, entity_name=entity_name).data()
+            )
+        paper_num = result[0]["PaperCount"]
         return paper_num
 
+    def get_entities_related_paper_num(self, entity_names):
+        query = """
+            UNWIND $entity_names AS entity_name
+            MATCH (e:Entity {name: entity_name})-[:RELATED_TO]->(p:Paper)
+            WITH entity_name, COUNT(p) AS PaperCount
+            RETURN entity_name, PaperCount
+        """
+
+        with self.driver.session() as session:
+            result = session.execute_read(
+                lambda tx: tx.run(query, entity_names=entity_names).data()
+            )
+        # 将查询结果转化为字典形式：实体名称 -> 论文数量
+        entity_paper_count = {
+            record["entity_name"]: record["PaperCount"] for record in result
+        }
+        return entity_paper_count
+
     def get_entity_text(self):
         query = """
             MATCH (e:Entity)-[:RELATED_TO]->(p:Paper)
@@ -472,11 +726,13 @@ class PaperClient:
         """
         with self.driver.session() as session:
             result = session.execute_read(lambda tx: tx.run(query).data())
-        text_list = [record['entity_text'] for record in result]
+        text_list = [record["entity_text"] for record in result]
         return text_list
-    
+
     def get_entity_combinations(self, venue_name, year):
-        def process_paper_relationships(session, entity_name_1, entity_name_2, abstract):
+        def process_paper_relationships(
+            session, entity_name_1, entity_name_2, abstract
+        ):
             if entity_name_2 < entity_name_1:
                 entity_name_1, entity_name_2 = entity_name_2, entity_name_1
             query = """
@@ -486,13 +742,17 @@ class PaperClient:
                 ON CREATE SET r.strength = 1
                 ON MATCH SET r.strength = r.strength + 1
             """
-            sentences = re.split(r'(?<!\w\.\w.)(?<![A-Z][a-z]\.)(?<=\.|\?)\s', abstract)
+            sentences = re.split(r"(?<!\w\.\w.)(?<![A-Z][a-z]\.)(?<=\.|\?)\s", abstract)
             for sentence in sentences:
                 sentence = sentence.lower()
                 if entity_name_1 in sentence and entity_name_2 in sentence:
                     # 如果两个实体在同一句话中出现过，则创建或更新 CONNECT 关系
                     session.execute_write(
-                        lambda tx: tx.run(query, entity_name_1=entity_name_1, entity_name_2=entity_name_2).data()
+                        lambda tx: tx.run(
+                            query,
+                            entity_name_1=entity_name_1,
+                            entity_name_2=entity_name_2,
+                        ).data()
                     )
                     # logger.debug(f"CONNECT relation created or updated between {entity_name_1} and {entity_name_2} for Paper ID {paper_id}")
                     break  # 如果找到一次出现就可以退出循环
@@ -506,13 +766,17 @@ class PaperClient:
             RETURN p.hash_id AS hash_id, entities[i].name AS entity_name_1, entities[j].name AS entity_name_2
         """
         with self.driver.session() as session:
-            result = session.execute_read(lambda tx: tx.run(query, venue_name=venue_name, year=year).data())
+            result = session.execute_read(
+                lambda tx: tx.run(query, venue_name=venue_name, year=year).data()
+            )
             for record in tqdm(result):
                 paper_id = record["hash_id"]
-                entity_name_1 = record['entity_name_1']
-                entity_name_2 = record['entity_name_2']
+                entity_name_1 = record["entity_name_1"]
+                entity_name_2 = record["entity_name_2"]
                 abstract = self.get_paper_attribute(paper_id, "abstract")
-                process_paper_relationships(session, entity_name_1, entity_name_2, abstract)
+                process_paper_relationships(
+                    session, entity_name_1, entity_name_2, abstract
+                )
 
     def build_citemap(self):
         citemap = defaultdict(set)
@@ -523,8 +787,8 @@ class PaperClient:
         with self.driver.session() as session:
             results = session.execute_read(lambda tx: tx.run(query).data())
         for result in results:
-            hash_id = result['hash_id']
-            cite_id_list = result['cite_id_list']
+            hash_id = result["hash_id"]
+            cite_id_list = result["cite_id_list"]
             if cite_id_list:
                 for cited_id in cite_id_list:
                     citemap[hash_id].add(cited_id)
@@ -537,12 +801,17 @@ class PaperClient:
         AUTH = (NEO4J_USERNAME, NEO4J_PASSWD)
         graph = Graph(URI, auth=AUTH)
         # 创建一个字典来保存数据
+        # 定义批次大小
         data = {"nodes": [], "relationships": []}
-        query = """
+        # 计算数据的总数（例如查询节点总数）
+        total_papers_query = "MATCH (e:Entity)-[:RELATED_TO]->(p:Paper) RETURN COUNT(DISTINCT p) AS count"
+        total_papers = graph.run(total_papers_query).evaluate()
+        print(f"total paper: {total_papers}")
+        query = f"""
             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)
         # 处理查询结果
         for record in tqdm(results):
@@ -550,39 +819,46 @@ class PaperClient:
             entity_node = record["e"]
             relationship = record["r"]
             # 将节点数据加入字典
-            data["nodes"].append({
-                "id": paper_node.identity,
-                "label": "Paper",
-                "properties": dict(paper_node)
-            })
-            data["nodes"].append({
-                "id": entity_node.identity,
-                "label": "Entity",
-                "properties": dict(entity_node)
-            })
+            data["nodes"].append(
+                {
+                    "id": paper_node.identity,
+                    "label": "Paper",
+                    "properties": dict(paper_node),
+                }
+            )
+            data["nodes"].append(
+                {
+                    "id": entity_node.identity,
+                    "label": "Entity",
+                    "properties": dict(entity_node),
+                }
+            )
             # 将关系数据加入字典
-            data["relationships"].append({
-                "start_node": entity_node.identity,
-                "end_node": paper_node.identity,
-                "type": "RELATED_TO",
-                "properties": dict(relationship)
-            })
+            data["relationships"].append(
+                {
+                    "start_node": entity_node.identity,
+                    "end_node": paper_node.identity,
+                    "type": "RELATED_TO",
+                    "properties": dict(relationship),
+                }
+            )
+        """
         query = """
             MATCH (p:Paper)
             WHERE p.venue_name='acl' and p.year='2024'
             RETURN p
         """
-        """
         results = graph.run(query)
         for record in tqdm(results):
             paper_node = record["p"]
             # 将节点数据加入字典
-            data["nodes"].append({
-                "id": paper_node.identity,
-                "label": "Paper",
-                "properties": dict(paper_node)
-            })
-        """
+            data["nodes"].append(
+                {
+                    "id": paper_node.identity,
+                    "label": "Paper",
+                    "properties": dict(paper_node),
+                }
+            )
         # 去除重复节点
         # data["nodes"] = [dict(t) for t in {tuple(d.items()) for d in data["nodes"]}]
         unique_nodes = []
@@ -595,9 +871,11 @@ class PaperClient:
                 unique_nodes.append(node)
         data["nodes"] = unique_nodes
         # 将数据保存为 JSON 文件
-        with open("./assets/data/scipip_neo4j_clean_backup.json", "w", encoding="utf-8") as f:
+        with open(
+            "./assets/data/scipip_neo4j_clean_backup.json", "w", encoding="utf-8"
+        ) as f:
             json.dump(data, f, ensure_ascii=False, indent=4)
-    
+
     def neo4j_import_data(self):
         # clear_database() # 清空数据库，谨慎执行
         URI = os.environ["NEO4J_URL"]
@@ -606,7 +884,9 @@ class PaperClient:
         AUTH = (NEO4J_USERNAME, NEO4J_PASSWD)
         graph = Graph(URI, auth=AUTH)
         # 从 JSON 文件中读取数据
-        with open("./assets/data/scipip_neo4j_clean_backup.json", "r", encoding="utf-8") as f:
+        with open(
+            "./assets/data/scipip_neo4j_clean_backup.json", "r", encoding="utf-8"
+        ) as f:
             data = json.load(f)
         # 创建节点
         nodes = {}

src/utils/paper_retriever.py

@@ -59,6 +59,7 @@ class CoCite:
 
     def __init__(self) -> None:
         if not self._initialized:
+            logger.debug("init co-cite map begin...")
             self.paper_client = PaperClient()
             citemap = self.paper_client.build_citemap()
             self.comap = defaultdict(lambda: defaultdict(int))
@@ -101,20 +102,16 @@ class Retriever(object):
 
     def retrieve_entities_by_enties(self, entities):
         # TODO: KG
-        expand_entities = []
-        for entity in entities:
-            expand_entities += self.paper_client.find_related_entities_by_entity(
-                entity,
-                n=self.config.RETRIEVE.kg_jump_num,
-                k=self.config.RETRIEVE.kg_cover_num,
-                relation_name=self.config.RETRIEVE.relation_name,
-            )
+        expand_entities = self.paper_client.find_related_entities_by_entity_list(
+            entities,
+            n=self.config.RETRIEVE.kg_jump_num,
+            k=self.config.RETRIEVE.kg_cover_num,
+            relation_name=self.config.RETRIEVE.relation_name,
+        )
         expand_entities = list(set(entities + expand_entities))
-        entity_paper_num_dict = {}
-        for entity in expand_entities:
-            entity_paper_num_dict[entity] = (
-                self.paper_client.get_entity_related_paper_num(entity)
-            )
+        entity_paper_num_dict = self.paper_client.get_entities_related_paper_num(
+            expand_entities
+        )
         new_entities = []
         entity_paper_num_dict = {
             k: v for k, v in entity_paper_num_dict.items() if v != 0
@@ -142,11 +139,7 @@ class Retriever(object):
         Return:
             related_paper: list(dict)
         """
-        related_paper = []
-        for paper_id in paper_id_list:
-            paper = {"hash_id": paper_id}
-            self.paper_client.update_paper_from_client(paper)
-            related_paper.append(paper)
+        related_paper = self.paper_client.update_papers_from_client(paper_id_list)
         return related_paper
 
     def calculate_similarity(self, entities, related_entities_list, use_weight=False):
@@ -333,7 +326,6 @@ class Retriever(object):
         similarity_threshold = self.config.RETRIEVE.similarity_threshold
         similarity_matrix = np.dot(target_paper_embedding, target_paper_embedding.T)
         target_labels = self.cluster_algorithm(target_paper_id_list, similarity_matrix)
-        # target_labels = list(range(0, len(target_paper_id_list)))
         target_paper_label_dict = dict(zip(target_paper_id_list, target_labels))
         logger.debug("Target paper cluster result: {}".format(target_paper_label_dict))
         logger.debug(
@@ -672,8 +664,7 @@ class SNKGRetriever(Retriever):
         )
         related_paper = set()
         related_paper.update(sn_paper_id_list)
-        for paper_id in sn_paper_id_list:
-            sn_entities += self.paper_client.find_entities_by_paper(paper_id)
+        sn_entities += self.paper_client.find_entities_by_paper_list(sn_paper_id_list)
         logger.debug("SN entities for retriever: {}".format(sn_entities))
         entities = list(set(entities + sn_entities))
         new_entities = self.retrieve_entities_by_enties(entities)