Add references generation.

app.py

@@ -3,6 +3,7 @@ import os
 import openai
 from auto_backgrounds import generate_backgrounds, generate_draft
 from utils.file_operations import hash_name
+from references_generator import generate_top_k_references
 
 # note: App白屏bug：允许第三方cookie
 # todo:
@@ -48,6 +49,9 @@ else:
 def clear_inputs(*args):
     return "", ""
 
+def clear_inputs_refs(*args):
+    return "", 5
+
 
 def wrapped_generator(paper_title, paper_description, openai_api_key=None,
                       paper_template="ICLR2022", tldr=True, max_num_refs=50, selected_sections=None, bib_refs=None, model="gpt-4",
@@ -91,6 +95,11 @@ def wrapped_generator(paper_title, paper_description, openai_api_key=None,
         return output
 
 
+def wrapped_references_generator(paper_title, num_refs):
+    return generate_top_k_references(paper_title, top_k=num_refs)
+
+
+
 theme = gr.themes.Default(font=gr.themes.GoogleFont("Questrial"))
 # .set(
 #     background_fill_primary='#E5E4E2',
@@ -105,6 +114,14 @@ ACADEMIC_PAPER = """## 一键生成论文初稿
 3. 在右侧下载.zip格式的输出，在Overleaf上编译浏览.  
 """
 
+
+REFERENCES = """## 一键搜索相关论文
+
+1. 在Title文本框中输入想要搜索文献的论文（比如Playing Atari with Deep Reinforcement Learning). 
+2. 点击Submit. 等待大概十分钟. 
+3. 在右侧JSON处会显示相关文献.  
+"""
+
 with gr.Blocks(theme=theme) as demo:
     gr.Markdown('''
     # Auto-Draft: 文献整理辅助工具
@@ -176,23 +193,22 @@ with gr.Blocks(theme=theme) as demo:
                     clear_button_pp = gr.Button("Clear")
                     submit_button_pp = gr.Button("Submit", variant="primary")
 
-            with gr.Tab("文献综述"):
-                gr.Markdown('''
-                <h1  style="text-align: center;">Coming soon!</h1>
-                ''')
-                # topic = gr.Textbox(value="Deep Reinforcement Learning", lines=1, max_lines=1,
-                #                    label="Topic", info="文献主题")
-                # with gr.Accordion("Advanced Setting"):
-                #     description_lr = gr.Textbox(lines=5, label="Description (Optional)", visible=True,
-                #                              info="对希望生成的综述的一些描述. 包括这篇论文的创新点, 主要贡献, 等.")
-                # with gr.Row():
-                #     clear_button_lr = gr.Button("Clear")
-                #     submit_button_lr = gr.Button("Submit", variant="primary")
-            with gr.Tab("论文润色"):
+            with gr.Tab("文献搜索 (NEW!)"):
+                gr.Markdown(REFERENCES)
+
+                title_refs = gr.Textbox(value="Playing Atari with Deep Reinforcement Learning", lines=1, max_lines=1,
+                                   label="Title", info="论文标题")
+                slider_refs = gr.Slider(minimum=1, maximum=100, value=5, step=1,
+                                   interactive=True, label="最相关的参考文献数目")
+                with gr.Row():
+                    clear_button_refs = gr.Button("Clear")
+                    submit_button_refs = gr.Button("Submit", variant="primary")
+
+            with gr.Tab("文献综述 (Coming soon!)"):
                 gr.Markdown('''
                 <h1  style="text-align: center;">Coming soon!</h1>
                 ''')
-            with gr.Tab("帮我想想该写什么论文!"):
+            with gr.Tab("Github文档 (Coming soon!)"):
                 gr.Markdown('''
                 <h1  style="text-align: center;">Coming soon!</h1>
                 ''')
@@ -207,13 +223,16 @@ with gr.Blocks(theme=theme) as demo:
              当`Cache`显示AVAILABLE的时候, 所有的输入和输出会被备份到我的云储存中. 显示NOT AVAILABLE的时候不影响实际使用. 
             `OpenAI API`: <span style="{style_mapping[IS_OPENAI_API_KEY_AVAILABLE]}">{availability_mapping[IS_OPENAI_API_KEY_AVAILABLE]}</span>.  `Cache`: <span style="{style_mapping[IS_CACHE_AVAILABLE]}">{availability_mapping[IS_CACHE_AVAILABLE]}</span>.''')
             file_output = gr.File(label="Output")
+            json_output = gr.JSON(label="References")
 
     clear_button_pp.click(fn=clear_inputs, inputs=[title, description_pp], outputs=[title, description_pp])
-    # submit_button_pp.click(fn=wrapped_generator,
-    # inputs=[title, description_pp, key, template, tldr, slider, sections, bibtex_file], outputs=file_output)
     submit_button_pp.click(fn=wrapped_generator,
                            inputs=[title, description_pp, key, template, tldr_checkbox, slider, sections, bibtex_file,
                                    model_selection], outputs=file_output)
 
+    clear_button_refs.click(fn=clear_inputs_refs, inputs=[title_refs, slider_refs], outputs=[title_refs, slider_refs])
+    submit_button_refs.click(fn=wrapped_references_generator,
+                           inputs=[title_refs, slider_refs], outputs=json_output)
+
 demo.queue(concurrency_count=1, max_size=5, api_open=False)
 demo.launch()

latex_templates/ICLR2022/template.tex

@@ -7,7 +7,7 @@
 \usepackage{hyperref}
 \usepackage{url} 
 \usepackage{algorithm}
-\usepackage{algorithmic}
+\usepackage{algpseudocode}
 
 \title{TITLE} 
 \author{GPT-4}

references_generator.py

@@ -0,0 +1,73 @@
+import os.path
+import json
+from utils.references import References
+from section_generator import section_generation_bg, keywords_generation, figures_generation, section_generation
+import itertools
+from gradio_client import Client
+
+def generate_raw_references(title, description="",
+                            bib_refs=None, tldr=False, max_kw_refs=10,  save_to="ref.bib"):
+    # load pre-provided references
+    ref = References(title, bib_refs)
+
+    # generate multiple keywords for searching
+    input_dict = {"title": title, "description": description}
+    keywords, usage = keywords_generation(input_dict)
+    keywords = list(keywords)
+    comb_keywords = list(itertools.combinations(keywords, 2))
+    for comb_keyword in comb_keywords:
+        keywords.append(" ".join(comb_keyword))
+    keywords = {keyword:max_kw_refs for keyword in keywords}
+    print(f"keywords: {keywords}\n\n")
+
+    ref.collect_papers(keywords, tldr=tldr)
+    paper_json = ref.to_json()
+
+    with open(save_to, "w") as f:
+        json.dump(paper_json, f)
+
+    return save_to, paper_json
+
+def generate_top_k_references(title, description="",
+                            bib_refs=None, tldr=False, max_kw_refs=10,  save_to="ref.bib", top_k=5):
+    json_path, json_content = generate_raw_references(title, description, bib_refs, tldr, max_kw_refs,  save_to)
+
+    client = Client("https://shaocongma-evaluate-specter-embeddings.hf.space/")
+    result = client.predict(
+        title,  # str  in 'Title' Textbox component
+        json_path,  # str (filepath or URL to file) in 'Papers JSON (as string)' File component
+        top_k,  # int | float (numeric value between 1 and 50) in 'Top-k Relevant Papers' Slider component
+        api_name="/get_k_relevant_papers"
+    )
+    with open(result) as f:
+        result = json.load(f)
+    return result
+
+if __name__ == "__main__":
+    import openai
+    openai.api_key = os.getenv("OPENAI_API_KEY")
+
+    title = "Using interpretable boosting algorithms for modeling environmental and agricultural data"
+    description = ""
+    save_to = "paper.json"
+    save_to, paper_json = generate_raw_references(title, description, save_to=save_to)
+
+    print("`paper.json` has been generated. Now evaluating its similarity...")
+
+    k = 5
+    client = Client("https://shaocongma-evaluate-specter-embeddings.hf.space/")
+    result = client.predict(
+        title,  # str  in 'Title' Textbox component
+        save_to,  # str (filepath or URL to file) in 'Papers JSON (as string)' File component
+        k,  # int | float (numeric value between 1 and 50) in 'Top-k Relevant Papers' Slider component
+        api_name="/get_k_relevant_papers"
+    )
+
+    with open(result) as f:
+        result = json.load(f)
+
+    print(result)
+
+    save_to = "paper2.json"
+    with open(save_to, "w") as f:
+        json.dump(result, f)

section_generator.py

@@ -90,7 +90,7 @@ def keywords_generation(input_dict):
     attempts_count = 0
     while attempts_count < max_attempts:
         try:
-            keywords, usage= get_gpt_responses(KEYWORDS_SYSTEM.format(min_refs_num=3, max_refs_num=5), title,
+            keywords, usage= get_gpt_responses(KEYWORDS_SYSTEM.format(min_refs_num=1, max_refs_num=10), title,
                                      model="gpt-3.5-turbo", temperature=0.4)
             print(keywords)
             output = json.loads(keywords)