init commit

.gitignore

@@ -0,0 +1,22 @@
+**/*.pdf
+**/*.json
+!assets/**/*.json
+!configs/**/*.json
+**/__pycache__
+**/result
+**/datasets
+**/*~HEAD
+**/mlruns
+**/wandb
+**/*_QLoRA
+**/dqz_*
+datasets
+**/*.log
+assets/data/scipip_neo4j_clean_backup.json
+assets/paper/
+tmp
+test/
+# configs
+
+# vs code
+.history

configs/config.py

@@ -0,0 +1,212 @@
+r"""_summary_
+-*- coding: utf-8 -*-
+
+Module : configs.config
+
+File Name : config.py
+
+Description : Load the config file, which supports referencing other configuration files. If a circular reference occurs, an exception will be thrown
+   
+Creation Date : 2024-08-18
+
+Author : Frank Kang([email protected])
+"""
+import pathlib
+import json
+
+import os
+import warnings
+
+from typing import Union, Any, IO
+from omegaconf import OmegaConf, DictConfig, ListConfig
+
+from .utils import get_dir
+
+INCLUDE_KEY = 'include'
+
+
+def get_api_aliases(llms_api, sum_api, gen_api):
+    if sum_api is None:
+        if llms_api is not None:
+            sum_api = llms_api
+        else:
+            sum_api = 'ZhipuAI'
+
+    if gen_api is None:
+        if llms_api is not None:
+            gen_api = llms_api
+        else:
+            gen_api = 'OpenAI'
+
+    return sum_api, gen_api
+
+
+def check_api_alias(config, api):
+    api = api.lower()
+    for k in config.keys():
+        if k.lower() == api:
+            return k
+    return None
+
+
+def update_config_with_api_aliases(config, llms_api, sum_api, gen_api):
+    sum_api, gen_api = get_api_aliases(llms_api, sum_api, gen_api)
+    sum_api_found = check_api_alias(config, sum_api)
+    if sum_api_found is None:
+        raise KeyError('{} cannot match any llms api in config'.format(sum_api))
+    gen_api_found = check_api_alias(config, gen_api)
+    if gen_api_found is None:
+        raise KeyError('{} cannot match any llms api in config'.format(gen_api))
+    config.used_llms_apis = {'summarization': sum_api_found, 'generation': gen_api_found}
+
+
+class ConfigReader:
+    """_summary_
+    Load the config file, which supports referencing other configuration files. If a circular reference occurs, an exception will be thrown
+
+    for examples:
+    ```
+    config = ConfigReader.load(file)
+    ```
+    """
+
+    def __init__(
+        self,
+        file_: Union[str, pathlib.Path, IO[Any]],
+        included: set | None = None
+    ) -> None:
+        """_summary_
+
+        Args:
+            file_ (Union[str, pathlib.Path, IO[Any]]): config
+            included (set | None, optional): Include config file. Defaults to None.
+
+        Raises:
+            FileNotFoundError: If the configuration file cannot be found
+            RecursionError: If there is a loop include
+        """
+        fname = ''
+        self.included = included if included is not None else set()
+        if isinstance(file_, str):
+            fname = file_
+            if not os.path.exists(fname):
+                template_path = '{}.template'.format(fname)
+                if os.path.exists(template_path):
+                    with open(fname, 'w', encoding='utf8') as wf:
+                        with open(template_path, 'r', encoding='utf8') as rf:
+                            wf.write(rf.read())
+                    warnings.warn(
+                        'cannot find file {}. Auto generate from {}'.format(
+                            fname, template_path))
+                else:
+                    raise FileNotFoundError(
+                        'cannot find file {}'.format(fname))
+        else:
+            fname = file_.name
+
+        suffix = fname.split('.')[-1]
+        if suffix == 'yaml':
+            config = OmegaConf.load(fname)
+        elif suffix == 'json':
+            if isinstance(file_, (str, IO[Any])):
+                with open(file_, 'r', encoding='utf8') as f:
+                    config = json.load(f)
+            else:
+                config = json.load(file_)
+            config = DictConfig(config)
+        if fname not in self.included:
+            self.included.add(fname)
+        else:
+            raise RecursionError()
+        self.__config = config
+        self.complied = False
+
+    def complie(self, config: DictConfig | None = None):
+        """_summary_
+
+        Resolve config to make include effective
+
+        Args:
+            config (DictConfig | None, optional): dict config. Defaults to None.
+
+        Raises:
+            RecursionError: If there is a loop include
+        """
+        modify_flag = False
+        if config is None:
+            config = self.__config
+            modify_flag = True
+
+        include_item = None
+
+        if INCLUDE_KEY in config.keys():
+            include_value = config.get(INCLUDE_KEY)
+            if isinstance(include_value, (list, ListConfig)):
+                include_item = [get_dir(p) for p in include_value]
+            else:
+                include_item = get_dir(include_value)
+        for key in config.keys():
+            value = config.get(key)
+            if isinstance(value, DictConfig):
+                self.complie(value)
+
+        if include_item is not None:
+            if isinstance(include_item, str):
+                included = self.included.copy()
+                if include_item in included:
+                    print(include_item, included)
+                    raise RecursionError()
+                included.add(include_item)
+                config.merge_with(ConfigReader.load(include_item, included))
+
+            else:
+                for item in include_item:
+                    included = self.included.copy()
+                    if item in included:
+                        print(include_item, included)
+                        raise RecursionError()
+                    config.merge_with(ConfigReader.load(item, included))
+                    included.add(item)
+
+        if modify_flag:
+            self.complied = True
+
+    @property
+    def config(self) -> DictConfig:
+        """_summary_
+
+        Obtain parsed dict config
+
+        Returns:
+            DictConfig: parsed dict config
+        """
+        if not self.complied:
+            self.complie()
+        return self.__config
+
+    @staticmethod
+    def load(
+        file_: Union[str, pathlib.Path, IO[Any]],
+        included: set | None = None,
+        **kwargs
+    ) -> DictConfig:
+        """_summary_
+
+        Class method loading configuration file
+
+        Args:
+            file_ (Union[str, pathlib.Path, IO[Any]]): config
+            included (set | None, optional): Include config file. Defaults to None.
+
+        Returns:
+            DictConfig: parsed dict config
+        """
+        config = ConfigReader(file_, included).config
+        if 'llms_api' in kwargs and 'sum_api' in kwargs and 'gen_api' in kwargs:
+            update_config_with_api_aliases(config, kwargs['llms_api'], kwargs['sum_api'], kwargs['gen_api'])
+            del kwargs['llms_api']
+            del kwargs['sum_api']
+            del kwargs['gen_api']
+        for k, v in kwargs.items():
+            config[k] = v
+        return config

configs/datasets.yaml

@@ -0,0 +1,38 @@
+DEFAULT:
+    pdf_cached: /data/llms/data/scipip-data/pdf_cached
+    ignore_paper_id_list: ./assets/data/ignore_paper_id_list.json
+    log_level: "DEBUG"
+    log_dir: ./log
+    embedding: "sentence-transformers/all-MiniLM-L6-v2"
+    device: "cpu" # "cpu"
+
+ARTICLE:
+    summarizing_prompt: ./prompt/summarizing.xml
+
+RETRIEVE:
+    cite_type: "all_cite_id_list"
+    limit_num: 100  # 限制entity对应的paper数量
+    sn_num_for_entity: 5 # SN搜索的文章数量，扩充entity
+    kg_jump_num: 1   # 跳数
+    kg_cover_num: 3  # entity重合数量
+    max_paper_num_after_filter: 10 # 过滤后最多保留的论文数量
+    min_paper_num_after_filter: 5 # 过滤后最多保留的论文数量
+    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
+    relation_name: "related" # "connect"
+    top_p_list: [0.1, 0.2, 0.3, 0.4, 0.5]
+    top_k_list: [10, 20, 30, 40, 50]
+    s_bg: 0
+    s_contribution: 0.5
+    s_summary: 0.5
+    similarity_threshold: 0.55
+
+used_llms_apis:
+    summarization: ZhipuAI
+    generation: OpenAI

configs/utils.py

@@ -0,0 +1,13 @@
+#!/usr/bin/env python
+# coding: utf-8
+# Author: Frank Kang
+# Data: 13 July 2024
+import os
+ROOT = os.path.dirname(os.path.dirname(__file__))
+
+
+def get_dir(config_dir):
+    if config_dir.startswith('.'):
+        return os.path.realpath(os.path.join(ROOT, config_dir))
+    else:
+        return os.path.realpath(config_dir)

requirements.txt

@@ -0,0 +1,155 @@
+altair==5.4.1
+annotated-types==0.7.0
+antlr4-python3-runtime==4.9.3
+anyio==4.6.2.post1
+arrow==1.3.0
+asttokens==2.4.1
+attrs==23.2.0
+beautifulsoup4==4.12.3
+bibtexparser==1.4.2
+blinker==1.8.2
+blis==0.7.11
+Brotli==1.0.9
+cachetools==5.5.0
+catalogue==2.0.10
+certifi==2024.8.30
+cffi==1.17.1
+charset-normalizer==3.3.2
+click==8.1.7
+cloudpathlib==0.16.0
+comm==0.2.2
+confection==0.1.5
+cryptography==43.0.0
+cymem==2.0.8
+debugpy==1.8.7
+decorator==5.1.1
+Deprecated==1.2.14
+distro==1.9.0
+docker==7.0.0
+dockerpty==0.4.1
+docopt==0.6.2
+exceptiongroup==1.2.2
+executing==2.1.0
+filelock==3.16.1
+free-proxy==1.1.2
+fsspec==2024.10.0
+gitdb==4.0.11
+GitPython==3.1.43
+gmpy2==2.1.2
+h11==0.14.0
+httpcore==1.0.6
+httpx==0.27.2
+idna==3.7
+interchange==2021.0.4
+ipykernel==6.29.5
+ipython==8.29.0
+jedi==0.19.1
+Jinja2==3.1.4
+joblib==1.4.2
+jsonschema==3.2.0
+jupyter_client==8.6.3
+jupyter_core==5.7.2
+langcodes==3.4.1
+language_data==1.2.0
+loguru==0.7.2
+lxml==5.3.0
+marisa-trie==1.2.1
+markdown-it-py==3.0.0
+MarkupSafe==2.1.3
+matplotlib-inline==0.1.7
+mdurl==0.1.2
+monotonic==1.6
+mpmath==1.3.0
+murmurhash==1.0.10
+narwhals==1.13.1
+neo4j==5.21.0
+nest-asyncio==1.6.0
+networkx==3.3
+numpy==1.26.0
+omegaconf==2.3.0
+openai==1.12.0
+outcome==1.3.0.post0
+packaging==23.2
+pandas==2.2.3
+pansi==2020.7.3
+parso==0.8.4
+pexpect==4.9.0
+pillow==10.4.0
+pip==24.2
+platformdirs==4.3.6
+preshed==3.0.9
+prompt_toolkit==3.0.48
+protobuf==5.28.3
+psutil==6.1.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+py2neo==2021.2.4
+pyarrow==18.0.0
+pycparser==2.21
+pydantic==2.9.2
+pydantic_core==2.23.4
+pydeck==0.9.1
+Pygments==2.18.0
+PyJWT==2.8.0
+pyOpenSSL==24.2.1
+pyparsing==3.2.0
+pyphen==0.17.0
+pyrsistent==0.20.0
+PySocks==1.7.1
+python-dateutil==2.9.0.post0
+python-dotenv==0.21.1
+pytz==2024.2
+PyYAML==6.0
+pyzmq==26.2.0
+regex==2024.9.11
+requests==2.31.0
+rich==13.9.4
+safetensors==0.4.5
+scikit-learn==1.5.2
+scipy==1.14.1
+selenium==4.25.0
+sentence-transformers==3.0.1
+setuptools==68.0.0
+shellingham==1.5.4
+six==1.16.0
+smart-open==6.4.0
+smmap==5.0.1
+sniffio==1.3.1
+sortedcontainers==2.4.0
+soupsieve==2.6
+spacy==3.7.4
+spacy-legacy==3.0.12
+spacy-loggers==1.0.5
+srsly==2.4.8
+stack-data==0.6.3
+streamlit==1.39.0
+sympy==1.13.2
+tenacity==9.0.0
+textstat==0.7.4
+texttable==1.7.0
+thinc==8.2.5
+threadpoolctl==3.5.0
+tokenizers==0.19.1
+toml==0.10.2
+torch==2.1.0
+tornado==6.4.1
+tqdm==4.66.6
+traitlets==5.14.3
+transformers==4.44.0
+trio==0.27.0
+trio-websocket==0.11.1
+triton==2.1.0
+typer==0.9.4
+types-python-dateutil==2.9.0.20241003
+typing_extensions==4.12.2
+tzdata==2024.2
+urllib3==1.26.19
+wasabi==1.1.3
+watchdog==5.0.3
+wcwidth==0.2.13
+weasel==0.3.4
+websocket-client==1.8.0
+wheel==0.44.0
+wrapt==1.16.0
+wsproto==1.2.0
+zhipuai==2.1.5.20230904

src/ai_scientist_idea.py

@@ -0,0 +1,124 @@
+# ai scientist 生成 idea
+# Reference: https://github.com/SakanaAI/AI-Scientist/blob/main/ai_scientist/generate_ideas.py
+from utils.paper_retriever import RetrieverFactory
+from utils.llms_api import APIHelper
+from utils.header import ConfigReader
+from omegaconf import OmegaConf
+import click
+import json
+from loguru import logger
+import warnings
+import time
+warnings.filterwarnings('ignore')
+
+
+class AiScientistIdeaGenerator():
+    def __init__(self, config) -> None:
+        self.api_helper = APIHelper(config)
+
+    def generate(message_input):
+        # @LuoYunxiang 
+        return
+
+@click.group()
+@click.pass_context
+def main(ctx):
+    """
+    Training and evaluation
+    """
+    print("Mode:", ctx.invoked_subcommand)
+
+@main.command()
+@click.option(
+    "-c",
+    "--config-path",
+    default='../configs/datasets.yaml',
+    type=click.File(),
+    required=True,
+    help="Dataset configuration file in YAML",
+)
+@click.option(
+    "--ids-path",
+    default='../assets/data/test_acl_2024.json',
+    type=click.File(),
+    required=True,
+    help="Dataset configuration file in YAML",
+)
+@click.option(
+    "-r",
+    "--retriever-name",
+    default='SNKG',
+    type=str,
+    required=True,
+    help="Retrieve method",
+)
+@click.option(
+    "--llms-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API alias used. If you do not have separate APIs for summarization and generation, you can use this unified setting. This option is ignored when setting the API to be used by summarization and generation separately",
+)
+@click.option(
+    "--sum-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API aliases used for summarization. When used, it will invalidate --llms-api",
+)
+@click.option(
+    "--gen-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API aliases used for generation. When used, it will invalidate --llms-api",
+)
+def generate(config_path, ids_path, retriever_name, **kwargs):
+    logger.add("ai_scientist_generate_{}.log".format(retriever_name), level="DEBUG")
+    logger.info("Retrieve name: {}".format(retriever_name))
+    # Configuration
+    config = ConfigReader.load(config_path, **kwargs)
+    api_helper = APIHelper(config)
+    eval_data = []
+    num = 0
+    for line in ids_path:
+        # Parse each line's JSON data
+        background = json.loads(line)
+        bg = background["background"]
+        entities = api_helper.generate_entity_list(bg)
+        logger.debug("Original entities from background: {}".format(entities))
+        rt = RetrieverFactory.get_retriever_factory().create_retriever(
+            retriever_name, 
+            config, 
+            use_cocite=config.RETRIEVE.use_cocite,
+            use_cluster_to_filter=config.RETRIEVE.use_cluster_to_filter
+        )
+        result = rt.retrieve(bg, entities, need_evaluate=False, target_paper_id_list=[], top_k=5)
+        related_paper = result["related_paper"]
+        logger.info("Find {} related papers...".format(len(related_paper)))
+        title_list = [paper["title"] for paper in related_paper]
+        contribution_list = [paper["summary"] for paper in related_paper]
+        message_input =  {
+          "Name": ",".join(entities),
+          "Title": ",".join(title_list),
+          "Experiment": ",".join(contribution_list)
+        }
+        print(message_input)
+        exit()
+        idea_generator = AiScientistIdeaGenerator(config)
+        # idea list(str)
+        idea = idea_generator.generate(message_input)
+        eval_data.append({
+            "background": bg,
+            "input": message_input,
+            "pred": idea
+        })
+        num += 1
+        if num >= 1:
+            break
+    logger.info("=== Finish ===")
+    with open("ai_scientist_output_new_idea.json", "w", encoding="utf-8") as f:
+        json.dump(eval_data, f, ensure_ascii=False, indent=4)
+
+if __name__ == "__main__":
+    main()

src/generator.py

@@ -0,0 +1,671 @@
+from utils.paper_retriever import RetrieverFactory
+from utils.paper_client import PaperClient
+from utils.llms_api import APIHelper
+from utils.header import ConfigReader
+from omegaconf import OmegaConf
+import click
+import json
+from loguru import logger
+import warnings
+import time
+import os
+
+warnings.filterwarnings("ignore")
+
+
+def extract_problem(problem, background):
+    start_keyword = "**Research Problem**"
+    end_keyword = "**Rationales**"
+    start_index = problem.find(start_keyword)
+    end_index = problem.find(end_keyword)
+    if start_index != -1 and end_index != -1:
+        research_problem = problem[start_index:end_index].strip()
+    else:
+        research_problem = background
+    return research_problem
+
+class IdeaGenerator:
+    def __init__(
+        self, config, paper_list: list[dict], cue_words: list = None, brainstorm: str = None
+    ) -> None:
+        self.api_helper = APIHelper(config)
+        self.paper_list = paper_list
+        self.cue_words = cue_words
+        self.brainstorm = brainstorm
+
+    def generate_with_cue_words(self, background: str):
+        problem, message_input = self.api_helper.generate_problem_with_cue_words(
+            background, self.paper_list, self.cue_words
+        )
+        idea = self.api_helper.generate_idea_with_cue_words(
+            problem, self.paper_list, self.cue_words
+        )
+        idea_filtered = self.api_helper.filter_idea(idea, background)
+        return message_input, problem, idea, idea_filtered
+
+    def generate_without_cue_words(self, background: str):
+        problem, message_input = self.api_helper.generate_problem(
+            background, self.paper_list
+        )
+        idea = self.api_helper.generate_idea(problem, self.paper_list)
+        idea_filtered = self.api_helper.filter_idea(idea, background)
+        return message_input, problem, idea, idea_filtered
+
+    def generate_with_cue_words_bs(self, background: str):
+        problem, message_input = self.api_helper.generate_problem_with_cue_words(
+            background, self.paper_list, self.cue_words
+        )
+        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)
+        return message_input, problem, idea, idea_filtered
+
+    def generate_without_cue_words_bs(self, background: str):
+        problem, message_input = self.api_helper.generate_problem(
+            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)
+        return message_input, problem, idea, idea_filtered
+
+    def generate_with_cue_words_ins(self, background: str):
+        problem, message_input = self.api_helper.generate_problem_with_cue_words(
+            background, self.paper_list, self.cue_words
+        )
+        research_problem = extract_problem(problem, background)
+        inspirations = []
+        for paper in self.paper_list:
+            inspiration = self.api_helper.generate_inspiration_with_cue_words(
+                research_problem, paper, self.cue_words
+            )
+            inspirations.append(inspiration)
+        idea = self.api_helper.generate_idea_by_inspiration_with_cue_words(
+            problem, inspirations, self.cue_words
+        )
+        idea_filtered = self.api_helper.filter_idea(idea, background)
+        return message_input, problem, inspirations, idea, idea_filtered
+
+    def generate_without_cue_words_ins(self, background: str):
+        problem, message_input = self.api_helper.generate_problem(
+            background, self.paper_list
+        )
+        research_problem = extract_problem(problem, background)
+        inspirations = []
+        for paper in self.paper_list:
+            inspiration = self.api_helper.generate_inspiration(
+                research_problem, paper
+            )
+            inspirations.append(inspiration)
+        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
+        )
+        research_problem = extract_problem(problem, background)
+        inspirations = []
+        for paper in self.paper_list:
+            inspiration = self.api_helper.generate_inspiration_with_cue_words(
+                research_problem, paper, self.cue_words
+            )
+            inspirations.append(inspiration)
+        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)
+        return message_input, problem, inspirations, idea, idea_filtered
+
+    def generate_without_cue_words_ins_bs(self, background: str):
+        problem, message_input = self.api_helper.generate_problem(
+            background, self.paper_list
+        )
+        research_problem = extract_problem(problem, background)
+        inspirations = []
+        for paper in self.paper_list:
+            inspiration = self.api_helper.generate_inspiration(
+                research_problem, paper
+            )
+            inspirations.append(inspiration)
+        idea = self.api_helper.generate_idea_by_inspiration(
+            problem, inspirations
+        )
+        idea_filtered = self.api_helper.integrate_idea(background, self.brainstorm, idea)
+        return message_input, problem, inspirations, idea, idea_filtered
+
+    def generate(
+        self,
+        background: str,
+        mode: str,
+        bs_mode: str = None,
+        use_cue_words: bool = False,
+    ):
+        mode_name = None
+        if mode == "backtracking":
+            mode_name = "Backtrack"
+        elif mode == "new_idea":
+            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
+                ) = (
+                    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
+                ) = (
+                    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
+                ) = (
+                    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
+                ) = (
+                    self.generate_without_cue_words_bs(background)
+                )
+
+        idea_modified = self.api_helper.modify_idea(background, idea_filtered)
+        median = {
+            "problem": problem,
+            "initial_idea": idea,
+            "filtered_idea": idea_filtered,
+        }
+        print("=====")
+        print(idea_modified)
+        print("=====")
+        exit()
+        return message_input, idea_modified, median
+
+    def generate_by_inspiration(
+        self,
+        background: str,
+        mode: str,
+        bs_mode: str = None,
+        use_cue_words: bool = False,
+    ):
+        mode_name = None
+        if mode == "backtracking":
+            mode_name = "Backtrack"
+        elif mode == "new_idea":
+            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
+                ) = (
+                    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
+                ) = (
+                    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
+                ) = (
+                    self.generate_with_cue_words_ins_bs(background)
+                )
+            else:
+                logger.info("{} using brainstorm_{} with cue words.".format(mode_name, bs_mode))
+                (
+                    message_input, 
+                    problem, 
+                    inspirations,
+                    idea, 
+                    idea_filtered
+                ) = (
+                    self.generate_without_cue_words_ins_bs(background)
+                )
+
+        idea_modified = self.api_helper.modify_idea(background, idea_filtered)
+        median = {
+            "problem": problem,
+            "inspirations": inspirations,
+            "initial_idea": idea,
+            "filtered_idea": idea_filtered,
+        }
+        return message_input, idea_modified, median
+
+
+@click.group()
+@click.pass_context
+def main(ctx):
+    """
+    Training and evaluation
+    """
+    print("Mode:", ctx.invoked_subcommand)
+
+
+@main.command()
+@click.option(
+    "-c",
+    "--config-path",
+    default="./configs/datasets.yaml",
+    type=click.File(),
+    required=True,
+    help="Dataset configuration file in YAML",
+)
+@click.option(
+    "--ids-path",
+    default="./assets/data/test_acl_2024.json",
+    type=click.File(),
+    required=True,
+    help="Dataset configuration file in YAML",
+)
+@click.option(
+    "-r",
+    "--retriever-name",
+    default="SNKG",
+    type=str,
+    required=True,
+    help="Retrieve method",
+)
+@click.option(
+    "--brainstorm-mode",
+    default="mode_a",
+    type=str,
+    required=True,
+    help="Choose your brainstorm mode (mode_a: no brainstorm, mode_b: brainstorm for idea generation, mode_c: brainstorm for idea generation and retrival)",
+)
+@click.option(
+    "--use-cue-words",
+    default=True,
+    type=bool,
+    required=True,
+    help="Use cue words in generation",
+)
+@click.option(
+    "--use-inspiration",
+    default=False,
+    type=bool,
+    required=True,
+    help="Use inspiration in generation",
+)
+@click.option(
+    "--llms-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API alias used. If you do not have separate APIs for summarization and generation, you can use this unified setting. This option is ignored when setting the API to be used by summarization and generation separately",
+)
+@click.option(
+    "--sum-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API aliases used for summarization. When used, it will invalidate --llms-api",
+)
+@click.option(
+    "--gen-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API aliases used for generation. When used, it will invalidate --llms-api",
+)
+@click.option(
+    "--num",
+    default=100,
+    type=int,
+    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):
+    # Configuration
+    config = ConfigReader.load(config_path, **kwargs)
+    logger.add(
+        "log/generate_{}_{}.log".format(time.time(), retriever_name),
+        level=config.DEFAULT.log_level,
+    )
+    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)
+    eval_data = []
+    processed_ids = set()
+    cur_num = 0
+    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")
+    if os.path.exists(output_file):
+        with open(output_file, "r", encoding="utf-8") as f:
+            try:
+                eval_data = json.load(f)
+                processed_ids = {paper["hash_id"] for paper in eval_data}
+                cur_num = len(eval_data)
+            except json.JSONDecodeError:
+                print("Failed to decode JSON, initializing eval_data as an empty list.")
+    print(f"{cur_num} papers have been processed.")
+    for line in ids_path:
+        # 解析每行的JSON数据
+        paper = json.loads(line)
+        if paper["hash_id"] in processed_ids:
+            print(f"Skipping already processed paper: {paper_id}")
+            continue
+        logger.info("\nbegin generate paper hash id {}".format(paper["hash_id"]))
+        # if "entities" in paper.keys():
+        #     entities = paper["entities"]
+        # else:
+        # 1. 获取背景信息
+        paper = paper_client.get_paper_by_id(paper["hash_id"])
+        if "motivation" in paper.keys():
+            bg = paper["motivation"]
+        else:
+            print(f"Paper hash_id {paper['hash_id']} doesn't have background...")
+            continue
+        if brainstorm_mode == "mode_b" or brainstorm_mode == "mode_c":
+            brainstorm = api_helper.generate_brainstorm(bg)
+        else:
+            brainstorm = None
+        if "entities" in paper.keys():
+            entities = paper["entities"]
+        else:
+            entities = api_helper.generate_entity_list(bg)
+        logger.debug("Original entities from background: {}".format(entities))
+        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))
+        else:
+            entities_bs = None
+            entities_all = entities
+        # 2. 获取真实引用文章 (用于评估)
+        cite_type = "cite_id_list"
+        # cite_type = config.RETRIEVE.cite_type
+        if cite_type in paper and len(paper[cite_type]) >= 5:
+            target_paper_id_list = paper[cite_type]
+        else:
+            logger.warning(
+                "Hash ID {} cited paper num less than 5...".format(paper["hash_id"])
+            )
+            continue
+        # 3. 检索相关论文
+        rt = RetrieverFactory.get_retriever_factory().create_retriever(
+            retriever_name, 
+            config, 
+            use_cocite=True, 
+            use_cluster_to_filter=True
+        )
+        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"]
+        # 4. 生成IDEA
+        if use_cue_words:
+            if "contribution" in paper.keys():
+                cue_words = api_helper.generate_entity_list(paper["contribution"])
+            else:
+                print(f"Paper hash_id {paper['hash_id']} doesn't have contribution...")
+                cue_words = None
+        else:
+            cue_words = None
+        idea_generator = IdeaGenerator(config, related_paper, cue_words, brainstorm)
+        if not use_inspiration:
+            message_input, idea_modified, median = idea_generator.generate(
+                bg, "backtracking", brainstorm_mode, use_cue_words
+            )
+        else:
+            message_input, idea_modified, median = (
+                idea_generator.generate_by_inspiration(
+                    bg, "backtracking", brainstorm_mode, use_cue_words
+                )
+            )
+        eval_data.append(
+            {
+                "hash_id": paper["hash_id"],
+                "background": bg,
+                "entities_bg": entities,
+                "brainstorm" : brainstorm,
+                "entities_bs": entities_bs,
+                "entities_rt": entities_rt,
+                "related_paper": [p["hash_id"] for p in related_paper],
+                "input": message_input,
+                "cue_words": cue_words,
+                "median": median,
+                "pred": idea_modified,
+                "ground_truth": paper["ground_truth"],
+            }
+        )
+        cur_num += 1
+        if cur_num % batch_size == 0:
+            with open(
+                output_file,
+                "w",
+                encoding="utf-8",
+            ) as f:
+                json.dump(eval_data, f, ensure_ascii=False, indent=4)
+        if cur_num >= num:
+            break
+    logger.info("=== Finish ===")
+    with open(
+        output_file,
+        "w",
+        encoding="utf-8",
+    ) as f:
+        json.dump(eval_data, f, ensure_ascii=False, indent=4)
+
+@main.command()
+@click.option(
+    "-c",
+    "--config-path",
+    default="./configs/datasets.yaml",
+    type=click.File(),
+    required=True,
+    help="Dataset configuration file in YAML",
+)
+@click.option(
+    "--ids-path",
+    default="./assets/data/test_background.json",
+    type=click.File(),
+    required=True,
+    help="Dataset configuration file in YAML",
+)
+@click.option(
+    "-r",
+    "--retriever-name",
+    default="SNKG",
+    type=str,
+    required=True,
+    help="Retrieve method",
+)
+@click.option(
+    "--brainstorm-mode",
+    default="mode_a",
+    type=str,
+    required=True,
+    help="Choose your brainstorm mode (mode_a: no brainstorm, mode_b: brainstorm for idea generation, mode_c: brainstorm for idea generation and retrival)",
+)
+@click.option(
+    "--use-inspiration",
+    default=False,
+    type=bool,
+    required=True,
+    help="Use inspiration in generation",
+)
+@click.option(
+    "--llms-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API alias used. If you do not have separate APIs for summarization and generation, you can use this unified setting. This option is ignored when setting the API to be used by summarization and generation separately",
+)
+@click.option(
+    "--sum-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API aliases used for summarization. When used, it will invalidate --llms-api",
+)
+@click.option(
+    "--gen-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API aliases used for generation. When used, it will invalidate --llms-api",
+)
+@click.option(
+    "--num",
+    default=100,
+    type=int,
+    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):
+    logger.add(
+        "log/generate_{}_{}.log".format(time.time(), retriever_name), level="DEBUG"
+    )  # 添加文件输出
+    logger.info("Retrieve name: {}".format(retriever_name))
+    # Configuration
+    config = ConfigReader.load(config_path, **kwargs)
+    api_helper = APIHelper(config)
+    eval_data = []
+    cur_num = 0
+    data_num = 0
+    batch_size = 2
+    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")
+    if os.path.exists(output_file):
+        with open(output_file, "r", encoding="utf-8") as f:
+            try:
+                eval_data = json.load(f)
+                bg_ids = {data["background"] for data in eval_data}
+                cur_num = len(eval_data)
+            except json.JSONDecodeError:
+                eval_data = []
+    print(f"{cur_num} datas have been processed.")
+    for line in ids_path:
+        # 解析每行的JSON数据
+        data = json.loads(line)
+        # 1. 获取背景信息
+        if "background" in data.keys():
+            bg = data["background"]
+        else:
+            data_num += 1
+            print(f"This data doesn't have background...")
+            continue
+        if bg in bg_ids:
+            data_num += 1
+            print(f"Skipping already processed data_{data_num}.")
+            continue
+        if brainstorm_mode == "mode_b" or brainstorm_mode == "mode_c":
+            brainstorm = api_helper.generate_brainstorm(bg)
+        else:
+            brainstorm = None
+        if "cue_words" in data.keys():
+            use_cue_words = True
+            cue_words = data["cue_words"]
+        else:
+            use_cue_words = False
+            cue_words = None
+        entities = api_helper.generate_entity_list(bg)
+        logger.debug("Original entities from background: {}".format(entities))
+        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))
+        else:
+            entities_bs = None
+            entities_all = entities
+        # 2. 检索相关论文
+        rt = RetrieverFactory.get_retriever_factory().create_retriever(
+            retriever_name,
+            config,
+            use_cocite=config.RETRIEVE.use_cocite,
+            use_cluster_to_filter=config.RETRIEVE.use_cluster_to_filter,
+        )
+        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"]
+        # 3. 生成IDEA
+        idea_generator = IdeaGenerator(config, related_paper, cue_words, brainstorm)
+        if not use_inspiration:
+            message_input, idea_modified, median = idea_generator.generate(
+                bg, "new_idea", brainstorm_mode, use_cue_words
+            )
+        else:
+            message_input, idea_modified, median = (
+                idea_generator.generate_by_inspiration(
+                    bg, "new_idea", brainstorm_mode, use_cue_words
+                )
+            )
+        eval_data.append(
+            {
+                "background": bg,
+                "entities_bg": entities,
+                "brainstorm" : brainstorm,
+                "entities_bs": entities_bs,
+                "entities_rt": entities_rt,
+                "related_paper": [p["hash_id"] for p in related_paper],
+                "input": message_input,
+                "cue_words": cue_words,
+                "median": median,
+                "pred": idea_modified,
+            }
+        )
+        cur_num += 1
+        if cur_num % batch_size == 0:
+            with open(
+                output_file,
+                "w",
+                encoding="utf-8",
+            ) as f:
+                json.dump(eval_data, f, ensure_ascii=False, indent=4)
+        if cur_num >= num:
+            break
+    logger.info("=== Finish ===")
+    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/pages/app_gradio_backup.py

@@ -0,0 +1,89 @@
+import gradio as gr
+from button_interface import Backend
+from generator import APIHelper
+from utils.header import ConfigReader
+
+DEBUG_MODE = False
+
+def generate_page(backend):
+    with gr.Blocks(title="Scientific Paper Idea Proposer") as demo:
+        ## Background, keywords parts
+        gr.Markdown(
+            """
+    # Scientific Paper Idea Proposer
+    """
+        )
+        # with gr.Blocks(theme="earneleh/paris") as d:
+        with gr.Blocks() as d:
+            with gr.Tab("Keywords"):
+                key_words = gr.Textbox(placeholder="Interested key words", label="Keywords (Provide at least 1 keyword)")
+            with gr.Tab("Background"):
+                background = gr.Textbox(placeholder="Background", label="Background")
+            if DEBUG_MODE:
+                with gr.Tab("Json"):
+                    json_file = gr.File()
+                    json_background = gr.Textbox(placeholder="Background", label="Background")
+                    json_strs = gr.Textbox(visible=False)
+                    json_file.upload(backend.upload_json_callback, inputs=[json_file], outputs=[json_background])
+            else:
+                json_strs = None
+
+        ## brainstorm ideas parts
+        # background2brainstorm = gr.Button("Continue (background2brainstorm)")
+        with gr.Row(equal_height=True):
+            gr.ClearButton(value="🆑 Clear", components=[background], scale=1)
+            background2brainstorm = gr.Button("😈 Brainstorm", scale=1)
+        # @gr.render(inputs=None, triggers=[background2brainstorm.click])
+        # def show_brainstorm():
+        # with gr.Accordion("Braining Ideas", open=False) as a1:
+        with gr.Row(equal_height=True):
+            brainstorm_txt = gr.Textbox(placeholder="Generated brainstorm ideas", label="Brainstorm ideas", info="Feel free to improve them before next step", max_lines=500)
+            brainstorm_md = gr.Markdown(label="Brainstorm ideas")
+
+        ## Expanded key words parts
+        # brainstorm2entities = gr.Button("Continue (brainstorm2entities)")
+        with gr.Row(equal_height=True):
+            gr.ClearButton(value="🆑 Clear", components=[brainstorm_txt], scale=1)
+            brainstorm2entities = gr.Button("Extract Entities", scale=1)
+        entities = gr.CheckboxGroup([], label="Expanded key words", visible=True)
+        entities2literature = gr.Button("📖 Retrieve Literature")
+        literature_intact = gr.State()
+        # entities2literature = gr.Button("Continue (retrieve literature)")
+
+        ## Retrieved literature parts
+        retrieved_literature = gr.Textbox(placeholder="Retrieved literature", label="Retrieved related works", info="", max_lines=500)
+        # literature2initial_ideas = gr.Button("Continue (generate initial ideas)")
+        with gr.Row(equal_height=True):
+            gr.ClearButton(value="🆑 Clear", components=[retrieved_literature], scale=1)
+            literature2initial_ideas = gr.Button("🤖 Generate Initial ideas", scale=1)
+
+
+        ## Initial ideas parts
+        with gr.Row():
+            initial_ideas_txt = gr.Textbox(placeholder="Initial ideas", label="Initial ideas", info="Feel free to improve them before next step", max_lines=500)
+            initial_ideas_md = gr.Markdown(label="Initial ideas")
+        # initial2final = gr.Button("Continue (generate final ideas)")
+        with gr.Row(equal_height=True):
+            gr.ClearButton(value="🆑 Clear", components=[initial_ideas_txt], scale=1)
+            initial2final = gr.Button("🔥 Refine Ideas")
+
+        ## Final ideas parts
+        with gr.Row():
+            final_ideas_txt = gr.Textbox(placeholder="Final ideas", label="Final ideas", info="", max_lines=500)
+            final_ideas_md = gr.Markdown(label="Final ideas")
+
+        # register callback
+        background2brainstorm.click(backend.background2brainstorm_callback, inputs=[background], outputs=[brainstorm_txt])
+        brainstorm2entities.click(backend.brainstorm2entities_callback, inputs=[background, brainstorm_txt], outputs=[entities])
+        brainstorm_txt.change(lambda input: input, inputs=brainstorm_txt, outputs=brainstorm_md)
+        initial_ideas_txt.change(lambda input: input, inputs=initial_ideas_txt, outputs=initial_ideas_md)
+        final_ideas_txt.change(lambda input: input, inputs=final_ideas_txt, outputs=final_ideas_md)
+        entities2literature.click(backend.entities2literature_callback, inputs=[background, entities], outputs=[retrieved_literature, literature_intact])
+        literature2initial_ideas.click(backend.literature2initial_ideas_callback, inputs=[background, literature_intact], outputs=[initial_ideas_txt, final_ideas_txt])
+        initial2final.click(backend.initial2final_callback, inputs=[initial_ideas_txt], outputs=[final_ideas_txt])
+    return demo
+
+if __name__ == "__main__":
+    backend = Backend()
+    demo = generate_page(backend)
+    demo.launch(server_name="0.0.0.0", share=True)

src/pages/button_interface.py

@@ -0,0 +1,109 @@
+import json
+from utils.paper_retriever import RetrieverFactory
+from utils.llms_api import APIHelper
+from utils.header import ConfigReader
+from generator import IdeaGenerator
+
+class Backend(object):
+    def __init__(self) -> None:
+        CONFIG_PATH = "./configs/datasets.yaml"
+        RETRIEVER_NAME = "SNKG"
+        USE_INSPIRATION = True
+        BRAINSTORM_MODE = "mode_c"
+
+        self.config = ConfigReader.load(CONFIG_PATH)
+        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.idea_generator = IdeaGenerator(self.config, None)
+        self.use_inspiration = USE_INSPIRATION
+        self.brainstorm_mode = BRAINSTORM_MODE
+
+    def background2brainstorm_callback(self, background, json_strs=None):
+        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
+            json_contents = json.loads(json_strs)
+            entities_bg = json_contents["entities_bg"]
+            entities_bs = json_contents["entities_bs"]
+            entities_all = entities_bg + entities_bs
+            # return gr.CheckboxGroup(choices=entities, value=entities, label="Expanded key words", visible=True)
+            return entities_all
+        else:
+            entities_bg = self.api_helper.generate_entity_list(background)
+            entities_bs = self.api_helper.generate_entity_list(brainstorm, 10)
+            entities_all = list(set(entities_bg) | set(entities_bs))
+            # return extracted_entities
+            # return gr.CheckboxGroup(choices=entities_all, value=entities_all, label="Expanded key words", visible=True)
+            return entities_all
+
+    def upload_json_callback(self, input):
+        # print(type(input))
+        # print(len(input))
+        # print(input) # temp file path
+        with open(input, "r") as json_file:
+            contents = json_file.read()
+            json_contents = json.loads(contents)
+        return [json_contents["background"], contents]
+
+    def entities2literature_callback(self, background, entities, json_strs=None):
+        if json_strs is not None:
+            json_contents = json.loads(json_strs)
+            res = ""
+            for i, p in enumerate(json_contents["related_paper"]):
+                res += "%d. " % (i + 1) + str(p)
+                if i < len(json_contents["related_paper"]) - 1:
+                    res += "\n"
+            return res, res
+        else:
+            result = self.retriever_factory.retrieve(background, entities, need_evaluate=False, target_paper_id_list=[])
+            res = ""
+            for i, p in enumerate(result["related_paper"]):
+                res += "%d. " % (i + 1) + str(p["title"])
+                if i < len(result["related_paper"]) - 1:
+                    res += "\n"
+            return res, result["related_paper"]
+
+    def literature2initial_ideas_callback(self, background, retrieved_literature, json_strs=None):
+        if json_strs is not None:
+            json_contents = json.loads(json_strs)
+            return json_contents["median"]["filtered_idea"]
+        else:
+            self.idea_generator.paper_list = retrieved_literature
+            if self.use_inspiration:
+                message_input, idea_modified, median = (
+                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["filtered_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)
+            return json_contents["median"]["modified_idea"]
+        else:
+            return final_ideas
+
+    def get_demo_i(self, i):
+        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 "
+    "how they make decisions becomes increasingly difficult. Challenge: 1 The complexity of model "
+    "interpretation: The billions of parameters and nonlinear decision paths within large-scale language "
+    "models make it very difficult to track and interpret specific outputs. The existing interpretation "
+    "methods usually only provide a local perspective and are difficult to systematize. 2. Transparency "
+    "and Fairness: In specific scenarios, models may exhibit biased or discriminatory behavior. Ensuring "
+    "the transparency of these models, reducing bias, and providing credible explanations is one of the current challenges.")

src/pages/one_click_generation.py

@@ -0,0 +1,140 @@
+import streamlit as st
+
+# st.set_page_config(layout="wide", page_title="🦜🔗 Generate Idea Step-by-step")
+
+## Pipeline global state
+# 1.0: Input background is in progress
+# 2.0: Brainstorming is in progress
+#  2.5 Brainstorming is finished
+# 3.0: Extracting entities is in progress
+#  3.5 Extracting entities is finished
+# 4.0: Retrieving literature is in progress
+#  4.5 Retrieving ideas is finished
+# 5.0: Generating initial ideas is in progress
+#  5.5 Generating initial ideas is finished
+# 6.0: Generating final ideas is in progress
+#  6.5 Generating final ideas is finished
+if "global_state_one_click" not in st.session_state:
+    st.session_state["global_state_one_click"] = 1.0
+
+def generate_sidebar():
+    st.sidebar.header("SciPIP", divider="rainbow")
+    st.sidebar.markdown(
+        ("SciPIP will generate ideas in one click. The generation pipeline is the same as "
+        "step-by-step generation, but you are free from caring about intermediate outputs.")
+    )
+
+    pipeline_list = ["1. Input Background", "2. Brainstorming", "3. Extracting Entities", "4. Retrieving Related Works", 
+                     "5. Generate Initial Ideas", "6. Generate Final Ideas"]
+    st.sidebar.header("Pipeline", divider="red")
+    for i in range(6):
+        st.sidebar.markdown(f"<font color='black'>{pipeline_list[i]}</font>", unsafe_allow_html=True)
+
+    st.sidebar.header("Supported Fields", divider="orange")
+    st.sidebar.caption("The supported fields are temporarily limited because we only collect literature "
+               "from ICML, ICLR, NeurIPS, ACL, and EMNLP. Support for other fields are in progress.")
+    st.sidebar.checkbox("Natural Language Processing (NLP)", value=True, disabled=True)
+    st.sidebar.checkbox("Computer Vision (CV)", value=False, disabled=True)
+    st.sidebar.checkbox("[Partial] Multimodal", value=True, disabled=True)
+    st.sidebar.checkbox("Incoming Other Fields", value=False, disabled=True)
+
+    st.sidebar.header("Help Us To Improve", divider="green")
+    st.sidebar.markdown("https://forms.gle/YpLUrhqs1ahyCAe99", unsafe_allow_html=True)
+
+
+def genrate_mainpage(backend):
+    st.title('💧 Generate Idea in One-click')
+    # st.markdown("# 🐳 Background")
+    # st.markdown("Available soon...")
+
+    if "messages" not in st.session_state:
+        st.session_state["messages"] = [{"role": "assistant", "content": "Please give me some key words or a background"}]
+    if "intermediate_output" not in st.session_state:
+        st.session_state["intermediate_output"] = {}
+
+    for msg in st.session_state.messages:
+        st.chat_message(msg["role"]).write(msg["content"])
+
+    def disable_submit():
+        st.session_state["enable_submmit"] = False
+
+    if prompt := st.chat_input(disabled=not st.session_state.get("enable_submmit", True), on_submit=disable_submit):
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        st.chat_message("user").write(prompt)
+        generate_ideas(backend, prompt)
+    elif st.session_state.get("use_demo_input", False):
+        generate_ideas(backend, st.session_state.get("demo_input"))
+        st.session_state["use_demo_input"] = False
+        del(st.session_state["demo_input"])
+
+    def get_demo_n(i):
+        demo_input = backend.get_demo_i(i)
+        st.session_state["enable_submmit"] = False
+        st.session_state.messages.append({"role": "user", "content": demo_input})
+        st.session_state["use_demo_input"] = True
+        st.session_state["demo_input"] = demo_input
+
+    cols = st.columns([2, 2])
+    cols[0].button("Example 1", on_click=get_demo_n, args=(1,), use_container_width=True, disabled=not st.session_state.get("enable_submmit", True))
+    cols[1].button("Example 2", on_click=get_demo_n, args=(2,), use_container_width=True, disabled=not st.session_state.get("enable_submmit", True))
+
+    def check_intermediate_outputs(id="brainstorms"):
+        msg = st.session_state["intermediate_output"].get(id, None)
+        if msg is not None:
+            st.session_state.messages.append(msg)
+        else:
+            st.toast(f"No {id} now!")
+
+    def reset():
+        del(st.session_state["messages"])
+        st.session_state["enable_submmit"] = True
+        st.session_state["global_state_one_click"] = 1.0
+        st.toast(f"The chat has been reset!")
+
+    cols = st.columns([1, 1, 1, 1])
+    cols[0].button("Check Brainstorms", on_click=check_intermediate_outputs, args=("brainstorms",), use_container_width=True)
+    cols[1].button("Check Entities", on_click=check_intermediate_outputs, args=("entities",), use_container_width=True)
+    cols[2].button("Check Retrieved Papers", on_click=check_intermediate_outputs, args=("related_works",), use_container_width=True)
+    cols[3].button("Reset Chat", on_click=reset, use_container_width=True, type="primary")
+
+def generate_ideas(backend, background):
+    with st.spinner(text="Brainstorming..."):
+        brainstorms = backend.background2brainstorm_callback(background)
+        st.session_state["intermediate_output"]["brainstorms"] = {"role": "assistant", "content": brainstorms}
+        # st.chat_message("assistant").write(brainstorms)
+        st.session_state["global_state_one_click"] = 2.5
+
+    with st.spinner(text="Extracting entities..."):
+        entities = backend.brainstorm2entities_callback(background, brainstorms)
+        st.session_state["intermediate_output"]["entities"] = {"role": "assistant", "content": entities}
+        # st.chat_message("assistant").write(entities)
+        st.session_state["global_state_one_click"] = 3.5
+
+    with st.spinner(text="Retrieving related works..."):
+        msg = "My initial ideas are:"
+        related_works, related_works_intact = backend.entities2literature_callback(background, entities)
+        st.session_state["intermediate_output"]["related_works"] = {"role": "assistant", "content": related_works}
+        # st.chat_message("assistant").write(related_works)
+        st.session_state["global_state_one_click"] = 4.5
+
+    with st.spinner(text="Generating initial ideas..."):
+        msg = "My initial ideas are:"
+        initial_ideas, final_ideas = backend.literature2initial_ideas_callback(background, related_works_intact)
+        st.session_state.messages.append({"role": "assistant", "content": msg})
+        st.chat_message("assistant").write(msg)
+        st.session_state.messages.append({"role": "assistant", "content": initial_ideas})
+        st.chat_message("assistant").write(initial_ideas)
+        st.session_state["global_state_one_click"] = 5.5
+
+    with st.spinner(text="Generating final ideas..."):
+        msg = "My final ideas after refinement are:"
+        final_ideas = backend.initial2final_callback(initial_ideas, final_ideas)
+        st.session_state.messages.append({"role": "assistant", "content": msg})
+        st.chat_message("assistant").write(msg)
+        st.session_state.messages.append({"role": "assistant", "content": final_ideas})
+        st.chat_message("assistant").write(final_ideas)
+        st.session_state["global_state_one_click"] = 6.5
+
+def one_click_generation(backend):
+    generate_sidebar()
+    genrate_mainpage(backend)

src/pages/step_by_step_generation.py

@@ -0,0 +1,192 @@
+import ast
+import streamlit as st
+
+def generate_sidebar():
+    st.sidebar.header("About", divider="rainbow")
+    st.sidebar.markdown(
+        ("SciPIP will generate ideas step by step. The generation pipeline is the same as "
+        "one-click generation, While you can improve each part manually after SciPIP providing the manuscript.")
+    )
+
+    DONE_COLOR = "black"
+    UNDONE_COLOR = "gray"
+    # INPROGRESS_COLOR = "#4d9ee6"
+    INPROGRESS_COLOR = "black"
+    color_list = []
+    pipeline_list = ["1. Input Background", "2. Brainstorming", "3. Extracting Entities", "4. Retrieving Related Works", 
+                     "5. Generate Initial Ideas", "6. Generate Final Ideas"]
+    for i in range(1, 8):
+        if st.session_state["global_state_step"] < i:
+            color_list.append(UNDONE_COLOR)
+        elif st.session_state["global_state_step"] == i:
+            color_list.append(INPROGRESS_COLOR)
+        elif st.session_state["global_state_step"] > i:
+            color_list.append(DONE_COLOR)
+    st.sidebar.header("Pipeline", divider="red")
+    for i in range(6):
+        st.sidebar.markdown(f"<font color='{color_list[i]}'>{pipeline_list[i]}</font>", unsafe_allow_html=True)
+        # if st.session_state["global_state_step"] == i + 1:
+        #     st.sidebar.progress(50, text=None)
+
+    st.sidebar.header("Supported Fields", divider="orange")
+    st.sidebar.caption("The supported fields are temporarily limited because we only collect literature "
+               "from ICML, ICLR, NeurIPS, ACL, and EMNLP. Support for other fields are in progress.")
+    st.sidebar.checkbox("Natural Language Processing (NLP)", value=True, disabled=True)
+    st.sidebar.checkbox("Computer Vision (CV)", value=False, disabled=True)
+    st.sidebar.checkbox("[Partial] Multimodal", value=True, disabled=True)
+    st.sidebar.checkbox("Incoming Other Fields", value=False, disabled=True)
+
+    st.sidebar.header("Help Us To Improve", divider="green")
+    st.sidebar.markdown("https://forms.gle/YpLUrhqs1ahyCAe99", unsafe_allow_html=True)
+
+def get_textarea_height(text_content):
+    if text_content is None:
+        return 100
+    lines = text_content.split("\n")
+    count = len(lines)
+    for line in lines:
+        count += len(line) // 96
+    return count * 23 + 20 # 23 is a magic number
+
+def genrate_mainpage(backend):
+    # print("refresh mainpage")
+    st.title('💦 Generate Idea Step-by-step')
+    st.markdown("# 🐳 Background")
+    with st.form('background_form') as bg_form:
+        background = st.session_state.get("background", "")
+        background = st.text_area("Input your field background", background, placeholder="Input your field background", height=200, label_visibility="collapsed")
+
+        cols = st.columns(2)
+        def click_demo_i(i):
+            st.session_state["background"] = backend.get_demo_i(i)
+        for i, col in enumerate(cols):
+            col.form_submit_button(f"Example {i + 1}", use_container_width=True, on_click=click_demo_i, args=(i,))
+
+        col1, col2 = st.columns([2, 30])
+        submitted = col1.form_submit_button('Submit', type="primary")
+        if submitted:
+            st.session_state["global_state_step"] = 2.0
+            with st.spinner(text="Brainstorming..."):
+                st.session_state["brainstorms"] = backend.background2brainstorm_callback(background)
+            # st.session_state["brainstorms"] = "Test text"
+            st.session_state["brainstorms_expand"] = True
+            st.session_state["global_state_step"] = 2.5
+        # st.warning('Please enter your OpenAI API key!', icon='⚠')
+
+    ## Brainstorms
+    st.markdown("# 👻 Brainstorms")
+    with st.expander("Here is the generated brainstorms", expanded=st.session_state.get("brainstorms_expand", False)):
+        # st.write("<div class='myclass'>")
+        col1, col2 = st.columns(2)
+        widget_height = get_textarea_height(st.session_state.get("brainstorms", ""))
+        brainstorms = col1.text_area(label="brainstorms", value=st.session_state.get("brainstorms", ""), 
+                                     label_visibility="collapsed", height=widget_height)
+        st.session_state["brainstorms"] = brainstorms
+        if brainstorms:
+            col2.markdown(f"{brainstorms}")
+        else:
+            col2.markdown(f"Please input the brainstorms on the left.")
+        # st.write("</div>")
+        col1, col2 = st.columns([2, 30])
+        submitted = col1.button('Submit')
+        if submitted:
+            st.session_state["global_state_step"] = 3.0
+            with st.spinner(text="Extracting entities..."):
+                st.session_state["entities"] = backend.brainstorm2entities_callback(background, brainstorms)
+            # st.session_state["entities"] = "entities"
+            st.session_state["global_state_step"] = 3.5
+            st.session_state["entities_expand"] = True
+
+    ## Entities
+    st.markdown("# 🐱 Extracted Entities")
+    with st.expander("Here is the extracted entities", expanded=st.session_state.get("entities_expand", False)):
+        col1, col2 = st.columns(2, )
+        entities = col1.text_area(label="entities", value=st.session_state.get("entities", "[]"), label_visibility="collapsed")
+        entities = ast.literal_eval(entities)
+        st.session_state["entities"] = entities
+        if entities:
+            col2.markdown(f"{entities}")
+        else:
+            col2.markdown(f"Please input the entities on the left.")
+        submitted = col1.button('Submit', key="entities_button")
+        if submitted:
+            st.session_state["global_state_step"] = 4.0
+            with st.spinner(text="Retrieving related works..."):
+                st.session_state["related_works"], st.session_state["related_works_intact"] = backend.entities2literature_callback(background, entities)
+            # st.session_state["related_works"] = "related works"
+            st.session_state["global_state_step"] = 4.5
+            st.session_state["related_works_expand"] = True
+
+    ## Retrieved related works
+    st.markdown("# 📖 Retrieved Related Works")
+    with st.expander("Here is the retrieved related works", expanded=st.session_state.get("related_works_expand", False)):
+        col1, col2 = st.columns(2, )
+        widget_height = get_textarea_height(st.session_state.get("related_works", ""))
+        related_works_title = col1.text_area(label="related_works", value=st.session_state.get("related_works", ""), 
+                                             label_visibility="collapsed", height=widget_height)
+        if related_works_title:
+            col2.markdown(f"{related_works_title}")
+        else:
+            col2.markdown(f"Please input the related works on the left.")
+        submitted = col1.button('Submit', key="related_works_button")
+        if submitted:
+            st.session_state["global_state_step"] = 5.0
+            with st.spinner(text="Generating initial ideas..."):
+                res = backend.literature2initial_ideas_callback(background, st.session_state["related_works_intact"])
+                st.session_state["initial_ideas"] = res[0]
+                st.session_state["final_ideas"] = res[1]
+            # st.session_state["initial_ideas"] = "initial ideas"
+            st.session_state["global_state_step"] = 5.5
+            st.session_state["initial_ideas_expand"] = True
+
+    ## Initial ideas
+    st.markdown("# 😼 Generated Initial Ideas")
+    with st.expander("Here is the generated initial ideas", expanded=st.session_state.get("initial_ideas_expand", False)):
+        col1, col2 = st.columns(2, )
+        widget_height = get_textarea_height(st.session_state.get("initial_ideas", ""))
+        initial_ideas = col1.text_area(label="initial_ideas", value=st.session_state.get("initial_ideas", ""),
+                                       label_visibility="collapsed", height=widget_height)
+        if initial_ideas:
+            col2.markdown(f"{initial_ideas}")
+        else:
+            col2.markdown(f"Please input the initial ideas on the left.")
+        submitted = col1.button('Submit', key="initial_ideas_button")
+        if submitted:
+            st.session_state["global_state_step"] = 6.0
+            with st.spinner(text="Generating final ideas..."):
+                st.session_state["final_ideas"] = backend.initial2final_callback(initial_ideas, st.session_state["final_ideas"])
+            # st.session_state["final_ideas"] = "final ideas"
+            st.session_state["global_state_step"] = 6.5
+            st.session_state["final_ideas_expand"] = True
+
+    ## Final ideas
+    st.markdown("# 😸 Generated Final Ideas")
+    with st.expander("Here is the generated final ideas", expanded=st.session_state.get("final_ideas_expand", False)):
+        col1, col2 = st.columns(2, )
+        widget_height = get_textarea_height(st.session_state.get("final_ideas", ""))
+        user_input = col1.text_area(label="final_ideas", value=st.session_state.get("final_ideas", ""),
+                                    label_visibility="collapsed", height=widget_height)
+        if user_input:
+            col2.markdown(f"{user_input}")
+        else:
+            col2.markdown(f"Please input the final ideas on the left.")
+        submitted = col1.button('Submit', key="final_ideas_button")
+
+def step_by_step_generation(backend):
+    ## Pipeline global state
+    # 1.0: Input background is in progress
+    # 2.0: Brainstorming is in progress
+    #  2.5 Brainstorming is finished
+    # 3.0: Extracting entities is in progress
+    #  3.5 Extracting entities is finished
+    # 4.0: Retrieving literature is in progress
+    #  4.5 Retrieving ideas is finished
+    # 5.0: Generating initial ideas is in progress
+    #  5.5 Generating initial ideas is finished
+    # 6.0: Generating final ideas is in progress
+    #  6.5 Generating final ideas is finished
+    if "global_state_step" not in st.session_state:
+        st.session_state["global_state_step"] = 1.0
+    # backend = button_interface.Backend()
+    genrate_mainpage(backend)
+    generate_sidebar()

src/paper_manager.py

@@ -0,0 +1,862 @@
+import os
+import json
+import re
+from sentence_transformers import SentenceTransformer
+from tqdm import tqdm
+from utils.paper_crawling import PaperCrawling
+from utils.paper_client import PaperClient
+from utils.hash import generate_hash_id
+from collections import defaultdict
+from utils.header import get_dir, ConfigReader
+from utils.llms_api import APIHelper
+from utils.paper_retriever import Retriever
+from utils import scipdf
+import click
+from collections import Counter
+from loguru import logger
+import warnings
+
+warnings.filterwarnings("ignore")
+
+unicode_pattern = r"\u00c0-\u00ff\u0100-\u017f\u0180-\u024f\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\u31f0-\u31ff"
+
+
+def find_methodology(article_dict):
+    def find_section_index(keywords):
+        for i, section in enumerate(article_dict["sections"], 1):
+            heading = section["heading"].lower()
+            text = section["text"].lower()
+            if any(keyword in heading for keyword in keywords):
+                return i - 1
+            i = -1
+        if i == -1:
+            for i, section in enumerate(article_dict["sections"], 1):
+                heading = section["heading"].lower()
+                text = section["text"].lower()
+                if any(
+                    keyword in re.split(r"(?<=[.!?])\s+", text)[-1]
+                    for keyword in keywords
+                ):
+                    return i
+        return -1
+
+    index = find_section_index(["experiment", "evaluation"])
+    if index == -1:
+        experiments_index = next(
+            (
+                i
+                for i, section in enumerate(article_dict["sections"])
+                if "experiment" in section["heading"].lower()
+                or "evaluation" in section["heading"].lower()
+            ),
+            5,
+        )
+        experiments_index = min(experiments_index, len(article_dict["sections"]))
+        texts = [
+            section["text"] for section in article_dict["sections"][1:experiments_index]
+        ]
+        methodology = " ".join(texts)
+        return methodology
+    texts = [
+        section["text"]
+        for section in article_dict["sections"][1:index]
+        if not any(
+            keyword in section["heading"].lower()
+            for keyword in ["relate", "previous", "background"]
+        )
+    ]
+    methodology = " ".join(texts)
+    return methodology
+
+
+def count_sb_pairs(text):
+    return len(re.findall(r"\[.*?\]", text))
+
+
+def count_rb_pairs(text):
+    return len(re.findall(r"\(.*?\)", text))
+
+
+def find_cite_paper(introduction, methodology, references):
+    """
+    Count the number of times []/() appear in the introduction,
+    and determine which one is the reference ()/[]
+    """
+    text = introduction + methodology
+    rb_count = count_rb_pairs(introduction)
+    sb_count = count_sb_pairs(introduction)
+    pattern = (
+        r"\b[A-Z"
+        + unicode_pattern
+        + r"][a-zA-Z"
+        + unicode_pattern
+        + r"]+(?: and [A-Z"
+        + unicode_pattern
+        + r"][a-zA-Z"
+        + unicode_pattern
+        + r"]+)?(?: et al\.)?, \d{4}[a-z]?\b"
+    )
+    pattern = (
+        r"\b[A-Z"
+        + unicode_pattern
+        + r"][a-zA-Z"
+        + unicode_pattern
+        + r"]+(?: and [A-Z"
+        + unicode_pattern
+        + r"][a-zA-Z"
+        + unicode_pattern
+        + r"]+)?(?: et al\.)?, \d{4}[a-z]?\b"
+    )
+    temp_list = re.findall(pattern, text)
+    ref_list = []
+    ref_title = []
+    if len(temp_list) > 0:
+        pattern = (
+            r"\b([A-Z"
+            + unicode_pattern
+            + r"][a-zA-Z"
+            + unicode_pattern
+            + r"]+)(?: and [A-Z"
+            + unicode_pattern
+            + r"][a-zA-Z"
+            + unicode_pattern
+            + r"]+)?(?: et al\.)?, (\d{4})[a-z]?\b"
+        )
+        for temp in temp_list:
+            match = re.search(pattern, temp)
+            ref_list.append({"authors": match.group(1), "year": match.group(2)})
+        for i, ref in enumerate(ref_list):
+            for j, r in enumerate(references):
+                if r["year"] == ref["year"] and ref["authors"] in r["authors"]:
+                    ref_title.append(r["title"])
+    if len(ref_title) <= 1:
+        ref_list = []
+        ref_title = []
+        if rb_count < sb_count:
+            pattern = r"\[\d+(?:,\s*\d+)*\]"
+        else:
+            pattern = r"\(\d+(?:,\s*\d+)*\)"
+        ref_list = re.findall(pattern, text)
+        # ref: ['[15, 16]', '[5]', '[2, 3, 8]']
+        combined_ref_list = []
+        for ref in ref_list:
+            numbers = re.findall(r"\d+", ref)
+            combined_ref_list.extend(map(int, numbers))
+        # Sort
+        ref_counts = Counter(combined_ref_list)
+        ref_counts = dict(sorted(ref_counts.items()))
+        ref_list = list(ref_counts.keys())
+        for idx in ref_list:
+            if idx < len(references):
+                ref_title.append(references[idx]["title"])
+    return ref_title
+
+
+class PaperManager:
+    def __init__(self, config, venue_name="acl", year="2013") -> None:
+        log_dir = config.DEFAULT.log_dir
+        if not os.path.exists(log_dir):
+            os.makedirs(log_dir)
+            print(f"Created log directory: {log_dir}")
+        log_file = os.path.join(log_dir, "paper_manager.log")
+        logger.add(log_file, level=config.DEFAULT.log_level)
+        self.venue_name = venue_name
+        self.year = year
+        self.data_type = "train"
+        self.paper_client = PaperClient(config)
+        self.paper_crawling = PaperCrawling(config, data_type=self.data_type)
+        self.embedding_model = SentenceTransformer(
+            model_name_or_path=get_dir(config.DEFAULT.embedding), device=self.config.DEFAULT.device
+        )
+        self.api_helper = APIHelper(config)
+        self.retriever = Retriever(config)
+        self.paper_id_map = defaultdict()
+        self.citemap = defaultdict(set)
+        self.year_list = [
+            "2013",
+            "2014",
+            "2015",
+            "2016",
+            "2017",
+            "2018",
+            "2019",
+            "2020",
+            "2021",
+            "2022",
+            "2023",
+            "2024",
+        ]
+        self.config = config
+        with open(config.DEFAULT.ignore_paper_id_list, "r", encoding="utf-8") as f:
+            try:
+                self.ignore_paper_pdf_url = [dic["pdf_url"] for dic in json.load(f)]
+            except:
+                self.ignore_paper_pdf_url = []
+
+    def create_vector_index(self):
+        index_exists = self.paper_client.check_index_exists()
+        if not index_exists:
+            print("Create vector index paper-embeddings")
+            self.paper_client.create_vector_index()
+
+    def clean_entity(self, entity):
+        if entity is None:
+            return None
+        cleaned_entity = re.sub(r"\([^)]*\)", "", entity)
+        cleaned_entity = re.sub(r"[^\w\s]", "", cleaned_entity)
+        cleaned_entity = re.sub(r"_", " ", cleaned_entity)
+        cleaned_entity = re.sub(r"\s+", " ", cleaned_entity).strip()
+        return cleaned_entity
+
+    def clean_text(self, text):
+        return text.replace(", , ", ", ")
+
+    def check_parse(self, paper):
+        # Required keys
+        required_keys = [
+            "abstract",
+            "introduction",
+            "reference",
+            "methodology",
+            "reference_filter",
+        ]
+        # Check for missing keys or None values
+        for key in required_keys:
+            if key not in paper or paper[key] is None:
+                logger.error(
+                    f"hash_id: {paper.get('hash_id')} pdf_url: {paper.get('pdf_url')} : "
+                    f"Missing or None '{key}' in paper."
+                )
+                return False
+        return True
+
+    def update_paper(
+        self,
+        paper,
+        need_download=False,
+        need_parse=False,
+        need_summary=False,
+        need_get_entities=False,
+        need_ground_truth=False,
+    ):
+        if paper["pdf_url"] in self.ignore_paper_pdf_url:
+            logger.warning(
+                "hash_id: {}, pdf_url: {} ignore".format(
+                    paper["hash_id"], paper["pdf_url"]
+                )
+            )
+            return
+        self.paper_client.update_paper_from_client(paper)
+        if need_download:
+            if not self.paper_crawling.download_paper(paper):
+                print(f"download paper {paper['pdf_url']} failed!")
+                return
+        if need_parse:
+            if not self.check_parse(paper):
+                logger.debug(f"begin to parse {paper['hash_id']}")
+                if not self.paper_crawling.download_paper(paper):
+                    logger.error(f"download paper {paper['pdf_url']} failed!")
+                    return
+                try:
+                    article_dict = scipdf.parse_pdf_to_dict(paper["pdf_path"])
+                    if "title" not in paper.keys() or paper["title"] is None:
+                        paper["title"] = article_dict["title"]
+                    paper["abstract"] = article_dict["abstract"]
+                    paper["introduction"] = article_dict["sections"][0]["text"]
+                    paper["methodology"] = find_methodology(article_dict)
+                    reference = []
+                    for ref in article_dict["references"]:
+                        reference.append(ref["title"])
+                    paper["reference"] = reference
+                    paper["reference_filter"] = find_cite_paper(
+                        paper["introduction"],
+                        paper["methodology"],
+                        article_dict["references"],
+                    )
+                    logger.info(f"{paper['hash_id']} parse success")
+                except Exception:
+                    logger.error(
+                        f"{paper['hash_id']}: {paper['pdf_url']}  parse error!"
+                    )
+
+        if need_summary:
+            if not self.check_parse(paper):
+                logger.error(f"paper {paper['hash_id']} need parse first...")
+            elif "summary" not in paper.keys():
+                result = self.api_helper(
+                    paper["title"], paper["abstract"], paper["introduction"]
+                )
+                if result is not None:
+                    paper["summary"] = result["summary"]
+                    paper["motivation"] = result["motivation"]
+                    paper["contribution"] = result["contribution"]
+                    logger.info(f"paper {paper['hash_id']} summary success...")
+                else:
+                    logger.warning(
+                        "hash_id: {}, pdf_url: {} summary failed...".format(
+                            paper["hash_id"], paper["pdf_url"]
+                        )
+                    )
+            if need_ground_truth:
+                if "ground_truth" not in paper.keys():
+                    if (
+                        "abstract" in paper.keys()
+                        and "contribution" in paper.keys()
+                        and "methodology" in paper.keys()
+                    ):
+                        paper["ground_truth"] = self.api_helper.generate_ground_truth(
+                            abstract=paper["abstract"],
+                            contribution=paper["contribution"],
+                            text=paper["methodology"],
+                        )
+                        logger.info(f"paper {paper['hash_id']} ground truth success...")
+                    else:
+                        logger.error("Can't get ground truth...please check")
+
+        # insert paper in database
+        if self.check_parse(paper):
+            self.paper_client.add_paper_node(paper)
+        else:
+            return
+
+        if need_get_entities and self.paper_client.check_entity_node_count(
+            paper["hash_id"]
+        ):
+            if (
+                paper["abstract"] is None
+                or paper["introduction"] is None
+                or paper["reference"] is None
+            ):
+                logger.error(f"paper need parse first")
+            entities = self.api_helper.generate_entity_list(paper["abstract"])
+            logger.info("hash_id {}, Entities: {}".format(paper["hash_id"], entities))
+            if entities is not None:
+                self.paper_client.add_entity_node(paper["hash_id"], entities)
+            else:
+                logger.warning(
+                    "hash_id: {}, pdf_url: {} entities None...".format(
+                        paper["hash_id"], paper["pdf_url"]
+                    )
+                )
+
+    def update_paper_local(
+        self,
+        paper,
+        need_download=False,
+        need_parse=False,
+        need_summary=False,
+        need_get_entities=False,
+        need_ground_truth=False,
+    ):
+        if paper["pdf_url"] in self.ignore_paper_pdf_url:
+            logger.warning(
+                "hash_id: {}, pdf_url: {} ignore".format(
+                    paper["hash_id"], paper["pdf_url"]
+                )
+            )
+            return
+        # keep the content of the paper node consistent with the database
+        self.paper_client.update_paper_from_client(paper)
+        if need_download:
+            if not self.paper_crawling.download_paper(paper):
+                print(f"download paper {paper['pdf_url']} failed!")
+                return
+        if need_parse:
+            if not self.check_parse(paper):  # haven't parse
+                logger.debug(f"begin to parse {paper['hash_id']}")
+                if not self.paper_crawling.download_paper(paper):
+                    logger.error(f"download paper {paper['pdf_url']} failed!")
+                    return
+                try:
+                    article_dict = scipdf.parse_pdf_to_dict(paper["pdf_path"])
+                    if "title" not in paper.keys() or paper["title"] is None:
+                        paper["title"] = article_dict["title"]
+                    paper["abstract"] = article_dict["abstract"]
+                    paper["introduction"] = article_dict["sections"][0]["text"]
+                    paper["methodology"] = find_methodology(article_dict)
+                    reference = []
+                    for ref in article_dict["references"]:
+                        reference.append(ref["title"])
+                    paper["reference"] = reference
+                    paper["reference_filter"] = find_cite_paper(
+                        paper["introduction"],
+                        paper["methodology"],
+                        article_dict["references"],
+                    )
+                    logger.info(f"{paper['hash_id']} parse success")
+                except Exception:
+                    logger.error(
+                        f"{paper['hash_id']}: {paper['pdf_url']}  parse error!"
+                    )
+
+        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"]
+            )
+            if result is not None:
+                paper["summary"] = result["summary"]
+                paper["motivation"] = result["motivation"]
+                paper["contribution"] = result["contribution"]
+                logger.info(f"paper {paper['hash_id']} summary success...")
+            else:
+                logger.warning(
+                    "hash_id: {}, pdf_url: {} summary failed...".format(
+                        paper["hash_id"], paper["pdf_url"]
+                    )
+                )
+
+            if need_ground_truth:
+                if (
+                    "abstract" in paper.keys()
+                    and "contribution" in paper.keys()
+                    and "methodology" in paper.keys()
+                ):
+                    paper["ground_truth"] = self.api_helper.generate_ground_truth(
+                        abstract=paper["abstract"],
+                        contribution=paper["contribution"],
+                        text=paper["methodology"],
+                    )
+                    logger.info(f"paper {paper['hash_id']} ground truth success...")
+                else:
+                    logger.error("Can't get ground truth...please check")
+
+        if need_get_entities and self.paper_client.check_entity_node_count(
+            paper["hash_id"]
+        ):
+            if (
+                paper["abstract"] is None
+                or paper["introduction"] is None
+                or paper["reference"] is None
+            ):
+                logger.error(f"paper need parse first")
+            entities = self.api_helper.generate_entity_list(paper["abstract"])
+            logger.info("hash_id {}, Entities: {}".format(paper["hash_id"], entities))
+            if entities is not None:
+                self.paper_client.add_entity_node(paper["hash_id"], entities)
+            else:
+                logger.warning(
+                    "hash_id: {}, pdf_url: {} entities None...".format(
+                        paper["hash_id"], paper["pdf_url"]
+                    )
+                )
+
+        with open(
+            self.config.output_path.replace(
+                ".json", "_{}.json".format(paper["hash_id"])
+            ),
+            "w",
+            encoding="utf8",
+        ) as f:
+            json.dump(paper, f)
+        return paper
+
+    def update_paper_from_json(
+        self,
+        need_download=True,
+        need_parse=False,
+        need_summary=False,
+        need_get_entities=False,
+        need_ground_truth=False,
+    ):
+        if self.year != "all":
+            logger.info(
+                "=== year {}, venue name {} ===".format(self.year, self.venue_name)
+            )
+            with open(
+                f"./assets/paper/{self.venue_name}/{self.venue_name}_{self.year}_paper_list.json",
+                "r",
+                encoding="utf8",
+            ) as f:
+                paper_list = json.load(f)
+            for paper in tqdm(paper_list):
+                self.update_paper(
+                    paper,
+                    need_download=need_download,
+                    need_parse=need_parse,
+                    need_summary=need_summary,
+                    need_get_entities=need_get_entities,
+                    need_ground_truth=need_ground_truth,
+                )
+        else:
+            if self.venue_name == "iccv":
+                self.year_list = ["2013", "2015", "2017", "2019", "2021", "2023"]
+            elif self.venue_name == "eccv":
+                self.year_list = ["2018", "2020", "2022", "2024"]
+            for year in self.year_list:
+                with open(
+                    f"./assets/paper/{self.venue_name}/{self.venue_name}_{year}_paper_list.json",
+                    "r",
+                    encoding="utf8",
+                ) as f:
+                    paper_list = json.load(f)
+                logger.info(
+                    "=== year {}, venue name {} ===".format(year, self.venue_name)
+                )
+                for paper in tqdm(paper_list):
+                    self.update_paper(
+                        paper,
+                        need_download=need_download,
+                        need_parse=need_parse,
+                        need_summary=need_summary,
+                        need_get_entities=need_get_entities,
+                        need_ground_truth=need_ground_truth,
+                    )
+
+    def update_paper_from_json_to_json(
+        self,
+        need_download=True,
+        need_parse=False,
+        need_summary=False,
+        need_get_entities=False,
+        need_ground_truth=False,
+    ):
+        result = []
+        if self.year != "all":
+            logger.info(
+                "=== year {}, venue name {} ===".format(self.year, self.venue_name)
+            )
+            with open(
+                f"./assets/paper/{self.venue_name}/{self.venue_name}_{self.year}_paper_list.json",
+                "r",
+                encoding="utf8",
+            ) as f:
+                paper_list = json.load(f)
+            result = [
+                self.update_paper_local(
+                    paper,
+                    need_download=need_download,
+                    need_parse=need_parse,
+                    need_summary=need_summary,
+                    need_get_entities=need_get_entities,
+                    need_ground_truth=need_ground_truth,
+                )
+                for paper in tqdm(paper_list)
+            ]
+
+        else:
+            if self.venue_name == "iccv":
+                self.year_list = ["2013", "2015", "2017", "2019", "2021", "2023"]
+            elif self.venue_name == "eccv":
+                self.year_list = ["2018", "2020", "2022", "2024"]
+            for year in self.year_list:
+                with open(
+                    f"./assets/paper/{self.venue_name}/{self.venue_name}_{year}_paper_list.json",
+                    "r",
+                    encoding="utf8",
+                ) as f:
+                    paper_list = json.load(f)
+                logger.info(
+                    "=== year {}, venue name {} ===".format(year, self.venue_name)
+                )
+                subresult = [
+                    self.update_paper_local(
+                        paper,
+                        need_download=need_download,
+                        need_parse=need_parse,
+                        need_summary=need_summary,
+                        need_get_entities=need_get_entities,
+                        need_ground_truth=need_ground_truth,
+                    )
+                    for paper in tqdm(paper_list)
+                ]
+                result += subresult
+
+        with open(self.config.output_path, "w", encoding="utf8") as f:
+            json.dump(result, f)
+
+    def insert_citation(self):
+        if self.year != "all":
+            year_list = [self.year]
+        else:
+            year_list = self.year_list
+        for year in year_list:
+            paper_list = self.paper_client.select_paper(self.venue_name, year)
+            for paper in tqdm(paper_list):
+                if (
+                    self.check_parse(paper)
+                    and len(paper["reference"]) > 0
+                    and "motivation" in paper.keys()
+                    and paper["motivation"] is not None
+                ):
+                    paper["cite_id_list"] = [
+                        generate_hash_id(ref_title)
+                        for ref_title in paper["reference_filter"]
+                    ]
+                    paper["cite_id_list"] = self.paper_client.filter_paper_id_list(
+                        paper["cite_id_list"], year=year
+                    )
+                    paper["all_cite_id_list"] = [
+                        generate_hash_id(ref_title) for ref_title in paper["reference"]
+                    ]
+                    paper["all_cite_id_list"] = self.paper_client.filter_paper_id_list(
+                        paper["all_cite_id_list"], year=year
+                    )
+                    if "entities" not in paper.keys() or len(paper["entities"]) < 3:
+                        paper["entities"] = self.api_helper.generate_entity_list(
+                            paper["abstract"]
+                        )
+                        logger.debug(
+                            "get entity from context: {}".format(paper["entities"])
+                        )
+                    logger.debug(
+                        "paper hash_id {}, cite_id_list {}, all_cite_id_list {}".format(
+                            paper["hash_id"],
+                            paper["cite_id_list"],
+                            paper["all_cite_id_list"],
+                        )
+                    )
+                else:
+                    paper["cite_id_list"] = []
+                    paper["all_cite_id_list"] = []
+                if (
+                    "entities" in paper.keys()
+                    and "cite_id_list" in paper.keys()
+                    and "all_cite_id_list" in paper.keys()
+                ):
+                    self.paper_client.add_paper_citation(paper)
+
+    def insert_entity_combinations(self):
+        if self.year != "all":
+            year_list = [self.year]
+        else:
+            year_list = self.year_list
+        for year in year_list:
+            self.paper_client.get_entity_combinations(self.venue_name, year)
+
+    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)
+
+    def cosine_similarity_search(self, data_type, context, k=1):
+        """
+        return related paper: list
+        """
+        embedding = self.embedding_model.encode(context)
+        result = self.paper_client.cosine_similarity_search(data_type, embedding, k)
+        return result
+
+    def generate_paper_list(self):
+        folder_path = f"./assets/paper/{self.venue_name}"
+        if not os.path.exists(folder_path):
+            os.makedirs(folder_path)
+        if self.year != "all":
+            logger.info(
+                "=== year {}, venue name {} ===".format(self.year, self.venue_name)
+            )
+            paper_list = self.paper_crawling.crawling(self.year, self.venue_name)
+            with open(
+                f"{folder_path}/{self.venue_name}_{self.year}_paper_list.json",
+                "w",
+            ) as f:
+                json.dump(paper_list, f, indent=4, ensure_ascii=False)
+        else:
+            for year in self.year_list:
+                logger.info(
+                    "=== year {}, venue name {} ===".format(year, self.venue_name)
+                )
+                paper_list = self.paper_crawling.crawling(year, self.venue_name)
+                with open(
+                    f"{folder_path}/{self.venue_name}_{year}_paper_list.json",
+                    "w",
+                ) as f:
+                    json.dump(paper_list, f, indent=4, ensure_ascii=False)
+
+
+@click.group()
+@click.pass_context
+def main(ctx):
+    """
+    Training and evaluation
+    """
+    print("Mode:", ctx.invoked_subcommand)
+
+
+@main.command()
+@click.option(
+    "-c",
+    "--config-path",
+    default=get_dir("./configs/datasets.yaml"),
+    type=click.File(),
+    required=True,
+    help="Dataset configuration file in YAML",
+)
+@click.option(
+    "--year",
+    default="2013",
+    type=str,
+    required=True,
+    help="Venue year",
+)
+@click.option(
+    "--venue-name",
+    default="acl",
+    type=str,
+    required=True,
+    help="Venue name",
+)
+@click.option(
+    "--llms-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API alias used. If you do not have separate APIs for summarization and generation, you can use this unified setting. This option is ignored when setting the API to be used by summarization and generation separately",
+)
+@click.option(
+    "--sum-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API aliases used for summarization. When used, it will invalidate --llms-api",
+)
+@click.option(
+    "--gen-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API aliases used for generation. When used, it will invalidate --llms-api",
+)
+def crawling(config_path, year, venue_name, **kwargs):
+    # Configuration
+    config = ConfigReader.load(config_path, **kwargs)
+    pm = PaperManager(config, venue_name, year)
+    pm.generate_paper_list()
+
+
+@main.command()
+@click.option(
+    "-c",
+    "--config-path",
+    default=get_dir("./configs/datasets.yaml"),
+    type=click.File(),
+    required=True,
+    help="Dataset configuration file in YAML",
+)
+@click.option(
+    "--year",
+    default="2013",
+    type=str,
+    required=True,
+    help="Venue year",
+)
+@click.option(
+    "--venue-name",
+    default="acl",
+    type=str,
+    required=True,
+    help="Venue name",
+)
+@click.option(
+    "--llms-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API alias used. If you do not have separate APIs for summarization and generation, you can use this unified setting. This option is ignored when setting the API to be used by summarization and generation separately",
+)
+@click.option(
+    "--sum-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API aliases used for summarization. When used, it will invalidate --llms-api",
+)
+@click.option(
+    "--gen-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API aliases used for generation. When used, it will invalidate --llms-api",
+)
+def update(config_path, year, venue_name, **kwargs):
+    # Configuration
+    config = ConfigReader.load(config_path, **kwargs)
+    pm = PaperManager(config, venue_name, year)
+    pm.update_paper_from_json(need_download=True)
+
+
+@main.command()
+@click.option(
+    "-c",
+    "--config-path",
+    default=get_dir("./configs/datasets.yaml"),
+    type=click.File(),
+    required=True,
+    help="Dataset configuration file in YAML",
+)
+@click.option(
+    "--year",
+    default="2013",
+    type=str,
+    required=True,
+    help="Venue year",
+)
+@click.option(
+    "--venue-name",
+    default="acl",
+    type=str,
+    required=True,
+    help="Venue name",
+)
+@click.option(
+    "--llms-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API alias used. If you do not have separate APIs for summarization and generation, you can use this unified setting. This option is ignored when setting the API to be used by summarization and generation separately",
+)
+@click.option(
+    "--sum-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API aliases used for summarization. When used, it will invalidate --llms-api",
+)
+@click.option(
+    "--gen-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API aliases used for generation. When used, it will invalidate --llms-api",
+)
+@click.option(
+    "-o",
+    "--output",
+    default=get_dir("./output/out.json"),
+    type=click.File("wb"),
+    required=True,
+    help="Dataset configuration file in YAML",
+)
+def local(config_path, year, venue_name, output, **kwargs):
+    # Configuration
+    output_path = output.name
+    if not os.path.exists(os.path.dirname(output_path)):
+        os.makedirs(os.path.dirname(output_path))
+    config = ConfigReader.load(config_path, output_path=output_path, **kwargs)
+    pm = PaperManager(config, venue_name, year)
+    pm.update_paper_from_json_to_json(
+        need_download=True, need_parse=True, need_summary=True, need_ground_truth=True
+    )
+
+
+@main.command()
+@click.option(
+    "-c",
+    "--config-path",
+    default=get_dir("./configs/datasets.yaml"),
+    type=click.File(),
+    required=True,
+    help="Dataset configuration file in YAML",
+)
+def embedding(config_path):
+    # Configuration
+    config = ConfigReader.load(config_path)
+    PaperManager(config).insert_embedding()
+
+
+if __name__ == "__main__":
+    main()

src/retriever.py

@@ -0,0 +1,179 @@
+import os
+from utils.paper_retriever import RetrieverFactory
+from utils.llms_api import APIHelper
+from utils.paper_client import PaperClient
+from utils.header import ConfigReader
+from omegaconf import OmegaConf
+import click
+import json
+from loguru import logger
+import warnings
+
+warnings.filterwarnings("ignore")
+
+
+@click.group()
+@click.pass_context
+def main(ctx):
+    """
+    Evaluate Retriever SN/KG/SNKG
+    """
+    print("Mode:", ctx.invoked_subcommand)
+
+
+@main.command()
+@click.option(
+    "-c",
+    "--config-path",
+    default="../configs/datasets.yaml",
+    type=click.File(),
+    required=True,
+    help="Dataset configuration file in YAML",
+)
+@click.option(
+    "--ids-path",
+    default="assets/data/test_acl_2024.json",
+    type=click.File(),
+    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",
+)
+@click.option(
+    "--llms-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API alias used. If you do not have separate APIs for summarization and generation, you can use this unified setting. This option is ignored when setting the API to be used by summarization and generation separately",
+)
+@click.option(
+    "--sum-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API aliases used for summarization. When used, it will invalidate --llms-api",
+)
+@click.option(
+    "--gen-api",
+    default=None,
+    type=str,
+    required=False,
+    help="The LLMS API aliases used for generation. When used, it will invalidate --llms-api",
+)
+def retrieve(
+    config_path, ids_path, retriever_name, co_cite, cluster_to_filter, **kwargs
+):
+    config = ConfigReader.load(config_path, **kwargs)
+    log_dir = config.DEFAULT.log_dir
+    if not os.path.exists(log_dir):
+        os.makedirs(log_dir)
+        print(f"Created log directory: {log_dir}")
+    log_file = os.path.join(
+        log_dir,
+        "retriever_eval_{}_cocite-{}_cluster-{}.log".format(
+            retriever_name, co_cite, cluster_to_filter
+        ),
+    )
+    logger.add(log_file, level=config.DEFAULT.log_level)
+    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)
+    precision = 0
+    filtered_precision = 0
+    recall = 0
+    filtered_recall = 0
+    num = 0
+    gt_reference_num = 0
+    retrieve_paper_num = 0
+    label_num = 0
+    top_k_precision = {p: 0 for p in config.RETRIEVE.top_k_list}
+    top_k_recall = {p: 0 for p in config.RETRIEVE.top_k_list}
+    # Init Retriever
+    rt = RetrieverFactory.get_retriever_factory().create_retriever(
+        retriever_name,
+        config,
+        use_cocite=co_cite,
+        use_cluster_to_filter=cluster_to_filter,
+    )
+    for line in ids_path:
+        paper = json.loads(line)
+        logger.info("\nbegin generate paper hash id {}".format(paper["hash_id"]))
+        # 1. Get Background
+        paper = paper_client.get_paper_by_id(paper["hash_id"])
+        if "motivation" in paper.keys():
+            bg = paper["motivation"]
+        else:
+            logger.error(f"paper hash_id {paper['hash_id']} doesn't have background...")
+            continue
+        if "entities" in paper.keys():
+            entities = paper["entities"]
+        else:
+            entities = api_helper.generate_entity_list(bg)
+        logger.info("origin 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]
+        else:
+            logger.warning(
+                "hash_id {} cite paper num less than 5 ...".format(paper["hash_id"])
+            )
+            continue
+        # 2. Retrieve
+        result = rt.retrieve(
+            bg, entities, need_evaluate=True, target_paper_id_list=target_paper_id_list
+        )
+        filtered_precision += result["filtered_precision"]
+        precision += result["precision"]
+        filtered_recall += result["filtered_recall"]
+        gt_reference_num += result["gt_reference_num"]
+        retrieve_paper_num += result["retrieve_paper_num"]
+        recall += result["recall"]
+        label_num += result["label_num"]
+        for k, v in result["top_k_matrix"].items():
+            top_k_recall[k] += v["recall"]
+            top_k_precision[k] += v["precision"]
+        num += 1
+        if num >= 100:
+            break
+        continue
+    logger.info("=== Finish Report ===")
+    logger.info(f"{'Test Paper Num:':<25} {num}")
+    logger.info(f"{'Average Precision:':<25} {precision/num:.3f}")
+    logger.info(f"{'Average Recall:':<25} {recall/num:.3f}")
+    logger.info(f"{'Average GT Ref Paper Num:':<25} {gt_reference_num/num:.3f}")
+    logger.info(f"{'Average Retrieve Paper Num:':<25} {retrieve_paper_num/num:.3f}")
+    logger.info(f"{'Average Label Num:':<25} {label_num/num:.3f}")
+    # Print Eval Result
+    logger.info("=== Top-K Metrics ===")
+    logger.info(
+        f"=== USE_COCIT: {co_cite}, USE_CLUSTER_TO_FILTER: {cluster_to_filter} ==="
+    )
+    logger.info("| Top K  | Recall | Precision |")
+    logger.info("|--------|--------|-----------|")
+    for k in config.RETRIEVE.top_k_list:
+        if k <= retrieve_paper_num / num:
+            logger.info(
+                f"| {k:<5} | {top_k_recall[k]/num:.3f}  | {top_k_precision[k]/num:.3f}    |"
+            )
+    logger.info("=" * 40)
+
+
+if __name__ == "__main__":
+    main()

src/utils/api/__init__.py

@@ -0,0 +1,29 @@
+r"""_summary_
+-*- coding: utf-8 -*-
+
+Module : data.utils.api
+
+File Name : __init__.py
+
+Description : API helper automatic registration, using HelperCompany can directly reflect the corresponding helper
+    If you need to add a new AI helper, please add the Python source file in the same level path, for example:
+    ```
+    @register_helper('name')
+    class CustomerHelper(BaseHelper):
+        ...
+    ```
+    Then import it into this file, for example
+    ```
+    from .customer_helper import CustomerHelper # noqa: F401, ensure autoregister
+    ```
+
+   
+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
+
+__all__ = ["HelperCompany"]

src/utils/api/base_helper.py

@@ -0,0 +1,207 @@
+r"""_summary_
+-*- coding: utf-8 -*-
+
+Module : data.utils.api.base_helper
+
+File Name : base_helper.py
+
+Description : API helper automatic registration, using HelperCompany can directly reflect the corresponding helper
+
+Creation Date : 2024-10-29
+
+Author : Frank Kang([email protected])
+"""
+
+from typing import Union, List, Optional
+from abc import ABCMeta
+from typing_extensions import Literal, override
+from ..base_company import BaseCompany
+from typing import Union
+from typing_extensions import Literal, override
+
+
+class NotGiven:
+    """
+    Copy from OpenAI
+
+    A sentinel singleton class used to distinguish omitted keyword arguments
+    from those passed in with the value None (which may have different behavior).
+
+    For example:
+
+    ```py
+    def get(timeout: Union[int, NotGiven, None] = NotGiven()) -> Response:
+        ...
+
+
+    get(timeout=1)  # 1s timeout
+    get(timeout=None)  # No timeout
+    get()  # Default timeout behavior, which may not be statically known at the method definition.
+    ```
+    """
+
+    def __bool__(self) -> Literal[False]:
+        return False
+
+    @override
+    def __repr__(self) -> str:
+        return "NOT_GIVEN"
+
+
+class HelperCompany(BaseCompany):
+    """_summary_
+
+    AI helper factory, inheriting BaseCompany
+
+    For example:
+    ```
+    helper_company = HelperCompany.get()
+
+    # Of course, you can also obtain the singleton using the following methods
+    helper_company = HelperCompany()
+
+    helper = helper_company[helper_name]
+    ```
+
+    @see data.utils.base_company.BaseCompany
+    """
+
+    @override
+    def __repr__(self) -> str:
+        return "HelperCompany"
+
+
+class register_helper:
+    """_summary_
+
+    Automatically register helper annotation classes
+    """
+
+    def __init__(self, helper_type, *args, **kwds):
+        self.helper_type = helper_type
+        self.init_args = args
+        self.init_kwds = kwds
+
+    def __call__(self, helper_cls, *args, **kwds):
+        helper_name = helper_cls.__name__
+        if HelperCompany.get().register(self.helper_type, helper_cls):
+
+            def _method(obj):
+                return helper_name
+
+            helper_cls.name = _method
+            return helper_cls
+        else:
+            raise KeyError()
+
+
+class BaseHelper:
+    """_summary_
+
+    Base class for API helper
+    """
+
+    __metaclass__ = ABCMeta
+
+    def __init__(self, api_key, model, base_url) -> None:
+        super(BaseHelper, self).__init__()
+        self.api_key = api_key
+        self.model = model
+        self.base_url = base_url
+        self.client = None
+
+    def create(
+        self,
+        *args,
+        messages: Union[str, List[str], List[int], object, None],
+        stream: Optional[Literal[False]] | Literal[True] | NotGiven = None,
+        temperature: Optional[float] | NotGiven = None,
+        top_p: Optional[float] | NotGiven = None,
+        max_tokens: int | NotGiven = None,
+        seed: int | NotGiven = None,
+        stop: Optional[Union[str, List[str], None]] | NotGiven = None,
+        tools: Optional[object] | NotGiven = None,
+        tool_choice: str | NotGiven = None,
+        extra_headers: None | NotGiven = None,
+        extra_body: None | NotGiven = None,
+        timeout: float | None | NotGiven = None,
+        **kwargs
+    ):
+        """
+        Creates a model response for the given chat conversation.
+
+        Args:
+            messages: A list of messages comprising the conversation so far.
+                [Example Python code](https://cookbook.openai.com/examples/how_to_format_inputs_to_chatgpt_models).
+
+            stream: If set, partial message deltas will be sent, like in ChatGPT. Tokens will be sent as data-only
+                [server-sent events](https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#Event_stream_format)
+                as they become available, with the stream terminated by a `data: [DONE]`
+                message.
+                [Example Python code](https://cookbook.openai.com/examples/how_to_stream_completions).
+
+            temperature: What sampling temperature to use, between 0 and 2. Higher values like 0.8 will
+                make the output more random, while lower values like 0.2 will make it more
+                focused and deterministic.
+
+                We generally recommend altering this or `top_p` but not both.
+
+            top_p: An alternative to sampling with temperature, called nucleus sampling, where the
+                model considers the results of the tokens with top_p probability mass. So 0.1
+                means only the tokens comprising the top 10% probability mass are considered.
+
+                We generally recommend altering this or `temperature` but not both.
+
+            max_tokens: The maximum number of [tokens](/tokenizer) that can be generated in the chat
+                completion.
+
+                The total length of input tokens and generated tokens is limited by the model's
+                context length.
+                [Example Python code](https://cookbook.openai.com/examples/how_to_count_tokens_with_tiktoken)
+                for counting tokens.
+
+            seed: This feature is in Beta. If specified, our system will make a best effort to
+                sample deterministically, such that repeated requests with the same `seed` and
+                parameters should return the same result. Determinism is not guaranteed, and you
+                should refer to the `system_fingerprint` response parameter to monitor changes
+                in the backend.
+
+            stop: Up to 4 sequences where the API will stop generating further tokens.
+
+            tools: A list of tools the model may call. Currently, only functions are supported as a
+                tool. Use this to provide a list of functions the model may generate JSON inputs
+                for. A max of 128 functions are supported.
+
+            tool_choice: Controls which (if any) tool is called by the model. `none` means the model will
+                not call any tool and instead generates a message. `auto` means the model can
+                pick between generating a message or calling one or more tools. `required` means
+                the model must call one or more tools. Specifying a particular tool via
+                `{"type": "function", "function": {"name": "my_function"}}` forces the model to
+                call that tool.
+
+                `none` is the default when no tools are present. `auto` is the default if tools
+                are present.
+
+            extra_headers: Send extra headers
+
+            extra_body: Add additional JSON properties to the request
+
+            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
+        )

src/utils/api/openai_helper.py

@@ -0,0 +1,40 @@
+r"""_summary_
+-*- coding: utf-8 -*-
+
+Module : data.utils.api.openai_helper
+
+File Name : openai_helper.py
+
+Description : Helper class for openai interface, generally not used directly.
+    For example:
+    ```
+    from data.utils.api import HelperCompany
+    helper = HelperCompany.get()['OpenAI']
+    ...
+    ```
+   
+Creation Date : 2024-10-29
+
+Author : Frank Kang([email protected])
+"""
+from openai import OpenAI
+from .base_helper import register_helper, BaseHelper
+
+
+@register_helper('OpenAI')
+class OpenAIHelper(BaseHelper):
+    """_summary_
+
+    Helper class for openai interface, generally not used directly.
+
+    For example:
+    ```
+    from data.utils.api import HelperCompany
+    helper = HelperCompany.get()['OpenAI']
+    ...
+    ```
+    """
+
+    def __init__(self, api_key, model, base_url=None, timeout=None):
+        super().__init__(api_key, model, base_url)
+        self.client = OpenAI(api_key=api_key, base_url=base_url, timeout=timeout)

src/utils/api/zhipuai_helper.py

@@ -0,0 +1,40 @@
+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-10-29
+
+Author : Frank Kang([email protected])
+"""
+from zhipuai import ZhipuAI
+from .base_helper import register_helper, BaseHelper
+
+
+@register_helper('ZhipuAI')
+class ZhipuAIHelper(BaseHelper):
+    """_summary_
+
+    Helper class for ZhipuAI interface, generally not used directly.
+
+    For example:
+    ```
+    from data.utils.api import HelperCompany
+    helper = HelperCompany.get()['ZhipuAI']
+    ...
+    ```
+    """
+
+    def __init__(self, api_key, model, base_url=None, timeout=None):
+        super().__init__(api_key, model, base_url)
+        self.client = ZhipuAI(api_key=api_key, base_url=base_url, timeout=timeout)

src/utils/base_company.py

@@ -0,0 +1,118 @@
+r"""_summary_
+-*- coding: utf-8 -*-
+
+Module : data.utils.base_company
+
+File Name : base_company.py
+
+Description : The base class of the factory class, used to register and reflect specific classes
+
+Creation Date : 2024-10-29
+
+Author : Frank Kang([email protected])
+"""
+import threading
+from typing import Any
+from typing_extensions import override
+
+
+class BaseCompany(object):
+    """_summary_
+
+    The base class of the factory class, used to register and reflect specific classes. Use singleton mode, so it is necessary to maintain consistency in the path when importing and changing classes
+
+    For example:
+    ```
+    base_company = BaseCompany.get()
+
+    # Of course, you can also obtain the singleton using the following methods
+    base_company = BaseCompany()
+
+    entity = base_company[registered_name]
+    ```
+    """
+    _instance = None
+    _lock = threading.Lock()
+
+    def __new__(cls, *args, **kwargs):
+        with cls._lock:
+            if cls._instance is None:
+                cls._instance = super(BaseCompany, cls).__new__(
+                    cls, *args, **kwargs)
+                cls._instance.init_factory()
+        return cls._instance
+
+    def __init__(self):
+        self.entities = {}
+
+    def init_factory(self):
+        """_summary_
+
+        Used for initializing singleton
+        """
+        self.entities = {}
+
+    @staticmethod
+    def get():
+        """_summary_
+
+        Method for obtaining singleton classes
+
+        For example:
+        ```
+        base_company = BaseCompany.get()
+        entity = base_company[registered_name]
+        ```
+
+        Returns:
+            BaseCompany: singleton
+        """
+        if BaseCompany._instance is None:
+            BaseCompany._instance = BaseCompany()
+        return BaseCompany._instance
+
+    def register(self, entity_name: str, entity: Any) -> bool:
+        """_summary_
+
+        Register the entity, which is called by the automatic registrar. Please do not call it yourself. Each name can only be registered once
+
+        Args:
+            entity_name (str): Name used for registration
+            entity (Any): Registered entity
+
+        Returns:
+            bool: Registration success returns true, failure returns false
+        """
+        if entity_name not in self.entities:
+            self.entities[entity_name] = entity
+            return True
+        else:
+            return False
+
+    def delete(self, entity_name: str) -> bool:
+        """_summary_
+
+        Remove registered entities, please use with caution
+
+        Args:
+            entity_name (str): The registered name of the registered entity
+
+        Returns:
+            bool: Success in deletion returns true, failure returns false
+        """
+        if entity_name in self.entities:
+            self.entities[entity_name] = None
+            del self.entities[entity_name]
+            return True
+        else:
+            return False
+
+    def __getitem__(self, key):
+        return self.entities[key]
+
+    def __len__(self):
+        return len(self.entities)
+
+    @override
+    def __repr__(self) -> str:
+        return "BaseCompany"

src/utils/hash.py

@@ -0,0 +1,76 @@
+import re
+import os
+import hashlib
+import struct
+from collections import Counter
+
+
+def check_env():
+    env_name_list = [
+        "NEO4J_URL",
+        "NEO4J_USERNAME",
+        "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...")
+
+
+def generate_hash_id(input_string):
+    if input_string is None:
+        return None
+    sha1_hash = hashlib.sha256(input_string.lower().encode("utf-8")).hexdigest()
+    binary_hash = bytes.fromhex(sha1_hash)
+    int64_hash = struct.unpack(">q", binary_hash[:8])[0]
+    return abs(int64_hash)
+
+
+def extract_ref_id(text, references):
+    """
+    references: paper["references"]
+    """
+    # 正则表达式模式，用于匹配[数字, 数字]格式
+    pattern = r"\[\d+(?:,\s*\d+)*\]"
+    # 提取所有匹配的内容
+    ref_list = re.findall(pattern, text)
+    # ref ['[15, 16]', '[5]', '[2, 3, 8]']
+    combined_ref_list = []
+    if len(ref_list) > 0:
+        # 说明是pattern 0
+        for ref in ref_list:
+            # 移除方括号并分割数字
+            numbers = re.findall(r"\d+", ref)
+            # 将字符串数字转换为整数并加入到列表中
+            combined_ref_list.extend(map(int, numbers))
+        # 去重并排序
+        ref_counts = Counter(combined_ref_list)
+        ref_counts = dict(sorted(ref_counts.items()))
+        # 对多个，只保留引用最多的一个
+        for ref in ref_list:
+            # 移除方括号并分割数字
+            numbers = re.findall(r"\d+", ref)
+            # 找到只引用了一次的
+            temp_list = []
+            for num in numbers:
+                num = int(num)
+                if ref_counts[num] == 1:
+                    temp_list.append(num)
+            if len(temp_list) == len(numbers):
+                temp_list = temp_list[1:]
+            for num in temp_list:
+                del ref_counts[num]
+    hash_id_list = []
+    for idx in ref_counts.keys():
+        hash_id_list.append(generate_hash_id(references[idx]))
+    return hash_id_list
+
+
+if __name__ == "__main__":
+    # 示例用法
+    input_string = "example_string"
+    hash_id = generate_hash_id(input_string)
+    print("INT64 Hash ID:", hash_id)

src/utils/header.py

@@ -0,0 +1,11 @@
+import sys
+import os
+
+sys.path.append(
+    os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
+)
+from configs.utils import get_dir
+from configs.config import ConfigReader
+from prompt.prompt_reader import Prompt, AssistantCreateQuery, MessageQuery
+
+__all__ = ["get_dir", "ConfigReader", "Prompt", "AssistantCreateQuery", "MessageQuery"]

src/utils/llms_api.py

@@ -0,0 +1,1354 @@
+from .api import HelperCompany
+import openai
+import re
+from .header import get_dir, Prompt, ConfigReader
+import traceback
+
+TAG_moti = "Motivations:"
+TAG_contr = "Details:"
+
+
+def clean_text(text):
+    cleaned_text = re.sub(r"-\s*\n", "", text)
+    cleaned_text = re.sub(r"\s*\n\s*", " ", cleaned_text)
+    return cleaned_text.strip()
+
+
+def clean_entities(input_string):
+    # 取出括号中的内容
+    cleaned_text = re.sub(r"\([^)]*\)", "", input_string)
+    # 使用正则表达式删除非字母字符
+    cleaned = re.sub(r"[^a-zA-Z\s]", "", input_string)
+    # 将多个空格替换为一个空格
+    cleaned = re.sub(r"\s+", " ", cleaned)
+    # 删除首尾空格
+    cleaned = cleaned.strip().lower()
+    return cleaned
+
+
+class APIHelper(object):
+
+    def __init__(self, config) -> None:
+        super(APIHelper, self).__init__()
+        self.config = config
+        self.__checkout_config__()
+        self.summarizer = self.get_helper(config, config.used_llms_apis.summarization)
+        self.generator = self.get_helper(config, config.used_llms_apis.generation)
+        self.prompt = Prompt(get_dir(config.ARTICLE.summarizing_prompt))
+
+    def get_helper(self, config, alias):
+        api_config = config[alias]
+        return HelperCompany.get()[api_config.type](
+            api_config.api_key, api_config.model, api_config.base_url, timeout=None
+        )
+
+    def __checkout_config__(self):
+        pass
+
+    def __call__(self, title: str, abstract: str, introduction: str) -> dict:
+        if title is None or abstract is None or introduction is None:
+            return None
+        try:
+            message = [
+                self.prompt.queries[0][0](
+                    title=title, abstract=abstract, introduction=introduction
+                )
+            ]
+            response1 = self.summarizer.create(
+                messages=message,
+            )
+            summary = clean_text(response1.choices[0].message.content)
+            message.append({"role": "assistant", "content": summary})
+            message.append(self.prompt.queries[1][0]())
+            response2 = self.summarizer.create(
+                messages=message,
+            )
+            detail = response2.choices[0].message.content
+            motivation = clean_text(detail.split(TAG_moti)[1].split(TAG_contr)[0])
+            contribution = clean_text(detail.split(TAG_contr)[1])
+            result = {
+                "summary": summary,
+                "motivation": motivation,
+                "contribution": contribution,
+            }
+        except Exception:
+            traceback.print_exc()
+            return None
+        return result
+
+    def generate_entity_list(self, abstract: str, max_num: int = 5) -> list:
+        common_examples = [
+            {
+                "content": "This paper presents a novel approach to automatic captioning of geo-tagged images by summarizing multiple webdocuments that contain information related to an image's location. The summarizer is biased by dependency pattern models towards sentences which contain features typically provided for different scene types such as those of churches, bridges, etc. Our results show that summaries biased by dependency pattern models lead to significantly higher ROUGE scores than both n-gram language models reported in previous work and also Wikipedia baseline summaries. Summaries generated using dependency patterns also lead to more readable summaries than those generated without dependency patterns.",
+                "entities": "dependency pattern models, automatic captioning",
+            },
+            {
+                "content": "In this paper, we describe the 2015 iteration of the SemEval shared task on Sentiment Analysis in Twitter. This was the most popular sentiment analysis shared task to date with more than 40 teams participating in each of the last three years. This year's shared task competition consisted of five sentiment prediction subtasks. Two were reruns from previous years: (A) sentiment expressed by a phrase in the context of a tweet, and (B) overall sentiment of a tweet. We further included three new subtasks asking to predict (C) the sentiment towards a topic in a single tweet, (D) the overall sentiment towards a topic in a set of tweets, and (E) the degree of prior polarity of a phrase.",
+                "entities": "sentiment analysis, shared task",
+            },
+            {
+                "content": 'This paper presents two different tools which may be used as a support of speech recognition. The tool "transc" is the first one and it generates the phonetic transcription (pronunciation) of given utterance. It is based mainly on fixed rules which can be defined for Czech pronunciation but it can work also with specified list of exceptions which is defined on lexicon basis. It allows the usage of "transc" for unknown text with high probability of correct phonetic transcription generation. The second part is devoted to lexicon management tool "lexedit" which may be useful in the phase of generation of pronunciation lexicon for collected corpora. The presented tool allows editing of pronunciation, playing examples of pronunciation, comparison with reference lexicon, updating of reference lexicon, etc.',
+                "entities": "speech recognition, phonetic transcription, lexicon management",
+            },
+            {
+                "content": "Previous research applying kernel methods to natural language parsing have focussed on proposing kernels over parse trees, which are hand-crafted based on domain knowledge and computational considerations. In this paper we propose a method for defining kernels in terms of a probabilistic model of parsing. This model is then trained, so that the parameters of the probabilistic model reflect the generalizations in the training data. The method we propose then uses these trained parameters to define a kernel for reranking parse trees. In experiments, we use a neural network based statistical parser as the probabilistic model, and use the resulting kernel with the Voted Perceptron algorithm to rerank the top 20 parses from the probabilistic model. This method achieves a significant improvement over the accuracy of the probabilistic model.",
+                "entities": "parse trees, probabilistic model, natural language parsing",
+            },
+        ]
+
+        few_shot_examples = []
+        for example in common_examples:
+            few_shot_examples.append(example)
+        
+        prompt_template_entity = '''
+        ### Task Description:
+        You are an AI researcher tasked with extracting the key entities from a given research paper content. These entities should represent the most important keywords or phrases that summarize the main topics or concepts discussed in the content.
+
+        ### Information Provided:
+        **Content**: Focus on this content, and extract entities that serve as concrete manifestations of the main themes and topics within it.
+
+        ### Approach:
+        Your entity extraction should be systematic:
+        - **Step 1**: Carefully read through the content to fully understand its main themes and topics.
+        - **Step 2**: Identify and list key entities central to the content, ensuring each entity is relevant, meaningful, and accurately represents the content.
+
+        ### Entity Guidelines:
+        - Each entity should be no longer than 5 words and contain at least 2 words.
+        - The entities should be nouns or noun phrases.
+        - The total number of entities should be less than or equal to {max_num}.
+
+        ### Examples:
+        {examples}
+
+        ### Specific information:
+        I will provide you with specific information now, please use them according to the instructions above:
+        **Content**: {content}
+
+        ### Format for Your Response:
+        Please just give me the entities and spilt them by ",":
+        <entity 1>,<entity2>,...
+        '''
+
+        if abstract is None:
+            return None
+        try:
+            examples_str = "\n".join(
+                f"[content]: {example['content']}\n[entity]: {example['entities']}\n###\n"
+                for example in few_shot_examples
+            )
+            system_input = "Now you are an expert in extracting key entities from research contents. You are good at identifying the most important keywords or phrases that summarize the main topics or concepts discussed in the content."
+            message = []
+            message.append({"role": "system", "content": system_input})
+            message_input = prompt_template_entity.format(
+                examples=examples_str, content=abstract, max_num=str(max_num)
+            )
+            message.append({"role": "user", "content": message_input})
+            response = self.summarizer.create(
+                messages=message,
+            )
+            entities = response.choices[0].message.content
+            entity_list = entities.strip().split(", ")
+            clean_entity_list = []
+            for entity in entity_list:
+                entity = clean_entities(entity)
+                if len(entity.split()) <= 20:
+                    clean_entity_list.append(entity)
+
+            if "entity" not in abstract.lower() and "entities" not in abstract.lower():
+                clean_entity_list = [
+                    re.sub(
+                        r"\bentity\b|\bentities\b", "", e, flags=re.IGNORECASE
+                    ).strip()
+                    for e in clean_entity_list
+                ]
+                clean_entity_list = [e for e in clean_entity_list if e]
+                clean_entity_list = [clean_entities(e) for e in clean_entity_list]
+        except Exception:
+            traceback.print_exc()
+            return None
+        return clean_entity_list
+
+    def generate_brainstorm(self, background: str) -> str:
+        prompt_template_brainstorming = """
+        ### Task Description:
+        You are an AI researcher tasked with brainstorming initial, innovative ideas to address a given research problem in AI. Focus on generating diverse and creative approaches rather than finalized methods. The ideas can be rough and in their infancy but should cover a range of possible directions that could be explored further.
+
+        ### Information Provided:
+        - **Research Background**: {background}
+
+        ### Approach:
+        Your brainstorming should be systematic:
+        - **Step 1**: Thoroughly understand the research background.
+        - **Step 2**: Generate a list of 4 to 6 high-level ideas or directions that could potentially solve problems in the given background. Be creative, think outside the box, and avoid merely rephrasing existing methods.
+
+        ### Format for Your Response:
+        Please present 4 to 6 ideas in the following format:
+        **Idea 1**: [Brief description of the first idea]
+        **Idea 2**: [Brief description of the second idea]
+        ...
+        """
+
+        if background is None:
+            return None
+        try:
+            # Initial brainstorming to generate raw ideas
+            brainstorming_prompt = prompt_template_brainstorming.format(
+                background=background
+            )
+            message = []
+            prompt_first = "Now you are a researcher in the field of AI with innovative and pioneering abilities. You are good at generating creative and original ideas."
+            message.append({"role": "system", "content": prompt_first})
+            message.append({"role": "user", "content": brainstorming_prompt})
+
+            # Call the API to generate brainstorming ideas
+            response_brainstorming = self.generator.create(
+                messages=message,
+            )
+            brainstorming_ideas = response_brainstorming.choices[0].message.content
+
+        except Exception:
+            traceback.print_exc()
+            return None
+
+        return brainstorming_ideas
+
+    def generate_problem(self, background: str, related_papers: list[dict]):
+        prompt_template_problem = """
+        ### Task Description:
+        You will receive a research background along with summaries, backgrounds, and contributions (methods) of several related papers. Your task is to carefully analyze this information and propose a research problem that is original, clear, feasible, relevant, and significant to its field. Additionally, provide the rationales behind the proposed problem.
+
+        ### Information Provided:
+        1. **Research Background**: This is your primary focus. The research problem you propose should be a direct reflection of this background.
+        2. **Related Papers**: These papers offer studies directly related to the primary research topic, providing additional insights and knowledge that will inform your proposed problem.
+
+        ### Approach:
+        Your approach should be systematic:
+        - **Step 1**: Begin by thoroughly understanding the core focus of the research background.
+        - **Step 2**: Review the summaries, backgrounds, and contributions (methods) of the related papers to gain broader insights into the primary research topic.
+        - **Step 3**: Based on the provided information, propose a research problem that meets the criteria of being original, clear, feasible, relevant, and significant. Support your problem statement with clear rationales.
+
+        ### Specific information:
+        I will provide you with specific information now, please use them according to the instructions above:
+        1. **Research Background**: {background}
+        2. **Related Papers**: {related_papers_information}
+
+        ### Format for Your Response:
+        **Research Problem**: [your problem]
+        - **Rationales**: [the rationale behind your problem]
+        """
+        if background is None or related_papers is None:
+            return None
+        try:
+            related_papers_information = ""
+            for i, paper in enumerate(related_papers):
+                related_papers_information += (
+                    "Related paper {i}: ".format(i=i + 1) + paper["title"]
+                )
+                related_papers_information += "\nSummary: " + paper["summary"]
+                related_papers_information += "\nBackgrounds: " + paper["motivation"]
+                related_papers_information += (
+                    "\nContributions: " + paper["contribution"] + "\n \n"
+                )
+            message = []
+            system_input = "Now you are a researcher in the field of AI with innovative and pioneering abilities. You are good at proposing novel and valuable questions based on research background."
+            message.append({"role": "system", "content": system_input})
+            message_input = prompt_template_problem.format(
+                background=background,
+                related_papers_information=related_papers_information,
+            )
+            message.append({"role": "user", "content": message_input})
+            response = self.generator.create(
+                messages=message,
+            )
+            problem = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+            return None
+        return problem, message_input
+
+    def generate_problem_with_cue_words(
+        self, background: str, related_papers: list[dict], cue_words: list
+    ):
+        prompt_template_problem = """
+        ### Task Description:
+        You will receive a research background and some cue words along with summaries, backgrounds, and contributions (methods) of several related papers. Your task is to carefully analyze this information and propose a research problem that is original, clear, feasible, relevant, and significant to its field. Additionally, provide the rationales behind the proposed problem.
+
+        ### Information Provided:
+        1. **Research Background**: This is your primary focus. The research problem you propose should be a direct reflection of this background.
+        2. **Cue Words**: Some of these words can provide direction and ideas for you to ask questions. They should be the focus of your attention.
+        3. **Related Papers**: These papers offer studies directly related to the primary research topic, providing additional insights and knowledge that will inform your proposed problem.
+
+        ### Approach:
+        Your approach should be systematic:
+        - **Step 1**: Begin by thoroughly understanding the core focus of the research background.
+        - **Step 2**: Read the cue words and then determine the approximate direction of your problem.
+        - **Step 3**: Review the summaries, backgrounds, and contributions (methods) of the related papers to gain broader insights into the primary research topic.
+        - **Step 4**: Based on the provided information, propose a research problem that meets the criteria of being original, clear, feasible, relevant, and significant. Support your problem statement with clear rationales.
+
+        ### Specific information:
+        I will provide you with specific information now, please use them according to the instructions above:
+        1. **Research Background**: {background}
+        2. **Cue Words**: {cue_words}
+        3. **Related Papers**: {related_papers_information}
+
+        ### Format for Your Response:
+        **Research Problem**: [your problem]
+        - **Rationales**: [the rationale behind your problem]
+        """
+        if background is None or related_papers is None or cue_words is None:
+            return None
+        try:
+            related_papers_information = ""
+            for i, paper in enumerate(related_papers):
+                related_papers_information += (
+                    "Related paper {i}: ".format(i=i + 1) + paper["title"]
+                )
+                related_papers_information += "\nSummary: " + paper["summary"]
+                related_papers_information += "\nBackgrounds: " + paper["motivation"]
+                related_papers_information += (
+                    "\nContributions: " + paper["contribution"] + "\n \n"
+                )
+            message = []
+            system_input = "Now you are a researcher in the field of AI with innovative and pioneering abilities. You are good at proposing novel and valuable questions based on research background."
+            message.append({"role": "system", "content": system_input})
+            message_input = prompt_template_problem.format(
+                background=background,
+                related_papers_information=related_papers_information,
+                cue_words=cue_words,
+            )
+            message.append({"role": "user", "content": message_input})
+            response = self.generator.create(
+                messages=message,
+            )
+            problem = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+            return None
+        return problem, message_input
+
+    def generate_inspiration(self, problem: str, related_paper: dict):
+        prompt_inspiration = """
+        ### Task Description:
+        You will be provided with a research problem, as well as the summary, backgrounds and contributions (methods) of a related paper. Your task is to extract a novel, effective, and specific inspiration from the related paper that can help addressing the research problem, and provide a brief rationale for this inspiration.
+
+        ### Information Provided:
+        1. **Research problem**: The key issues or aspects of the problem that need to be addressed. These will serve as the foundation for generating your inspiration.
+        2. **Related paper**: Draw insights from the abstract, background, and methods of the related paper. Delve deeply into these methods, understand the motivations behind them, and critically assess how they might contribute to solving the research problem. Avoid merely replicating the methods; instead, synthesize relevant aspects with your own insights to derive a novel inspiration.
+
+        ### Approach:
+        Your approach should be systematic:
+        - **Step 1**: Thoroughly read the research problem to clearly understand the primary focus.
+        - **Step 2**: Review the summary, background, and contributions (methods) of the related paper. Evaluate whether the methods proposed in the paper can provide solutions or insights relevant to the research problem.
+        - **Step 3**: Based on the information provided in the paper and your own analysis, propose a novel, effective, and specific inspiration. Include a rationale explaining how this inspiration helps addressing the research problem.
+
+        ### Specific Information:
+        I will provide you with specific information now, please use them according to the instructions above:
+        1. **Research problem**: {problem}
+        2. **Related paper**: {related_paper_information}
+
+        ### Format for Your Response:
+        Your output should be around 200 words and follow the format:
+        **Inspiration**: [Your novel, effective, and specific idea derived from the related paper]
+        - **Rationale**: [The brief reasoning behind how this inspiration help addressing the research problem]
+        """
+        if problem is None or related_paper is None:
+            return None
+        try:
+            related_paper_information = ""
+            related_paper_information += "Related paper : " + related_paper["title"]
+            related_paper_information += "\nSummary: " + related_paper["summary"]
+            related_paper_information += "\nBackgrounds: " + related_paper["motivation"]
+            related_paper_information += (
+                "\nContributions: " + related_paper["contribution"] + "\n \n"
+            )
+            message = []
+            system_input = "Now you are a researcher in the field of AI with innovative and pioneering abilities. You are good at extracting novel and valuable inspirations from papers."
+            message.append({"role": "system", "content": system_input})
+            message_input = prompt_inspiration.format(
+                problem=problem, related_paper_information=related_paper_information
+            )
+            message.append({"role": "user", "content": message_input})
+            response = self.generator.create(
+                messages=message,
+            )
+            inspiration = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+            return None
+        return inspiration
+
+    def generate_inspiration_with_cue_words(
+        self, problem: str, related_paper: dict, cue_words: list
+    ):
+        prompt_inspiration = """
+        ### Task Description:
+        You will be provided with a research problem, some cue words as well as the summary, backgrounds and contributions (methods) of a related paper. Your task is to extract a novel, effective, and specific inspiration from the related paper that can help addressing the research problem, and provide a brief rationale for this inspiration.
+
+        ### Information Provided:
+        1. **Research problem**: The key issues or aspects of the problem that need to be addressed. These will serve as the foundation for generating your inspiration.
+        2. **Cue Words**: Some of these words can provide direction for you to extract inspiration. They should be the focus of your attention.
+        3. **Related paper**: Draw insights from the abstract, background, and methods of the related paper. Delve deeply into these methods, understand the motivations behind them, and critically assess how they might contribute to solving the research problem. Avoid merely replicating the methods; instead, synthesize relevant aspects with your own insights to derive a novel inspiration.
+
+        ### Approach:
+        Your approach should be systematic:
+        - **Step 1**: Thoroughly read the research problem to clearly understand the primary focus.
+        - **Step 2**: Review the summary, background, and contributions (methods) of the related paper. Evaluate whether the methods proposed in the paper can provide solutions or insights relevant to the research problem.
+        - **Step 3**: Read the cue words and consider whether these words can be combined with information of the related paper to  help providing inspiration.
+        - **Step 4**: Based on the information provided in the paper and your own analysis, propose a novel, effective, and specific inspiration. Include a rationale explaining how this inspiration helps addressing the research problem.
+
+        ### Specific Information:
+        I will provide you with specific information now, please use them according to the instructions above:
+        1. **Research problem**: {problem}
+        2. **Cue Words**: {cue_words}
+        3. **Related paper**: {related_paper_information}
+
+        ### Format for Your Response:
+        Your output should be around 200 words and follow the format:
+        **Inspiration**: [Your novel, effective, and specific idea derived from the related paper]
+        - **Rationale**: [The brief reasoning behind how this inspiration help addressing the research problem]
+        """
+        if problem is None or related_paper is None or cue_words is None:
+            return None
+        try:
+            related_paper_information = ""
+            related_paper_information += "Related paper : " + related_paper["title"]
+            related_paper_information += "\nSummary: " + related_paper["summary"]
+            related_paper_information += "\nBackgrounds: " + related_paper["motivation"]
+            related_paper_information += (
+                "\nContributions: " + related_paper["contribution"] + "\n \n"
+            )
+            message = []
+            system_input = "Now you are a researcher in the field of AI with innovative and pioneering abilities. You are good at extracting novel and valuable inspirations from papers."
+            message.append({"role": "system", "content": system_input})
+            message_input = prompt_inspiration.format(
+                problem=problem,
+                related_paper_information=related_paper_information,
+                cue_words=cue_words,
+            )
+            message.append({"role": "user", "content": message_input})
+            response = self.generator.create(
+                messages=message,
+            )
+            inspiration = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+            return None
+        return inspiration
+
+    def generate_idea(self, problem: str, related_papers: list[dict]) -> str:
+        prompt_template_idea = """
+        ### Task Description:
+        You will be provided with a research problem along with its rationales. Your task is to brainstorm some ideas that are clear, innovative, valid, and comprehensive to address the problem. Additionally, some cue words along with summaries, backgrounds, and contributions (methods) of related papers will be provided as sources of inspiration for generating novel ideas.
+
+        ### Information Provided:
+        1. **Research Problem & Rationales**: The key issues or aspects of the problem that need to be addressed. These will form the foundation for generating your ideas.
+        2. **Related Papers**: Draw inspiration from the abstracts, backgrounds, and methods of these papers. Delve deeply into these methods, understand the motivations behind them, and think critically about how they might inform your approach. Avoid merely stacking existing methods; instead, integrate relevant aspects with your own insights to create original solutions.
+
+        ### Approach:
+        Your approach should be systematic:
+        - **Step 1**: Thoroughly read the research problem to understand your primary focus.
+        - **Step 2**: Review the summaries, backgrounds, and contributions (methods) of the related papers to gain a broader perspective and insights relevant to the problem.
+        - **Step 3**: Based on the provided information, propose some ideas that are clear, innovative, valid, and comprehensive.
+
+        ### Specific Information:
+        I will provide you with specific information now, please use them according to the instructions above:
+        1. **Research Problem & Rationales**: {problem}
+        2. **Related Papers**: {related_papers_information}
+
+        ### Format for Your Response:
+        Please ensure that your final ideas include about 10 entries, presented in the following format:
+        **Idea 1**: [The first method idea]
+        **Idea 2**: [The second method idea]
+        **Idea 3**: [The third method idea]
+        ...
+        """
+        if problem is None or related_papers is None:
+            return None
+        try:
+            related_papers_information = ""
+            for i, dict in enumerate(related_papers):
+                related_papers_information += (
+                    "Related paper {i}: ".format(i=i + 1) + dict["title"]
+                )
+                related_papers_information += "\nSummary: " + dict["summary"]
+                related_papers_information += "\nBackgrounds: " + dict["motivation"]
+                related_papers_information += (
+                    "\nContributions: " + dict["contribution"] + "\n \n"
+                )
+            message = []
+            system_input = "Now you are a researcher in the field of AI with innovative and pioneering abilities. You are good at using innovative and original methods to solve cutting-edge problems in the field of AI."
+            message.append({"role": "system", "content": system_input})
+            message_input = prompt_template_idea.format(
+                problem=problem, related_papers_information=related_papers_information
+            )
+            message.append({"role": "user", "content": message_input})
+            response = self.generator.create(
+                messages=message,
+            )
+            idea = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+            return None
+        return idea
+
+    def generate_idea_with_cue_words(
+        self, problem: str, related_papers: list[dict], cue_words: list
+    ) -> str:
+        prompt_template_idea = """
+        ### Task Description:
+        You will be provided with a research problem along with its rationales. Your task is to brainstorm some ideas that are clear, innovative, valid, and comprehensive to address the problem. Additionally, some cue words along with summaries, backgrounds, and contributions (methods) of related papers will be provided as sources of inspiration for generating novel ideas.
+
+        ### Information Provided:
+        1. **Research Problem & Rationales**: The key issues or aspects of the problem that need to be addressed. These will form the foundation for generating your ideas.
+        2. **Cue Words**: Some of these words can inspire or provide direction for you to generate ideas. You can try to think deeply in these directions and perspectives.
+        3. **Related Papers**: Draw inspiration from the abstracts, backgrounds, and methods of these papers. Delve deeply into these methods, understand the motivations behind them, and think critically about how they might inform your approach. Avoid merely stacking existing methods; instead, integrate relevant aspects with your own insights to create original solutions.
+
+        ### Approach:
+        Your approach should be systematic:
+        - **Step 1**: Thoroughly read the research problem to understand your primary focus.
+        - **Step 2**: Read the cue words and think about whether these words can inspire or provide direction for you to come up with ideas.
+        - **Step 3**: Review the summaries, backgrounds, and contributions (methods) of the related papers to gain a broader perspective and insights relevant to the problem.
+        - **Step 4**: Based on the provided information, propose some ideas that are clear, innovative, valid, and comprehensive.
+
+        ### Specific Information:
+        I will provide you with specific information now, please use them according to the instructions above:
+        1. **Research Problem & Rationales**: {problem}
+        2. **Cue Words**: {cue_words}
+        3. **Related Papers**: {related_papers_information}
+
+        ### Format for Your Response:
+        Please ensure that your final ideas include about 10 entries, presented in the following format:
+        **Idea 1**: [The first method idea]
+        **Idea 2**: [The second method idea]
+        **Idea 3**: [The third method idea]
+        ...
+        """
+        if problem is None or related_papers is None or cue_words is None:
+            return None
+        try:
+            related_papers_information = ""
+            for i, dict in enumerate(related_papers):
+                related_papers_information += (
+                    "Related paper {i}: ".format(i=i + 1) + dict["title"]
+                )
+                related_papers_information += "\nSummary: " + dict["summary"]
+                related_papers_information += "\nBackgrounds: " + dict["motivation"]
+                related_papers_information += (
+                    "\nContributions: " + dict["contribution"] + "\n \n"
+                )
+            message = []
+            system_input = "Now you are a researcher in the field of AI with innovative and pioneering abilities. You are good at using innovative and original methods to solve cutting-edge problems in the field of AI."
+            message.append({"role": "system", "content": system_input})
+            message_input = prompt_template_idea.format(
+                problem=problem,
+                related_papers_information=related_papers_information,
+                cue_words=cue_words,
+            )
+            message.append({"role": "user", "content": message_input})
+            response = self.generator.create(
+                messages=message,
+            )
+            idea = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+            return None
+        return idea
+
+    def generate_idea_by_inspiration(self, problem: str, inspirations: list[str]):
+        prompt_template_idea = """
+        ### Task Description:
+        You will be provided with a research problem and its rationales, along with inspirations and their rationales extracted from related papers. Your task is to brainstorm some ideas that are clear, innovative, valid, and comprehensive to address the problem.
+
+        ### Information Provided:
+        1. **Research problem & Rationales**: The key issues or aspects of the problem that need to be addressed. These will form the foundation for generating your ideas.
+        2. **Inspirations**: Insights and ideas extracted from related papers that may provide valuable perspectives or techniques applicable to the research problem.
+
+        ### Approach:
+        Your approach should be systematic:
+        - **Step 1**: Thoroughly read and understand the research problem to identify your primary focus.
+        - **Step 2**: Review the inspirations extracted from the related papers to gain a broader perspective and insights relevant to the research topic.
+        - **Step 3**: Based on the provided information, propose some ideas that are clear, innovative, valid, and comprehensive.
+        
+        ### Specific Information:
+        I will provide you with specific information now, please use them according to the instructions above:
+        1. **Research problem & Rationales**: {problem}
+        2. **Inspirations**: {inspirations_text}
+
+        ### Format for Your Response:
+        Please ensure that your final ideas include about 10 entries, presented in the following format:
+        **Idea 1**: [The first method idea]
+        **Idea 2**: [The second method idea]
+        **Idea 3**: [The third method idea]
+        ...
+        """
+        if problem is None or inspirations is None:
+            return None
+        try:
+            inspirations_text = ""
+            for i, inspiration in enumerate(inspirations):
+                inspirations_text += (
+                    "Inspiration {i}: ".format(i=i + 1) + "\n" + inspiration + "\n \n"
+                )
+            message = []
+            system_input = "Now you are a researcher in the field of AI with innovative and pioneering abilities. You are good at using innovative and original methods to solve cutting-edge problems in the field of AI."
+            message.append({"role": "system", "content": system_input})
+            message_input = prompt_template_idea.format(
+                problem=problem, inspirations_text=inspirations_text
+            )
+            message.append({"role": "user", "content": message_input})
+            response = self.generator.create(
+                messages=message,
+            )
+            idea = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+            return None
+        return idea
+
+    def generate_idea_by_inspiration_with_cue_words(
+        self, problem: str, inspirations: list[str], cue_words: list
+    ):
+        prompt_template_idea = """
+        ### Task Description:
+        You will be provided with a research problem, its rationales and some cue words, along with inspirations and their rationales extracted from related papers. Your task is to brainstorm some ideas that are clear, innovative, valid, and comprehensive to address the problem.
+
+        ### Information Provided:
+        1. **Research problem & Rationales**: The key issues or aspects of the problem that need to be addressed. These will form the foundation for generating your ideas.
+        2. **Cue Words**: Some of these words can inspire or provide direction for you to generate ideas. You can try to think deeply in these directions and perspectives.
+        3. **Inspirations**: Insights and ideas extracted from related papers that may provide valuable perspectives or techniques applicable to the research problem.
+
+        ### Approach:
+        Your approach should be systematic:
+        - **Step 1**: Thoroughly read and understand the research problem to identify your primary focus.
+        - **Step 2**: Read the cue words and think about whether these words can inspire or provide direction for you to come up with ideas.
+        - **Step 3**: Review the inspirations extracted from the related papers to gain a broader perspective and insights relevant to the research topic.
+        - **Step 4**: Based on the provided information, propose some ideas that are clear, innovative, valid, and comprehensive.
+        
+        ### Specific Information:
+        I will provide you with specific information now, please use them according to the instructions above:
+        1. **Research problem & Rationales**: {problem}
+        2. **Cue Words**: {cue_words}
+        3. **Inspirations**: {inspirations_text}
+
+        ### Format for Your Response:
+        Please ensure that your final ideas include about 10 entries, presented in the following format:
+        **Idea 1**: [The first method idea]
+        **Idea 2**: [The second method idea]
+        **Idea 3**: [The third method idea]
+        ...
+        """
+        if problem is None or inspirations is None or cue_words is None:
+            return None
+        try:
+            inspirations_text = ""
+            for i, inspiration in enumerate(inspirations):
+                inspirations_text += (
+                    "Inspiration {i}: ".format(i=i + 1) + "\n" + inspiration + "\n \n"
+                )
+            message = []
+            system_input = "Now you are a researcher in the field of AI with innovative and pioneering abilities. You are good at using innovative and original methods to solve cutting-edge problems in the field of AI."
+            message.append({"role": "system", "content": system_input})
+            message_input = prompt_template_idea.format(
+                problem=problem,
+                inspirations_text=inspirations_text,
+                cue_words=cue_words,
+            )
+            message.append({"role": "user", "content": message_input})
+            response = self.generator.create(
+                messages=message,
+            )
+            idea = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+            return None
+        return idea
+    
+    def integrate_idea(self, background: str, brainstorm: str, idea: str) -> str:
+        prompt_template_idea = """
+        Task Description: 
+        You will be provided with research background information along with a set of ideas you generated previously from with the related paper information, and a set of brainstorming ideas concerning the same research topic. Your task is to combine these ideas and generate new ones, the new ideas you generate should base on the ideas you generated previously, and integrate creative parts of the brainstorming ideas. Consider the background thoroughly, taking into account the novelty and practicability of each idea. If you think an idea you generate is reasonable and valuable, feel free to retain it. 
+
+        ### Information Provided:
+        1. **Research Background**: The starting point for idea generation based on the research context.
+        2. **Brainstorming Ideas**: These ideas were generated purely from the research background, focusing on innovation and may not be directly related to the problem.
+        3. **Generated Ideas**: These are the ideas you previously generated by considering both the research background and related papers.
+
+        ### Approach:
+        - **Step 1**: Review the research background and original ideas to understand the foundation of the problem.
+        - **Step 2**: Consider the brainstorming ideas and original ideas together. Combine, improve, or expand upon them, integrating insights from the related papers.
+        - **Step 3**: Propose new ideas that are innovative and practical, ensuring they align with the research background.
+
+        ### Specific Information:
+        1. **Research Background**: {background}
+        2. **Brainstorming Ideas**: {brainstorm}
+        3. **Generated Ideas**: {idea}
+
+        ### Format for Your Response:
+        Please ensure that your final ideas include 5-6 entries and present the integrated ideas in the following format:
+        **Idea 1**: [The first method idea]
+        **Idea 2**: [The second method idea]
+        **Idea 3**: [The third method idea]
+        ...
+        """
+        if background is None or brainstorm is None or idea is None:
+            return None
+        try:
+            message = []
+            system_input = "Now you are a researcher in the field of AI with innovative and pioneering abilities. You are good at generating innovative and original ideas to solve cutting-edge problems in the field of AI."
+            message.append({"role": "system", "content": system_input})
+            message_input = prompt_template_idea.format(
+                background=background, brainstorm=brainstorm, idea=idea)
+            message.append({"role": "user", "content": message_input})
+            response = self.generator.create(
+                messages=message,
+            )
+            idea = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+            return None
+        return idea
+
+    def filter_idea(self, idea: str, background: str) -> str:
+        prompt_template_filter = """
+        ### Task Description: 
+        You will be provided with some ideas you previously generated, and a research background. Your task is to select 5-6 ideas that best address the problems described in the research background (priority) and ideas that are relatively novel and feasible (secondary), and then record the ideas and their content in given format. Remember that the content of idea includes everything about the idea.
+
+        ### Information Provided:
+        1. **Ideas**: These are the ideas you previously generated based on the research background and several related papers.
+        2. **Research Background**: This document describes specific problems and challenges that need to be addressed.
+
+        ### Approach:
+        Your approach should be systematic:
+        - **Step 1**: Analyze the research background to understand the specific problems that need solutions.
+        - **Step 2**: Critically review the ideas, selecting 5-6 ideas that are most effective in solving the problems in the research background (priority) and that are also relatively novel and feasible (secondary).
+
+        ### Specific Information:
+        I will provide you with specific information now; please use them according to the instructions above:
+        1. **Ideas**: {idea}
+        2. **Research Background**: {background}
+
+        ### Format for Your Response:
+        Please ensure that your final ideas include 5-6 entries, whose content has not been modified. Don't generate any explanation and just present the filtered ideas as well as their content in the following format:
+        **Idea 1**: [The first method idea]  
+        **Idea 2**: [The second method idea]  
+        **Idea 3**: [The third method idea]  
+        ...
+        """
+        if background is None or idea is None:
+            return None
+        try:
+            message = []
+            system_input = "Now you are a researcher in the field of AI. You are good at selecting the ideas that meet the requirements."
+            message.append({"role": "system", "content": system_input})
+            message_input = prompt_template_filter.format(
+                idea=idea, background=background
+            )
+            message.append({"role": "user", "content": message_input})
+            response = self.generator.create(
+                messages=message,
+            )
+            idea_filtered = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+            return None
+        return idea_filtered
+
+    def modify_idea(self, background: str, idea: str) -> str:
+        prompt_template_modify = """
+        ### Task Description:
+        You will be provided with the research background and the original ideas you previously generated. Your task is to refine these original ideas by filtering out those with low feasibility and insufficient novelty while enhancing the most critical and relevant ideas to make them more novel, feasible, targeted, and specific. If applicable, you may include formulas or algorithms to support the ideas. Additionally, please adhere to the following requirements:
+        1. Do not generate ideas that are repetitive or contradictory.
+        2. Ensure that the generated ideas are coherent and form a cohesive whole.
+
+        ### Information Provided:
+        1. **Research background**: This is the starting point of the original idea and the basis for analyzing whether the idea should be filtered.
+        2. **Original ideas**: These are the ideas you previously generated based on research background and several related papers.
+
+        ### Approach:
+        Your approach should be systematic:
+        - **Step 1**: Thoroughly review the research background to understand the context and objectives.
+        - **Step 2**: Analyze the original ideas critically, identifying aspects with low feasibility or insufficient novelty, and then filter out them.
+        - **Step 3**: Enhance the most critical and relevant ideas by making them more novel, feasible, targeted, and specific. Incorporate formulas or algorithms if they strengthen the ideas.
+
+        ### Specific Information:
+        I will provide you with specific information now, please use them according to the instructions above:
+        1. **Research background**: {background}
+        2. **Original idea**: {idea}
+
+        ### Format for Your Response:
+        Please ensure that your response only includes the final ideas, which include 2 to 4 entries, presented in the following format:
+        **Idea 1**: [The first method idea]
+        - **Details**: [Details of the first idea]
+        **Idea 2**: [The second method idea]
+        - **Details**: [Details of the second idea]
+        ...
+        """
+        if background is None or idea is None:
+            return None
+        try:
+            message = []
+            system_input = "Now you are a researcher in the field of AI with innovative and pioneering abilities. You are good at using innovative and original methods to solve cutting-edge problems in the field of AI."
+            message.append({"role": "system", "content": system_input})
+            message_input = prompt_template_modify.format(
+                background=background, idea=idea
+            )
+            message.append({"role": "user", "content": message_input})
+            response = self.generator.create(
+                messages=message,
+            )
+            idea_modified = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+            return None
+        return idea_modified
+
+    def generate_ground_truth(self, abstract: str, contribution: str, text: str) -> str:
+        prompt_method = """
+        ### Task Description:
+        You will be provided with the abstract and a text extracted from a paper and three contributions of the paper. Your task is to filter, refine, and revise the content of the contributions through the text provided to you.
+
+        ### Information Provided:
+        1. **Abstract**: It's the abstract directly extracted from the paper.
+        2. **Contributions**: These are the contributions (methods) we have summarized based on the abstract and introduction of the paper.
+        3. **Text**: It's the text directly extracted from the paper, containing the methodology of the paper.
+
+        ### Approach:
+        Your approach should be systematic:
+        - **Step 1**: Start by reading the abstract and contributions, to understand the main work of this paper.
+        - **Step 2**: Then, read the text, to find information related to the contributions and ignore other information. If you think there is missing content in the contributions section, you can add one. On the contrary, if you think there is content duplication, merge or delete one. Please ensure that the final contributions have 2 to 4 entries.
+        - **Step 3**: Finally, provide specific details for each contribution as detailed and comprehensive as possible based on the content in the text. If applicable, you may include formulas or algorithms to support the ideas.
+
+        ### Specific Information:
+        I will provide you with specific information now, please use them according to the instructions above:
+        1. **Abstract**: {abstract}
+        2. **Contribution**: {contribution}
+        3. **Text**: {text}
+
+        ### Format for Your Response:
+        Your output should follow the format, and please note that your subject should not be 'the paper' but 'this method' or the specific method name:
+        **Idea 1**: [The first method idea]
+        - **Details**: [Details of the first idea]
+        **Idea 2**: [The second method idea]
+        - **Details**: [Details of the second idea]
+        ...
+        """
+        ground_truth = None
+        if abstract is None or contribution is None or text is None:
+            return None
+        try:
+            message = []
+            prompt = prompt_method.format(
+                abstract=abstract, contribution=contribution, text=text
+            )
+            message.append({"role": "user", "content": prompt})
+            response = self.summarizer.create(
+                messages=message,
+            )
+            ground_truth = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+        return ground_truth
+    
+    def transfer_form(self, idea: str):
+        prompt_template_transfer = """
+        ### Task Description:
+        I will give you some ideas, please standardize the output format of the ideas without simplifying or modifying their specific content. Note that the content of each idea includes everything about the idea。
+
+        ### Specific Information:
+        I will provide you with specific information now, please use them according to the instructions above:
+        **Ideas**:
+        '''
+        {idea}
+        '''
+        ### Format for Your Response:
+        Please ensure that your final answer is strictly presented in the following format:
+        **1.**<The content of the first idea>. 
+        **2.**<The content of the second idea>. 
+        ...
+        """
+        if idea is None:
+            return None
+        try:
+            message = []
+            message_input = prompt_template_transfer.format(idea=idea)
+            message.append({"role": "user", "content": message_input})
+            response = self.generator.create(
+                messages=message,
+            )
+            idea_norm = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+            return None
+        return idea_norm
+
+    def select_contribution(self, idea: str, contribution: list[str]) -> str:
+        prompt_template_select = """
+        ### Task Description:
+        You will be provided with an idea you previously generated, and some reference ideas. Your task is to select the idea that is most similar to the one you generated from the reference ideas.
+
+        ### Information Provided:
+        1. **Generated Idea**: This is the idea you previously generated based on research background and several related papers.
+        2. **Reference Ideas**: These are the ideas that you should select from.
+
+        ### Approach:
+        Your approach should be systematic:
+        - **Step 1**: Analyze the generated idea to understand the methods it describes.
+        - **Step 2**: Critically review the reference ideas, selecting the idea that is most similar to the methods in the generated idea.
+
+        ### Specific Information:
+        I will provide you with specific information now, please use them according to the instructions above:
+        1. **Idea**: {idea}
+        2. **Reference Ideas**: {reference_ideas}
+
+        ### Format for Your Response:
+        Your answer can only have one number (strating from 1), indicating the number of the most similar idea, and cannot contain any other content.
+        """
+        if idea is None or contribution is None:
+            return None
+        try:
+            message = []
+            reference_ideas = ""
+            for i, idea in enumerate(contribution):
+                reference_ideas += "Idea {i}: ".format(i=i + 1) + "\n" + idea + "\n \n"
+            message_input = prompt_template_select.format(
+                idea=idea, reference_ideas=reference_ideas
+            )
+            message.append({"role": "user", "content": message_input})
+            response = self.generator.create(
+                messages=message,
+                max_tokens=10,
+            )
+            index = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+            return None
+        return index
+
+    def get_similarity_score(self, idea: str, contribution: str) -> str:
+        prompt_template_select = """
+        ### Task Description:
+        You will be provided with an idea you previously generated, and a reference idea. Your task is to determine the similarity between the generated idea and the reference idea and give a score from 0 to 5.
+
+        ### Information Provided:
+        1. **Generated Idea**: This is the idea you previously generated based on research background and several related papers.
+        2. **Reference Idea**: This is the idea we provide you with that you need to compare with the generated idea.
+
+        ### Approach:
+        You should follow the following scoring criteria:
+        - **0**: The generated idea and reference idea are completely unrelated with no discernible similarities.
+        - **1**: The generated idea and reference idea have a vague connection, but differ significantly in their main concepts or approach.
+        - **2**: The generated idea and reference idea share a general concept but differ in most key aspects such as methodology or application.
+        - **3**: The generated idea and reference idea are similar in several areas, including general concept and some aspects of methodology, but differ in details or specific approaches.
+        - **4**: The generated idea and reference idea are largely similar in concept, methodology, and approach, with only minor differences in specifics.
+        - **5**: The generated idea and reference idea are nearly identical in all key aspects, including concept, methodology, and approach.
+
+        ### Specific Information:
+        I will provide you with specific information now, please use them according to the instructions above:
+        1. **Generated Idea**: {idea}
+        2. **Reference Idea**: {reference_idea}
+
+        ### Format for Your Response:
+        Your answer can only have one number (from 0 to 5), indicating the similarity score, and cannot contain any other content.
+        """
+        if idea is None or contribution is None:
+            return None
+        try:
+            message = []
+            reference_ideas = ""
+            message_input = prompt_template_select.format(
+                idea=idea, reference_idea=contribution
+            )
+            message.append({"role": "user", "content": message_input})
+            response = self.generator.create(
+                messages=message,
+                max_tokens=10,
+            )
+            score = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+            return None
+        return score
+
+    def novelty_eval(
+        self, current_round: int, num_rounds: int, max_num_iterations: int, idea: str, last_query_results: str, msg_history: list
+    ):
+        novelty_system_msg = """You are an ambitious AI PhD student who is looking to publish a paper that will contribute significantly to the field.
+        You have an idea and you want to check if it is novel or not. I.e., not overlapping significantly with existing literature or already well explored.
+        Be a harsh critic for novelty, ensure there is a sufficient contribution in the idea for a new conference or workshop paper.
+        You will be given access to the Semantic Scholar API, which you may use to survey the literature and find relevant papers to help you make your decision.
+        The top 10 results for any search query will be presented to you with the abstracts.
+
+        You will be given {num_rounds} rounds to decide on the paper.
+        At any round, compare the provided idea with the information found in the article and provide a novelty score from 0 to 10.
+        In each search round, you should give a query and a novelty score based on the information in the relevant papers. 
+        If there are no relevant papers, give a novelty score based on your own feelings.
+        """
+
+        novelty_prompt = '''Round {current_round}/{num_rounds}.
+        You have this idea:
+
+        """
+        {idea}
+        """
+
+        The results of the last query are (empty on first round):
+        """
+        {last_query_results}
+        """
+
+        Respond in the following format:
+
+        THOUGHT:
+        <THOUGHT>
+
+        RESPONSE:
+        ```json
+        <JSON>
+        ```
+
+        In <THOUGHT>, first briefly reason over the idea and identify any query that could help you suggest a score based on its novelty. Then give your perceived novelty score.
+
+        In <JSON>, respond in JSON format with ONLY the following field:
+        - "Query": An optional search query to search the literature (e.g. attention is all you need). You must make a query if you have not decided this round.
+        - "Novelty Score": A novelty score from 0 to 10.
+
+        A query will work best if you are able to recall the exact name of the paper you are looking for, or the authors.
+        This JSON will be automatically parsed, so ensure the format is precise. (the JSON MUST contain the "Query" and the "Novelty Score")
+        In the last round, you should assign a "" value to the "Query" even if you don't need to generate it.'''
+        msg=novelty_prompt.format(
+            current_round=current_round,
+            num_rounds=max_num_iterations,
+            idea=idea,
+            last_query_results=last_query_results,
+        )
+        system_message=novelty_system_msg.format(
+            num_rounds=max_num_iterations,
+        )
+        if msg_history is None:
+            msg_history = []
+        try:
+            new_msg_history = msg_history + [{"role": "user", "content": msg}]
+            response = self.generator.create(
+                messages=[
+                    {"role": "system", "content": system_message},
+                    *new_msg_history,
+                ],
+                temperature=0.75,
+                max_tokens=3000,
+                n=1,
+                stop=None,
+                seed=0,
+            )
+            content = response.choices[0].message.content
+            new_msg_history = new_msg_history + [{"role": "assistant", "content": content}]
+
+        except Exception:
+            traceback.print_exc()
+            return None
+        return content, new_msg_history
+
+    def compare_same(self, idea1:str, idea2:str, idea3:str, idea4:str, idea5:str) -> str:
+        system_input = """
+        You are an artificial intelligence researcher with extensive knowledge in this field, and now you need to make a comprehensive comparison among five ideas.
+        You will obtain a comparison standard, compare every point on the standard.
+        """
+        input_message = '''
+        ### Comparison Standard:
+        """
+        **Clarity**: It evaluates whether the method is articulated in a straightforward and coherent manner, facilitating a comprehensive understanding for both practitioners and researchers, thus enabling effective application and potential adaptation in similar studies.
+        **Novelty**: It assesses the degree to which the method presents novel ideas or transformative strategies that challenge conventional practices, fostering advancements in the field and inspiring future research directions.
+        **Feasibility**: It examines the practicality and implementability of the method, ensuring that the required resources, time, and expertise are realistically available for its execution within the constraints of the study environment.
+        **Generalizability**: It determines how broadly the method can be extended or adapted to various contexts, populations, or situations, evaluating its applicability beyond the specific conditions of the study while maintaining relevance and effectiveness.
+        """
+
+        ### You should compare these five ideas:
+        """IDEA1
+        {idea1}
+        """
+        """IDEA2
+        {idea2}
+        """
+        """IDEA3
+        {idea3}
+        """
+        """IDEA4
+        {idea4}
+        """
+        """IDEA5
+        {idea5}
+        """
+
+        ### Respond in the following format:
+
+        THOUGHT:
+        ```thought
+        <THOUGHT>
+        ```
+
+        RESPONSE:
+        ```json
+        <JSON>
+        ```
+
+        In <THOUGHT>, You can record your reasoning process to make your comparison more organized..
+
+        In <JSON>, respond in JSON format with ONLY the following field:
+        - "Clarity": Provide an array consisting of 1-5, representing each idea separately, with the better idea placed at the beginning (e.g. [4, 5, 3, 2, 1])
+        - "Novelty": Same as above.
+        - "Feasibility": Same as above.
+        - "Generalizability": Same as above.
+        - "Overall Ranking": Same as above.
+
+        This JSON will be automatically parsed, so ensure the format is precise.
+        '''
+        if idea1 is None or idea2 is None or idea3 is None or idea4 is None or idea5 is None:
+            return None
+        try:
+            message = []
+            message.append({"role": "system", "content": system_input})
+            message_input = input_message.format(
+                idea1=idea1, idea2=idea2, idea3=idea3, idea4=idea4, idea5=idea5)
+            message.append({"role": "user", "content": message_input})
+            response = self.generator.create(
+                messages=message,
+            )
+            result = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+            return None
+        return result
+
+    def compare_all(self, idea1:str, idea2:str) -> str:
+        system_input = """
+        You are an artificial intelligence researcher with extensive knowledge in this field, and now you need to make a comprehensive comparison among five ideas.
+        You will obtain a comparison standard, compare every point on the standard, and make a overall ranking at the end.
+        """
+        input_message = '''
+        ### Comparison Standard:
+        """
+        **Novelty**: It assesses the degree to which the method presents novel ideas or transformative strategies that challenge conventional practices, fostering advancements in the field and inspiring future research directions.
+        **Feasibility**: It examines the practicality and implementability of the method, ensuring that the required resources, time, and expertise are realistically available for its execution within the constraints of the study environment.
+        **Clarity**: It evaluates whether the method is articulated in a straightforward and coherent manner, facilitating a comprehensive understanding for both practitioners and researchers, thus enabling effective application and potential adaptation in similar studies.
+        **Generalizability**: It determines how broadly the method can be extended or adapted to various contexts, populations, or situations, evaluating its applicability beyond the specific conditions of the study while maintaining relevance and effectiveness.
+        """
+
+        ### You should compare these five ideas:
+        """IDEA1
+        {idea1}
+        """
+        """IDEA2
+        {idea2}
+        """
+
+        ### Respond in the following format:
+
+        THOUGHT:
+        ```thought
+        <THOUGHT>
+        ```
+
+        RESPONSE:
+        ```json
+        <JSON>
+        ```
+
+        In <THOUGHT>, You can record your reasoning process and explain why you think the idea is better in each aspect in detail to make your comparison more organized.
+
+        In <JSON>, respond in JSON format with ONLY the following field:
+        - "Novelty": Provide an array consisting of 1 and 2, representing each idea separately, with the better idea placed at the beginning (e.g. [1, 2]).
+        - "Feasibility": Same as above.
+        - "clarity": Same as above.
+        - "Generalizability": Same as above.
+        - "Overall Ranking": Same as above.
+
+
+        This THOUGHT and JSON will be automatically parsed, so ensure the format is precise.
+        '''
+        if idea1 is None or idea2 is None:
+            return None
+        try:
+            message = []
+            message.append({"role": "system", "content": system_input})
+            message_input = input_message.format(
+                idea1=idea1, idea2=idea2)
+            message.append({"role": "user", "content": message_input})
+            response = self.generator.create(
+                messages=message,
+            )
+            result = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+            return None
+        return result
+
+    def compare_novelty_and_feasibility(self, idea1:str, idea2:str) -> str:
+        system_input = """
+        You are an artificial intelligence researcher with extensive knowledge in this field, and now you need to make a comprehensive comparison between two ideas.
+        You will obtain a comparison standard, compare every point on the standard, and make a ranking at the end.
+        """
+        input_message = '''
+        ### Comparison Standard:
+        """
+        **Novelty**: It assesses the degree to which the method presents novel ideas or transformative strategies that challenge conventional practices, fostering advancements in the field and inspiring future research directions.
+        **Feasibility**: It examines the practicality and implementability of the method, ensuring that the required resources, time, and expertise are realistically available for its execution within the constraints of the study environment.
+        """
+
+        ### You should compare these five ideas:
+        """IDEA1
+        {idea1}
+        """
+        """IDEA2
+        {idea2}
+        """
+
+        ### Respond in the following format:
+
+        THOUGHT:
+        ```thought
+        <THOUGHT>
+        ```
+
+        RESPONSE:
+        ```json
+        <JSON>
+        ```
+
+        In <THOUGHT>, You can record your reasoning process and explain why you think the idea is better in each aspect in detail to make your comparison more organized.
+
+        In <JSON>, respond in JSON format with ONLY the following field:
+        - "Novelty": Provide an array consisting of 1 and 2, representing each idea separately, with the better idea placed at the beginning (e.g. [1, 2]).
+        - "Feasibility": Same as above.
+
+        This THOUGHT and JSON will be automatically parsed, so ensure the format is precise.
+        '''
+        if idea1 is None or idea2 is None:
+            return None
+        try:
+            message = []
+            message.append({"role": "system", "content": system_input})
+            message_input = input_message.format(
+                idea1=idea1, idea2=idea2)
+            message.append({"role": "user", "content": message_input})
+            response = self.generator.create(
+                messages=message,
+            )
+            result = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+            return None
+        return result
+    
+    def compare_novelty(self, idea1:str, idea2:str) -> str:
+        system_input = """
+        You are an artificial intelligence researcher with extensive knowledge in this field, and now you need to make a comparison between two ideas.
+        You will obtain a comparison standard, compare the novelty between the ideas, and make a ranking at the end.
+        """
+        input_message = '''
+        ### Comparison Standard:
+        """
+        **Novelty**: It assesses the degree to which the method presents novel ideas or transformative strategies that challenge conventional practices, fostering advancements in the field and inspiring future research directions.
+        """
+
+        ### You should compare these five ideas:
+        """IDEA1
+        {idea1}
+        """
+        """IDEA2
+        {idea2}
+        """
+
+        ### Respond in the following format:
+
+        THOUGHT:
+        ```thought
+        <THOUGHT>
+        ```
+
+        RESPONSE:
+        ```json
+        <JSON>
+        ```
+
+        In <THOUGHT>, You can record your reasoning process and explain why you think the idea is better in each aspect in detail to make your comparison more organized.
+
+        In <JSON>, respond in JSON format with ONLY the following field:
+        - "Novelty": Provide an array consisting of 1 and 2, representing each idea separately, with the better idea placed at the beginning (e.g. [1, 2]).
+
+        This THOUGHT and JSON will be automatically parsed, so ensure the format is precise and don't forget the label "Novelty".
+        '''
+        if idea1 is None or idea2 is None:
+            return None
+        try:
+            message = []
+            message.append({"role": "system", "content": system_input})
+            message_input = input_message.format(
+                idea1=idea1, idea2=idea2)
+            message.append({"role": "user", "content": message_input})
+            response = self.generator.create(
+                messages=message,
+            )
+            result = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+            return None
+        return result
+
+    def compare_feasibility(self, idea1:str, idea2:str) -> str:
+        system_input = """
+        You are an artificial intelligence researcher with extensive knowledge in this field, and now you need to make a comparison between two ideas.
+        You will obtain a comparison standard, compare the feasibility between the ideas, and make a ranking at the end.
+        """
+        input_message = '''
+        ### Comparison Standard:
+        """
+        **Feasibility**: It examines the practicality and implementability of the method, ensuring that the required resources, time, and expertise are realistically available for its execution within the constraints of the study environment.
+        """
+
+        ### You should compare these five ideas:
+        """IDEA1
+        {idea1}
+        """
+        """IDEA2
+        {idea2}
+        """
+
+        ### Respond in the following format:
+
+        THOUGHT:
+        ```thought
+        <THOUGHT>
+        ```
+
+        RESPONSE:
+        ```json
+        <JSON>
+        ```
+
+        In <THOUGHT>, You can record your reasoning process and explain why you think the idea is better in each aspect in detail to make your comparison more organized.
+
+        In <JSON>, respond in JSON format with ONLY the following field:
+        - "Feasibility": Provide an array consisting of 1 and 2, representing each idea separately, with the better idea placed at the beginning (e.g. [1, 2]).
+
+        This THOUGHT and JSON will be automatically parsed, so ensure the format is precise and don't forget the label "Feasibility".
+        '''
+        if idea1 is None or idea2 is None:
+            return None
+        try:
+            message = []
+            message.append({"role": "system", "content": system_input})
+            message_input = input_message.format(
+                idea1=idea1, idea2=idea2)
+            message.append({"role": "user", "content": message_input})
+            response = self.generator.create(
+                messages=message,
+            )
+            result = response.choices[0].message.content
+        except Exception:
+            traceback.print_exc()
+            return None
+        return result
+
+
+if __name__ == "__main__":
+    config = ConfigReader.load("/mnt/llms/data/scimon-plus-data/configs/datasets.yaml")
+    api_helper = APIHelper(config=config)

src/utils/paper_client.py

@@ -0,0 +1,630 @@
+import os
+import re
+import json
+from tqdm import tqdm
+from neo4j import GraphDatabase
+from collections import defaultdict, deque
+from py2neo import Graph, Node, Relationship
+from loguru import logger
+
+class PaperClient:
+    def __init__(self, config) -> None:
+        self.config = config
+        self.driver = self.get_neo4j_driver()
+        self.teb_model = None
+
+    def get_neo4j_driver(self):
+        # 配置信息
+        URI = os.environ["NEO4J_URL"]
+        NEO4J_USERNAME = os.environ["NEO4J_USERNAME"]
+        NEO4J_PASSWD = os.environ["NEO4J_PASSWD"]
+        AUTH = (NEO4J_USERNAME, NEO4J_PASSWD)
+        # 连接到 Neo4j 数据库
+        driver = GraphDatabase.driver(URI, auth=AUTH)
+        return driver
+
+    def update_paper_from_client(self, paper):
+        paper_id = paper["hash_id"]
+        if paper_id is None:
+            return None
+        query = f"""
+            MATCH (p:Paper {{hash_id: {paper_id}}})
+            RETURN p
+            """
+        with self.driver.session() as session:
+            result = session.execute_read(lambda tx: tx.run(query).data())
+        if result:
+            paper_from_client = result[0]['p']
+            if paper_from_client is not None:
+                paper.update(paper_from_client)
+    
+    def get_paper_attribute(self, paper_id, attribute_name):
+        query = f"""
+            MATCH (p:Paper {{hash_id: {paper_id}}})
+            RETURN p.{attribute_name} AS attributeValue
+            """
+        with self.driver.session() as session:
+            result = session.execute_read(lambda tx: tx.run(query).data())
+        if result:
+            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}'}})
+            RETURN p
+            """
+        with self.driver.session() as session:
+            result = session.execute_read(lambda tx: tx.run(query).data())
+        if result:
+            return result[0]['p']
+        else:
+            return None
+
+    def get_paper_from_term(self, entity):
+        if entity is None:
+            return None
+        query = """
+            MATCH (p:Paper)
+            WHERE p.entity = $entity
+            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())
+        if 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):
+        query = """
+            MATCH (e:Entity)-[:RELATED_TO]->(p:Paper {hash_id: $hash_id})
+            RETURN e.name AS entity_name
+        """
+        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 []
+
+    def find_paper_by_entity(self, entity_name):
+        query = """
+            MATCH (e1:Entity {name: $entity_name})-[:RELATED_TO]->(p:Paper)
+            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())
+        if 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 = """
+        MATCH (e:Entity {name: $entity_name})-[:RELATED_TO]->(p:Paper)
+        WHERE p.abstract CONTAINS $entity_name OR
+          p.introduction CONTAINS $entity_name OR
+          p.methodology CONTAINS $entity_name
+        RETURN p.abstract AS abstract, 
+            p.introduction AS introduction, 
+            p.methodology AS methodology, 
+            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())
+        for record in result:
+            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])
+
+        return sentences
+
+    def select_paper(self, venue_name, year):
+        query = """
+            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())
+        if result:
+            return [record['n'] for record in result]
+        else:
+            return []
+
+    def add_paper_node(self, paper: dict):
+        if "summary" not in paper.keys():
+            paper["summary"] = None
+        if "abstract" not in paper.keys():
+            paper["abstract"] = None
+        if "introduction" not in paper.keys():
+            paper["introduction"] = None
+        if "reference" not in paper.keys():
+            paper["reference"] = None
+        if "cite" not in paper.keys():
+            paper["cite"] = None
+        if "motivation" not in paper.keys():
+            paper["motivation"] = None
+        if "contribution" not in paper.keys():
+            paper["contribution"] = None
+        if "methodology" not in paper.keys():
+            paper["methodology"] = None
+        if "ground_truth" not in paper.keys():
+            paper["ground_truth"] = None
+        if "reference_filter" not in paper.keys():
+            paper["reference_filter"] = None
+        if "conclusions" not in paper.keys():
+            paper["conclusions"] = None
+        query = """
+            MERGE (p:Paper {hash_id: $hash_id})
+            ON CREATE SET p.venue_name = $venue_name, p.year = $year, p.title = $title, p.pdf_url = $pdf_url, p.abstract = $abstract, p.introduction = $introduction, p.reference = $reference, p.summary = $summary, p.motivation = $motivation, p.contribution = $contribution, p.methodology = $methodology, p.ground_truth = $ground_truth, p.reference_filter = $reference_filter, p.conclusions = $conclusions
+            ON MATCH SET p.venue_name = $venue_name, p.year = $year, p.title = $title, p.pdf_url = $pdf_url, p.abstract = $abstract, p.introduction = $introduction, p.reference = $reference, p.summary = $summary, p.motivation = $motivation, p.contribution = $contribution, p.methodology = $methodology, p.ground_truth = $ground_truth, p.reference_filter = $reference_filter, p.conclusions = $conclusions
+            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())
+
+    def check_entity_node_count(self, hash_id: int):
+        query_check_count = """
+            MATCH (e:Entity)-[:RELATED_TO]->(p:Paper {hash_id: $hash_id})
+            RETURN count(e) AS entity_count
+        """
+        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())
+            if result[0]["entity_count"] > 3:
+                return False
+        return True
+
+    def add_entity_node(self, hash_id: int, entities: list):
+        query = """
+            MERGE (e:Entity {name: $entity_name})
+            WITH e
+            MATCH (p:Paper {hash_id: $hash_id})
+            MERGE (e)-[:RELATED_TO]->(p)
+            RETURN e, p
+        """
+        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())
+        
+    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):
+        if hash_id is not None:
+            query = """
+            MATCH (p:Paper {hash_id: $hash_id})
+            WHERE p.abstract IS NOT NULL
+            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:
+            query = """
+            MATCH (p:Paper)
+            WHERE p.abstract IS NOT NULL
+            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).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.config.DEFAULT.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())
+
+    def add_paper_bg_embedding(self, embedding_model, hash_id=None):
+        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
+            """
+            with self.driver.session() as session:
+                results = session.execute_write(lambda tx: tx.run(query, hash_id=hash_id).data())
+        else:
+            query = """
+            MATCH (p:Paper)
+            WHERE p.motivation IS NOT NULL
+            RETURN p.motivation AS context, p.hash_id AS hash_id
+            """
+            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.config.DEFAULT.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()
+            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):
+        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
+            """
+            with self.driver.session() as session:
+                results = session.execute_write(lambda tx: tx.run(query, hash_id=hash_id).data())
+        else:
+            query = """
+            MATCH (p:Paper)
+            WHERE p.contribution IS NOT NULL
+            RETURN p.contribution AS context, p.hash_id AS hash_id
+            """
+            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.config.DEFAULT.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())
+    
+
+    def add_paper_summary_embedding(self, embedding_model, hash_id=None):
+        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
+            """
+            with self.driver.session() as session:
+                results = session.execute_write(lambda tx: tx.run(query, hash_id=hash_id).data())
+        else:
+            query = """
+            MATCH (p:Paper)
+            WHERE p.summary IS NOT NULL
+            RETURN p.summary AS context, p.hash_id AS hash_id
+            """
+            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.config.DEFAULT.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())
+        
+    def cosine_similarity_search(self, embedding, k=1, type_name="embedding"):
+        query = f"""
+            MATCH (paper:Paper)
+            WITH paper,
+                vector.similarity.cosine(paper.{type_name}, $embedding) AS score
+            WHERE score > 0
+            RETURN paper, score
+            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 = [] 
+        for result in results:
+            related_paper.append(result["paper"]["hash_id"])
+        return related_paper
+
+    def create_vector_index(self):
+        """
+        适用于Paper节点
+        针对Paper节点上的是属性 embedding 进行索引
+        索引向量的维度为384
+        适用余弦相似度作为计算相似度的方法
+        """
+        query = """
+            CREATE VECTOR INDEX `paper-embeddings`
+            FOR (n:Paper) ON (n.embedding)
+            OPTIONS {indexConfig: {
+            `vector.dimensions`: 384,
+            `vector.similarity_function`: 'cosine'
+            }}
+            """
+        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 []
+        # WHERE p.year < "2024" AND p.venue_name <> "acl"
+        query = """
+            UNWIND $paper_id_list AS hash_id
+            MATCH (p:Paper {hash_id: hash_id})
+            WHERE p.year < $year
+            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]
+        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:
+            result = session.execute_read(lambda tx: tx.run(query).data())
+        for record in result:
+            if record["name"] == "paper-embeddings":
+                return True
+        return False
+
+    def clear_database(self):
+        query = """
+        MATCH (n)
+        DETACH DELETE n
+        """
+        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)
+            WITH COUNT(p) AS PaperCount
+            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']
+        return paper_num
+
+    def get_entity_text(self):
+        query = """
+            MATCH (e:Entity)-[:RELATED_TO]->(p:Paper)
+            WHERE p.venue_name = $venue_name and p.year = $year
+            WITH p, collect(e.name) AS entity_names
+            RETURN p, reduce(text = '', name IN entity_names | text + ' ' + name) AS entity_text
+        """
+        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]
+        return text_list
+    
+    def get_entity_combinations(self, venue_name, year):
+        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 = """
+                MATCH (e1:Entity {name: $entity_name_1})
+                MATCH (e2:Entity {name: $entity_name_2})
+                MERGE (e1)-[r:CONNECT]->(e2)
+                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)
+            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()
+                    )
+                    # logger.debug(f"CONNECT relation created or updated between {entity_name_1} and {entity_name_2} for Paper ID {paper_id}")
+                    break  # 如果找到一次出现就可以退出循环
+
+        query = """
+            MATCH (e:Entity)-[:RELATED_TO]->(p:Paper)
+            WHERE p.venue_name=$venue_name and p.year=$year
+            WITH p, collect(e) as entities
+            UNWIND range(0, size(entities)-2) as i
+            UNWIND range(i+1, size(entities)-1) as j
+            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())
+            for record in tqdm(result):
+                paper_id = record["hash_id"]
+                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)
+
+    def build_citemap(self):
+        citemap = defaultdict(set)
+        query = """
+            MATCH (p:Paper)
+            RETURN p.hash_id AS hash_id, p.cite_id_list AS cite_id_list
+        """
+        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']
+            if cite_id_list:
+                for cited_id in cite_id_list:
+                    citemap[hash_id].add(cited_id)
+        return citemap
+
+    def neo4j_backup(self):
+        URI = os.environ["NEO4J_URL"]
+        NEO4J_USERNAME = os.environ["NEO4J_USERNAME"]
+        NEO4J_PASSWD = os.environ["NEO4J_PASSWD"]
+        AUTH = (NEO4J_USERNAME, NEO4J_PASSWD)
+        graph = Graph(URI, auth=AUTH)
+        query = """
+            MATCH (e:Entity)-[r:RELATED_TO]->(p:Paper)
+            RETURN p, e, r
+        """
+        results = graph.run(query)
+        # 创建一个字典来保存数据
+        data = {"nodes": [], "relationships": []}
+        # 处理查询结果
+        for record in tqdm(results):
+            paper_node = record["p"]
+            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["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"] = [dict(t) for t in {tuple(d.items()) for d in data["nodes"]}]
+        unique_nodes = []
+        seen = set()
+        for node in tqdm(data["nodes"]):
+            # 将字典项转换为不可变的元组，以便用于集合去重
+            node_tuple = str(tuple(sorted(node.items())))
+            if node_tuple not in seen:
+                seen.add(node_tuple)
+                unique_nodes.append(node)
+        data["nodes"] = unique_nodes
+        # 将数据保存为 JSON 文件
+        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"]
+        NEO4J_USERNAME = os.environ["NEO4J_USERNAME"]
+        NEO4J_PASSWD = os.environ["NEO4J_PASSWD"]
+        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:
+            data = json.load(f)
+        # 创建节点
+        nodes = {}
+        for node_data in data["nodes"]:
+            label = node_data["label"]
+            properties = node_data["properties"]
+            node = Node(label, **properties)
+            graph.create(node)
+            nodes[node_data["id"]] = node
+
+        # 创建关系
+        for relationship_data in data["relationships"]:
+            start_node = nodes[relationship_data["start_node"]]
+            end_node = nodes[relationship_data["end_node"]]
+            properties = relationship_data["properties"]
+            rel_type = relationship_data["type"]
+            relationship = Relationship(start_node, rel_type, end_node, **properties)
+            graph.create(relationship)
+
+    def get_paper_by_id(self, hash_id):
+        paper = {"hash_id": hash_id}
+        self.update_paper_from_client(paper)
+        return paper
+
+
+if __name__ == "__main__":
+    from header import get_dir, ConfigReader
+    config_path = get_dir("./configs/datasets.yaml")
+    config = ConfigReader.load(config_path)
+    paper_client = PaperClient(config)
+    # paper_client.neo4j_backup()
+    paper_client.neo4j_import_data()

src/utils/paper_crawling.py

@@ -0,0 +1,363 @@
+import os
+import re
+import random
+import requests
+from requests.exceptions import RequestException
+import re
+from bs4 import BeautifulSoup
+from .hash import generate_hash_id
+from .header import get_dir
+from loguru import logger
+
+with open(get_dir("./assets/data/user_agents.txt"), "r", encoding="utf8") as f:
+    user_agents = [l.rstrip() for l in f.readlines()]
+
+
+def extract_title_from_index(index_url):
+    try:
+        headers = {"User-Agent": random.choice(user_agents)}
+        response_title = requests.get(index_url, headers=headers)
+        response_title.raise_for_status()
+        soup = BeautifulSoup(response_title.content, "html.parser")
+        papers = soup.find_all("h2", id="title")
+        for paper in papers:
+            title = paper.text.strip()
+            return title
+    except RequestException as e:
+        logger.error(f"Failed to extract title from {index_url}: {e}")
+        return None
+
+
+def extract_year_from_index(index_url):
+    try:
+        headers = {"User-Agent": random.choice(user_agents)}
+        response_year = requests.get(index_url, headers=headers)
+        response_year.raise_for_status()
+        soup = BeautifulSoup(response_year.content, "html.parser")
+
+        year_tag = soup.find("dt", text="Year:")
+        if year_tag:
+            year_dd = year_tag.find_next_sibling("dd")
+            if year_dd:
+                year = year_dd.text.strip()
+                return year
+        else:
+            print(f"Year not found in {index_url}")
+            return None
+    except requests.RequestException as e:
+        print(f"Failed to extract year from {index_url}: {e}")
+        return None
+
+
+def extract_pdf_url_from_index(index_url, id):
+    try:
+        headers = {"User-Agent": random.choice(user_agents)}
+        response = requests.get(index_url, headers=headers)
+        response.raise_for_status()
+        soup = BeautifulSoup(response.content, "html.parser")
+        pdf_link = soup.find("a", href=True, string=re.compile(r"\bPDF\b", re.I))
+        if pdf_link:
+            pdf_url = pdf_link["href"]
+            return pdf_url
+        else:
+            logger.warning(f"No PDF link found on {index_url}")
+            return None
+    except RequestException as e:
+        logger.error(f"Failed to extract PDF URL from {index_url}: {e}")
+        return None
+
+
+class PaperCrawling:
+    def __init__(self, config, data_type="train") -> None:
+        self.base_url = "https://aclanthology.org/"
+        self.data_type = data_type
+        self.paper_pdf_folder = config.DEFAULT.pdf_cached
+        if not os.path.exists(self.paper_pdf_folder):
+            os.makedirs(self.paper_pdf_folder)
+            logger.info(f"Created directory '{self.paper_pdf_folder}'")
+
+    def need_to_parse(self, paper: dict):
+        if (
+            paper["abstract"] is None
+            or paper["introduction"] is None
+            or paper["reference"] is None
+        ):
+            return True
+        return False
+
+    def get_title(self, paper):
+        index_url = f"{self.base_url}{paper['id']}/"
+        title = extract_title_from_index(index_url)
+        return title
+
+    def get_year(self, paper):
+        index_url = f"{self.base_url}{paper['id']}/"
+        year = extract_year_from_index(index_url)
+        return year
+
+    def get_pdf_url(self, paper):
+        if "pdf_url" not in paper.keys() or paper["pdf_url"] is None:
+            index_url = f"{self.base_url}{paper['id']}/"
+            paper["pdf_url"] = extract_pdf_url_from_index(index_url, paper["id"])
+
+    def download_paper(self, paper):
+        headers = {"User-Agent": random.choice(user_agents)}
+        pdf_folder = os.path.join(
+            self.paper_pdf_folder, f"{paper['venue_name']}", f"{paper['year']}"
+        )
+        file_path = os.path.join(pdf_folder, f"{paper['hash_id']}.pdf")
+        paper["pdf_path"] = file_path
+        paper_url = paper["pdf_url"]
+        if not os.path.exists(pdf_folder):
+            os.makedirs(pdf_folder)
+        if os.path.exists(file_path):
+            # print("pdf file {} exist ...".format(file_path))
+            return True
+        try:
+            response = requests.get(paper_url, headers=headers, timeout=10)
+            response.raise_for_status()
+        except Exception:
+            print(f"download failed... {paper['pdf_url']}")
+            return False
+
+        if response.status_code == 200:
+            with open(file_path, "wb") as f:
+                f.write(response.content)
+            logger.info("download success {}".format(paper_url))
+            logger.info(f"save {file_path}")
+            return True
+        else:
+            print("download failed, status code: {}".format(response.status_code))
+            return False
+
+    def get_page(self, url):
+        headers = {"User-Agent": random.choice(user_agents)}
+        try:
+            response = requests.get(url, headers=headers)
+            if response.status_code == 200:
+                response.encoding = response.apparent_encoding
+                return response.text
+            return None
+        except RequestException as e:
+            print(e)
+
+    def crawling(self, year, venue_name):
+        paper_list = []
+        paper_html_list = []
+
+        def append_paper_to_list(pdf_url, title):
+            for paper in paper_html_list:
+                if paper["title"] == title:
+                    if paper["pdf_url"] != pdf_url:
+                        logger.warning(
+                            f"Different PDF URL found for the same title '{title}'."
+                        )
+                    return
+            paper_html_list.append({"pdf_url": pdf_url, "title": title})
+
+        if venue_name == "nips":
+            if year == "2024":
+                return []
+            base_url = "https://papers.nips.cc/paper_files/paper/{}"
+            target_url = base_url.format(year)
+            target_html = self.get_page(target_url)
+            soup = BeautifulSoup(target_html, "html.parser")
+            ids = soup.find("div", {"class": "container-fluid"}).find_all("li")
+            for id in ids:
+                a = id.find("a")
+                href = a.attrs.get("href")
+                pdf_url = "https://papers.nips.cc{}".format(
+                    href.replace("hash", "file")
+                    .replace("Abstract", "Paper")
+                    .replace("html", "pdf")
+                )
+                title = a.text
+                append_paper_to_list(pdf_url, title)
+            for paper_html in paper_html_list:
+                title = paper_html["title"]
+                pdf_url = paper_html["pdf_url"]
+                hash_id = generate_hash_id(title)
+                paper_list.append(
+                    {
+                        "hash_id": hash_id,
+                        "year": year,
+                        "venue_name": venue_name,
+                        "title": title,
+                        "pdf_url": pdf_url,
+                    }
+                )
+
+        elif venue_name == "cvpr":
+            base_url = "https://openaccess.thecvf.com/CVPR{}"
+            dict_cvpr = {
+                "2018": ["2018-06-19", "2018-06-20", "2018-06-21"],
+                "2019": ["2019-06-18", "2019-06-28", "2019-06-20"],
+                "2020": ["2020-06-16", "2020-06-17", "2020-06-18"],
+                "2021": ["all"],
+                "2022": ["all"],
+                "2023": ["all"],
+            }
+            if year in dict_cvpr.keys():
+                day_list = dict_cvpr[year]
+                target_url = [
+                    base_url.format(year) + "?day={}".format(day) for day in day_list
+                ]
+            else:
+                target_url = [base_url.format(year)]
+            print("paper list from {}".format(target_url))
+            for url in target_url:
+                target_html = self.get_page(url)
+                soup = BeautifulSoup(target_html, "html.parser")
+                dl_elements = soup.find("div", {"id": "content"}).find_all("dl")
+                for dl in dl_elements:
+                    dt_elements = dl.find_all("dt")
+                    dd_elements = dl.find_all("dd")
+                if year in dict_cvpr.keys():
+                    dd_elements.pop(0)
+                for idx in range(len(dt_elements)):
+                    title = dt_elements[idx].text
+                    href = dd_elements[idx * 2 + 1].find("a").attrs.get("href")
+                    pdf_url = "https://openaccess.thecvf.com/{}".format(href)
+                    hash_id = generate_hash_id(title)
+                    paper_list.append(
+                        {
+                            "hash_id": hash_id,
+                            "year": year,
+                            "venue_name": venue_name,
+                            "title": title,
+                            "pdf_url": pdf_url,
+                        }
+                    )
+
+        elif venue_name == "emnlp":
+            if year == "2024":
+                return []
+            if year not in ["2020", "2021", "2022", "2023"]:
+                dev_id = "main-container"
+            else:
+                dev_id = "{}emnlp-main".format(year)
+            base_url = "https://aclanthology.org/events/emnlp-{}"
+            target_url = base_url.format(year)
+            target_html = self.get_page(target_url)
+            soup = BeautifulSoup(target_html, "html.parser")
+            ids = soup.find("div", {"id": dev_id}).find_all("p")
+            for id in ids:
+                a = id.find("a")
+                pdf_url = a.attrs.get("href")
+                title = id.find("strong").get_text()
+                append_paper_to_list(pdf_url, title)
+            for paper_html in paper_html_list:
+                title = paper_html["title"]
+                hash_id = generate_hash_id(title)
+                pdf_url = paper_html["pdf_url"]
+                if "http" not in pdf_url:
+                    continue
+                paper_list.append(
+                    {
+                        "hash_id": hash_id,
+                        "year": year,
+                        "venue_name": venue_name,
+                        "title": title,
+                        "pdf_url": pdf_url,
+                    }
+                )
+
+        elif venue_name == "naacl":
+            # https://aclanthology.org/
+            if year in ["2023", "2020", "2017", "2014"]:
+                return []
+            dev_id = "main-container"
+            base_url = "https://aclanthology.org/events/naacl-{}/"
+            target_url = base_url.format(year)
+            target_html = self.get_page(target_url)
+            soup = BeautifulSoup(target_html, "html.parser")
+            ids = soup.find("div", {"id": dev_id}).find_all("p")
+            for id in ids:
+                a = id.find("a")
+                pdf_url = a.attrs.get("href")
+                title = id.find("strong").get_text()
+                append_paper_to_list(pdf_url, title)
+            for paper_html in paper_html_list:
+                title = paper_html["title"]
+                hash_id = generate_hash_id(title)
+                pdf_url = paper_html["pdf_url"]
+                paper_list.append(
+                    {
+                        "hash_id": hash_id,
+                        "year": year,
+                        "venue_name": venue_name,
+                        "title": title,
+                        "pdf_url": pdf_url,
+                    }
+                )
+
+        elif venue_name == "acl":
+            dev_id = "main-container"
+            base_url = "https://aclanthology.org/events/acl-{}/"
+            target_url = base_url.format(year)
+            target_html = self.get_page(target_url)
+            soup = BeautifulSoup(target_html, "html.parser")
+            ids = soup.find("div", {"id": dev_id}).find_all("p")
+            for id in ids:
+                a = id.find("a")
+                pdf_url = a.attrs.get("href")
+                title = id.find("strong").get_text()
+                append_paper_to_list(pdf_url, title)
+
+            for paper_html in paper_html_list:
+                title = paper_html["title"]
+                hash_id = generate_hash_id(title)
+                pdf_url = paper_html["pdf_url"]
+                if "http" not in pdf_url:
+                    continue
+                paper_list.append(
+                    {
+                        "hash_id": hash_id,
+                        "year": year,
+                        "venue_name": venue_name,
+                        "title": title,
+                        "pdf_url": pdf_url,
+                    }
+                )
+
+        elif venue_name == "icml":
+            hit = {
+                "2024": "v235",
+                "2023": "v202",
+                "2022": "v162",
+                "2021": "v139",
+                "2020": "v119",
+                "2019": "v97",
+                "2018": "v80",
+                "2017": "v70",
+                "2016": "v48",
+                "2015": "v37",
+                "2014": "v32",
+                "2013": "v28",
+            }
+            dev_id = "container"
+            base_url = "https://proceedings.mlr.press/{}/"
+            target_url = base_url.format(hit[year])
+            target_html = self.get_page(target_url)
+            soup = BeautifulSoup(target_html, "html.parser")
+            ids = soup.find("main", {"class": "page-content"}).find_all(
+                "div", {"class": "paper"}
+            )
+            for id in ids:
+                title = id.find("p", class_="title").text
+                pdf_url = id.find("a", text="Download PDF")["href"]
+                append_paper_to_list(pdf_url, title)
+            for paper_html in paper_html_list:
+                title = paper_html["title"]
+                hash_id = generate_hash_id(title)
+                pdf_url = paper_html["pdf_url"]
+                paper_list.append(
+                    {
+                        "hash_id": hash_id,
+                        "year": year,
+                        "venue_name": venue_name,
+                        "title": title,
+                        "pdf_url": pdf_url,
+                    }
+                )
+        return paper_list

src/utils/paper_retriever.py

@@ -0,0 +1,749 @@
+import torch
+import itertools
+import threading
+import numpy as np
+from sentence_transformers import SentenceTransformer
+from sklearn.feature_extraction.text import CountVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+from collections import Counter, defaultdict
+from loguru import logger
+from abc import ABCMeta, abstractmethod
+from .paper_client import PaperClient
+from .paper_crawling import PaperCrawling
+from .llms_api import APIHelper
+from .header import get_dir
+
+
+class UnionFind:
+    def __init__(self, n):
+        self.parent = list(range(n))
+        self.rank = [1] * n
+
+    def find(self, x):
+        if self.parent[x] != x:
+            self.parent[x] = self.find(self.parent[x])
+        return self.parent[x]
+
+    def union(self, x, y):
+        rootX = self.find(x)
+        rootY = self.find(y)
+        if rootX != rootY:
+            if self.rank[rootX] > self.rank[rootY]:
+                self.parent[rootY] = rootX
+            elif self.rank[rootX] < self.rank[rootY]:
+                self.parent[rootX] = rootY
+            else:
+                self.parent[rootY] = rootX
+                self.rank[rootX] += 1
+
+
+def can_merge(uf, similarity_matrix, i, j, threshold):
+    root_i = uf.find(i)
+    root_j = uf.find(j)
+    for k in range(len(similarity_matrix)):
+        if uf.find(k) == root_i or uf.find(k) == root_j:
+            if (
+                similarity_matrix[i][k] < threshold
+                or similarity_matrix[j][k] < threshold
+            ):
+                return False
+    return True
+
+
+class CoCite:
+    def __init__(self, config) -> None:
+        self.paper_client = PaperClient(config)
+        citemap = self.paper_client.build_citemap()
+        self.comap = defaultdict(
+            lambda: defaultdict(int)
+        )
+        for paper_id, cited_id in citemap.items():
+            for id0, id1 in itertools.combinations(cited_id, 2):
+                # ensure comap[id0][id1] == comap[id1][id0] 
+                self.comap[id0][id1] += 1
+                self.comap[id1][id0] += 1
+        logger.debug("init co-cite map success")
+
+    def get_cocite_ids(self, id_, k=1):
+        sorted_items = sorted(self.comap[id_].items(), key=lambda x: x[1], reverse=True)
+        top_k = sorted_items[:k]
+        paper_ids = []
+        for item in top_k:
+            paper_ids.append(item[0])
+        paper_ids = self.paper_client.filter_paper_id_list(paper_ids)
+        return paper_ids
+
+
+class Retriever(object):
+    __metaclass__ = ABCMeta
+    retriever_name = "BASE"
+
+    def __init__(self, config, use_cocite=False, use_cluster_to_filter=False):
+        self.config = config
+        self.use_cocite = use_cocite
+        self.use_cluster_to_filter = use_cluster_to_filter
+        self.paper_client = PaperClient(config)
+        self.cocite = CoCite(config)
+        self.api_helper = APIHelper(config=config)
+        self.embedding_model = SentenceTransformer(
+            model_name_or_path=get_dir(config.DEFAULT.embedding), device=self.config.DEFAULT.device
+        )
+        self.paper_crawling = PaperCrawling(config=config)
+        self.vectorizer = CountVectorizer()
+
+    @abstractmethod
+    def retrieve(self, bg, entities, use_evaluate):
+        pass
+
+    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 = 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)
+            )
+        new_entities = []
+        entity_paper_num_dict = {
+            k: v for k, v in entity_paper_num_dict.items() if v != 0
+        }
+        entity_paper_num_dict = dict(
+            sorted(entity_paper_num_dict.items(), key=lambda item: item[1])
+        )
+        sum_paper_num = 0
+        for key, value in entity_paper_num_dict.items():
+            if sum_paper_num <= 100:
+                sum_paper_num += value
+                new_entities.append(key)
+            elif (
+                value < self.config.RETRIEVE.limit_num
+                and sum_paper_num < self.config.RETRIEVE.sum_paper_num
+            ):
+                sum_paper_num += value
+                new_entities.append(key)
+        return new_entities
+
+    def update_related_paper(self, paper_id_list):
+        """
+        Args:
+            paper_id_list: list
+        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)
+        return related_paper
+
+    def calculate_similarity(self, entities, related_entities_list, use_weight=False):
+        if use_weight:
+            vec1 = self.vectorizer.transform([" ".join(entities)]).toarray()[0]
+            weighted_vec1 = np.array(
+                [
+                    vec1[i] * self.log_inverse_freq.get(word, 1)
+                    for i, word in enumerate(self.vectorizer.get_feature_names_out())
+                ]
+            )
+            vecs2 = self.vectorizer.transform(
+                [
+                    " ".join(related_entities)
+                    for related_entities in related_entities_list
+                ]
+            ).toarray()
+            weighted_vecs2 = np.array(
+                [
+                    [
+                        vec2[i] * self.log_inverse_freq.get(word, 1)
+                        for i, word in enumerate(
+                            self.vectorizer.get_feature_names_out()
+                        )
+                    ]
+                    for vec2 in vecs2
+                ]
+            )
+            similarity = cosine_similarity([weighted_vec1], weighted_vecs2)[0]
+        else:
+            vec1 = self.vectorizer.transform([" ".join(entities)])
+            vecs2 = self.vectorizer.transform(
+                [
+                    " ".join(related_entities)
+                    for related_entities in related_entities_list
+                ]
+            )
+            similarity = cosine_similarity(vec1, vecs2)[0]
+        return similarity
+
+    def cal_related_score(
+        self, context, related_paper_id_list, entities=None, type_name="motivation"
+    ):
+        score_1 = np.zeros((len(related_paper_id_list)))
+        score_2 = np.zeros((len(related_paper_id_list)))
+        if entities is None:
+            entities = self.api_helper.generate_entity_list(context)
+            logger.debug("get entity from context: {}".format(entities))
+        origin_vector = self.embedding_model.encode(
+            context, convert_to_tensor=True, device=self.config.DEFAULT.device
+        ).unsqueeze(0)
+        related_contexts = [
+            self.paper_client.get_paper_attribute(paper_id, type_name)
+            for paper_id in related_paper_id_list
+        ]
+        if len(related_contexts) > 0:
+            context_embeddings = self.embedding_model.encode(
+                related_contexts, batch_size=512, convert_to_tensor=True, device=self.config.DEFAULT.device
+            )
+            score_1 = torch.nn.functional.cosine_similarity(
+                origin_vector, context_embeddings
+            )
+            score_1 = score_1.cpu().numpy()
+            if self.config.RETRIEVE.need_normalize:
+                score_1 = score_1 / np.max(score_1)
+        # score_2 not enable 
+        # if self.config.RETRIEVE.beta != 0:
+        score_sn_dict = dict(zip(related_paper_id_list, score_1))
+        score_en_dict = dict(zip(related_paper_id_list, score_2))
+        score_all_dict = dict(
+            zip(
+                related_paper_id_list,
+                score_1 * self.config.RETRIEVE.alpha
+                + score_2 * self.config.RETRIEVE.beta,
+            )
+        )
+        return score_sn_dict, score_en_dict, score_all_dict
+
+    def filter_related_paper(self, score_dict, top_k):
+        if len(score_dict) <= top_k:
+            return list(score_dict.keys())
+        if not self.use_cluster_to_filter:
+            paper_id_list = (
+                list(score_dict.keys())[:top_k]
+                if len(score_dict) >= top_k
+                else list(score_dict.keys())
+            )
+            return paper_id_list
+        else:  
+            # clustering filter, ensure that each category the highest score save first
+            paper_id_list = list(score_dict.keys())
+            paper_embedding_list = [
+                self.paper_client.get_paper_attribute(paper_id, "embedding") for paper_id in paper_id_list
+            ]
+            paper_embedding = np.array(paper_embedding_list)
+            paper_embedding_list = [
+                self.paper_client.get_paper_attribute(paper_id, "contribution_embedding") for paper_id in paper_id_list
+            ]
+            paper_contribution_embedding = np.array(paper_embedding_list)
+            paper_embedding_list = [
+                self.paper_client.get_paper_attribute(paper_id, "summary_embedding") for paper_id in paper_id_list
+            ]
+            paper_summary_embedding = np.array(paper_embedding_list)
+            weight_embedding = self.config.RETRIEVE.s_bg
+            weight_contribution = self.config.RETRIEVE.s_contribution
+            weight_summary = self.config.RETRIEVE.s_summary
+            paper_embedding = (
+                weight_embedding * paper_embedding +
+                weight_contribution * paper_contribution_embedding +
+                weight_summary * paper_summary_embedding
+            )
+            similarity_matrix = np.dot(paper_embedding, paper_embedding.T)
+            related_labels = self.cluster_algorithm(paper_id_list, similarity_matrix)
+            related_paper_label_dict = dict(zip(paper_id_list, related_labels))
+            label_group = {}
+            for paper_id, label in related_paper_label_dict.items():
+                if label not in label_group:
+                    label_group[label] = []
+                label_group[label].append(paper_id)
+            paper_id_list = []
+            while len(paper_id_list) < top_k:
+                for label, papers in label_group.items():
+                    if papers:
+                        paper_id_list.append(papers.pop(0))
+                        if len(paper_id_list) >= top_k:
+                            break
+            return paper_id_list
+
+    def cosine_similarity_search(self, context, k=1, type_name="embedding"):
+        """
+        return related paper: list
+        """
+        embedding = self.embedding_model.encode(context)
+        result = self.paper_client.cosine_similarity_search(
+            embedding, k, type_name=type_name
+        )
+        # backtrack: first is itself
+        result = result[1:]
+        return result
+
+    def cluster_algorithm(self, paper_id_list, similarity_matrix):
+        threshold = self.config.RETRIEVE.similarity_threshold
+        uf = UnionFind(len(paper_id_list))
+        # merge
+        for i in range(len(similarity_matrix)):
+            for j in range(i + 1, len(similarity_matrix)):
+                if similarity_matrix[i][j] >= threshold:
+                    if can_merge(uf, similarity_matrix, i, j, threshold):
+                        uf.union(i, j)
+        cluster_labels = [uf.find(i) for i in range(len(similarity_matrix))]
+        return cluster_labels
+
+    def eval_related_paper_in_all(self, score_all_dict, target_paper_id_list):
+        score_all_dict = dict(
+            sorted(score_all_dict.items(), key=lambda item: item[1], reverse=True)
+        )
+        result = {}
+        related_paper_id_list = list(score_all_dict.keys())
+        if len(related_paper_id_list) == 0:
+            for k in self.config.RETRIEVE.top_k_list:
+                result[k] = {"recall": 0, "precision": 0}
+            return result, 0, 0, 0
+        all_paper_id_set = set(related_paper_id_list)
+        all_paper_id_set.update(target_paper_id_list)
+        all_paper_id_list = list(all_paper_id_set)
+        paper_embedding_list = [
+            self.paper_client.get_paper_attribute(paper_id, "embedding")
+            for paper_id in target_paper_id_list
+        ]
+        paper_embedding = np.array(paper_embedding_list)
+        paper_embedding_list = [
+            self.paper_client.get_paper_attribute(paper_id, "contribution_embedding")
+            for paper_id in target_paper_id_list
+        ]
+        paper_contribution_embedding = np.array(paper_embedding_list)
+        paper_embedding_list = [
+            self.paper_client.get_paper_attribute(paper_id, "summary_embedding")
+            for paper_id in target_paper_id_list
+        ]
+        paper_summary_embedding = np.array(paper_embedding_list)
+        weight_embedding = self.config.RETRIEVE.s_bg
+        weight_contribution = self.config.RETRIEVE.s_contribution
+        weight_summary = self.config.RETRIEVE.s_summary
+        target_paper_embedding = (
+            weight_embedding * paper_embedding
+            + weight_contribution * paper_contribution_embedding
+            + weight_summary * paper_summary_embedding
+        )
+        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(
+            {
+                paper_id: self.paper_client.get_paper_attribute(paper_id, "title")
+                for paper_id in target_paper_label_dict.keys()
+            }
+        )
+
+        all_labels = []
+        for paper_id in all_paper_id_list:
+            paper_bg_embedding = [
+                self.paper_client.get_paper_attribute(paper_id, "embedding")
+            ]
+            paper_bg_embedding = np.array(paper_bg_embedding)
+            paper_contribution_embedding = [
+                self.paper_client.get_paper_attribute(
+                    paper_id, "contribution_embedding"
+                )
+            ]
+            paper_contribution_embedding = np.array(paper_contribution_embedding)
+            paper_summary_embedding = [
+                self.paper_client.get_paper_attribute(paper_id, "summary_embedding")
+            ]
+            paper_summary_embedding = np.array(paper_summary_embedding)
+            paper_embedding = (
+                weight_embedding * paper_bg_embedding
+                + weight_contribution * paper_contribution_embedding
+                + weight_summary * paper_summary_embedding
+            )
+            similarities = cosine_similarity(paper_embedding, target_paper_embedding)[0]
+            if np.any(similarities >= similarity_threshold):
+                all_labels.append(target_labels[np.argmax(similarities)])
+            else:
+                all_labels.append(-1)  # other class: -1
+        all_paper_label_dict = dict(zip(all_paper_id_list, all_labels))
+        all_label_counts = Counter(all_paper_label_dict.values())
+        logger.debug(f"all label counts : {all_label_counts}")
+        target_label_counts = Counter(target_paper_label_dict.values())
+        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)
+        max_k_paper_id_list = self.filter_related_paper(score_all_dict, top_k=max_k)
+        for k in self.config.RETRIEVE.top_k_list:
+            # 前top k 的文章
+            top_k = min(k, len(max_k_paper_id_list))
+            top_k_paper_id_list = max_k_paper_id_list[:top_k]
+            top_k_paper_label_dict = {}
+            for paper_id in top_k_paper_id_list:
+                top_k_paper_label_dict[paper_id] = all_paper_label_dict[paper_id]
+            logger.debug(
+                "=== top k {} paper id list : {}".format(k, top_k_paper_label_dict)
+            )
+            logger.debug(
+                {
+                    paper_id: self.paper_client.get_paper_attribute(paper_id, "title")
+                    for paper_id in top_k_paper_label_dict.keys()
+                }
+            )
+            top_k_label_counts = Counter(top_k_paper_label_dict.values())
+            logger.debug(f"top K label counts : {top_k_label_counts}")
+            top_k_label_list = list(top_k_label_counts.keys())
+            match_label_list = list(set(target_label_list) & set(top_k_label_list))
+            logger.debug(f"match label list : {match_label_list}")
+            recall = 0
+            precision = 0
+            for label in match_label_list:
+                recall += target_label_counts[label]
+            for label in match_label_list:
+                precision += top_k_label_counts[label]
+            recall /= len(target_paper_id_list)
+            precision /= len(top_k_paper_id_list)
+            result[k] = {"recall": recall, "precision": precision}
+
+        related_paper_id_list = list(score_all_dict.keys())
+        related_paper_label_dict = {}
+        for paper_id in related_paper_id_list:
+            related_paper_label_dict[paper_id] = all_paper_label_dict[paper_id]
+        related_label_counts = Counter(related_paper_label_dict.values())
+        logger.debug(f"top K label counts : {related_label_counts}")
+        related_label_list = list(related_label_counts.keys())
+        match_label_list = list(set(target_label_list) & set(related_label_list))
+        recall = 0
+        precision = 0
+        for label in match_label_list:
+            recall += target_label_counts[label]
+        for label in match_label_list:
+            precision += related_label_counts[label]
+        recall /= len(target_paper_id_list)
+        precision /= len(related_paper_id_list)
+        logger.debug(result)
+        return result, len(target_label_counts), recall, precision
+
+
+class RetrieverFactory(object):
+    _instance = None
+    _lock = threading.Lock()
+
+    def __new__(cls, *args, **kwargs):
+        with cls._lock:
+            if cls._instance is None:
+                cls._instance = super(RetrieverFactory, cls).__new__(
+                    cls, *args, **kwargs
+                )
+                cls._instance.init_factory()
+        return cls._instance
+
+    def init_factory(self):
+        self.retriever_classes = {}
+
+    @staticmethod
+    def get_retriever_factory():
+        if RetrieverFactory._instance is None:
+            RetrieverFactory._instance = RetrieverFactory()
+        return RetrieverFactory._instance
+
+    def register_retriever(self, retriever_name, retriever_class) -> bool:
+        if retriever_name not in self.retriever_classes:
+            self.retriever_classes[retriever_name] = retriever_class
+            return True
+        else:
+            return False
+
+    def delete_retriever(self, retriever_name) -> bool:
+        if retriever_name in self.retriever_classes:
+            self.retriever_classes[retriever_name] = None
+            del self.retriever_classes[retriever_name]
+            return True
+        else:
+            return False
+
+    def __getitem__(self, key):
+        return self.retriever_classes[key]
+
+    def __len__(self):
+        return len(self.retriever_classes)
+
+    def create_retriever(self, retriever_name, *args, **kwargs) -> Retriever:
+        if retriever_name not in self.retriever_classes:
+            raise ValueError(f"Unknown retriever type: {retriever_name}")
+        else:
+            return self.retriever_classes[retriever_name](*args, **kwargs)
+
+
+class autoregister:
+    def __init__(self, retriever_name, *args, **kwds):
+        self.retriever_name = retriever_name
+
+    def __call__(self, cls, *args, **kwds):
+        if RetrieverFactory.get_retriever_factory().register_retriever(
+            self.retriever_name, cls
+        ):
+            cls.retriever_name = self.retriever_name
+            return cls
+        else:
+            raise KeyError()
+
+
+@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 retrieve_paper(self, bg):
+        entities = []
+        sn_paper_id_list = self.cosine_similarity_search(
+            context=bg,
+            k=self.config.RETRIEVE.sn_retrieve_paper_num,
+        )
+        related_paper = set()
+        related_paper.update(sn_paper_id_list)
+        cocite_id_set = set()
+        if self.use_cocite:
+            for paper_id in related_paper:
+                cocite_id_set.update(
+                    self.cocite.get_cocite_ids(
+                        paper_id, k=self.config.RETRIEVE.cocite_top_k
+                    )
+                )
+            related_paper = related_paper.union(cocite_id_set)
+        related_paper = list(related_paper)
+        logger.debug(f"paper num before filter: {len(related_paper)}")
+        result = {
+            "paper": related_paper,
+            "entities": entities,
+            "cocite_paper": list(cocite_id_set),
+        }
+        return result
+
+    def retrieve(self, bg, entities, need_evaluate=True, target_paper_id_list=[]):
+        """
+        Args:
+            context: string
+        Return:
+            list(dict)
+        """
+        if need_evaluate:
+            if target_paper_id_list is None or len(target_paper_id_list) == 0:
+                logger.error(
+                    "If you need evaluate retriever, please input target paper is list..."
+                )
+            else:
+                target_paper_id_list = list(set(target_paper_id_list))
+        retrieve_result = self.retrieve_paper(bg)
+        related_paper_id_list = retrieve_result["paper"]
+        retrieve_paper_num = len(related_paper_id_list)
+        _, _, score_all_dict = self.cal_related_score(
+            bg,
+            related_paper_id_list=related_paper_id_list,
+            entities=entities
+        )
+        top_k_matrix = {}
+        recall = 0
+        precision = 0
+        filtered_recall = 0
+        filtered_precision = 0
+        if need_evaluate:
+            top_k_matrix, label_num, recall, precision = self.eval_related_paper_in_all(
+                score_all_dict, target_paper_id_list
+            )
+            logger.debug("Top P matrix:{}".format(top_k_matrix))
+            logger.debug("before filter:")
+            logger.debug(f"Recall: {recall:.3f}")
+            logger.debug(f"Precision: {precision:.3f}")
+        related_paper = self.filter_related_paper(score_all_dict, top_k=10)
+        related_paper = self.update_related_paper(related_paper)
+        result = {
+            "recall": recall,
+            "precision": precision,
+            "filtered_recall": filtered_recall,
+            "filtered_precision": filtered_precision,
+            "related_paper": related_paper,
+            "related_paper_id_list": related_paper_id_list,
+            "cocite_paper_id_list": retrieve_result["cocite_paper"],
+            "entities": retrieve_result["entities"],
+            "top_k_matrix": top_k_matrix,
+            "gt_reference_num": len(target_paper_id_list),
+            "retrieve_paper_num": retrieve_paper_num,
+            "label_num": label_num,
+        }
+        return result
+
+
+@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 retrieve_paper(self, entities):
+        new_entities = self.retrieve_entities_by_enties(entities)
+        logger.debug("KG entities for retriever: {}".format(new_entities))
+        related_paper = set()
+        for entity in new_entities:
+            paper_id_set = set(self.paper_client.find_paper_by_entity(entity))
+            related_paper = related_paper.union(paper_id_set)
+        cocite_id_set = set()
+        if self.use_cocite:
+            for paper_id in related_paper:
+                cocite_id_set.update(self.cocite.get_cocite_ids(paper_id))
+            related_paper = related_paper.union(cocite_id_set)
+        related_paper = list(related_paper)
+        logger.debug(f"paper num before filter: {len(related_paper)}")
+        result = {
+            "paper": related_paper,
+            "entities": entities,
+            "cocite_paper": list(cocite_id_set),
+        }
+        return result
+
+    def retrieve(self, bg, entities, need_evaluate=True, target_paper_id_list=[]):
+        """
+        Args:
+            context: string
+        Return:
+            list(dict)
+        """
+        if need_evaluate:
+            if target_paper_id_list is None or len(target_paper_id_list) == 0:
+                logger.error(
+                    "If you need evaluate retriever, please input target paper is list..."
+                )
+            else:
+                target_paper_id_list = list(set(target_paper_id_list))
+                logger.debug(f"target paper id list: {target_paper_id_list}")
+        retrieve_result = self.retrieve_paper(entities)
+        related_paper_id_list = retrieve_result["paper"]
+        retrieve_paper_num = len(related_paper_id_list)
+        _, _, score_all_dict = self.cal_related_score(
+            bg, related_paper_id_list=related_paper_id_list, entities=entities
+        )
+        top_k_matrix = {}
+        recall = 0
+        precision = 0
+        filtered_recall = 0
+        filtered_precision = 0
+        if need_evaluate:
+            top_k_matrix, label_num, recall, precision = self.eval_related_paper_in_all(
+                score_all_dict, target_paper_id_list
+            )
+            logger.debug("Top P ACC:{}".format(top_k_matrix))
+            logger.debug("before filter:")
+            logger.debug(f"Recall: {recall:.3f}")
+            logger.debug(f"Precision: {precision:.3f}")
+        related_paper = self.filter_related_paper(score_all_dict, top_k=10)
+        related_paper = self.update_related_paper(related_paper)
+        result = {
+            "recall": recall,
+            "precision": precision,
+            "filtered_recall": filtered_recall,
+            "filtered_precision": filtered_precision,
+            "related_paper": related_paper,
+            "related_paper_id_list": related_paper_id_list,
+            "cocite_paper_id_list": retrieve_result["cocite_paper"],
+            "entities": retrieve_result["entities"],
+            "top_k_matrix": top_k_matrix,
+            "gt_reference_num": len(target_paper_id_list),
+            "retrieve_paper_num": retrieve_paper_num,
+            "label_num": label_num,
+        }
+        return result
+
+
+@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 retrieve_paper(self, bg, entities):
+        sn_entities = []
+        sn_paper_id_list = self.cosine_similarity_search(
+            context=bg, k=self.config.RETRIEVE.sn_num_for_entity
+        )
+        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)
+        logger.debug("SN entities for retriever: {}".format(sn_entities))
+        entities = list(set(entities + sn_entities))
+        new_entities = self.retrieve_entities_by_enties(entities)
+        logger.debug("SNKG entities for retriever: {}".format(new_entities))
+        for entity in new_entities:
+            paper_id_set = set(self.paper_client.find_paper_by_entity(entity))
+            related_paper = related_paper.union(paper_id_set)
+        cocite_id_set = set()
+        if self.use_cocite:
+            for paper_id in related_paper:
+                cocite_id_set.update(self.cocite.get_cocite_ids(paper_id))
+            related_paper = related_paper.union(cocite_id_set)
+        related_paper = list(related_paper)
+        result = {
+            "paper": related_paper,
+            "entities": entities,
+            "cocite_paper": list(cocite_id_set),
+        }
+        return result
+
+    def retrieve(
+        self, bg, entities, need_evaluate=True, target_paper_id_list=[], top_k=10
+    ):
+        """
+        Args:
+            context: string
+        Return:
+            list(dict)
+        """
+        if need_evaluate:
+            if target_paper_id_list is None or len(target_paper_id_list) == 0:
+                logger.error(
+                    "If you need evaluate retriever, please input target paper is list..."
+                )
+            else:
+                target_paper_id_list = list(set(target_paper_id_list))
+                logger.debug(f"target paper id list: {target_paper_id_list}")
+        retrieve_result = self.retrieve_paper(bg, entities)
+        related_paper_id_list = retrieve_result["paper"]
+        retrieve_paper_num = len(related_paper_id_list)
+        _, _, score_all_dict = self.cal_related_score(
+            bg, related_paper_id_list=related_paper_id_list, entities=entities
+        )
+        top_k_matrix = {}
+        recall = 0
+        precision = 0
+        filtered_recall = 0
+        filtered_precision = 0
+        label_num = 0
+        if need_evaluate:
+            top_k_matrix, label_num, recall, precision = self.eval_related_paper_in_all(
+                score_all_dict, target_paper_id_list
+            )
+            logger.debug("Top K matrix:{}".format(top_k_matrix))
+            logger.debug("before filter:")
+            logger.debug(f"Recall: {recall:.3f}")
+            logger.debug(f"Precision: {precision:.3f}")
+        related_paper = self.filter_related_paper(score_all_dict, top_k)
+        related_paper = self.update_related_paper(related_paper)
+        result = {
+            "recall": recall,
+            "precision": precision,
+            "filtered_recall": filtered_recall,
+            "filtered_precision": filtered_precision,
+            "related_paper": related_paper,
+            "cocite_paper_id_list": retrieve_result["cocite_paper"],
+            "related_paper_id_list": related_paper_id_list,
+            "entities": retrieve_result["entities"],
+            "top_k_matrix": top_k_matrix,
+            "gt_reference_num": (
+                len(target_paper_id_list) if target_paper_id_list is not None else 0
+            ),
+            "retrieve_paper_num": retrieve_paper_num,
+            "label_num": label_num,
+        }
+        return result