fix knowledge database error.

support OPENAI_API_BASE.
backup all generated files when possible.

app.py

@@ -1,27 +1,44 @@
+import uuid
 import gradio as gr
 import os
 import openai
 from auto_backgrounds import generate_backgrounds, generate_draft
-from utils.file_operations import hash_name, list_folders
-from references_generator import generate_top_k_references
+from utils.file_operations import list_folders, urlify
+from huggingface_hub import snapshot_download
 
 # todo:
 #   6. get logs when the procedure is not completed. *
 #   7. 自己的文件库； 更多的prompts
 #   2. 实现别的功能
-#   3. Check API Key GPT-4 Support.
 # future:
 #   generation.log sometimes disappears (ignore this)
 #   1. Check if there are any duplicated citations
 #   2. Remove potential thebibliography and bibitem in .tex file
 
+#######################################################################################################################
+# Environment Variables
+#######################################################################################################################
+# OPENAI_API_KEY: OpenAI API key for GPT models
+# OPENAI_API_BASE: (Optional) Support alternative OpenAI minors
+# GPT4_ENABLE: (Optional) Set it to 1 to enable GPT-4 model.
+
+# AWS_ACCESS_KEY_ID: (Optional) Access AWS cloud storage (you need to edit `BUCKET_NAME` in `utils/storage.py` if you need to use this function)
+# AWS_SECRET_ACCESS_KEY: (Optional) Access AWS cloud storage (you need to edit `BUCKET_NAME` in `utils/storage.py` if you need to use this function)
+# KDB_REPO: (Optional) A Huggingface dataset hosting Knowledge Databases
+# HF_TOKEN: (Optional) Access to KDB_REPO
+
 #######################################################################################################################
 # Check if openai and cloud storage available
 #######################################################################################################################
 openai_key = os.getenv("OPENAI_API_KEY")
+openai_api_base = os.getenv("OPENAI_API_BASE")
+if openai_api_base is not None:
+    openai.api_base = openai_api_base
+GPT4_ENABLE = os.getenv("GPT4_ENABLE") # disable GPT-4 for public repo
+
 access_key_id = os.getenv('AWS_ACCESS_KEY_ID')
 secret_access_key = os.getenv('AWS_SECRET_ACCESS_KEY')
-GPT4_ENABLE = os.getenv("GPT4_ENABLE") # by default None.
+
 if access_key_id is None or secret_access_key is None:
     print("Access keys are not provided. Outputs cannot be saved to AWS Cloud Storage.\n")
     IS_CACHE_AVAILABLE = False
