Add webui (testing).

app.py

@@ -0,0 +1,37 @@
+import gradio as gr
+import openai
+from auto_backgrounds import generate_backgrounds
+
+# todo: 1. remove repeated entry in bibfile (go to references.py)
+#       2. (maybe) multiple commas error (see Overleaf)
+#       3. create a huggingface space. test it using multiple devices!
+#       4. further polish auto_backgrounds.py. Make backgrounds have multiple subsection.
+#       5. Design a good layout of huggingface space.
+def generate_backgrounds(t1, t2):
+    return "README.md"
+
+def clear_inputs(text1, text2):
+    return ("", "")
+
+with gr.Blocks() as demo:
+    gr.Markdown('''
+    # Auto-Draft: 论文结构辅助工具
+    
+    用法: 输入任意论文标题, 点击Submit, 等待大概十分钟, 下载output.zip.  
+    ''')
+    with gr.Row():
+        with gr.Column():
+            title = gr.Textbox(value="Playing Atari Game with Deep Reinforcement Learning", lines=1, max_lines=1, label="Title")
+            description = gr.Textbox(lines=5, label="Description (Optional)")
+
+            with gr.Row():
+                clear_button = gr.Button("Clear")
+                submit_button = gr.Button("Submit")
+        with gr.Column():
+            file_output = gr.outputs.File()
+
+    clear_button.click(fn=clear_inputs, inputs=[title, description], outputs=[title, description])
+    submit_button.click(fn=generate_backgrounds, inputs=[title, description], outputs=file_output)
+
+demo.queue(concurrency_count=1, max_size=5, api_open=False)
+demo.launch()

auto_backgrounds.py

@@ -6,6 +6,7 @@ import datetime
 import shutil
 import time
 import logging
+import os
 
 TOTAL_TOKENS = 0
 TOTAL_PROMPTS_TOKENS = 0
@@ -31,6 +32,16 @@ def log_usage(usage, generating_target, print_out=True):
         print(message)
     logging.info(message)
 
+def make_archive(source, destination):
+    base = os.path.basename(destination)
+    name = base.split('.')[0]
+    format = base.split('.')[1]
+    archive_from = os.path.dirname(source)
+    archive_to = os.path.basename(source.strip(os.sep))
+    shutil.make_archive(name, format, archive_from, archive_to)
+    shutil.move('%s.%s'%(name,format), destination)
+    return destination
+
 def pipeline(paper, section, save_to_path, model):
     """
     The main pipeline of generating a section.
@@ -55,10 +66,10 @@ def pipeline(paper, section, save_to_path, model):
             f.write(r"\end{abstract}")
     else:
         with open(tex_file, "w") as f:
-            f.write(f"\section{{{section}}}\n")
+            f.write(f"\section{{{section.upper()}}}\n")
         with open(tex_file, "a") as f:
             f.write(output)
-    time.sleep(20)
+    time.sleep(5)
     print(f"{section} has been generated. Saved to {tex_file}.")
     return usage
 
@@ -106,6 +117,8 @@ def generate_backgrounds(title, description="", template="ICLR2022", model="gpt-
         except Exception as e:
             print(f"Failed to generate {section} due to the error: {e}")
     print(f"The paper {title} has been generated. Saved to {save_to_path}.")
+    # shutil.make_archive("output.zip", 'zip', save_to_path)
+    return make_archive(save_to_path, save_to_path+"output.zip")
 
 if __name__ == "__main__":
     title = "Reinforcement Learning"

auto_draft.py

@@ -7,11 +7,22 @@ import datetime
 import shutil
 import time
 import logging
+import os
 
 TOTAL_TOKENS = 0
 TOTAL_PROMPTS_TOKENS = 0
 TOTAL_COMPLETION_TOKENS = 0
 
+def make_archive(source, destination):
+    base = os.path.basename(destination)
+    name = base.split('.')[0]
+    format = base.split('.')[1]
+    archive_from = os.path.dirname(source)
+    archive_to = os.path.basename(source.strip(os.sep))
+    shutil.make_archive(name, format, archive_from, archive_to)
+    shutil.move('%s.%s'%(name,format), destination)
+    return destination
+
 
 def log_usage(usage, generating_target, print_out=True):
     global TOTAL_TOKENS
@@ -59,7 +70,7 @@ def pipeline(paper, section, save_to_path, model):
             f.write(f"\section{{{section}}}\n")
         with open(tex_file, "a") as f:
             f.write(output)
-    time.sleep(20)
+    time.sleep(5)
     print(f"{section} has been generated. Saved to {tex_file}.")
     return usage
 
@@ -121,6 +132,7 @@ def generate_draft(title, description="", template="ICLR2022", model="gpt-4"):
         except Exception as e:
             print(f"Failed to generate {section} due to the error: {e}")
     print(f"The paper {title} has been generated. Saved to {save_to_path}.")
+    return make_archive(save_to_path, save_to_path+"output.zip")
 
 if __name__ == "__main__":
     # title = "Training Adversarial Generative Neural Network with Adaptive Dropout Rate"

utils/references.py

@@ -49,7 +49,14 @@ def _collect_papers_arxiv(keyword, counts=3):
             # Extract the year
             published = entry.find(f"{namespace}published").text
             year = published.split("-")[0]
-            arxiv_id = re.search(r'\d+\.\d+', link).group(0)
+
+            founds = re.search(r'\d+\.\d+', link)
+            if founds is None:
+                # some links are not standard; such as "https://arxiv.org/abs/cs/0603127v1".
+                # will be solved in the future.
+                continue
+            else:
+                arxiv_id = founds.group(0)
             journal = f"arXiv preprint arXiv:{arxiv_id}"
             result = {
                 "paper_id": arxiv_id,
@@ -93,6 +100,17 @@ class References:
         for key, counts in keywords_dict.items():
             self.papers = self.papers + process(key, counts)
 
+        # TODO: remove repeated entries
+        # test this
+        seen = set()
+        papers = []
+        for paper in self.papers:
+            paper_id = paper["paper_id"]
+            if paper_id not in seen:
+                seen.add(paper_id)
+                papers.append(paper)
+        self.papers = papers
+
     def to_bibtex(self, path_to_bibtex="ref.bib"):
         """
         Turn the saved paper list into bibtex file "ref.bib". Return a list of all `paper_id`.