update idea generate by bs

.gitignore

@@ -13,7 +13,9 @@
 datasets
 **/*.log
 assets/data/scipip_neo4j_clean_backup.json
-assets/paper/
+src/utils/scipdf/pdf/pdffigures2/
+assets/paper
+assets/model
 tmp
 test/
 # configs

src/generator.py

@@ -9,7 +9,7 @@ from loguru import logger
 import warnings
 import time
 import os
-
+from utils.hash import check_env, check_embedding
 warnings.filterwarnings("ignore")
 
 
@@ -198,10 +198,6 @@ class IdeaGenerator:
             "initial_idea": idea,
             "filtered_idea": idea_filtered,
         }
-        print("=====")
-        print(idea_modified)
-        print("=====")
-        exit()
         return message_input, idea_modified, median
 
     def generate_by_inspiration(
@@ -252,7 +248,7 @@ class IdeaGenerator:
                     self.generate_with_cue_words_ins_bs(background)
                 )
             else:
-                logger.info("{} using brainstorm_{} with cue words.".format(mode_name, bs_mode))
+                logger.info("{} using brainstorm_{} without cue words.".format(mode_name, bs_mode))
                 (
                     message_input, 
                     problem, 
@@ -308,14 +304,14 @@ def main(ctx):
 )
 @click.option(
     "--brainstorm-mode",
-    default="mode_a",
+    default="mode_c",
     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,
+    default=False,
     type=bool,
     required=True,
     help="Use cue words in generation",
@@ -327,27 +323,6 @@ def main(ctx):
     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,
@@ -356,6 +331,8 @@ def main(ctx):
     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):
+    check_env()
+    check_embedding() 
     # Configuration
     config = ConfigReader.load(config_path, **kwargs)
     logger.add(
@@ -518,7 +495,7 @@ def backtracking(config_path, ids_path, retriever_name, brainstorm_mode, use_cue
 )
 @click.option(
     "--brainstorm-mode",
-    default="mode_a",
+    default="mode_c",
     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)",
@@ -530,27 +507,6 @@ def backtracking(config_path, ids_path, retriever_name, brainstorm_mode, use_cue
     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,
@@ -559,6 +515,8 @@ def backtracking(config_path, ids_path, retriever_name, brainstorm_mode, use_cue
     help="The number of data you want to process",
 )
 def new_idea(config_path, ids_path, retriever_name, brainstorm_mode, use_inspiration, num, **kwargs):
+    check_env()
+    check_embedding()
     logger.add(
         "log/generate_{}_{}.log".format(time.time(), retriever_name), level="DEBUG"
     )  # 添加文件输出

src/pages/button_interface.py

@@ -73,12 +73,13 @@ class Backend(object):
                     res += "\n"
             return res, result["related_paper"]
 
-    def literature2initial_ideas_callback(self, background, retrieved_literature, json_strs=None):
+    def literature2initial_ideas_callback(self, background, brainstorms, retrieved_literature, json_strs=None):
         if json_strs is not None:
             json_contents = json.loads(json_strs)
             return json_contents["median"]["filtered_idea"]
         else:
             self.idea_generator.paper_list = retrieved_literature
+            self.idea_generator.brainstorm = brainstorms
             if self.use_inspiration:
                 message_input, idea_modified, median = (
                 self.idea_generator.generate_by_inspiration(

src/pages/one_click_generation.py

@@ -119,7 +119,7 @@ def generate_ideas(backend, background):
 
     with st.spinner(text="Generating initial ideas..."):
         msg = "My initial ideas are:"
-        initial_ideas, final_ideas = backend.literature2initial_ideas_callback(background, related_works_intact)
+        initial_ideas, final_ideas = backend.literature2initial_ideas_callback(background, brainstorms, 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})

src/pages/step_by_step_generation.py

@@ -132,7 +132,7 @@ def genrate_mainpage(backend):
         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"])
+                res = backend.literature2initial_ideas_callback(background, brainstorms, 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"

src/retriever.py

@@ -8,6 +8,7 @@ import click
 import json
 from loguru import logger
 import warnings
+from utils.hash import check_env, check_embedding
 
 warnings.filterwarnings("ignore")
 
@@ -55,30 +56,11 @@ def main(ctx):
     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
 ):
+    check_env()
+    check_embedding()
     config = ConfigReader.load(config_path, **kwargs)
     log_dir = config.DEFAULT.log_dir
     if not os.path.exists(log_dir):

src/utils/hash.py

@@ -5,33 +5,42 @@ import struct
 from collections import Counter
 from huggingface_hub import hf_hub_download
 
+ENV_CHECKED = False
+EMBEDDING_CHECKED = False
+
 def check_embedding():
-    # Define the repository and files to download
-    repo_id = "sentence-transformers/all-MiniLM-L6-v2"  # "BAAI/bge-small-en-v1.5"
-    files_to_download = [
-        "config.json",
-        "pytorch_model.bin",
-        "tokenizer_config.json",
-        "vocab.txt",
-    ]
-    # Download each file and save it to the /model/bge directory
-    for file_name in files_to_download:
-        print("Checking for file: ", file_name)
-        hf_hub_download(repo_id=repo_id, filename=file_name, local_dir=f"./assets/model/{repo_id}")
+    global EMBEDDING_CHECKED
+    if not EMBEDDING_CHECKED:
+        # Define the repository and files to download
+        repo_id = "sentence-transformers/all-MiniLM-L6-v2"  # "BAAI/bge-small-en-v1.5"
+        files_to_download = [
+            "config.json",
+            "pytorch_model.bin",
+            "tokenizer_config.json",
+            "vocab.txt",
+        ]
+        # Download each file and save it to the /model/bge directory
+        for file_name in files_to_download:
+            print("Checking for file: ", file_name)
+            hf_hub_download(repo_id=repo_id, filename=file_name, local_dir=f"./assets/model/{repo_id}")
+        EMBEDDING_CHECKED = True
 
 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...")
+    global ENV_CHECKED
+    if not ENV_CHECKED:
+        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...")
+        ENV_CHECKED = True
 
 
 def generate_hash_id(input_string):

src/utils/llms_api.py

@@ -49,6 +49,14 @@ class APIHelper(object):
         pass
 
     def __call__(self, title: str, abstract: str, introduction: str) -> dict:
+        if os.environ["MODEL_NAME"] not in [
+            "glm4",
+            "glm4-air",
+            "qwen-max",
+            "qwen-plus",
+        ]:
+            raise ValueError(f"Check model name...")
+
         if title is None or abstract is None or introduction is None:
             return None
         try:
@@ -102,8 +110,8 @@ class APIHelper(object):
         few_shot_examples = []
         for example in common_examples:
             few_shot_examples.append(example)
-        
-        prompt_template_entity = '''
+
+        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.
 
@@ -130,7 +138,7 @@ class APIHelper(object):
         ### Format for Your Response:
         Please just give me the entities and spilt them by ",":
         <entity 1>,<entity2>,...
-        '''
+        """
 
         if abstract is None:
             return None
@@ -192,6 +200,7 @@ class APIHelper(object):
         """
 
         if background is None:
+            print("Input background is empty ...")
             return None
         try:
             # Initial brainstorming to generate raw ideas
@@ -202,7 +211,6 @@ class APIHelper(object):
             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,
@@ -665,7 +673,7 @@ class APIHelper(object):
             traceback.print_exc()
             return None
         return idea
-    
+
     def integrate_idea(self, background: str, brainstorm: str, idea: str) -> str:
         prompt_template_idea = """
         Task Description: 
@@ -700,7 +708,8 @@ class APIHelper(object):
             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)
+                background=background, brainstorm=brainstorm, idea=idea
+            )
             message.append({"role": "user", "content": message_input})
             response = self.generator.create(
                 messages=message,
@@ -851,7 +860,7 @@ class APIHelper(object):
         except Exception:
             traceback.print_exc()
         return ground_truth
-    
+
     def transfer_form(self, idea: str):
         prompt_template_transfer = """
         ### Task Description:
@@ -973,7 +982,13 @@ class APIHelper(object):
         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
+        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.
@@ -1018,13 +1033,13 @@ class APIHelper(object):
         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(
+        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(
+        system_message = novelty_system_msg.format(
             num_rounds=max_num_iterations,
         )
         if msg_history is None:
@@ -1043,14 +1058,18 @@ class APIHelper(object):
                 seed=0,
             )
             content = response.choices[0].message.content
-            new_msg_history = new_msg_history + [{"role": "assistant", "content": 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:
+    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.
@@ -1104,13 +1123,20 @@ class APIHelper(object):
 
         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:
+        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)
+                idea1=idea1, idea2=idea2, idea3=idea3, idea4=idea4, idea5=idea5
+            )
             message.append({"role": "user", "content": message_input})
             response = self.generator.create(
                 messages=message,
@@ -1121,7 +1147,7 @@ class APIHelper(object):
             return None
         return result
 
-    def compare_all(self, idea1:str, idea2:str) -> str:
+    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.
@@ -1172,8 +1198,7 @@ class APIHelper(object):
         try:
             message = []
             message.append({"role": "system", "content": system_input})
-            message_input = input_message.format(
-                idea1=idea1, idea2=idea2)
+            message_input = input_message.format(idea1=idea1, idea2=idea2)
             message.append({"role": "user", "content": message_input})
             response = self.generator.create(
                 messages=message,
@@ -1184,7 +1209,7 @@ class APIHelper(object):
             return None
         return result
 
-    def compare_novelty_and_feasibility(self, idea1:str, idea2:str) -> str:
+    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.
@@ -1229,8 +1254,7 @@ class APIHelper(object):
         try:
             message = []
             message.append({"role": "system", "content": system_input})
-            message_input = input_message.format(
-                idea1=idea1, idea2=idea2)
+            message_input = input_message.format(idea1=idea1, idea2=idea2)
             message.append({"role": "user", "content": message_input})
             response = self.generator.create(
                 messages=message,
@@ -1240,8 +1264,8 @@ class APIHelper(object):
             traceback.print_exc()
             return None
         return result
-    
-    def compare_novelty(self, idea1:str, idea2:str) -> str:
+
+    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.
@@ -1284,8 +1308,7 @@ class APIHelper(object):
         try:
             message = []
             message.append({"role": "system", "content": system_input})
-            message_input = input_message.format(
-                idea1=idea1, idea2=idea2)
+            message_input = input_message.format(idea1=idea1, idea2=idea2)
             message.append({"role": "user", "content": message_input})
             response = self.generator.create(
                 messages=message,
@@ -1296,7 +1319,7 @@ class APIHelper(object):
             return None
         return result
 
-    def compare_feasibility(self, idea1:str, idea2:str) -> str:
+    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.
@@ -1339,8 +1362,7 @@ class APIHelper(object):
         try:
             message = []
             message.append({"role": "system", "content": system_input})
-            message_input = input_message.format(
-                idea1=idea1, idea2=idea2)
+            message_input = input_message.format(idea1=idea1, idea2=idea2)
             message.append({"role": "user", "content": message_input})
             response = self.generator.create(
                 messages=message,