@@ -48,6 +65,13 @@ DEFAULT_SECTIONS = ["introduction", "related works", "backgrounds", "methodology
 
 MODEL_LIST = ['gpt-4', 'gpt-3.5-turbo', 'gpt-3.5-turbo-16k']
 
+HF_TOKEN = os.getenv("HF_TOKEN")
+REPO_ID = os.getenv("KDB_REPO")
+if HF_TOKEN is not None and REPO_ID is not None:
+    snapshot_download(REPO_ID, repo_type="dataset", local_dir="knowledge_databases/",
+                      local_dir_use_symlinks=False, token=HF_TOKEN)
+    KDB_LIST = ["(None)"] + list_folders("knowledge_databases")
+
 #######################################################################################################################
 # Load the list of templates & knowledge databases
 #######################################################################################################################
@@ -140,7 +164,6 @@ def clear_inputs(*args):
 def clear_inputs_refs(*args):
     return "", 5
 
-
 def wrapped_generator(
         paper_title, paper_description,  # main input
         openai_api_key=None, openai_url=None,  # key
@@ -149,10 +172,8 @@ def wrapped_generator(
         paper_template="ICLR2022", selected_sections=None, model="gpt-4", prompts_mode=False,  # outputs parameters
         cache_mode=IS_CACHE_AVAILABLE  # handle cache mode
 ):
-    # if `cache_mode` is True, then follow the following steps:
-    #        check if "title"+"description" have been generated before
-    #        if so, download from the cloud storage, return it
-    #        if not, generate the result.
+    # if `cache_mode` is True, then always upload the generated content to my S3.
+    file_name_upload = urlify(paper_title) + "_" + uuid.uuid1().hex + ".zip"
     if bib_refs is not None:
         bib_refs = bib_refs.name
     if openai_api_key is not None:
@@ -161,19 +182,6 @@ def wrapped_generator(
             openai.Model.list()
         except Exception as e:
             raise gr.Error(f"Key错误. Error: {e}")
-
-    if cache_mode:
-        from utils.storage import list_all_files, download_file
-        # check if "title"+"description" have been generated before
-        input_dict = {"title": paper_title, "description": paper_description,
-                      "generator": "generate_draft"}
-        file_name = hash_name(input_dict) + ".zip"
-        file_list = list_all_files()
-        # print(f"{file_name} will be generated. Check the file list {file_list}")
-        if file_name in file_list:
-            # download from the cloud storage, return it
-            download_file(file_name)
-            return file_name
     try:
         output = generate_draft(
             paper_title, description=paper_description, # main input
@@ -183,19 +191,12 @@ def wrapped_generator(
            )
         if cache_mode:
             from utils.storage import upload_file
-            upload_file(output)
+            upload_file(output, target_name=file_name_upload)
     except Exception as e:
         raise gr.Error(f"生成失败. Error: {e}")
     return output
 
 
-def wrapped_references_generator(paper_title, num_refs, openai_api_key=None):
-    if openai_api_key is not None:
-        openai.api_key = openai_api_key
-        openai.Model.list()
-    return generate_top_k_references(paper_title, top_k=num_refs)
-
-
 with gr.Blocks(theme=theme) as demo:
     gr.Markdown(ANNOUNCEMENT)
 
@@ -271,9 +272,6 @@ with gr.Blocks(theme=theme) as demo:
                             max_tokens_kd_slider = gr.Slider(minimum=256, maximum=8192, value=2048, step=2,
                                                       interactive=True, label="MAX_TOKENS",
                                                       info="知识库内容占用Prompts中的Token数")
-                            # template = gr.Dropdown(label="Template", choices=ALL_TEMPLATES, value="Default",
-                            #                        interactive=True,
-                            #                        info="生成论文的参考模板.")
                             domain_knowledge = gr.Dropdown(label="预载知识库",
                                                            choices=ALL_DATABASES,
                                                            value="(None)",
@@ -283,18 +281,6 @@ with gr.Blocks(theme=theme) as demo:
                 with gr.Row():
                     clear_button_pp = gr.Button("Clear")
                     submit_button_pp = gr.Button("Submit", variant="primary")
-
-            # with gr.Tab("文献搜索"):
-            #     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>
@@ -308,16 +294,6 @@ with gr.Blocks(theme=theme) as demo:
             gr.Markdown(STATUS)
             file_output = gr.File(label="Output")
             json_output = gr.JSON(label="References")
-
-
-    # def wrapped_generator(
-    #         paper_title, paper_description,  # main input
-    #         openai_api_key=None, openai_url=None,  # key
-    #         tldr=True, max_kw_refs=10, bib_refs=None, max_tokens_ref=2048,  # references
-    #         knowledge_database=None, max_tokens_kd=2048, query_counts=10,  # domain knowledge
-    #         paper_template="ICLR2022", selected_sections=None, model="gpt-4", prompts_mode=False,  # outputs parameters
-    #         cache_mode=IS_CACHE_AVAILABLE  # handle cache mode
-    # ):
     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, url,
@@ -325,9 +301,5 @@ with gr.Blocks(theme=theme) as demo:
                                    domain_knowledge, max_tokens_kd_slider, query_counts_slider,
                                    template, sections, model_selection, prompts_mode], 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, key], outputs=json_output)
-
 demo.queue(concurrency_count=1, max_size=5, api_open=False)
 demo.launch(show_error=True)

auto_backgrounds.py

@@ -295,6 +295,7 @@ if __name__ == "__main__":
     import openai
 
     openai.api_key = os.getenv("OPENAI_API_KEY")
+    openai.api_base = os.getenv("OPENAI_API_BASE")
 
     target_title = "Playing Atari with Decentralized Reinforcement Learning"
     output = generate_draft(target_title, knowledge_database="ml_textbook_test")

kdb_test.py

@@ -0,0 +1,62 @@
+from utils.knowledge import Knowledge
+from langchain.vectorstores import FAISS
+from utils.file_operations import list_folders
+from huggingface_hub import snapshot_download
+import gradio as gr
+import os
+import json
+from models import EMBEDDINGS
+
+HF_TOKEN = os.getenv("HF_TOKEN")
+REPO_ID = os.getenv("KDB_REPO")
+
+snapshot_download(REPO_ID, repo_type="dataset", local_dir="knowledge_databases/",
+                  local_dir_use_symlinks=False, token=HF_TOKEN)
+ALL_KDB = ["(None)"] + list_folders("knowledge_databases")
+
+
+
+def query_from_kdb(input, kdb, query_counts):
+    if kdb == "(None)":
+        return {"knowledge_database": "(None)", "input": input, "output": ""}, ""
+
+    db_path = f"knowledge_databases/{kdb}"
+    db_config_path = os.path.join(db_path, "db_meta.json")
+    db_index_path = os.path.join(db_path, "faiss_index")
+    if os.path.isdir(db_path):
+        # load configuration file
+        with open(db_config_path, "r", encoding="utf-8") as f:
+            db_config = json.load(f)
+        model_name = db_config["embedding_model"]
+        embeddings = EMBEDDINGS[model_name]
+        db = FAISS.load_local(db_index_path, embeddings)
+        knowledge = Knowledge(db=db)
+        knowledge.collect_knowledge({input: query_counts}, max_query=query_counts)
+        domain_knowledge = knowledge.to_json()
+    else:
+        raise RuntimeError(f"Failed to query from FAISS.")
+    return domain_knowledge, ""
+
+ANNOUNCEMENT = """"""
+
+with gr.Blocks() as demo:
+    gr.HTML(ANNOUNCEMENT)
+    with gr.Row():
+        with gr.Column():
+            kdb_dropdown = gr.Dropdown(choices=ALL_KDB, value="(None)")
+            user_input = gr.Textbox(label="Input")
+            button_retrieval = gr.Button("Query", variant="primary")
+
+            with gr.Accordion("Advanced Setting", open=False):
+                query_counts_slider = gr.Slider(minimum=1, maximum=20, value=10, step=1,
+                                                       interactive=True, label="QUERY_COUNTS",
+                                                       info="从知识库内检索多少条内容")
+
+        retrieval_output = gr.JSON(label="Output")
+
+
+    button_retrieval.click(fn=query_from_kdb, inputs=[user_input, kdb_dropdown, query_counts_slider], outputs=[retrieval_output, user_input])
+demo.queue(concurrency_count=1, max_size=5, api_open=False)
+demo.launch(show_error=True)
+
+

models/embeddings.py

@@ -1,5 +1,12 @@
 from langchain.embeddings import HuggingFaceEmbeddings
+import os
 
+openai_api_key = os.getenv("OPENAI_API_KEY")
+if openai_api_key is not None:
+    from langchain.embeddings.openai import OpenAIEmbeddings
+    openai_embedding = OpenAIEmbeddings(model="text-embedding-ada-002", openai_api_key=openai_api_key)
+else:
+    openai_embedding = None
 
 model_name = 'sentence-transformers/all-MiniLM-L6-v2'
 model_kwargs = {'device': 'cpu'}
@@ -11,4 +18,4 @@ all_minilm_l6_v2 = HuggingFaceEmbeddings(
     encode_kwargs=encode_kwargs)
 
 
-EMBEDDINGS = {"all-MiniLM-L6-v2": all_minilm_l6_v2}
+EMBEDDINGS = {"text-embedding-ada-002": openai_embedding, "all-MiniLM-L6-v2": all_minilm_l6_v2}

utils/file_operations.py

@@ -2,6 +2,14 @@ import hashlib
 import os, shutil
 import datetime
 from utils.tex_processing import replace_title
+import re
+
+def urlify(s):
+    # Remove all non-word characters (everything except numbers and letters)
+    s = re.sub(r"[^\w\s]", '', s)
+    # Replace all runs of whitespace with a single dash
+    s = re.sub(r"\s+", '_', s)
+    return s
 
 def hash_name(input_dict):
     '''

utils/knowledge.py

@@ -44,4 +44,18 @@ class Knowledge:
                 break
             else:
                 prompts.append(prompt)
-        return "".join(prompts)
+        return "".join(prompts)
+
+    def to_json(self):
+        if len(self.contents) == 0:
+            return {}
+        output = {}
+        for idx, content in enumerate(self.contents):
+            output[str(idx)] = {
+                "content": content["content"],
+                "score": str(content["score"])
+            }
+        print(output)
+        return output
+
+