update new version

README.md

@@ -36,15 +36,16 @@ SciPIP is a scientific paper idea generation tool powered by a large language mo
 ![SciPIP](./assets/pic/demo.png)
 
 
-🤗 Try it on the Hugging Face (Coming Soon... You can deploy it at your own computer now.)
+🤗 Try it on the Hugging Face: https://huggingface.co/spaces/lihuigu/SciPIP (The demo uses the old version code temporally and will be updated soon.)
 
 ## Updates
 
 - [x] Idea generation in a GUI enviroment (web app).
-- [x] Idea generation for the NLP and multimodal (partial) field.
-- [ ] Idea generation for the CV field.
+- [x] Idea generation for the NLP and multimodal field.
+- [x] Idea generation for the CV field.
 - [ ] Idea generation for other fields.
-- [ ] Release the Huggingface demo.
+- [x] Release the Huggingface demo.
+- [x] Support DeepSeek-v3 as an backend, now. 🎉 🎉 🎉 
 
 ## Prerequisites
 
@@ -64,11 +65,11 @@ The following enviroments are tested under Ubuntu 22.04 with python>=3.10.3.
    ## Install Neo4j database
    sudo apt install -y openjdk-17-jre # Install Neo4j required JDK
    # cd ~/Downloads # or /your/path/to/download/Neo4j
-   wget http://dist.neo4j.org/neo4j-community-5.20.0-unix.tar.gz
-   tar -xvf neo4j-community-5.20.0-unix.tar.gz
+   wget http://dist.neo4j.org/neo4j-community-5.25.1-unix.tar.gz
+   tar -xvf neo4j-community-5.25.1-unix.tar.gz
    
    ## Start Neo4j
-   cd ./neo4j-community-5.20.0
+   cd ./neo4j-community-5.25.1
    # Uncomment server.default_listen_address=0.0.0.0 in conf/neo4j.conf to visit Neo4j through a browser
    sed -i 's/# server.default_listen_address=0.0.0.0/server.default_listen_address=0.0.0.0/g' ./conf/neo4j.conf
    ./bin/neo4j start
@@ -93,17 +94,14 @@ The following enviroments are tested under Ubuntu 22.04 with python>=3.10.3.
    ```
 3. **Prepare the literature database**
    
-   1. Download the literature data from [this link](https://drive.google.com/file/d/1NZTDpxKo7bmxwXPI03dgikEemKGLkwne/view?usp=sharing) and save it to `assets/data/scipip_neo4j_clean_backup.json`.
-   2. Then, run the following command to load the literature into Neo4j database (It may 40-60 minutes):
-   ```
-   python src/utils/paper_client.py
-   ```
+   1. Download the literature data from [google_drive](https://drive.google.com/file/d/1kZmJff8am-JGegZZQx0qxlC7o7YgBURg/view?usp=sharing) or [baidu disk](https://pan.baidu.com/s/1S22Evi5ReL0MvahFoQ-ipA?pwd=scip). Replace the `/your/path/neo4j-community-5.25.1/data` folder with our provided `data` folder, which contains literature of CV, NLP, ML, *etc.*
+   2. [Optional] Prepare the embedding model. Our algorithm uses **jina-embedding v3** and will automatically download it from Huggingface the first time the program is run. However, if you're concerned about potential download failures due to network issues, you can download it in advance and place it in the specified directory.
    
-4. **[Optional] Prepare the embedding model**. Our algorithm uses SentenceBERT and **will automatically download** it from Huggingface the first time the program is run. However, if you're concerned about potential download failures due to network issues, you can download it in advance and place it in the specified directory.
    ```bash
-   cd /root/path/of/SciPIP && mkdir -p assets/model/sentence-transformers
-   git clone https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2 assets/model/sentence-transformers/all-MiniLM-L6-v2 assets/model/sentence-transformers
+   cd /root/path/of/SciPIP && mkdir -p assets/model/
+   git clone https://huggingface.co/jinaai/jina-embeddings-v3 assets/model
    ```
+   
 
 ## Run In a Browser (Recommended)
 
@@ -116,36 +114,20 @@ Then, visit `http://localhost:8501` in your browser with an interactive envirome
 
 ## Run In a Terminal
 
-**1. BackTracking of ACL 2024**
-
-```
-python src/generator.py backtracking --brainstorm-mode mode_c --use-cue-words True --use-inspiration True --num 1
-```
-
-Results dump in `assets/output_idea/output_backtracking_mode_c_cue_True_ins_True.json`.
+**1. Generate new idea**
 
-**2. Generate new idea**
-
-Input your backgound and cue words in `assets/data/test_background.json`
+Input your backgound in `assets/data/test_background.json`
 
 ```
 python src/generator.py new-idea --brainstorm-mode mode_c --use-inspiration True --num 2
 ```
 
-Results dump in `assets/output_idea/output_new_idea_mode_c_ins_True.json`.
+Results dump in `assets/output_idea/output-file.json`.
 
 ## Others
 
-### Retrieve Eval
-
-Generate retrieve eval log result in `./log`.
-
-```
-bash scripts/retriever_eval.sh
-```
-
 ### Database Construction
-SciPIP uses Neo4j as its database. You can directly import the provided data or add your own research papers.
+SciPIP uses Neo4j as its database. You can directly import the provided data or add your own research papers.ddddddsfasdfldsafkldsjfdkls
 ```
 wget https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.7.1/en_core_web_sm-3.7.1-py3-none-any.whl
 pip install en_core_web_sm-3.7.1-py3-none-any.whl

app.py

@@ -22,11 +22,9 @@ if __name__ == "__main__":
     def fn2():
         step_by_step_generation.step_by_step_generation(backend)
 
-    pg = st.navigation(
-        [
-            st.Page(homepage.home_page, title=_("🏠️ Homepage")),
-            st.Page(fn1, title=_("💧 One-click Generation")),
-            st.Page(fn2, title=_("💦 Step-by-step Generation")),
-        ]
-    )
-    pg.run()
+    pg = st.navigation([
+        st.Page(homepage.home_page, title=_("🏠️ Homepage")),
+        st.Page(fn1, title=_("💧 One-click Generation")),
+        st.Page(fn2, title=_("💦 Step-by-step Generation")),
+    ])
+    pg.run()

assets/data/test_background.json

@@ -1,2 +1,6 @@
-{"background": "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."}
-{"background": "Multimodal learning is committed to integrating multiple information sources such as text, images, audio, and video to create more powerful and universal AI models. The research on unified representation aims to find representation methods that can generalize across modalities. Challenge: 1 Modal alignment: There is heterogeneity between different modalities, and how to achieve semantic alignment of these modalities to ensure that the model can comprehensively understand different types of data is a core challenge. 2. Data sparsity and imbalance: There is usually an imbalance in the amount of data in different modalities, such as video and audio data being relatively scarce, while text data is relatively abundant. How to effectively utilize and fuse these modalities to avoid overfitting or underfitting remains a research difficulty.", "cue_words": ["cross-modal embedding", "dat augmentation", "modality-aware fusion", "heterogeneous data integration"]}
+[{"background": "1. The limitations of existing methods in leveraging nonverbal information for discerning complex semantics in unsupervised scenarios. \n2. The recognition that non-verbal modalities (video and audio) play a critical role in performing unsupervised clustering and can provide useful cues for semantics discovery."},
+{"background": "1. The need to reduce the memory footprint and computational requirements of Large Language Models (LLMs) for deployment on resource-constrained devices. \n2. The challenge of preserving model performance at sub-4-bit quantization levels, where existing methods significantly degrade the fidelity of model weights."},
+{"background": "1. The need to trace back and understand which training data influences specific generations in large language models for tasks like risk assessment, model retraining, and improving explainability. \n2. The challenge of scaling existing influence estimation methods to large language models due to their massive size and the computational expense of processing and storing gradients."},
+{"background": "1. The need for a unifying framework to integrate the multiple dimensions of lexical semantic change detected by historical linguists. \n2. The desire to align theoretical insights from linguistics with the methodological sophistication of natural language processing to better understand social and cultural change."},
+{"background": "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."},
+{"background": "Multimodal learning is committed to integrating multiple information sources such as text, images, audio, and video to create more powerful and universal AI models. The research on unified representation aims to find representation methods that can generalize across modalities. Challenge: 1 Modal alignment: There is heterogeneity between different modalities, and how to achieve semantic alignment of these modalities to ensure that the model can comprehensively understand different types of data is a core challenge. 2. Data sparsity and imbalance: There is usually an imbalance in the amount of data in different modalities, such as video and audio data being relatively scarce, while text data is relatively abundant. How to effectively utilize and fuse these modalities to avoid overfitting or underfitting remains a research difficulty.", "cue_words": ["cross-modal embedding", "dat augmentation", "modality-aware fusion", "heterogeneous data integration"]}]

assets/prompt/expand_background.xml

@@ -0,0 +1,44 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!DOCTYPE body [
+  <!ENTITY warning "Warning: Something bad happened... please refresh and try again.">
+]>
+<body>
+<query rank="0">
+<title>System Message</title>
+<text>
+You are a teacher in the field of AI, skilled at clearly explaining AI concepts to students. Your student is an undergraduate in AI with a basic understanding of deep learning.
+</text>
+</query>
+<query rank="1">
+<title>User Message</title>
+<text>
+# Task Description:
+You are teaching your undergraduate about a specific subfield of AI research. You have a brief description of the research background, and now you need to explain its meaning and purpose in detail to your undergraduate. Keep in mind that your undergraduate may be completely unfamiliar with the technical terms in the research background. I will give you an example. The example begins with "# Example 1" and includes a Brief Research Background, several Technical Terms, and the corresponding Detailed Research Background. Then, your task starts with "# Your Task", containing "Your Brief Research Background" and "Your Technical Terms". Your job is to expand Your Brief Research Background into a Detailed Research Background by referring to Example 1. Note that the research background in Example 1 are unrelated to yours, so the key focus should be on the relationship between the Brief Research Background and the Detailed Research Background. You should directly start with your response and do not start with a section title like "## Detailed Background". 
+
+# Example 1
+
+## Brief Research Background
+
+During the inference process of large language models, the KV cache grows with the context and the length of the generated content, occupying an increasing amount of GPU memory. How can we minimize the memory usage of the KV cache as much as possible to extend the text length that large language models can handle without increasing the cache size?
+
+## Technical Terms
+
+large language models, kv cache, gpu memory
+
+## Detailed Research Background
+
+Large language models use a Transformer-based architecture and generate text autoregressively, outputting one token at a time during inference. Within the Transformer, there is a self-attention module where each token is associated with three vectors: Q (query), K (key), and V (value). For each newly generated token, its Q needs to be computed with the K and V of all previously generated tokens. To avoid recalculating K and V repeatedly, the K and V of all tokens are stored in GPU memory, which is referred to as KV-cache. During the inference process of large language models, the KV cache grows with the context and the length of the generated content, occupying an increasing amount of GPU memory. How can we minimize the memory usage of the KV cache as much as possible to extend the text length that large language models can handle without increasing the cache size?
+
+# Your Task
+
+## Your Research Background
+
+{brief_background}
+
+## Your Technical Terms
+
+{keywords}
+
+</text>
+</query>
+</body>

assets/prompt/expand_idea.xml

@@ -0,0 +1,75 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!DOCTYPE body [
+  <!ENTITY warning "Warning: Something bad happened... please refresh and try again.">
+]>
+<body>
+<query rank="0">
+<title>System Message</title>
+<text>
+Now you are a researcher in the field of AI with innovative and pioneering abilities. You are good at transforming a brief scientific idea into a concrete algorithm.
+</text>
+</query>
+<query rank="1">
+<title>User Message</title>
+<text>
+# Task Description:
+You are an AI researcher conducting studies in a specific domain. Someone has provided you with a brief scientific idea, and your task is to transform it into a detailed, feasible, and concrete algorithm. If necessary, you may incorporate formulas to elaborate on the algorithm in the Latex format. I will give you an example. The example begins with "# Example 1" and includes a Brief Scientific Idea and its corresponding Detailed Scientific Idea. Then, your task starts with "# Your Task", containing "Your Brief Scientific Idea". Your job is to expand Your Brief Scientific Idea into a Detailed Scientific Idea by referring to Example 1. Note that the ideas in Example 1 are unrelated to your idea, so the key focus should be on the relationship between the Brief Scientific Idea and the Detailed Scientific Idea. You should directly start with your response and do not start with a section title like "## Detailed Scientific Idea". 
+
+# Example 1
+
+## Example Brief Scientific Idea
+
+1. The use of a dual-stage process that first quantizes the LLM's model weights into 4-bit and then introduces a side network that leverages downsampled outputs and hidden states from the quantized LLM to make task-specific predictions.
+2. The innovative application of several low-rank adapter methods, such as MaxPooling and AvgPooling, within the side network to perform downsampling and significantly reduce the number of trainable parameters and the memory footprint of optimizer states.
+3. The aggregation of hidden states from the quantized LLM and the side network using a learnable parameter, which allows for efficient parallel computation of the LLM and side network without increasing inference latency.
+
+## Example Detailed Scientific Idea
+
+Building on the concept of Parameter-Efficient Fine-Tuning (PEFT), which aims to adapt Large Language Models (LLMs) to specific tasks without the full computational cost of training all parameters, the following algorithm integrates a memory-efficient strategy enhanced through quantization and an auxiliary side network. This allows for efficient fine-tuning and inference on large models with reduced memory requirements and computational demands.
+
+1. **Quantization of LLM to 4-bit Precision:**
+   1. Utilize 4-bit quantization to reduce the memory footprint of the LLM's weights. Each floating-point parameter in the LLM is converted to a 4-bit representation for efficient memory utilization.
+   2. Begin with the conversion of 16-bit parameters to 4-bit using the relation:
+      $$
+      X_{{4bit}} = \text{{round}}\left(\frac{{M_{{4bit}}}}{{\text{{Absmax}}(X_{{16bit}})}} \cdot X_{{16bit}}\right)
+      $$
+      where $M_{{4bit}}$ is the maximum value representable in 4 bits, ensuring quantization minimizes precision loss by managing outliers through block-wise separate quantization.
+
+2. **Side Network for Memory-Efficient Tuning:**
+   1. Implement a side network $g$ with dimensions reduced by a factor $r$ relative to the original model $f$. The side network processes information more economically, storing less data and reducing computational load during training.
+   2. Define the hidden state transformation at the $i$-th layer as:
+      $$
+      h_{{gi}}^{{16bit}} = (1 - \beta_i) * \text{{downsample}}_i(h_{{fi}}^{{16bit}}) + \beta_i * h_{{gi-1}}^{{16bit}}
+      $$
+      where $\beta_i$ is a learnable gating parameter and $\text{{downsample}}_i$ reduces dimensionality.
+
+3. **Low-Memory Gradient Calculation:**
+   1. Perform backpropagation limited to the side network $g$, excluding the calculation of gradients for the quantized weights in $f$, leveraging the pre-trained knowledge while focusing computational resources on the task-specific adaptation.
+   2. The gradient computation is detached from the main LLM, avoiding the costly backpropagation through large transformer layers and focusing updates through efficient gradient paths within $g$.
+
+4. **Combining Outputs for Inference:**
+   1. At inference, blend the outputs from the LLM and side network as a weighted sum:
+      $$
+      h_N^{{16bit}} = \alpha h_{{fN}}^{{16bit}} + (1-\alpha) h_{{gN}}^{{16bit}}
+      $$
+      where $\alpha$ is a learnable parameter initialized to prioritize the pre-trained model influence gradually allowing task customization through tuning.
+
+5. **Optimized Training Procedure:**
+   1. Integrate efficient downsampling techniques, such as LoRA and Adapter models, reducing parameter size significantly without losing efficacy.
+   2. Maintain a 16-bit floating-point data type for computations in forward and backward passes to balance precision and performance, ensuring that the quantized network remains robust and generalizable.
+
+By synthesizing quantization with an auxiliary network, the algorithm achieves a robust parameter-efficient fine-tuning technique, significantly reducing memory overhead, improving inference speed, and maintaining high performance despite the minimal parameters being updated. This approach effectively supports large-scale models, facilitating application in environments with constrained computational resources.
+
+# Your Task
+
+## Your Research Background
+
+{background}
+
+## Your Brief Scientific Idea
+
+{brief_idea}
+
+</text>
+</query>
+</body>

assets/prompt/generate_brainstorm.xml

@@ -13,21 +13,21 @@ Now you are a researcher in the field of AI with innovative and pioneering abili
 <title>User Message</title>
 <text>
 ### Task Description:
-trunkYou 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.
+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. The ideas should cover a range of possible directions that could be explored further.
 
-trunk### Information Provided:
-trunk- **Research Background**: {background}
+### Information Provided:
+- **Research Background**: {background}
 
-trunk### Approach:
-trunkYour brainstorming should be systematic:
-trunk- **Step 1**: Thoroughly understand the research background.
-trunk- **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.
+### Approach:
+Your brainstorming should be systematic:
+- **Step 1**: Thoroughly understand the research background.
+- **Step 2**: Generate a list of 3 to 4 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.
 
-trunk### Format for Your Response:
-trunkPlease present 4 to 6 ideas in the following format:
-trunk**Idea 1**: [Brief description of the first idea]
-trunk**Idea 2**: [Brief description of the second idea]
-trunk...
+### Format for Your Response:
+Please present 3 to 4 ideas in the following format:
+**Idea 1**: [Brief description of the first idea]
+**Idea 2**: [Brief description of the second idea]
+...
 </text>
 </query>
 </body>

assets/prompt/generate_concise_method.xml

@@ -0,0 +1,67 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!DOCTYPE body [
+  <!ENTITY warning "Warning: Something bad happened... please refresh and try again.">
+]>
+<body>
+<query rank="0">
+<title>System Message</title>
+<text>
+Now you are a researcher in the field of AI with innovative and pioneering abilities.
+</text>
+</query>
+<query rank="1">
+<title>User Message</title>
+<text>
+# Task Description:
+You are an AI researcher conducting studies in a specific domain. Someone has provided you with a methodology section, and your task is to transform it into another style. I will give you an example. The example begins with "# Example 1" and includes a Example Summarized Methods. Then, your task starts with "# Your Task", containing "Your Methodology Section". Your job is to transform Your Methodology Section into a Summarized Methods by referring to Example 1. Note that the ideas in Example 1 are unrelated to your idea, so the key focus should be on the style of Example Summarized Methods. You should directly start with your response and do not start with a section title like "## Your Summarized Methods".
+
+# Example 1
+
+## Example Summarized Methods
+
+Building on the concept of Parameter-Efficient Fine-Tuning (PEFT), which aims to adapt Large Language Models (LLMs) to specific tasks without the full computational cost of training all parameters, the following algorithm integrates a memory-efficient strategy enhanced through quantization and an auxiliary side network. This allows for efficient fine-tuning and inference on large models with reduced memory requirements and computational demands.
+
+1. **Quantization of LLM to 4-bit Precision:**
+   1. Utilize 4-bit quantization to reduce the memory footprint of the LLM's weights. Each floating-point parameter in the LLM is converted to a 4-bit representation for efficient memory utilization.
+   2. Begin with the conversion of 16-bit parameters to 4-bit using the relation:
+      \[
+      X_{{4bit}} = \text{{round}}\left(\frac{{M_{{4bit}}}}{{\text{{Absmax}}(X_{{16bit}})}} \cdot X_{{16bit}}\right)
+      \]
+      where \(M_{{4bit}}\) is the maximum value representable in 4 bits, ensuring quantization minimizes precision loss by managing outliers through block-wise separate quantization.
+
+2. **Side Network for Memory-Efficient Tuning:**
+   1. Implement a side network \(g\) with dimensions reduced by a factor \(r\) relative to the original model \(f\). The side network processes information more economically, storing less data and reducing computational load during training.
+   2. Define the hidden state transformation at the \(i\)-th layer as:
+      \[
+      h_{{gi}}^{{16bit}} = (1 - \beta_i) * \text{{downsample}}_i(h_{{fi}}^{{16bit}}) + \beta_i * h_{{gi-1}}^{{16bit}}
+      \]
+      where \(\beta_i\) is a learnable gating parameter and \(\text{{downsample}}_i\) reduces dimensionality.
+
+3. **Low-Memory Gradient Calculation:**
+   1. Perform backpropagation limited to the side network \(g\), excluding the calculation of gradients for the quantized weights in \)f\), leveraging the pre-trained knowledge while focusing computational resources on the task-specific adaptation.
+   2. The gradient computation is detached from the main LLM, avoiding the costly backpropagation through large transformer layers and focusing updates through efficient gradient paths within \(g\).
+
+4. **Combining Outputs for Inference:**
+   1. At inference, blend the outputs from the LLM and side network as a weighted sum:
+      \[
+      h_N^{{16bit}} = \alpha h_{{fN}}^{{16bit}} + (1-\alpha) h_{{gN}}^{{16bit}}
+      \]
+      where \(\alpha\) is a learnable parameter initialized to prioritize the pre-trained model influence gradually allowing task customization through tuning.
+
+5. **Optimized Training Procedure:**
+   1. Integrate efficient downsampling techniques, such as LoRA and Adapter models, reducing parameter size significantly without losing efficacy.
+   2. Maintain a 16-bit floating-point data type for computations in forward and backward passes to balance precision and performance, ensuring that the quantized network remains robust and generalizable.
+
+By synthesizing quantization with an auxiliary network, the algorithm achieves a robust parameter-efficient fine-tuning technique, significantly reducing memory overhead, improving inference speed, and maintaining high performance despite the minimal parameters being updated. This approach effectively supports large-scale models, facilitating application in environments with constrained computational resources.
+
+# Your Task
+
+## Your Methodology Section
+
+{methodology}
+
+## Your Summarized Methods
+
+</text>
+</query>
+</body>

assets/prompt/generate_idea_by_inspiration.xml

@@ -12,30 +12,19 @@ Now you are a researcher in the field of AI with innovative and pioneering abili
 <query rank="1">
 <title>User Message</title>
 <text>
-### 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.
+# Task Description:
+You will be provided with a research problem, along with inspirations extracted from related papers. Your task is to identify and combine these inspirations to propose 3 to 4 different ideas to solve the research problem. The ideas should be innovative and try to avoid using the same inspirations repeatedly. Each idea should starts with "**Idea **:".
 
-### Information Provided:
-1. **Research problem &amp; 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.
+# Information Provided:
 
-### 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.
+## Research problem
 
-### Specific Information:
-I will provide you with specific information now, please use them according to the instructions above:
-1. **Research problem &amp; Rationales**: {problem}
-2. **Inspirations**: {inspirations_text}
+{background}
+
+## Inspirations
+
+{inspirations}
 
-### 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]
-...
 </text>
 </query>
 </body>

assets/prompt/generate_inspiration_with_detail_method.xml

@@ -0,0 +1,77 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!DOCTYPE body [
+  <!ENTITY warning "Warning: Something bad happened... please refresh and try again.">
+]>
+<body>
+<query rank="0">
+<title>System Message</title>
+<text>
+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.
+</text>
+</query>
+<query rank="1">
+<title>User Message</title>
+<text>
+# Task Description:
+You will be provided with a research problem, as well as some reference materials. Your task is to extract a novel, effective, and specific inspiration from the materials that can help addressing the research problem. I will give you an example. The example begins with "# Example 1" and includes a Example Problem, Example Materials, and Example Inspiration. Your task is to read Your Problem and Your Materials, and extract Inspirations by referring to Example 1. Note that the contents in Example 1 are unrelated to yours, so the key focus should be on the relationship among the Background, Materials, and Inspiration. Only output the three most significant inspirations and each inspiration should be concluded with three sentences. You should directly start with your response and do not start with a section title like "## Your Inspirations". Further, if you believe the materials do not contribute to solving the problem described in the Background, you may simply reply with "None" and provide no further response.
+
+# Example 1
+
+## Example Problem
+
+1. The need to ensure that chain-of-thought (CoT) rationales generated by large language models are consistent with their predictions and faithfully justify those decisions. 
+2. The desire to distill the reasoning capabilities of large LMs into smaller models without losing the quality and faithfulness of the rationales.
+
+## Example Materials
+
+In addressing the challenges of generating faithful Chain-Of-Thought (CoT) rationales and consistent student outputs in knowledge distillation, we introduce the Self-Consistent Chain-Of-Thought Distillation (SCOTT) method. SCOTT is designed to enhance consistency in rationale generation and counter the pitfalls of hallucination and reasoning shortcuts in language models. This approach involves training a smaller student model to produce rationales that align with its predictions, learning from a larger teacher model to achieve this. Our approach leverages contrastive decoding and counterfactual reasoning to improve the quality and faithfulness of rationales.
+
+1. **Contrastive Decoding for Teacher Model:**
+   1. Employ contrastive decoding to generate more relevant and answer-grounded rationales by the teacher model, which mitigates issues related to hallucination common in language models.
+   2. This is achieved by introducing perturbed answers and evaluating the plausibility shift of each token to ensure rationales support the intended answers more distinctly:
+      $$
+      G(t_i | a^*) = \log \frac{{P(t_i | p, q, A, a^*, t_{{&lt;i}})}}{{P(t_i | p, q, A, a', t_{{&lt;i}})}}
+      $$
+      Here, $a'$ represents a perturbed answer, used to fine-tune the rationale's specificity towards the correct answer.
+
+2. **Counterfactual Reasoning for Student Model:**
+   1. Train the student model to validate its rationale against its predictions through counterfactual reasoning, requiring the student to adjust predictions when confronted with altered context or rationale.
+   2. Implement this by incorporating variabilities in rationales that lead to different answers and ensuring the model understands and adjusts accordingly:
+      $$
+      L_{{\text{{counterfactual}}}} = - \sum \log P(t_i | q, r', t_{{&lt;i}})
+      $$
+      where $r'$ is a rationale leading to a perturbed answer, encouraging the student model to reflect such dependencies in its decision-making process.
+
+3. **Holistic Training Approach:**
+   1. Integrate the contrastive decoding outputs and counterfactual reasoning objective into the student’s training to simultaneously focus on consistency in rationale generation and alignment with the predictions.
+   2. By incorporating more on-topic rationale-answer pairs and utilizing both factual and counterfactual losses, the student model's faithfulness and performance are improved:
+      $$
+      L_{{\text{{total}}}} = L_{{\text{{factual}}}} + L_{{\text{{counterfactual}}}}
+      $$
+
+4. **Experimentation and Validation:**
+   1. Conduct experiments on open-domain QA tasks where knowledge-intensive reasoning is essential, assessing both rationale consistency and the student's alignment between rationale and prediction.
+   2. Results indicate the proposed SCOTT method leads to improved student faithfulness, maintaining competitive performance with additional advantages in rationale justification consistency compared to baseline models.
+   3. Additional ablation studies reveal robustness across various student model sizes, ensuring consistent rationale fidelity irrespective of model capacity.
+
+Through the integration of contrastive decoding and counterfactual reasoning, SCOTT offers a novel and robust approach to improve rationale consistency and model faithfulness in natural language processing tasks, enhancing interpretability and performance.
+
+## Example Inspiration
+
+1. When prompting large language models to generate rationales, the faithfulness of the rationale to the prediction can be enhanced using contrastive decoding. Specifically, for a given prediction, the model-generated rationale should differ as much as possible from the rationale generated for other predictions.
+
+2. The chain-of-thought (CoT) reasoning ability of smaller language models can be improved through chain-of-thought distillation. During distillation, for the same question, when the chain-of-thought content differs, the model's predictions should also differ.
+
+# You Task
+
+## Your Problem
+
+{background}
+
+## Your Materials
+
+{detail_method}
+
+</text>
+</query>
+</body>

configs/datasets.yaml

@@ -3,30 +3,40 @@ DEFAULT:
     ignore_paper_id_list: ./assets/data/ignore_paper_id_list.json
     log_level: "DEBUG"
     log_dir: ./log
-    embedding: sentence-transformers/all-MiniLM-L6-v2
+    # embedding: sentence-transformers/all-MiniLM-L6-v2
+    # embedding: BAAI/llm-embedder
+    embedding: jina-embeddings-v3
+    embedding_task: text-matching # ONLY FOR JINA_v3, retrieval.passage, text-matching, retrieval.query
+    embedding_database: text-matching # ONLY FOR JINA_v3, retrieval.passage, text-matching, retrieval.query
+
 
 ARTICLE:
     summarizing_prompt: ./assets/prompt/summarizing.xml
 
 RETRIEVE:
     retriever_name: "SNKG"
-    use_cocite: False
-    use_cluster_to_filter: False # 过滤器中使用聚类算法
-    cite_type: "all_cite_id_list"
+    # retriever_name: "SN"
+    SN_field_name: "background"
+    use_cocite: True
+    use_cluster_to_filter: True # 过滤器中使用聚类算法
+    cite_type: "cite_id_list"
     limit_num: 100  # 限制entity对应的paper数量
-    sn_num_for_entity: 3 # SN搜索的文章数量，扩充entity
-    kg_jump_num: 1   # 跳数
-    kg_cover_num: 3  # entity重合数量
-    sum_paper_num: 50  # 最多检索到的paper数量
-    sn_retrieve_paper_num: 55 # 通过SN检索到的文章
+    sn_num_for_entity: 5 # SN搜索的文章数量，扩充entity
+    kg_jump_num: 1   # 跳数，这个参数是不用的，就默认一次
+    kg_cover_num: 7  # entity重合数量，就是两个entity共同同时出现在了kg_cover_num篇文章中
+    sum_paper_num: 100  # 最多检索到的paper数量，指的是通过entity检索到的paper数量
+    sn_retrieve_paper_num: 100 # 通过SN检索到的文章
+    all_retrieve_paper_num: 10
     cocite_top_k: 1
     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
+    top_k_list: [10, 20, 30, 40, 50, 60, 80, 100, 120, 150]
+    s_bg: 1.0
+    s_contribution: 0.0
+    s_summary: 0.0
+    s_abstract: 0.0
+    similarity_threshold: 0.95
+    # similarity_threshold: 0.55

scripts/env.sh

@@ -1,6 +1,13 @@
 export NEO4J_URL="bolt://127.0.0.1:7687"
 export NEO4J_USERNAME="neo4j" # default neo4j
 export NEO4J_PASSWD="****" # your passwd
+
+## Use GPT-4o
+# export MODEL_NAME="gpt-4o"
+# export MODEL_TYPE="OpenAI"
+# export MODEL_API_KEY="sk-************************************************"
+# export BASE_URL="https://********************************"
+
 ## Use Qwen
 export MODEL_NAME="qwen-turbo"
 export MODEL_TYPE="OpenAI"

scripts/retriever_eval.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-python src/retriever.py retrieve \
+CUDA_VISIBLE_DEVICES=0 python src/retriever.py retrieve \
   -c configs/datasets.yaml \
   --ids-path assets/data/test_acl_2024.json
 

src/app_pages/button_interface.py

@@ -4,6 +4,7 @@ from utils.llms_api import APIHelper
 from utils.header import ConfigReader
 from utils.hash import check_env, check_embedding
 from generator import IdeaGenerator
+import functools
 
 
 class Backend(object):
@@ -36,94 +37,53 @@ class Backend(object):
         except (FileNotFoundError, json.JSONDecodeError) as e:
             print(f"Error loading examples from {path}: {e}")
             return []
+        
+    def background2entities_callback(self, background):
+        return self.api_helper.generate_entity_list(background)
 
-    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 background2expandedbackground_callback(self, background, entities):
+        keywords_str = functools.reduce(lambda x, y: f"{x}, {y}", entities)
+        expanded_background = self.api_helper.expand_background(background, keywords_str)
+        return expanded_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 background2brainstorm_callback(self, expanded_background):
+        return self.api_helper.generate_brainstorm(expanded_background)
+
+    def brainstorm2entities_callback(self, brainstorm, entities):
+        entities_bs = self.api_helper.generate_entity_list(brainstorm, 10)
+        entities_all = list(set(entities) | set(entities_bs))
+        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:
-            result = json.loads(json_strs)
-            res = []
-            for i, p in enumerate(result["related_paper"]):
-                res.append(str(p))
-        else:
-            result = self.retriever_factory.retrieve(
-                background, entities, need_evaluate=False, target_paper_id_list=[]
-            )
-            res = []
-            for i, p in enumerate(result["related_paper"]):
-                res.append(f'{p["title"]}. {p["venue_name"].upper()} {p["year"]}.')
+    def entities2literature_callback(self, expanded_background, entities):
+        result = self.retriever_factory.retrieve(
+            expanded_background, entities, need_evaluate=False, target_paper_id_list=[]
+        )
+        res = []
+        for i, p in enumerate(result["related_paper"]):
+            res.append(f'{p["title"]}. {p["venue_name"].upper()} {p["year"]}.')
         return res, result["related_paper"]
 
     def literature2initial_ideas_callback(
-        self, background, brainstorms, retrieved_literature, json_strs=None
+        self, expanded_background, brainstorms, retrieved_literature
     ):
-        if json_strs is not None:
-            json_contents = json.loads(json_strs)
-            return json_contents["median"]["initial_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(
-                        background, "new_idea", self.brainstorm_mode, False
-                    )
-                )
-            else:
-                message_input, idea_modified, median = self.idea_generator.generate(
-                    background, "new_idea", self.brainstorm_mode, False
-                )
-            return median["initial_idea"], idea_modified
+        self.idea_generator.paper_list = retrieved_literature
+        self.idea_generator.brainstorm = brainstorms
+        _, _, inspirations, initial_ideas, idea_filtered, final_ideas = (
+            self.idea_generator.generate_ins_bs(expanded_background)
+        )
+        return idea_filtered, final_ideas
 
-    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 initial2final_callback(self, initial_ideas, final_ideas):
+        return final_ideas
 
     def get_demo_i(self, i):
         if 0 <= i < len(self.examples):
             return self.examples[i].get("background", "Background not found.")
         else:
             return "Example not found. Please select a valid index."
-
-    #     return ("The application scope of large-scale language models such as GPT-4 and LLaMA "
-    # "has rapidly expanded, demonstrating powerful capabilities in natural language processing "
-    # "and multimodal tasks. However, as the size and complexity of the models increase, understanding "
-    # "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/app_pages/homepage.py

@@ -40,7 +40,7 @@ def generate_mainpage():
         st.header("Resources")
         st.markdown("Our paper: [https://arxiv.org/abs/2410.23166](https://arxiv.org/abs/2410.23166)")
         st.markdown("Our github repository: [https://github.com/cheerss/SciPIP](https://github.com/cheerss/SciPIP)")
-        st.markdown("Our Huggingface demo: Coming soon...")
+        st.markdown("Our Huggingface demo: [https://huggingface.co/spaces/lihuigu/SciPIP](https://huggingface.co/spaces/lihuigu/SciPIP)")
         # st.page_link("https://arxiv.org/abs/2410.23166", label="Our paper: https://arxiv.org/abs/2410.23166", icon=None)
         # st.page_link("https://github.com/cheerss/SciPIP", label="Our github repository: https://github.com/cheerss/SciPIP", icon=None)
 
@@ -71,7 +71,7 @@ def generate_mainpage():
         st.header("相关资源")
         st.markdown("论文: [https://arxiv.org/abs/2410.23166](https://arxiv.org/abs/2410.23166)")
         st.markdown("Github仓库: [https://github.com/cheerss/SciPIP](https://github.com/cheerss/SciPIP)")
-        st.markdown("Huggingface演示: 敬请期待...")
+        st.markdown("Huggingface演示: [https://huggingface.co/spaces/lihuigu/SciPIP](https://huggingface.co/spaces/lihuigu/SciPIP)")
         # st.page_link("https://arxiv.org/abs/2410.23166", label="Our paper: https://arxiv.org/abs/2410.23166", icon=None)
         # st.page_link("https://github.com/cheerss/SciPIP", label="Our github repository: https://github.com/cheerss/SciPIP", icon=None)
 

src/app_pages/one_click_generation.py

@@ -1,9 +1,8 @@
 import streamlit as st
+from loguru import logger
 from .locale import _
 from .sidebar_components import get_sidebar_header, get_sidebar_supported_fields, get_help_us_improve, get_language_select
 
-# 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
@@ -37,7 +36,7 @@ def generate_sidebar():
     get_help_us_improve()
 
 def generate_mainpage(backend):
-    st.title(_("💧 One-click Generation"))
+    st.title(_("One-click Generation"))
 
     if "messages" not in st.session_state:
         st.session_state["messages"] = [{"role": "assistant", "content": "Please give me some key words or a background"}]
@@ -95,42 +94,47 @@ def generate_mainpage(backend):
     cols[3].button(_("Reset Chat"), on_click=reset, use_container_width=True, type="primary")
 
 def generate_ideas(backend, background):
+    with st.spinner(text=("Extracting entities from the user's input...")):
+        entities_bg = backend.background2entities_callback(background)
+
+    with st.spinner(text=("Understanding the user's input...")):
+        expanded_background = backend.background2expandedbackground_callback(background, entities_bg)
+        st.session_state["intermediate_output"]["expanded_background"] = {"role": "assistant", "content": expanded_background}
+
     with st.spinner(text=("Brainstorming...")):
-        brainstorms = backend.background2brainstorm_callback(background)
+        brainstorms = backend.background2brainstorm_callback(expanded_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
+        st.chat_message("assistant").write("I have the following thoughts, but I'll search the literature to further consolidate and improve the ideas.")
+        st.chat_message("assistant").write(brainstorms)
 
-    with st.spinner(text=("Extracting entities...")):
-        entities = backend.brainstorm2entities_callback(background, brainstorms)
-        st.session_state["intermediate_output"]["entities"] = {"role": "assistant", "content": entities}
+    with st.spinner(text=("Extracting entities for literature retrieval...")):
+        entities_all = backend.brainstorm2entities_callback(brainstorms, entities_bg)
+        st.session_state["intermediate_output"]["entities"] = {"role": "assistant", "content": entities_all}
         # 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)
+        msg = "The retrieved works include:"
+        related_works, related_works_intact = backend.entities2literature_callback(expanded_background, entities_all)
         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:"
+    with st.spinner(text="Generating ideas... (This may take up to 5 minutes)"):
         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})
-        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})
+        logger.info(f"Num of initial ideas: {len(initial_ideas)}, num of final ideas: {len(final_ideas)}")
+        # assert len(initial_ideas) == len(final_ideas)
+        msg = f"I have {len(initial_ideas)} ideas:"
         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
+        for i in range(len(initial_ideas)):
+            output = f"""### Concise Idea
+{initial_ideas[i]}
+
+### Idea in Detail:
+
+{final_ideas[i]}
+
+"""
+            st.session_state.messages.append({"role": "assistant", "content": output})
+            st.chat_message("assistant").write(output)
 
 def one_click_generation(backend):
     generate_sidebar()

src/app_pages/sidebar_components.py

@@ -11,9 +11,9 @@ def get_sidebar_supported_fields():
     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(_("Computer Vision (CV)"), value=True, disabled=True)
 
-    st.sidebar.checkbox(_("[Partial] Multimodal"), value=True, disabled=True)
+    st.sidebar.checkbox(_("Multimodal"), value=True, disabled=True)
     st.sidebar.checkbox(_("Incoming Other Fields"), value=False, disabled=True)
 
 def get_help_us_improve():

src/app_pages/step_by_step_generation.py

@@ -44,8 +44,8 @@ def get_textarea_height(text_content):
     return max(count * 23 + 20, 100) # 23 is a magic number
 
 def generate_mainpage(backend):
-    st.title(_("💦 Step-by-step Generation"))
-    st.header(_("🐳 Background"))
+    st.title(_("Step-by-step Generation"))
+    st.header(_("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")
@@ -60,112 +60,101 @@ def generate_mainpage(backend):
         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)
+            with st.spinner(text="Let me first brainstorm some ideas..."):
+                st.session_state["entities_bg"] = backend.background2entities_callback(background)
+                st.session_state["expanded_background"] = backend.background2expandedbackground_callback(
+                    background, st.session_state["entities_bg"]
+                )
+                st.session_state["brainstorms"] = backend.background2brainstorm_callback(
+                    st.session_state["expanded_background"]
+                )
             # st.session_state["brainstorms"] = "Test text"
             st.session_state["brainstorms_expand"] = True
             st.session_state["global_state_step"] = 2.5
 
     ## Brainstorms
-    st.header(_("👻 Brainstorms"))
-    with st.expander("", 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(_("Please input the brainstorms on the left."))
-        # st.write("</div>")
-        col1, col2 = st.columns([2, 20])
-        submitted = col1.button(_("Submit"), type="primary")
-        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
+    if st.session_state["global_state_step"] >= 2.5:
+        st.header(_("Brainstorms"))
+        with st.expander("", 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(_("Please input the brainstorms on the left."))
+            # st.write("</div>")
+            col1, col2 = st.columns([2, 20])
+            submitted = col1.button(_("Submit"), type="primary")
+            if submitted:
+                st.session_state["global_state_step"] = 3.0
+                with st.spinner(text="I'am extracting keywords in the background and brainstorming ideas"):
+                    st.session_state["entities"] = backend.brainstorm2entities_callback(brainstorms, st.session_state["entities_bg"])
+                # st.session_state["entities"] = "entities"
+                st.session_state["global_state_step"] = 3.5
+                st.session_state["entities_expand"] = True
 
     ## Entities
-    st.header(_("🐱 Extracted Entities"))
-    with st.expander("", expanded=st.session_state.get("entities_expand", False)):
-        ## pills
-        def update_entities():
-            return
-        ori_entities = st.session_state.get("entities", [])
-        entities_updated = st.session_state.get("entities_updated", ori_entities)
-        entities_updated = st.pills(label="entities", options=ori_entities, selection_mode="multi", 
-                            default=ori_entities, label_visibility="collapsed", on_change=update_entities)
-        st.session_state["entities_updated"] = entities_updated
+    if st.session_state["global_state_step"] >= 3.5:
+        st.header(_("Extracted Entities"))
+        with st.expander("", expanded=st.session_state.get("entities_expand", False)):
+            ## pills
+            def update_entities():
+                return
+            ori_entities = st.session_state.get("entities", [])
+            entities_updated = st.session_state.get("entities_updated", ori_entities)
+            entities_updated = st.pills(label="entities", options=ori_entities, selection_mode="multi", 
+                                default=ori_entities, label_visibility="collapsed", on_change=update_entities)
+            st.session_state["entities_updated"] = entities_updated
 
-        submitted = st.button(_("Submit"), key="entities_button", type="primary")
-        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_updated)
-            st.session_state["related_works_use_state"] = [True] * len(st.session_state["related_works"])
-            st.session_state["global_state_step"] = 4.5
-            st.session_state["related_works_expand"] = True
+            submitted = st.button(_("Submit"), key="entities_button", type="primary")
+            if submitted:
+                st.session_state["global_state_step"] = 4.0
+                with st.spinner(text="I am retrieving related works for more ideas..."):
+                    st.session_state["related_works"], st.session_state["related_works_intact"] = \
+                        backend.entities2literature_callback(st.session_state["expanded_background"], entities_updated)
+                st.session_state["related_works_use_state"] = [True] * len(st.session_state["related_works"])
+                st.session_state["global_state_step"] = 4.5
+                st.session_state["related_works_expand"] = True
 
     ## Retrieved related works
-    st.header(_("📖 Retrieved Related Works"))
-    with st.expander("", expanded=st.session_state.get("related_works_expand", False)):
-        related_works = st.session_state.get("related_works", [])
-        for i, rw in enumerate(related_works):
-            checked = st.checkbox(rw, value=st.session_state.get("related_works_use_state")[i])
-            st.session_state.get("related_works_use_state")[i] = checked
+    if st.session_state["global_state_step"] >= 4.5:
+        st.header(_("Retrieved Related Works"))
+        with st.expander("", expanded=st.session_state.get("related_works_expand", False)):
+            related_works = st.session_state.get("related_works", [])
+            for i, rw in enumerate(related_works):
+                checked = st.checkbox(rw, value=st.session_state.get("related_works_use_state")[i])
+                st.session_state.get("related_works_use_state")[i] = checked
 
-        submitted = st.button(_("Submit"), key="related_works_button", type="primary")
-        if submitted:
-            st.session_state["global_state_step"] = 5.0
-            with st.spinner(text="Generating initial ideas..."):
-                st.session_state["selected_related_works_intact"] = []
-                for s, p in zip(st.session_state.get("related_works_use_state"), st.session_state["related_works_intact"]):
-                    if s:
-                        st.session_state["selected_related_works_intact"].append(p)
-                res = backend.literature2initial_ideas_callback(background, brainstorms, st.session_state["selected_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
+            submitted = st.button(_("Submit"), key="related_works_button", type="primary")
+            if submitted:
+                st.session_state["global_state_step"] = 5.0
+                with st.spinner(text="I am generating final ideas..."):
+                    st.session_state["selected_related_works_intact"] = []
+                    for s, p in zip(st.session_state.get("related_works_use_state"), st.session_state["related_works_intact"]):
+                        if s:
+                            st.session_state["selected_related_works_intact"].append(p)
+                    res = backend.literature2initial_ideas_callback(background, brainstorms, st.session_state["selected_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.header(_("😼 Generated Initial Ideas"))
-    with st.expander("", 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(_("Please input the initial ideas on the left."))
-        submitted = col1.button(_("Submit"), key="initial_ideas_button", type="primary")
-        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.header(_("😸 Generated Final Ideas"))
-    with st.expander("", 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(_("Please input the final ideas on the left."))
-        submitted = col1.button(_("Submit"), key="final_ideas_button", type="primary")
+    if st.session_state["global_state_step"] >= 5.5:
+        st.header(_("Generated Ideas"))
+        with st.expander("", expanded=st.session_state.get("initial_ideas_expand", False)):
+            for initial_idea, final_idea in zip(st.session_state.get("initial_ideas", ""), st.session_state.get("final_ideas", "")):
+                st.divider()
+                st.markdown("### Concise Idea")
+                st.markdown(initial_idea)
+                st.markdown("### Idea in Detail")
+                st.markdown(final_idea)
+                st.divider()
 
 def step_by_step_generation(backend):
     ## Pipeline global state

src/config/reader.py

@@ -159,5 +159,5 @@ class ConfigReader:
         """
         config = ConfigReader(file_, included).config
         for k, v in kwargs.items():
-            config[k] = v
+            config.get(k, {}).update(v)
         return config

src/generator.py

@@ -1,3 +1,4 @@
+import functools
 from utils.paper_retriever import RetrieverFactory
 from utils.paper_client import PaperClient
 from utils.llms_api import APIHelper
@@ -10,6 +11,7 @@ import warnings
 import time
 import os
 from utils.hash import check_env, check_embedding
+import threading
 
 warnings.filterwarnings("ignore")
 
@@ -26,30 +28,40 @@ def extract_problem(problem, background):
     return research_problem
 
 
+def extract_ideas(idea_str):
+    if idea_str is None:
+        return ""
+    ideas = []
+    for i in range(1, 100): # 100 is a magic number
+        start_word = f"**Idea {i}"
+        end_word = f"**Idea {i+1}"
+        start_index = idea_str.find(start_word)
+        end_index = idea_str.find(end_word)
+        if start_index != -1 and end_index != -1:
+            ideas.append(idea_str[start_index:end_index].strip())
+            # idea_str = idea_str[start_index+end_index+1:]
+        elif start_index != -1:
+            ideas.append(idea_str[start_index:].strip())
+            break
+        else:
+            break
+    return ideas if ideas else [idea_str]
+
+
 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):
+        """Generate ideas without cue words and brainstorm
+        """
         problem, message_input = self.api_helper.generate_problem(
             background, self.paper_list
         )
@@ -57,19 +69,9 @@ class IdeaGenerator:
         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):
+        """Generate ideas without cue words, but brainstorm
+        """
         problem, message_input = self.api_helper.generate_problem(
             background, self.paper_list
         )
@@ -79,24 +81,9 @@ class IdeaGenerator:
         )
         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):
+        """Generate ideas without cue words and brainstorm, but inspiration
+        """
         problem, message_input = self.api_helper.generate_problem(
             background, self.paper_list
         )
@@ -109,26 +96,9 @@ class IdeaGenerator:
         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):
+        """Generate ideas without cue words, but inspiration and brainstorm
+        """
         problem, message_input = self.api_helper.generate_problem(
             background, self.paper_list
         )
@@ -143,6 +113,52 @@ class IdeaGenerator:
         )
         return message_input, problem, inspirations, idea, idea_filtered
 
+    def generate_ins_bs(self, detail_background: str):
+        """Generate ideas with inspiration and brainstorm
+        """
+        inspirations = []
+
+        ## generate inspirations
+        processes = []
+        def generate_inspiration(paper, i):
+            detail_method = self.api_helper.generate_concise_method(paper["methodology"])
+            inspiration = self.api_helper.generate_inspiration_with_detail_method(detail_background, detail_method)
+            logger.info(f"Generate inspiration for related paper {i} succeed")
+            if not(inspiration.startswith("None") or (len(inspiration) < 100 and "None" in inspiration)):
+                inspirations.append(inspiration)
+
+        for i, paper in enumerate(self.paper_list):
+            p = threading.Thread(target=generate_inspiration, args=(paper, i))
+            processes.append(p)
+            p.start()
+        for p in processes:
+            p.join(120)
+        
+        ## generate ideas through all inspirations
+        logger.info("Generate inspirations for all related papers succeed")
+        idea = self.api_helper.generate_idea_by_inspiration(detail_background, inspirations)
+        initial_ideas = extract_ideas(idea)
+        logger.info("Generate ideas from inspirations succeed")
+        idea_filtered = self.api_helper.integrate_idea(detail_background, self.brainstorm, idea)
+        logger.info("Idea integration succeed")
+
+        ## expand ideas
+        ideas_filtered = extract_ideas(idea_filtered)
+        final_ideas = ["None"] * len(ideas_filtered)
+        def expand_idea(detail_background: str, idea: str, i):
+            final_ideas[i] = self.api_helper.expand_idea(detail_background, idea)
+            logger.info(f"Expand the {i}th idea succeed")
+        processes = []
+        for i, idea in enumerate(ideas_filtered):
+            p = threading.Thread(target=expand_idea, args=(detail_background, idea, i))
+            processes.append(p)
+            p.start()
+        for p in processes:
+            p.join(120)
+
+        ## reture
+        return None, None, inspirations, initial_ideas, ideas_filtered, final_ideas
+
     def generate(
         self,
         background: str,
@@ -156,37 +172,21 @@ class IdeaGenerator:
         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)
-                )
+            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)
+            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 = {
@@ -209,37 +209,21 @@ class IdeaGenerator:
         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)
-                )
+            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_{} without cue words.".format(
-                        mode_name, bs_mode
-                    )
-                )
-                (message_input, problem, inspirations, idea, idea_filtered) = (
-                    self.generate_without_cue_words_ins_bs(background)
+            logger.info(
+                "{} using brainstorm_{} without cue words.".format(
+                    mode_name, bs_mode
                 )
+            )
+            (message_input, problem, inspirations, idea, idea_filtered) = (
+                self.generate_without_cue_words_ins_bs(background)
+            )
 
         idea_modified = self.api_helper.modify_idea(background, idea_filtered)
         median = {
@@ -271,209 +255,22 @@ def main(ctx):
 )
 @click.option(
     "--ids-path",
-    default="./assets/data/test_acl_2024.json",
+    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_c",
+    "--out-path",
+    default="./assets/output_idea/",
     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=False,
-    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(
-    "--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,
-):
-    check_env()
-    check_embedding()
-    # 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()
-    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
-        )
-        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(),
+    "--out-file",
+    default="out-file.json",
+    type=str,
     required=True,
     help="Dataset configuration file in YAML",
 )
@@ -499,6 +296,12 @@ def backtracking(
     required=True,
     help="Use inspiration in generation",
 )
+@click.option(
+    "--expand-intermediate",
+    default=False,
+    type=bool,
+    help="The number of data you want to process",
+)
 @click.option(
     "--num",
     default=100,
@@ -509,9 +312,12 @@ def backtracking(
 def new_idea(
     config_path,
     ids_path,
+    out_path,
+    out_file,
     retriever_name,
     brainstorm_mode,
     use_inspiration,
+    expand_intermediate,
     num,
     **kwargs,
 ):
@@ -523,16 +329,16 @@ def new_idea(
     # Configuration
     config = ConfigReader.load(config_path, **kwargs)
     api_helper = APIHelper(config)
+    paper_client = PaperClient()
     check_embedding(config.DEFAULT.embedding)
     eval_data = []
     cur_num = 0
     data_num = 0
-    batch_size = 2
+    batch_size = 1
     bg_ids = set()
-    output_dir = "./assets/output_idea/"
-    os.makedirs(output_dir, exist_ok=True)
+    os.makedirs(out_path, exist_ok=True)
     output_file = os.path.join(
-        output_dir, f"output_new_idea_{brainstorm_mode}_ins_{use_inspiration}.json"
+        out_path, out_file
     )
     if os.path.exists(output_file):
         with open(output_file, "r", encoding="utf-8") as f:
@@ -543,10 +349,12 @@ def new_idea(
             except json.JSONDecodeError:
                 eval_data = []
     logger.debug(f"{cur_num} datas have been processed.")
-    for line in ids_path:
+    all_input = json.load(ids_path)
+    for line in all_input:
         # 解析每行的JSON数据
-        data = json.loads(line)
-        # 1. 获取背景信息
+        # data = json.loads(line)
+        data = line
+        ### 1. 获取背景信息
         if "background" in data.keys():
             bg = data["background"]
         else:
@@ -557,17 +365,28 @@ def new_idea(
             data_num += 1
             print(f"Skipping already processed data_{data_num}.")
             continue
+        
+        ## extract entities from background
+        entities = api_helper.generate_entity_list(bg)
+
+        ## expand background to a detailed version
+        keywords_str = functools.reduce(lambda x, y: f"{x}, {y}", entities)
+        expanded_background = api_helper.expand_background(bg, keywords_str)
+
+        ## brainstorm according to the background
         if brainstorm_mode == "mode_b" or brainstorm_mode == "mode_c":
-            brainstorm = api_helper.generate_brainstorm(bg)
+            brainstorm = api_helper.generate_brainstorm(expanded_background)
+            seperate_brainstorm = extract_ideas(brainstorm)
+            ## expand the brainstorms to a detailed version
+            expanded_brainstorms = []
+            if expand_intermediate:
+                for i, sb in enumerate(seperate_brainstorm):
+                    expanded_brainstorms.append(api_helper.expand_idea(expanded_background, sb))
+                    logger.info(f"Expand the {i}th brainstorm succeed")
         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)
+        
+        ## Extract entities from the brainstorm result
         logger.debug("Original entities from background: {}".format(entities))
         if brainstorm_mode == "mode_c":
             entities_bs = api_helper.generate_entity_list(brainstorm, 10)
@@ -576,49 +395,54 @@ def new_idea(
         else:
             entities_bs = None
             entities_all = entities
-        # 2. 检索相关论文
+
+        ### 2. 检索相关论文
         rt = RetrieverFactory.get_retriever_factory().create_retriever(
             retriever_name, config
         )
         result = rt.retrieve(
-            bg, entities_all, need_evaluate=False, target_paper_id_list=[]
+            expanded_background, 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
-                )
-            )
+        for paper in related_paper:
+            if not ("detail_method" in paper):
+                paper["detail_method"] = api_helper.generate_concise_method(paper["methodology"])
+                if isinstance(paper["detail_method"], str):
+                    paper_client.insert_new_field(paper["hash_id"], "detail_method", paper["detail_method"])
+                    logger.info(f"Add new field detail method to paper: {paper['hash_id']} succeed")
+        logger.info("Generate detail methods for all related papers succeed")
+
+        ### 3. 生成IDEA
+        idea_generator = IdeaGenerator(config, related_paper, brainstorm)
+        _, _, inspirations, initial_ideas, idea_filtered, final_ideas = idea_generator.generate_ins_bs(expanded_background)
+        expanded_initial_ideas = []
+        if expand_intermediate:
+            for i, initial_idea in enumerate(initial_ideas):
+                expanded_initial_ideas.append(api_helper.expand_idea(expanded_background, initial_idea))
+                logger.info(f"Expand the {i}th initial idea succeed")
         eval_data.append(
             {
                 "background": bg,
+                "expanded_background": expanded_background,
                 "entities_bg": entities,
                 "brainstorm": brainstorm,
+                "seperate_brainstorm": seperate_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,
+                "related_paper": [p["title"] for p in related_paper],
+                "inspirations": inspirations,
+                "initial_ideas": initial_ideas,
+                "filtered_ideas": idea_filtered,
+                "expanded_final_ideas": final_ideas,
+                "expanded_brainstorms": expanded_brainstorms,
+                "expanded_initial_ideas": expanded_initial_ideas,
             }
         )
         cur_num += 1
         if cur_num % batch_size == 0:
-            with open(
-                output_file,
-                "w",
-                encoding="utf-8",
-            ) as f:
+            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

src/paper_manager.py

@@ -22,6 +22,12 @@ unicode_pattern = r"\u00c0-\u00ff\u0100-\u017f\u0180-\u024f\u4e00-\u9fff\u3040-\
 
 
 def find_methodology(article_dict):
+    """For an article dict (representing an article), return the methodology part
+    Args:
+        article_dict
+    Returns:
+        methodology: str
+    """
     def find_section_index(keywords):
         for i, section in enumerate(article_dict["sections"], 1):
             heading = section["heading"].lower()
@@ -70,10 +76,14 @@ def find_methodology(article_dict):
 
 
 def count_sb_pairs(text):
+    """Find the number of square brackets (possible citations)
+    """
     return len(re.findall(r"\[.*?\]", text))
 
 
 def count_rb_pairs(text):
+    """Find the number of round brackets (possible citations)
+    """
     return len(re.findall(r"\(.*?\)", text))
 
 
@@ -85,17 +95,18 @@ def find_cite_paper(introduction, methodology, references):
     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"
-    )
+    ## Seems redudant, remove repeated definition of pattern
+    # 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
@@ -199,11 +210,22 @@ class PaperManager:
             self.paper_client.create_vector_index()
 
     def clean_entity(self, entity):
+        """The extracted entities may be noisy, remove all noisy characters
+        Args:
+            entity (str): an entity
+        Returns:
+            cleaned_entity (str): entity after cleaning
+        """
         if entity is None:
             return None
+        # remove all () and their contents
         cleaned_entity = re.sub(r"\([^)]*\)", "", entity)
+        # remove non-word characters (e.g., punctuations)
         cleaned_entity = re.sub(r"[^\w\s]", "", cleaned_entity)
+        # replace _ as a whitespace
         cleaned_entity = re.sub(r"_", " ", cleaned_entity)
+        # remove multiple continuous blanks, remove leading and trailing spaces 
+        # (\s means blank characters)
         cleaned_entity = re.sub(r"\s+", " ", cleaned_entity).strip()
         return cleaned_entity
 
@@ -347,6 +369,8 @@ class PaperManager:
         need_get_entities=False,
         need_ground_truth=False,
     ):
+        """Parse a paper, dump the result into a json file
+        """
         if paper["pdf_url"] in self.ignore_paper_pdf_url:
             logger.warning(
                 "hash_id: {}, pdf_url: {} ignore".format(
@@ -511,6 +535,8 @@ class PaperManager:
         need_get_entities=False,
         need_ground_truth=False,
     ):
+        """Parse a paper and dump into the json file
+        """
         result = []
         if self.year != "all":
             logger.info(
@@ -632,6 +658,45 @@ class PaperManager:
         # )
         # self.paper_client.add_paper_summary_embedding(self.embedding_model, hash_id)
 
+    def add_new_embedding(self, hash_id=None, to="all"):
+        """add new embeddings for abstract, background, contribution, and summary
+        """
+        postfix_set = {
+            "sentence-transformers/all-MiniLM-L6-v2": "",
+            "BAAI/llm-embedder": "_llm_embedder",
+            "jina-embeddings-v3": "_jina_v3"
+        }
+        postfix = postfix_set[self.config.DEFAULT.embedding]
+        if "jina" in postfix:
+            if self.config.DEFAULT.embedding_task == "text-matching":
+                postfix += "_text_matching"
+            elif self.config.DEFAULT.embedding_task == "retrieval.query":
+                postfix += "_query"
+            elif self.config.DEFAULT.embedding_task == "retrieval.passage":
+                postfix += "_passage"
+            else:
+                assert False
+        if to == "all" or to == "abstract":
+            self.paper_client.update_paper_embedding(
+                self.embedding_model, hash_id,
+                name="abstract", postfix=postfix
+            )
+        if to == "all" or to == "background":
+            self.paper_client.update_paper_embedding(
+                self.embedding_model, hash_id,
+                name="background", postfix=postfix
+            )
+        if to == "all" or to == "contribution":
+            self.paper_client.update_paper_embedding(
+                self.embedding_model, hash_id,
+                name="contribution", postfix=postfix
+            )
+        if to == "all" or to == "summary":
+            self.paper_client.update_paper_embedding(
+                self.embedding_model, hash_id,
+                name="summary", postfix=postfix
+            )
+
     def cosine_similarity_search(self, data_type, context, k=1):
         """
         return related paper: list
@@ -641,6 +706,12 @@ class PaperManager:
         return result
 
     def generate_paper_list(self):
+        """Dump paper list into a json file, the json is saved at "folder_path"
+        Args:
+            None
+        Return:
+            None
+        """
         folder_path = f"./assets/paper/{self.venue_name}"
         if not os.path.exists(folder_path):
             os.makedirs(folder_path)
@@ -721,6 +792,14 @@ def main(ctx):
     help="The LLMS API aliases used for generation. When used, it will invalidate --llms-api",
 )
 def crawling(config_path, year, venue_name, **kwargs):
+    """Download paper list in to a json file
+    Args:
+        config_path (str):
+        year (int): the paper's publication data
+        venue_name (str): CVPR, etc.
+    Resturns:
+        None
+    """
     # Configuration
     config = ConfigReader.load(config_path, **kwargs)
     pm = PaperManager(config, venue_name, year)
@@ -772,6 +851,9 @@ def crawling(config_path, year, venue_name, **kwargs):
     help="The LLMS API aliases used for generation. When used, it will invalidate --llms-api",
 )
 def update(config_path, year, venue_name, **kwargs):
+    """Read paper lists from assets/paper directory, insert them into database,
+    including downloading, parsing, etc., but not embedding
+    """
     # Configuration
     config = ConfigReader.load(config_path, **kwargs)
     pm = PaperManager(config, venue_name, year)
@@ -831,6 +913,8 @@ def update(config_path, year, venue_name, **kwargs):
     help="Dataset configuration file in YAML",
 )
 def local(config_path, year, venue_name, output, **kwargs):
+    """Parse papers and dump them into json files
+    """
     # Configuration
     output_path = output.name
     if not os.path.exists(os.path.dirname(output_path)):
@@ -842,7 +926,6 @@ def local(config_path, year, venue_name, output, **kwargs):
         need_download=True, need_parse=True, need_summary=True
     )
 
-
 @main.command()
 @click.option(
     "-c",
@@ -853,10 +936,93 @@ def local(config_path, year, venue_name, output, **kwargs):
     help="Dataset configuration file in YAML",
 )
 def embedding(config_path):
+    """Insert embedding for papers in the database
+    """
     # Configuration
     config = ConfigReader.load(config_path)
     PaperManager(config).insert_embedding()
 
+@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 add_new_embedding(config_path):
+    """Insert another new embedding for papers in the database
+    """
+    # Configuration
+    config = ConfigReader.load(config_path)
+    PaperManager(config).add_new_embedding(to="all")
+
+@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(
+    "--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(
+    "--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(
+    "-o",
+    "--output",
+    default=get_dir("./output/out.json"),
+    type=click.File("wb"),
+    required=True,
+    help="Dataset configuration file in YAML",
+)
+def parse_papers_to_json(config_path, venue_name, year, output, **kwargs):
+    """Read json files and download papers, then parse them and dump into jsons
+    """
+    # 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=venue_name, year=year)
+    pm.update_paper_from_json_to_json(
+        need_download=True, need_parse=True, need_summary=True
+    )
+
 
 if __name__ == "__main__":
     main()

src/retriever.py

@@ -22,7 +22,10 @@ def main(ctx):
     print("Mode:", ctx.invoked_subcommand)
 
 
-@main.command()
+@main.command(context_settings=dict(
+    ignore_unknown_options=True,
+    allow_extra_args=True,
+))
 @click.option(
     "-c",
     "--config-path",
@@ -38,9 +41,21 @@ def main(ctx):
     required=True,
     help="Dataset configuration file in YAML",
 )
-def retrieve(
-    config_path, ids_path, **kwargs
-):
+@click.pass_context
+def retrieve(ctx,
+    config_path, ids_path
+): 
+    initial_kwargs={ctx.args[i][2:]: ctx.args[i+1] for i in range(0, len(ctx.args), 2)}
+    kwargs = {"RETRIEVE": {}, "DEFAULT": {}}
+    for k, v in initial_kwargs.items():
+        if "num" in k:
+            kwargs["RETRIEVE"][k] = int(v)
+        elif "s_" in k:
+            kwargs["RETRIEVE"][k] = float(v)
+        elif "use_cocite" in k:
+            kwargs["RETRIEVE"][k] = bool(int(v))
+        else:
+            kwargs["RETRIEVE"][k] = v
     config = ConfigReader.load(config_path, **kwargs)
     check_embedding(config.DEFAULT.embedding)
     check_env()
@@ -82,8 +97,8 @@ def retrieve(
         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"]
+        if "background" in paper.keys():
+            bg = paper["background"]
         else:
             logger.error(f"paper hash_id {paper['hash_id']} doesn't have background...")
             continue

src/utils/api/base_helper.py

@@ -128,7 +128,7 @@ class BaseHelper:
                     response = r.json()["data"]["output"]
                     return response  # 或者根据需要返回其他内容
                 else:
-                    print("服务请求失败，响应状态码：", response.status_code)
+                    print("服务请求失败，响应状态码：", r.status_code)
             except RequestException as e:
                 print("请求发生错误：", e)
 

src/utils/hash.py

@@ -21,7 +21,7 @@ def check_embedding(repo_id):
         if repo_id in [
             "sentence-transformers/all-MiniLM-L6-v2",
             "BAAI/bge-small-en-v1.5",
-            "BAAAI/llm_embedder",
+            "BAAI/llm-embedder",
         ]:
             # repo_id = "sentence-transformers/all-MiniLM-L6-v2"
             # repo_id = "BAAI/bge-small-en-v1.5"
@@ -31,6 +31,18 @@ def check_embedding(repo_id):
                 "tokenizer_config.json",
                 "vocab.txt",
             ]
+        elif repo_id in [
+            "jina-embeddings-v3",
+        ]:
+            files_to_download = [
+                "model.safetensors",
+                "modules.json",
+                "tokenizer.json",
+                "config_sentence_transformers.json",
+                "tokenizer_config.json",
+                "1_Pooling/config.json",
+                "config.json",
+            ]
         elif repo_id in ["Alibaba-NLP/gte-base-en-v1.5"]:
             files_to_download = [
                 "config.json",
@@ -89,6 +101,8 @@ class EmbeddingModel:
                 device=device,
                 trust_remote_code=True,
             )
+            if "jina-embeddings-v3" in config.DEFAULT.embedding:
+                cls._instance.embedding_model[0].default_task = config.DEFAULT.embedding_task
             print(f"==== using device {device} ====")
         return cls._instance
 

src/utils/llms_api.py

@@ -62,15 +62,15 @@ 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",
-            "gpt-4o-mini",
-            "local",
-        ]:
-            raise ValueError(f"Check model name...")
+        # if os.environ["MODEL_NAME"] not in [
+        #     "glm4",
+        #     "glm4-air",
+        #     "qwen-max",
+        #     "qwen-plus",
+        #     "gpt-4o-mini",
+        #     "local",
+        # ]:
+        #     raise ValueError(f"Check model name...")
 
         if title is None or abstract is None or introduction is None:
             return None
@@ -102,6 +102,25 @@ class APIHelper(object):
             return None
         return result
 
+    def generate_concise_method(self, methodology: str):
+        prompt = get_prompt()
+        if methodology is None:
+            return None
+        try:
+            message = [
+                prompt[0][0](),
+                prompt[1][0](
+                    methodology=methodology
+                ),
+            ]
+            detail_method = self.generator.create(
+                messages=message,
+            )
+        except Exception:
+            traceback.print_exc()
+            return None
+        return detail_method
+
     def generate_entity_list(self, abstract: str, max_num: int = 5) -> list:
         prompt = get_prompt()
 
@@ -163,6 +182,46 @@ class APIHelper(object):
             return None
 
         return brainstorming_ideas
+    
+    def expand_idea(self, background: str, idea: str) -> str:
+        prompt = get_prompt()
+
+        if background is None:
+            print("Input background is empty ...")
+            return None
+        try:
+            # Initial brainstorming to generate raw ideas
+            message = [prompt[0][0](), prompt[1][0](background=background, brief_idea=idea)]
+            # Call the API to generate brainstorming ideas
+            detail_ideas = self.generator.create(
+                messages=message,
+            )
+
+        except Exception:
+            traceback.print_exc()
+            return None
+
+        return detail_ideas
+
+    def expand_background(self, brief_background: str, keywords: str) -> str:
+        prompt = get_prompt()
+
+        if brief_background is None:
+            print("Input brief background is empty ...")
+            return None
+        try:
+            # Initial brainstorming to generate raw ideas
+            message = [prompt[0][0](), prompt[1][0](brief_background=brief_background, keywords=keywords)]
+            # Call the API to generate brainstorming ideas
+            detail_background= self.generator.create(
+                messages=message,
+            )
+
+        except Exception:
+            traceback.print_exc()
+            return None
+
+        return detail_background
 
     def generate_problem(self, background: str, related_papers: list[dict]):
         prompt = get_prompt()
@@ -238,6 +297,27 @@ class APIHelper(object):
             traceback.print_exc()
             return None
         return inspiration
+    
+
+    def generate_inspiration_with_detail_method(self, background: str, detail_method: str):
+        prompt = get_prompt()
+        if background is None or detail_method is None:
+            return None
+        try:
+            message = [
+                prompt[0][0](),
+                prompt[1][0](
+                    background=background, detail_method=detail_method
+                ),
+            ]
+            response = self.generator.create(
+                messages=message,
+            )
+            inspiration = response
+        except Exception:
+            traceback.print_exc()
+            return None
+        return inspiration
 
     def generate_inspiration_with_cue_words(
         self, problem: str, related_paper: dict, cue_words: list
@@ -325,10 +405,10 @@ class APIHelper(object):
             return None
         return idea
 
-    def generate_idea_by_inspiration(self, problem: str, inspirations: list[str]):
+    def generate_idea_by_inspiration(self, background: str, inspirations: list[str]):
         prompt = get_prompt()
 
-        if problem is None or inspirations is None:
+        if background is None or inspirations is None:
             return None
         try:
             inspirations_text = "".join(
@@ -340,7 +420,7 @@ class APIHelper(object):
 
             message = [
                 prompt[0][0](),
-                prompt[1][0](problem=problem, inspirations_text=inspirations_text),
+                prompt[1][0](background=background, inspirations=inspirations_text),
             ]
             response = self.generator.create(
                 messages=message,

src/utils/paper_client.py

@@ -1,3 +1,18 @@
+r"""_summary_
+-*- coding: utf-8 -*-
+
+Module : prompt.utils
+
+File Name : utils.paper_client
+
+Description : paper client, all operations about neo4j database are in PaperClient
+
+Creation Date : 2024-11-09
+
+Modification Date : 2024-12-17
+
+Author : Lihui Gu (code), Wenxiao Wang (comments)
+"""
 import os
 import re
 import json
@@ -34,7 +49,13 @@ class PaperClient:
         return driver
 
     def update_paper_from_client(self, paper):
-        paper_id = paper["hash_id"]
+        """Read paper from the database (client), update it info into `paper`
+        Args:
+            paper (str): a paper's hash_id
+        Returns:
+            None
+        """
+        paper_id = paper.get("hash_id", None)
         if paper_id is None:
             return None
         query = f"""
@@ -49,6 +70,12 @@ class PaperClient:
                 paper.update(paper_from_client)
 
     def update_papers_from_client(self, paper_id_list):
+        """Read paper from the database (client)
+        Args:
+            paper_id_list (List of str)
+        Returns:
+            List of papers read from the database
+        """
         query = """
             UNWIND $papers AS paper
             MATCH (p:Paper {hash_id: paper.hash_id})
@@ -67,6 +94,13 @@ class PaperClient:
         return [r["result"] for r in result]
 
     def get_paper_attribute(self, paper_id, attribute_name):
+        """Get some attribute of a certain paper
+        Args:
+            paper_id (str):
+            attribute_name (str):
+        Returns:
+            The certain attribute
+        """
         query = f"""
             MATCH (p:Paper {{hash_id: {paper_id}}})
             RETURN p.{attribute_name} AS attributeValue
@@ -80,6 +114,13 @@ class PaperClient:
             return None
 
     def get_papers_attribute(self, paper_id_list, attribute_name):
+        """Get some attribute of a list of papers
+        Args:
+            paper_id (List of str):
+            attribute_name (str):
+        Returns:
+            List of certain attribute
+        """
         query = """
             UNWIND $paper_ids AS paper_id
             MATCH (p:Paper {hash_id: paper_id})
@@ -95,6 +136,13 @@ class PaperClient:
         return paper_attributes
 
     def get_paper_by_attribute(self, attribute_name, anttribute_value):
+        """Get some paper whose `attribute_name` is exactly equal to `anttribute_value`
+        Args:
+            anttribute_name
+            anttribute_value
+        Returns:
+            The first exact match paper object or None
+        """
         query = f"""
             MATCH (p:Paper {{{attribute_name}: '{anttribute_value}'}})
             RETURN p
@@ -107,6 +155,13 @@ class PaperClient:
             return None
 
     def get_paper_from_term(self, entity):
+        """Get paper from entity. The method is so strict that paper.entities must be 
+        exactly equal to entity. The method is not used now.
+        Args:
+            entity:
+        Returns:
+
+        """
         if entity is None:
             return None
         query = """
@@ -126,6 +181,14 @@ class PaperClient:
     def find_related_entities_by_entity_list(
         self, entity_names, n=1, k=3, relation_name="related"
     ):
+        """Find all entities related to an entity name
+        Args:
+            entity_names (List): list of entities
+            n: not used
+            k: entity a and b are related if they co-occure in at least `k` papers
+        Returns:
+            related_entities (List): list of entities who are related with any entity in `entity_names`
+        """
         related_entities = set()
         query = """
             UNWIND $batch_entities AS entity_name
@@ -145,6 +208,12 @@ class PaperClient:
         return list(related_entities)
 
     def find_entities_by_paper_list(self, hash_ids: list):
+        """Retrieve entities for a list of papers:
+        Args:
+            hash_ids (List of papers):
+        Returns:
+            entity_list (List of List of entities): each item is also a list, meaning all entities from a paper
+        """
         query = """
             UNWIND $hash_ids AS hash_id
             MATCH (e:Entity)-[:RELATED_TO]->(p:Paper {hash_id: hash_id})
@@ -161,6 +230,12 @@ class PaperClient:
         return entity_list
 
     def find_paper_by_entity(self, entity_name):
+        """Find all papers with `entity_name`
+        Args:
+            entity_name (str)
+        Returns:
+            res (List of hash_ids): papers with `entity_name`
+        """
         query = """
             MATCH (e1:Entity {name: $entity_name})-[:RELATED_TO]->(p:Paper)
             RETURN p.hash_id AS hash_id
@@ -181,6 +256,19 @@ class PaperClient:
         return []
 
     def find_sentences_by_entity(self, entity_name):
+        """Find all sentences with a certain `entity_name`
+        Args:
+            entity_name (str)
+        Return:
+            sentences (List of strs): One str concatenates all sentences with `entity_name` in a section
+                E.g. [
+                    "abstract sentence 1 from paper 1.abstract sentence 2 from paper 1",
+                    "introduction sentence 1 from paper 1.introduction sentence 2 from paper 1",
+                    "methodology sentence 1 from paper 1.",
+                    "abstract sentence 1 from paper 2.abstract sentence 2 from paper 2",
+                    ...
+                ]
+        """
         query = """
         MATCH (e:Entity {name: $entity_name})-[:RELATED_TO]->(p:Paper)
         WHERE p.abstract CONTAINS $entity_name OR
@@ -215,6 +303,13 @@ class PaperClient:
         return sentences
 
     def select_paper(self, venue_name, year):
+        """Retrieve a list of papers which published at the `venue_name` in `year`
+        Args:
+            venue_name (str)
+            year (int?)
+        Returns:
+            result (List of paper node)
+        """
         query = """
             MATCH (n:Paper) where n.year=$year and n.venue_name=$venue_name return n
         """
@@ -228,6 +323,12 @@ class PaperClient:
             return []
 
     def add_paper_node(self, paper: dict):
+        """Add a paper node
+        Args:
+            paper (Dict)
+        Returns:
+            None
+        """
         if "summary" not in paper.keys():
             paper["summary"] = None
         if "abstract" not in paper.keys():
@@ -279,6 +380,12 @@ class PaperClient:
             )
 
     def check_entity_node_count(self, hash_id: int):
+        """Whether a paper has more than `3` entities
+        Args:
+            hash_id: a paper's hash_id
+        Returns:
+            True if has <= 2 entitis, False otherwise
+        """
         query_check_count = """
             MATCH (e:Entity)-[:RELATED_TO]->(p:Paper {hash_id: $hash_id})
             RETURN count(e) AS entity_count
@@ -293,6 +400,13 @@ class PaperClient:
         return True
 
     def add_entity_node(self, hash_id: int, entities: list):
+        """Add a entity node, and link it to its paper
+        Args:
+            hash_id: a paper's id
+            entities: a paper's all entities
+        Returns:
+            None
+        """
         query = """
             MERGE (e:Entity {name: $entity_name})
             WITH e
@@ -309,6 +423,14 @@ class PaperClient:
                 )
 
     def add_paper_citation(self, paper: dict):
+        """Add citations for the paper node, set its cite_id_list, entities, and all_cite_id_list
+        `cite_id_list` means citations in the Introduction section
+        `all_cite_id_list` means all citations
+        Args:
+            paper (Dict of a paper)
+        Returns:
+            None
+        """
         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
             """
@@ -323,9 +445,122 @@ class PaperClient:
                 ).data()
             )
 
+    def insert_new_field(self, hash_id: str, field_name: str, content):
+        if hash_id is not None:
+            query = f"""
+                MATCH (n {{hash_id: $hash_id}}) 
+                SET n.{field_name} = $content
+                RETURN n
+            """
+            with self.driver.session() as session:
+                result = session.execute_write(
+                    lambda tx: tx.run(
+                        query, hash_id=hash_id, content=content
+                    ).data()
+                )
+            return result
+        else:
+            return None
+
+    def update_paper_embedding(
+        self, embedding_model, hash_id=None, batch_size=512, name="abstract", postfix=""
+    ):
+        """Extract paper embedding and store in the database
+        Args:
+            embedding_model (TODO: what model?): an pytorch embedding model
+            hash_id (str): add embedding for a paper if hash_id is not None.
+                Otherwise, all papers will be handled with a batch size of 512
+            batch_size: if hash_id is None, all papers will be processed with `batch_size`
+        """
+        if hash_id is not None:
+            query = f"""
+            MATCH (p:Paper {{hash_id: $hash_id}})
+            WHERE p.{name} IS NOT NULL
+            RETURN p.{name} 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()
+                )
+            # contexts = [result["title"] + result["context"] for result in results]
+            if name == "abstract":
+                contexts = [result["title"] + result["context"] for result in results]
+            else:
+                contexts = [result["context"] for result in results]
+            paper_ids = [result["hash_id"] for result in results]
+            context_embeddings = embedding_model.encode(
+                contexts, convert_to_tensor=True, device=self.device
+            )
+            query = f"""
+                MERGE (p:Paper {{hash_id: $hash_id}})
+                ON CREATE SET p.{name}_embedding{postfix} = $embedding
+                ON MATCH SET p.{name}_embedding{postfix} = $embedding
+            """
+            for idx, hash_id in tqdm(enumerate(paper_ids)):
+                embedding = (
+                    context_embeddings[idx].detach().cpu().numpy().flatten().tolist()
+                )
+                with self.driver.session() as session:
+                    results = session.execute_write(
+                        lambda tx: tx.run(
+                            query, hash_id=hash_id, embedding=embedding
+                        ).data()
+                    )
+            return
+        offset = 0
+        while True:
+            query = f"""
+            MATCH (p:Paper)
+            WHERE p.{name} IS NOT NULL
+            RETURN p.{name} AS context, p.hash_id AS hash_id, p.title AS title
+            SKIP $offset LIMIT $batch_size
+            """
+            with self.driver.session() as session:
+                results = session.execute_write(
+                    lambda tx: tx.run(
+                        query, offset=offset, batch_size=batch_size
+                    ).data()
+                )
+            if not results:
+                break
+            if name == "abstract":
+                contexts = [result["title"] + result["context"] for result in results]
+            else:
+                contexts = [result["context"] for result in results]
+            paper_ids = [result["hash_id"] for result in results]
+            context_embeddings = embedding_model.encode(
+                contexts,
+                batch_size=batch_size,
+                convert_to_tensor=True,
+                device=self.device,
+            )
+            write_query = f"""
+            UNWIND $data AS row
+            MERGE (p:Paper {{hash_id: row.hash_id}})
+            ON CREATE SET p.{name}_embedding{postfix} = row.embedding
+            ON MATCH SET p.{name}_embedding{postfix} = row.embedding
+            """
+            data_to_write = []
+            context_embeddings = context_embeddings.detach().cpu().tolist()
+            for idx, hash_id in enumerate(paper_ids):
+                data_to_write.append({"hash_id": hash_id, "embedding": context_embeddings[idx]})
+            with self.driver.session() as session:
+                session.execute_write(
+                    lambda tx: tx.run(write_query, data=data_to_write)
+                )
+            offset += batch_size
+            logger.info(f"== Processed batch starting at offset {offset} ==")
+
     def add_paper_abstract_embedding(
         self, embedding_model, hash_id=None, batch_size=512
     ):
+        """Extract paper abstract embedding and store in the database
+        Args:
+            embedding_model (TODO: what model?): an pytorch embedding model
+            hash_id (str): add abstract embedding for a paper if hash_id is not None.
+                Otherwise, all papers will be handled with a batch size of 512
+            batch_size: if hash_id is None, all papers will be processed with `batch_size`
+        """
         if hash_id is not None:
             query = """
             MATCH (p:Paper {hash_id: $hash_id})
@@ -401,6 +636,13 @@ class PaperClient:
             logger.info(f"== Processed batch starting at offset {offset} ==")
 
     def add_paper_bg_embedding(self, embedding_model, hash_id=None, batch_size=512):
+        """Extract paper background embedding and store in the database
+        Args:
+            embedding_model (TODO: what model?): an pytorch embedding model
+            hash_id (str): add background embedding for a paper if hash_id is not None.
+                Otherwise, all papers will be handled with a batch size of 512
+            batch_size: if hash_id is None, all papers will be processed with `batch_size`
+        """
         if hash_id is not None:
             query = """
             MATCH (p:Paper {hash_id: $hash_id})
@@ -478,6 +720,13 @@ class PaperClient:
     def add_paper_contribution_embedding(
         self, embedding_model, hash_id=None, batch_size=512
     ):
+        """Extract paper contribution embedding and store in the database
+        Args:
+            embedding_model (TODO: what model?): an pytorch embedding model
+            hash_id (str): add contribution embedding for a paper if hash_id is not None.
+                Otherwise, all papers will be handled with a batch size of 512
+            batch_size: if hash_id is None, all papers will be processed with `batch_size`
+        """
         if hash_id is not None:
             query = """
             MATCH (p:Paper {hash_id: $hash_id})
@@ -555,6 +804,13 @@ class PaperClient:
     def add_paper_summary_embedding(
         self, embedding_model, hash_id=None, batch_size=512
     ):
+        """Extract paper summary embedding and store in the database
+        Args:
+            embedding_model (TODO: what model?): an pytorch embedding model
+            hash_id (str): add summary embedding for a paper if hash_id is not None.
+                Otherwise, all papers will be handled with a batch size of 512
+            batch_size: if hash_id is None, all papers will be processed with `batch_size`
+        """
         if hash_id is not None:
             query = """
             MATCH (p:Paper {hash_id: $hash_id})
@@ -567,6 +823,7 @@ class PaperClient:
                 )
             contexts = [result["context"] for result in results]
             paper_ids = [result["hash_id"] for result in results]
+            # context_embeddings are pytorch.Tensor
             context_embeddings = embedding_model.encode(
                 contexts, convert_to_tensor=True, device=self.device
             )
@@ -630,6 +887,15 @@ class PaperClient:
             logger.info(f"== Processed batch starting at offset {offset} ==")
 
     def cosine_similarity_search(self, embedding, k=1, type_name="embedding"):
+        """Retrieve all papers whose `type_name` embedding is similar to `embedding`
+        (cosine_sim > 0 and return in a descending order)
+        Args:
+            embedding (TODO: type): the embedding to be checked
+            k: only return topk papers with highest similarities
+            type_name: "abstract_embedding", "summary_embedding", etc.
+        Returns:
+            related_paper (List of str): hash_id of retrieved papers
+        """
         query = f"""
             MATCH (paper:Paper)
             WITH paper,
@@ -649,7 +915,7 @@ class PaperClient:
 
     def create_vector_index(self):
         """
-        适用于Paper节点
+        适用于Paper节点，这里的语句应该是针对所有数据库里的paper都做索引
         针对Paper节点上的是属性 embedding 进行索引
         索引向量的维度为384
         适用余弦相似度作为计算相似度的方法
@@ -666,6 +932,13 @@ class PaperClient:
             session.execute_write(lambda tx: tx.run(query).data())
 
     def filter_paper_id_list(self, paper_id_list, year="2024"):
+        """Retrieve all papers' ids which released before "year" (not contained) and existed in the database
+        Args:
+            paper_id_list (List of str): a list of paper ids
+            year: the paper before
+        Returns:
+            existing_paper_ids (List of str): paper_ids that satisfy the conditions
+        """
         if not paper_id_list:
             return []
         # WHERE p.year < "2024" AND p.venue_name <> "acl"

src/utils/paper_crawling.py

@@ -142,6 +142,19 @@ class PaperCrawling:
             print(e)
 
     def crawling(self, year, venue_name):
+        """
+        Args:
+        Returns:
+            paper_list (List of Dict):[
+                {
+                    "hash_id": hash_id, hash id of the paper
+                    "year": year, published year
+                    "venue_name": venue_name, venue name
+                    "title": title, paper title
+                    "pdf_url": pdf_url, paper url
+                }
+            ]
+        """
         paper_list = []
         paper_html_list = []
 

src/utils/paper_retriever.py

@@ -36,6 +36,9 @@ class UnionFind:
 
 
 def can_merge(uf, similarity_matrix, i, j, threshold):
+    """Condition of i and j can be merged: After merging, the similarity of any two nodes 
+    from root_i and root_j are larger than threshold
+    """
     root_i = uf.find(i)
     root_j = uf.find(j)
     for k in range(len(similarity_matrix)):
@@ -72,6 +75,8 @@ class CoCite:
             CoCite._initialized = True
 
     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 = []
@@ -82,6 +87,12 @@ class CoCite:
 
 
 class Retriever(object):
+    """The superclass of all retrievers
+    Args:
+        config: 
+    Returns:
+        A Retriever instance
+    """
     __metaclass__ = ABCMeta
     retriever_name = "BASE"
 
@@ -95,12 +106,38 @@ class Retriever(object):
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.embedding_model = get_embedding_model(config)
         self.paper_crawling = PaperCrawling(config=config)
-
+        if self.config.DEFAULT.embedding == "sentence-transformers/all-MiniLM-L6-v2":
+            self.embedding_postfix = ""
+        elif self.config.DEFAULT.embedding == "BAAI/llm-embedder":
+            self.embedding_postfix = "_llm_embedder"
+        elif self.config.DEFAULT.embedding == "jina-embeddings-v3":
+            self.embedding_postfix = "_jina_v3"
+            if self.config.DEFAULT.embedding_database == "text-matching":
+                self.embedding_postfix += "_text_matching"
+            elif self.config.DEFAULT.embedding_database == "retrieval.query":
+                self.embedding_postfix += "_query"
+            elif self.config.DEFAULT.embedding_database == "retrieval.passage":
+                self.embedding_postfix += "_passage"
     @abstractmethod
     def retrieve(self, bg, entities, use_evaluate):
+        """Retrieve papers, should be implemented by the sub-class
+        Args:
+            None
+        Returns:
+            None
+        """
         pass
 
     def retrieve_entities_by_enties(self, entities):
+        """The method do three things:
+        1. Expand entities according to entities co-occurence
+        2. Count the number of papers related to each expanded entity. Sort entities in terms of their occurence times in ascending order
+        3. Initial new entities. Retrieve entities one by one until the number of related papers reach a threshold
+        Args:
+            entities: A List of entities, e.g., [str, str, ...]
+        Returns:
+            new_entities: A List of entities after expansion, e.g., [str, str, ...]
+        """
         # TODO: KG
         expand_entities = self.paper_client.find_related_entities_by_entity_list(
             entities,
@@ -135,14 +172,16 @@ class Retriever(object):
     def update_related_paper(self, paper_id_list):
         """
         Args:
-            paper_id_list: list
+            paper_id_list (List of hash_id): e.g., [1231214, 46345]
         Return:
-            related_paper: list(dict)
+            related_paper (List of dict):
         """
         related_paper = self.paper_client.update_papers_from_client(paper_id_list)
         return related_paper
 
     def calculate_similarity(self, entities, related_entities_list, use_weight=False):
+        """[Deprecated] Calculate the similarities between two lists of entities
+        """
         if use_weight:
             vec1 = self.vectorizer.transform([" ".join(entities)]).toarray()[0]
             weighted_vec1 = np.array(
@@ -181,8 +220,21 @@ class Retriever(object):
         return similarity
 
     def cal_related_score(
-        self, embedding, related_paper_id_list, type_name="embedding"
+        self, embedding, related_paper_id_list, type_name="background_embedding"
     ):
+        """Calculate the cosine similarity between the input background's embedding and
+        given list of papers
+        Args:
+            embedding: the embedding of the input background
+            related_paper_id_list (List of int): the paper ids in the database
+        Returns:
+            Empty dict: {}
+            Empty dict: {}
+            score_all_dict: 
+                paper_id1: score1,
+                paper_id2: score2,
+                ...
+        """
         score_1 = np.zeros((len(related_paper_id_list)))
         # score_2 = np.zeros((len(related_paper_id_list)))
         origin_vector = torch.tensor(embedding).to(self.device).unsqueeze(0)
@@ -211,6 +263,15 @@ class Retriever(object):
         return {}, {}, score_all_dict
 
     def filter_related_paper(self, score_dict, top_k):
+        """Pick top_k papers from all retrieved papers in terms of score_dict. If clustering
+        is not used, top_k papers with highest scores will be picked. If clustering is used,
+        we will pick papers from each cluster in turn util top_k papers are chosen.
+        Args:
+            score_dict (dict): dict of (paper_id, similarity with user input background)
+            top_k (int): pick top_k papers
+        Returns:
+
+        """
         if len(score_dict) <= top_k:
             return list(score_dict.keys())
         if not self.use_cluster_to_filter:
@@ -221,33 +282,49 @@ class Retriever(object):
             )
             return paper_id_list
         else:
+            ## Calculate the final embedding for each paper, which is the weighted average
+            ## background_embedding (embedding), contribution_embedding, and summary_embedding.
             # clustering filter, ensure that each category the highest score save first
+            # background embedding
             paper_id_list = list(score_dict.keys())
             paper_embedding_list = [
-                self.paper_client.get_paper_attribute(paper_id, "embedding")
+                self.paper_client.get_paper_attribute(paper_id, f"background_embedding{self.embedding_postfix}")
                 for paper_id in paper_id_list
             ]
             paper_embedding = np.array(paper_embedding_list)
+            # contribution embedding
             paper_embedding_list = [
                 self.paper_client.get_paper_attribute(
-                    paper_id, "contribution_embedding"
+                    paper_id, f"contribution_embedding{self.embedding_postfix}"
                 )
                 for paper_id in paper_id_list
             ]
             paper_contribution_embedding = np.array(paper_embedding_list)
+            # summary embedding
             paper_embedding_list = [
-                self.paper_client.get_paper_attribute(paper_id, "summary_embedding")
+                self.paper_client.get_paper_attribute(paper_id, f"summary_embedding{self.embedding_postfix}")
                 for paper_id in paper_id_list
             ]
             paper_summary_embedding = np.array(paper_embedding_list)
-            weight_embedding = self.config.RETRIEVE.s_bg
+            # abstract embedding
+            paper_embedding_list = [
+                self.paper_client.get_paper_attribute(paper_id, f"abstract_embedding{self.embedding_postfix}")
+                for paper_id in paper_id_list
+            ]
+            paper_abstract_embedding = np.array(paper_embedding_list)
+
+            weight_background = self.config.RETRIEVE.s_bg
             weight_contribution = self.config.RETRIEVE.s_contribution
             weight_summary = self.config.RETRIEVE.s_summary
+            weight_abstract = self.config.RETRIEVE.s_abstract
             paper_embedding = (
-                weight_embedding * paper_embedding
+                weight_background * paper_embedding
                 + weight_contribution * paper_contribution_embedding
                 + weight_summary * paper_summary_embedding
+                + weight_abstract * paper_abstract_embedding
             )
+
+            ## similarity_matrix of all retrieved papers
             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))
@@ -257,6 +334,7 @@ class Retriever(object):
                     label_group[label] = []
                 label_group[label].append(paper_id)
             paper_id_list = []
+            # randomly pick a paper from each cluster in turn until top_k papers are chosen
             while len(paper_id_list) < top_k:
                 for label, papers in label_group.items():
                     if papers:
@@ -265,9 +343,12 @@ class Retriever(object):
                             break
             return paper_id_list
 
-    def cosine_similarity_search(self, embedding, k=1, type_name="embedding"):
-        """
-        return related paper: list
+    def cosine_similarity_search(self, embedding, k=1, type_name="background_embedding"):
+        """Retrieve papers through embedding
+        Args:
+            embedding: the input embedding
+        Returns:
+            result (List of Papers): return related papers with the least embedding distance
         """
         result = self.paper_client.cosine_similarity_search(
             embedding, k, type_name=type_name
@@ -277,6 +358,8 @@ class Retriever(object):
         return result
 
     def cluster_algorithm(self, paper_id_list, similarity_matrix):
+        """
+        """
         threshold = self.config.RETRIEVE.similarity_threshold
         uf = UnionFind(len(paper_id_list))
         # merge
@@ -298,34 +381,51 @@ class Retriever(object):
             for k in self.config.RETRIEVE.top_k_list:
                 result[k] = {"recall": 0, "precision": 0}
             return result, 0, 0, 0
+        
+        ## merge retrieved papers and target papers and clustering
+        ## clustering according to the combination of background, contribution, and summary_embedding
         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)
+        # get all target papers' background_embedding
         paper_embedding_list = [
-            self.paper_client.get_paper_attribute(paper_id, "embedding")
+            self.paper_client.get_paper_attribute(paper_id, f"background_embedding{self.embedding_postfix}")
             for paper_id in target_paper_id_list
         ]
         paper_embedding = np.array(paper_embedding_list)
+        # get all target papers' contribution_embedding
         paper_embedding_list = [
-            self.paper_client.get_paper_attribute(paper_id, "contribution_embedding")
+            self.paper_client.get_paper_attribute(paper_id, f"contribution_embedding{self.embedding_postfix}")
             for paper_id in target_paper_id_list
         ]
         paper_contribution_embedding = np.array(paper_embedding_list)
+        # get all target papers' summary_embedding
+        paper_embedding_list = [
+            self.paper_client.get_paper_attribute(paper_id, f"summary_embedding{self.embedding_postfix}")
+            for paper_id in target_paper_id_list
+        ]
+        # abstract embedding
         paper_embedding_list = [
-            self.paper_client.get_paper_attribute(paper_id, "summary_embedding")
+            self.paper_client.get_paper_attribute(paper_id, f"abstract_embedding{self.embedding_postfix}")
             for paper_id in target_paper_id_list
         ]
+        paper_abstract_embedding = np.array(paper_embedding_list)
+
         paper_summary_embedding = np.array(paper_embedding_list)
-        weight_embedding = self.config.RETRIEVE.s_bg
+        weight_background = self.config.RETRIEVE.s_bg
         weight_contribution = self.config.RETRIEVE.s_contribution
         weight_summary = self.config.RETRIEVE.s_summary
+        weight_abstract = self.config.RETRIEVE.s_abstract
+        # 2D matrix of size [# of target papers, embedding dimension]
         target_paper_embedding = (
-            weight_embedding * paper_embedding
+            weight_background * paper_embedding
             + weight_contribution * paper_contribution_embedding
             + weight_summary * paper_summary_embedding
+            + weight_abstract * paper_abstract_embedding
         )
         similarity_threshold = self.config.RETRIEVE.similarity_threshold
         similarity_matrix = np.dot(target_paper_embedding, target_paper_embedding.T)
+        # return each target_paper's cluster label
         target_labels = self.cluster_algorithm(target_paper_id_list, similarity_matrix)
         target_paper_label_dict = dict(zip(target_paper_id_list, target_labels))
         logger.debug("Target paper cluster result: {}".format(target_paper_label_dict))
@@ -335,28 +435,41 @@ class Retriever(object):
                 for paper_id in target_paper_label_dict.keys()
             }
         )
-
+        
+        ## calculate the similarity between each two papers
         all_labels = []
         for paper_id in all_paper_id_list:
+            # for each paper, get its background_embedding
             paper_bg_embedding = [
-                self.paper_client.get_paper_attribute(paper_id, "embedding")
+                self.paper_client.get_paper_attribute(paper_id, f"background_embedding{self.embedding_postfix}")
             ]
             paper_bg_embedding = np.array(paper_bg_embedding)
+            # for each paper, get its contribution_embedding
             paper_contribution_embedding = [
                 self.paper_client.get_paper_attribute(
-                    paper_id, "contribution_embedding"
+                    paper_id, f"contribution_embedding{self.embedding_postfix}"
                 )
             ]
             paper_contribution_embedding = np.array(paper_contribution_embedding)
+            # for each paper, get its summary_embedding
             paper_summary_embedding = [
-                self.paper_client.get_paper_attribute(paper_id, "summary_embedding")
+                self.paper_client.get_paper_attribute(paper_id, f"summary_embedding{self.embedding_postfix}")
             ]
             paper_summary_embedding = np.array(paper_summary_embedding)
+            # for each paper, get its abstract_embedding
+            paper_abstract_embedding = [
+                self.paper_client.get_paper_attribute(paper_id, f"abstract_embedding{self.embedding_postfix}")
+            ]
+            paper_abstract_embedding = np.array(paper_abstract_embedding)
+            
             paper_embedding = (
-                weight_embedding * paper_bg_embedding
+                weight_background * paper_bg_embedding
                 + weight_contribution * paper_contribution_embedding
                 + weight_summary * paper_summary_embedding
+                + weight_abstract * paper_abstract_embedding
             )
+
+            # vector of size embedding dimension
             similarities = cosine_similarity(paper_embedding, target_paper_embedding)[0]
             if np.any(similarities >= similarity_threshold):
                 all_labels.append(target_labels[np.argmax(similarities)])
@@ -364,14 +477,15 @@ class Retriever(object):
                 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}")
+        logger.debug(f"All labels and the number of papers of each label: {all_label_counts}")
         target_label_counts = Counter(target_paper_label_dict.values())
-        logger.debug(f"target label counts : {target_label_counts}")
+        logger.debug(f"All labels and the number of target papers of each label : {target_label_counts}")
         target_label_list = list(target_label_counts.keys())
         max_k = max(self.config.RETRIEVE.top_k_list)
         logger.info("=== Begin filter related paper ===")
         max_k_paper_id_list = self.filter_related_paper(score_all_dict, top_k=max_k)
         logger.info("=== End filter related paper ===")
+        ## calculate recall and precision of first {10, 20, 30, ...} papers
         for k in self.config.RETRIEVE.top_k_list:
             # 前top k 的文章
             top_k = min(k, len(max_k_paper_id_list))
@@ -380,7 +494,7 @@ class Retriever(object):
             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)
+                "=== ideal top {}, real top {} paper id list : {}".format(k, top_k, top_k_paper_label_dict)
             )
             logger.debug(
                 {
@@ -389,7 +503,7 @@ class Retriever(object):
                 }
             )
             top_k_label_counts = Counter(top_k_paper_label_dict.values())
-            logger.debug(f"top K label counts : {top_k_label_counts}")
+            logger.debug(f"Retrieved {top_k} papers have K different label: {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}")
@@ -403,6 +517,7 @@ class Retriever(object):
             precision /= len(top_k_paper_id_list)
             result[k] = {"recall": recall, "precision": precision}
 
+        ## calculate recall and precision of all retrieved papers
         related_paper_id_list = list(score_all_dict.keys())
         related_paper_label_dict = {}
         for paper_id in related_paper_id_list:
@@ -419,11 +534,19 @@ class Retriever(object):
             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):
+    """RetrieverFactory is a singleton class, which will return cls._instance if it has been 
+    created, it saves all Retriever instances.
+    Args:
+        None
+    Returns:
+        The singleton instance of the RetrieverFactory
+    """
     _instance = None
     _lock = threading.Lock()
 
@@ -441,11 +564,24 @@ class RetrieverFactory(object):
 
     @staticmethod
     def get_retriever_factory():
+        """The method can also return the singleton instance of the RetrieverFactory
+        Args:
+            None
+        Returns:
+            The singleton instance of the RetrieverFactory
+        """
         if RetrieverFactory._instance is None:
             RetrieverFactory._instance = RetrieverFactory()
         return RetrieverFactory._instance
 
     def register_retriever(self, retriever_name, retriever_class) -> bool:
+        """Register a new retriever class (not instance) to the RetrieverFactory
+        Args:
+            retriever_name: str
+            retriever_class: a class object (not instance)
+        Returns:
+            True if add successfully, False otherwise
+        """
         if retriever_name not in self.retriever_classes:
             self.retriever_classes[retriever_name] = retriever_class
             return True
@@ -467,8 +603,14 @@ class RetrieverFactory(object):
         return len(self.retriever_classes)
 
     def create_retriever(self, retriever_name, *args, **kwargs) -> Retriever:
+        """Return a retriever instance
+        Args:
+            retriever_name: str
+        Returns:
+            The retriever
+        """
         if retriever_name not in self.retriever_classes:
-            raise ValueError(f"Unknown retriever type: {retriever_name}")
+            raise ValueError(f"Unknown retriever type: {retriever_name}. retriever_name should be one of {self.retriever_classes.keys()}")
         else:
             return self.retriever_classes[retriever_name](*args, **kwargs)
 
@@ -493,11 +635,24 @@ class SNRetriever(Retriever):
         super().__init__(config)
 
     def retrieve_paper(self, bg):
+        """Retrieve papers P (a set) according to embeddings' similarity between the input 
+        background and the backgrounds from the database. Optionally, you can also retrieve 
+        papers co-cited with P.
+        Args:
+            bg (str): the input background
+        Returns:
+            result (dict):
+                "background_embedding": embedding of the input background,
+                "paper" (List of int): all retrieved related_papers' ids,
+                "entities" (List): An empty list (TODO: remove),
+                "cocite_paper" (List of int): all papers cocited with embedding-retrieved papers
+        """
         entities = []
         embedding = self.embedding_model.encode(bg, device=self.device)
         sn_paper_id_list = self.cosine_similarity_search(
             embedding=embedding,
             k=self.config.RETRIEVE.sn_retrieve_paper_num,
+            type_name=f"{self.config.RETRIEVE.SN_field_name}_embedding{self.embedding_postfix}"
         )
         related_paper = set()
         related_paper.update(sn_paper_id_list)
@@ -513,7 +668,7 @@ class SNRetriever(Retriever):
         related_paper = list(related_paper)
         logger.debug(f"paper num before filter: {len(related_paper)}")
         result = {
-            "embedding": embedding,
+            f"background_embedding{self.embedding_postfix}": embedding,
             "paper": related_paper,
             "entities": entities,
             "cocite_paper": list(cocite_id_set),
@@ -523,9 +678,21 @@ class SNRetriever(Retriever):
     def retrieve(self, bg, entities, need_evaluate=True, target_paper_id_list=[]):
         """
         Args:
-            context: string
+            bg (str): The user input background
         Return:
-            list(dict)
+            result (dict):
+                "recall": recall of paper retrieval, 0 if need_evaluate==False,
+                "precision": precision of paper retrieval, 0 if need_evaluate==False,,
+                "filtered_recall": recall of paper retrieval after filtering, 0 if need_evaluate==False,,
+                "filtered_precision": precision of paper retrieval after filtering, 0 if need_evaluate==False,,
+                "related_paper": all retrieved related_papers. !!! [ The most important item ]
+                "related_paper_id_list": all retrieved related_papers' ids. !!! [ The most important item ]
+                "cocite_paper_id_list": retrieve_result["cocite_paper"],
+                "entities": retrieve_result["entities"], always empty
+                "top_k_matrix": top_k_matrix, 0 if need_evaluate==False
+                "gt_reference_num": len(target_paper_id_list)
+                "retrieve_paper_num": len(related_paper_id_list),
+                "label_num": TODO,
         """
         if need_evaluate:
             if target_paper_id_list is None or len(target_paper_id_list) == 0:
@@ -537,23 +704,27 @@ class SNRetriever(Retriever):
         retrieve_result = self.retrieve_paper(bg)
         related_paper_id_list = retrieve_result["paper"]
         retrieve_paper_num = len(related_paper_id_list)
+        # scores between the input background and all retrieved papers
         _, _, score_all_dict = self.cal_related_score(
-            retrieve_result["embedding"], related_paper_id_list=related_paper_id_list
+            retrieve_result[f"background_embedding{self.embedding_postfix}"], related_paper_id_list=related_paper_id_list,
+            type_name=f"{self.config.RETRIEVE.SN_field_name}_embedding{self.embedding_postfix}"
         )
         top_k_matrix = {}
         recall = 0
         precision = 0
         filtered_recall = 0
+        label_num = 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("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=10)
+        ## For idea generation, only top 10 papers will be used, which has no relations with retriveal evaluation
+        related_paper = self.filter_related_paper(score_all_dict, top_k=self.config.RETRIEVE.all_retrieve_paper_num)
         related_paper = self.update_related_paper(related_paper)
         result = {
             "recall": recall,
@@ -578,6 +749,8 @@ class KGRetriever(Retriever):
         super().__init__(config)
 
     def retrieve_paper(self, entities):
+        """Retrieve according to entities
+        """
         new_entities = self.retrieve_entities_by_enties(entities)
         logger.debug("KG entities for retriever: {}".format(new_entities))
         related_paper = set()
@@ -618,12 +791,14 @@ class KGRetriever(Retriever):
         retrieve_paper_num = len(related_paper_id_list)
         embedding = self.embedding_model.encode(bg, device=self.device)
         _, _, score_all_dict = self.cal_related_score(
-            embedding, related_paper_id_list=related_paper_id_list
+            embedding, related_paper_id_list=related_paper_id_list,
+            type_name=f"background_embedding{self.embedding_postfix}"
         )
         top_k_matrix = {}
         recall = 0
         precision = 0
         filtered_recall = 0
+        label_num = 0
         filtered_precision = 0
         if need_evaluate:
             top_k_matrix, label_num, recall, precision = self.eval_related_paper_in_all(
@@ -633,7 +808,7 @@ class KGRetriever(Retriever):
             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.filter_related_paper(score_all_dict, top_k=self.config.RETRIEVE.all_retrieve_paper_num)
         related_paper = self.update_related_paper(related_paper)
         result = {
             "recall": recall,
@@ -659,28 +834,44 @@ class SNKGRetriever(Retriever):
 
     def retrieve_paper(self, bg, entities):
         sn_entities = []
+        ## 1. Retrieve papers according to the embeddings of input background
         embedding = self.embedding_model.encode(bg, device=self.device)
         sn_paper_id_list = self.cosine_similarity_search(
-            embedding, k=self.config.RETRIEVE.sn_num_for_entity
+            embedding, k=self.config.RETRIEVE.sn_num_for_entity,
+            type_name=f"{self.config.RETRIEVE.SN_field_name}_embedding{self.embedding_postfix}"
         )
         related_paper = set()
         related_paper.update(sn_paper_id_list)
+        logger.debug(f"SN retrieve {len(related_paper)} papers")
+
+        ## 2. Retrieve papers according to entites
+        # Fetch all entities from embedding-retrieved papers
         sn_entities += self.paper_client.find_entities_by_paper_list(sn_paper_id_list)
         logger.debug("SN entities for retriever: {}".format(sn_entities))
         entities = list(set(entities + sn_entities))
+        # Expand entity list through synonyms
         new_entities = self.retrieve_entities_by_enties(entities)
         logger.debug("SNKG entities for retriever: {}".format(new_entities))
+        paper_id_set = 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)
+            paper_id_set.update(self.paper_client.find_paper_by_entity(entity))
+        related_paper = related_paper.union(paper_id_set)
+        logger.debug(f"Entity retrieve {len(paper_id_set)} papers")
+        logger.debug(f"SN+entity retrieve {len(related_paper)} papers")
+
+        ## 3. Retrieve papers according to citation co-occurrence
         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)
+        logger.debug(f"Cocite retrieve {len(cocite_id_set)} papers")
+        logger.debug(f"SN+entity+cocite retrieve {len(related_paper)} papers")
+
+        ## 4. Return retrieval results
         related_paper = list(related_paper)
         result = {
-            "embedding": embedding,
+            f"background_embedding{self.embedding_postfix}": embedding,
             "paper": related_paper,
             "entities": entities,
             "cocite_paper": list(cocite_id_set),
@@ -688,7 +879,7 @@ class SNKGRetriever(Retriever):
         return result
 
     def retrieve(
-        self, bg, entities, need_evaluate=True, target_paper_id_list=[], top_k=10
+        self, bg, entities, need_evaluate=True, target_paper_id_list=[]
     ):
         """
         Args:
@@ -709,7 +900,8 @@ class SNKGRetriever(Retriever):
         retrieve_paper_num = len(related_paper_id_list)
         logger.info("=== Begin cal related paper score ===")
         _, _, score_all_dict = self.cal_related_score(
-            retrieve_result["embedding"], related_paper_id_list=related_paper_id_list
+            retrieve_result[f"background_embedding{self.embedding_postfix}"], related_paper_id_list=related_paper_id_list,
+            type_name=f"background_embedding{self.embedding_postfix}"
         )
         logger.info("=== End cal related paper score ===")
         top_k_matrix = {}
@@ -727,7 +919,7 @@ class SNKGRetriever(Retriever):
             logger.debug(f"Recall: {recall:.3f}")
             logger.debug(f"Precision: {precision:.3f}")
         logger.info("=== Begin filter related paper score ===")
-        related_paper = self.filter_related_paper(score_all_dict, top_k)
+        related_paper = self.filter_related_paper(score_all_dict, self.config.RETRIEVE.all_retrieve_paper_num)
         logger.info("=== End filter related paper score ===")
         related_paper = self.update_related_paper(related_paper)
         result = {