初始化代码

README.md

@@ -1,12 +1,33 @@
 ---
-title: Semi-Annual Security Exam
-emoji: 👁
-colorFrom: blue
-colorTo: red
-sdk: gradio
-sdk_version: 3.35.2
+title: Faiss Chat
+emoji: 🐠
+colorFrom: indigo
+colorTo: purple
+sdk: gradio 
+sdk_version: 3.32.0
 app_file: app.py
 pinned: false
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# FAISS Chat: Chat with FAISS database
+
+Webui版本的Langchain-Chat. 目前支持两个功能:
+* 将本地PDF和TXT文件打包上传, 构建FAISS向量数据库. 
+* 直接上传本地的FAISS向量数据库. 
+
+
+## 更新日志
+
+
+* 2023-06-04: 
+  * 支持读取图片格式的图表数据(目前支持JPG, PNG)
+  
+* 2023-06-04: 
+  * 支持更多文件格式 (目前支持PDF, TXT, MD, TEX)
+  * 支持更多Embedding Models (目前支持[text-embedding-ada-002](https://openai.com/blog/new-and-improved-embedding-model), [text2vec-large-chinese](https://huggingface.co/GanymedeNil/text2vec-large-chinese), 和[distilbert-dot-tas_b-b256-msmarco](https://huggingface.co/sebastian-hofstaetter/distilbert-dot-tas_b-b256-msmarco) )
+  * 优化本地知识库文件结构. 
+
+## 体验地址
+[Huggingface Space](https://huggingface.co/spaces/shaocongma/faiss_chat)
+

app.py

@@ -0,0 +1,264 @@
+import json
+import os
+import time
+import uuid
+from datetime import datetime
+
+import gradio as gr
+import openai
+from huggingface_hub import HfApi
+from langchain.document_loaders import PyPDFLoader, \
+    UnstructuredPDFLoader, PyPDFium2Loader, PyMuPDFLoader, PDFPlumberLoader
+
+from knowledge.faiss_handler import create_faiss_index_from_zip, load_faiss_index_from_zip
+from knowledge.img_handler import process_image, add_markup
+from llms.chatbot import OpenAIChatBot
+from llms.embeddings import EMBEDDINGS_MAPPING
+from utils import make_archive
+
+UPLOAD_REPO_ID=os.getenv("UPLOAD_REPO_ID")
+HF_TOKEN=os.getenv("HF_TOKEN")
+openai.api_key = os.getenv("OPENAI_API_KEY")
+openai.api_base == os.getenv("OPENAI_API_BASE")
+hf_api = HfApi(token=HF_TOKEN)
+
+ALL_PDF_LOADERS = [PyPDFLoader, UnstructuredPDFLoader, PyPDFium2Loader, PyMuPDFLoader, PDFPlumberLoader]
+ALL_EMBEDDINGS = EMBEDDINGS_MAPPING.keys()
+PDF_LOADER_MAPPING = {loader.__name__: loader for loader in ALL_PDF_LOADERS}
+
+
+#######################################################################################################################
+# Host multiple vector database for use
+#######################################################################################################################
+# todo: add this feature in the future
+
+
+
+INSTRUCTIONS = '''# FAISS Chat: 和本地数据库聊天!
+
+***2023-06-06更新:*** 
+1. 支持读取图片格式的图表数据(目前支持JPG, PNG).
+2. 在"总结图表(Demo)"的标签页里提供了这个模块的测试.
+  
+***2023-06-04更新:*** 
+1. 支持更多的Embedding Model (目前支持[text-embedding-ada-002](https://openai.com/blog/new-and-improved-embedding-model), [text2vec-large-chinese](https://huggingface.co/GanymedeNil/text2vec-large-chinese), 和[distilbert-dot-tas_b-b256-msmarco](https://huggingface.co/sebastian-hofstaetter/distilbert-dot-tas_b-b256-msmarco) )
+2. 支持更多的文件格式(PDF, TXT, TEX, 和MD).
+3. 所有生成的数据库都可以在[这个数据集](https://huggingface.co/datasets/shaocongma/shared-faiss-vdb)里访问了！如果不希望文件被上传，可以在高级设置里关闭. 
+'''
+
+
+def load_zip_as_db(file_from_gradio,
+                   pdf_loader,
+                   embedding_model,
+                   chunk_size=300,
+                   chunk_overlap=20,
+                   upload_to_cloud=True):
+    if chunk_size <= chunk_overlap:
+        return "chunk_size小于chunk_overlap. 创建失败.", None, None
+    if file_from_gradio is None:
+        return "文件为空. 创建失败.", None, None
+    pdf_loader = PDF_LOADER_MAPPING[pdf_loader]
+    zip_file_path = file_from_gradio.name
+    project_name = uuid.uuid4().hex
+    db, project_name, db_meta = create_faiss_index_from_zip(zip_file_path, embeddings=embedding_model,
+                                                   pdf_loader=pdf_loader, chunk_size=chunk_size,
+                                                         chunk_overlap=chunk_overlap, project_name=project_name)
+    index_name = project_name + ".zip"
+    make_archive(project_name, index_name)
+    date = datetime.today().strftime('%Y-%m-%d')
+    if upload_to_cloud:
+        hf_api.upload_file(path_or_fileobj=index_name,
+                           path_in_repo=f"{date}/faiss_{index_name}.zip",
+                           repo_id=UPLOAD_REPO_ID,
+                           repo_type="dataset")
+    return "成功创建知识库. 可以开始聊天了!", index_name, db, db_meta
+
+
+def load_local_db(file_from_gradio):
+    if file_from_gradio is None:
+        return "文件为空. 创建失败.", None
+    zip_file_path = file_from_gradio.name
+    db = load_faiss_index_from_zip(zip_file_path)
+
+    return "成功读取知识库. 可以开始聊天了!", db
+
+
+def extract_image(image_path):
+    from PIL import Image
+    print("Image Path:", image_path)
+    im = Image.open(image_path)
+    table = process_image(im)
+    print(f"Success in processing the image. Table: {table}")
+    return table, add_markup(table)
+
+
+def describe(image):
+    table = add_markup(process_image(image))
+    _INSTRUCTION = 'Read the table below to answer the following questions.'
+    question = "Please refer to the above table, and write a summary of no less than 200 words based on it in Chinese, ensuring that your response is detailed and precise. "
+    prompt_0shot = _INSTRUCTION + "\n" + add_markup(table) + "\n" + "Q: " + question + "\n" + "A:"
+
+    messages = [{"role": "assistant", "content": prompt_0shot}]
+    response = openai.ChatCompletion.create(
+        model="gpt-3.5-turbo",
+        messages=messages,
+        temperature=0.7,
+        top_p=1,
+        frequency_penalty=0,
+        presence_penalty=0,
+    )
+    ret = response.choices[0].message['content']
+    return ret
+
+
+with gr.Blocks() as demo:
+    local_db = gr.State(None)
+
+    def get_augmented_message(message, local_db, query_count, preprocessing, meta):
+        print(f"Receiving message: {message}")
+
+        print("Detecting if the user need to read image from the local database...")
+        # read the db_meta.json from the local file
+        # read the images file list
+        files = meta["files"]
+        source_path = meta["source_path"]
+        # with open(meta.name, "r", encoding="utf-8") as f:
+        #     files = json.load(f)["files"]
+        img_files = []
+        for file in files:
+            if os.path.splitext(file)[1] in [".png", ".jpg"]:
+                img_files.append(file)
+
+        # scan user's input to see if it contains images' name
+        do_extract_image = False
+        target_file = None
+        for file in img_files:
+            img = os.path.splitext(file)[0]
+            if img in message:
+                do_extract_image = True
+                target_file = file
+                break
+
+        # extract image to tables
+        image_info = ""
+        if do_extract_image:
+            print("The user needs to read image from the local database. Extract image ... ")
+            target_file = os.path.join(source_path, target_file)
+            _, image_info = extract_image(target_file)
+        if len(image_info)>0:
+            image_content = {"content": image_info, "source": os.path.basename(target_file)}
+        else:
+            image_content = None
+
+        print("Querying references from the local database...")
+        contents = []
+        try:
+            if query_count > 0:
+                docs = local_db.similarity_search(message, k=query_count)
+                for i in range(query_count):
+                    # pre-processing each chunk
+                    content = docs[i].page_content.replace('\n', ' ')
+                    # pre-process meta data
+                    contents.append(content)
+        except:
+            print("Failed to query from the local database. ")
+        # generate augmented_message
+        print("Success in querying references: {}".format(contents))
+        if image_content is not None:
+            augmented_message =  f"{image_content}\n\n---\n\n" + "\n\n---\n\n".join(contents) + "\n\n-----\n\n"
+        else:
+            augmented_message =  "\n\n---\n\n".join(contents) + "\n\n-----\n\n"
+        return augmented_message + "\n\n" + f"'user_input': {message}"
+
+
+    def respond(message, local_db, chat_history, meta, query_count=5, test_mode=False, response_delay=5, preprocessing=False):
+        gpt_chatbot = OpenAIChatBot()
+        print("Chat History: ", chat_history)
+        print("Local DB: ", local_db is None)
+        for chat in chat_history:
+            gpt_chatbot.load_chat(chat)
+        if local_db is None or query_count == 0:
+            bot_message = gpt_chatbot(message)
+            print(bot_message)
+            print(message)
+            chat_history.append((message, bot_message))
+            return "", chat_history
+        else:
+            augmented_message = get_augmented_message(message, local_db, query_count, preprocessing, meta)
+            bot_message = gpt_chatbot(augmented_message, original_message=message)
+            print(message)
+            print(augmented_message)
+            print(bot_message)
+            if test_mode:
+                chat_history.append((augmented_message, bot_message))
+            else:
+                chat_history.append((message, bot_message))
+            time.sleep(response_delay)  # sleep 5 seconds to avoid freq. wall.
+            return "", chat_history
+
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown(INSTRUCTIONS)
+
+            with gr.Row():
+                with gr.Tab("从本地PDF文件创建知识库"):
+                    zip_file = gr.File(file_types=[".zip"], label="本地PDF文件(.zip)")
+                    create_db = gr.Button("创建知识库", variant="primary")
+                    with gr.Accordion("高级设置", open=False):
+                        embedding_selector = gr.Dropdown(ALL_EMBEDDINGS,
+                                                         value="distilbert-dot-tas_b-b256-msmarco",
+                                                         label="Embedding Models")
+                        pdf_loader_selector = gr.Dropdown([loader.__name__ for loader in ALL_PDF_LOADERS],
+                                                          value=PyPDFLoader.__name__, label="PDF Loader")
+                        chunk_size_slider = gr.Slider(minimum=50, maximum=2000, step=50, value=500,
+                                                      label="Chunk size (tokens)")
+                        chunk_overlap_slider = gr.Slider(minimum=0, maximum=500, step=1, value=50,
+                                                         label="Chunk overlap (tokens)")
+                        save_to_cloud_checkbox = gr.Checkbox(value=False, label="把数据库上传到云端")
+
+
+                    file_dp_output = gr.File(file_types=[".zip"], label="(输出)知识库文件(.zip)")
+                with gr.Tab("读取本地知识库文件"):
+                    file_local = gr.File(file_types=[".zip"], label="本地知识库文件(.zip)")
+                    load_db = gr.Button("读取已创建知识库", variant="primary")
+
+                with gr.Tab("总结图表(Demo)"):
+                    gr.Markdown(r"代码来源于: https://huggingface.co/spaces/fl399/deplot_plus_llm")
+                    input_image = gr.Image(label="Input Image", type="pil", interactive=True)
+                    extract = gr.Button("总结", variant="primary")
+
+                    output_text = gr.Textbox(lines=8, label="Output")
+
+
+
+
+        with gr.Column():
+            status = gr.Textbox(label="用来显示程序运行状态的Textbox")
+            chatbot = gr.Chatbot()
+
+            msg = gr.Textbox()
+            submit = gr.Button("Submit", variant="primary")
+            with gr.Accordion("高级设置", open=False):
+                json_output = gr.JSON()
+                with gr.Row():
+                    query_count_slider = gr.Slider(minimum=0, maximum=10, step=1, value=3,
+                                                  label="Query counts")
+                    test_mode_checkbox = gr.Checkbox(label="Test mode")
+
+
+    # def load_pdf_as_db(file_from_gradio,
+    #                    pdf_loader,
+    #                    embedding_model,
+    #                    chunk_size=300,
+    #                    chunk_overlap=20,
+    #                    upload_to_cloud=True):
+    msg.submit(respond, [msg, local_db, chatbot, json_output, query_count_slider, test_mode_checkbox], [msg, chatbot])
+    submit.click(respond, [msg, local_db, chatbot, json_output, query_count_slider, test_mode_checkbox], [msg, chatbot])
+
+    create_db.click(load_zip_as_db, [zip_file, pdf_loader_selector, embedding_selector, chunk_size_slider, chunk_overlap_slider, save_to_cloud_checkbox],
+                    [status, file_dp_output, local_db, json_output])
+    load_db.click(load_local_db, [file_local], [status, local_db])
+
+    extract.click(describe, [input_image], [output_text])
+
+demo.launch(show_api=False)

knowledge/faiss_handler.py

@@ -0,0 +1,203 @@
+import json
+import uuid
+from langchain.vectorstores import FAISS
+import os
+from tqdm.auto import tqdm
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.document_loaders import DirectoryLoader, TextLoader
+from llms.embeddings import EMBEDDINGS_MAPPING
+import tiktoken
+import zipfile
+import pickle
+
+tokenizer_name = tiktoken.encoding_for_model('gpt-4')
+tokenizer = tiktoken.get_encoding(tokenizer_name.name)
+EMBED_MODEL = "text-embedding-ada-002"
+EMBED_DIM = 1536
+METRIC = 'cosine'
+
+#######################################################################################################################
+# Files handler
+#######################################################################################################################
+def check_existence(path):
+    return os.path.isfile(path) or os.path.isdir(path)
+
+
+def list_files(directory, ext=".pdf"):
+    # List all files in the directory
+    files_in_directory = os.listdir(directory)
+    # Filter the list to only include PDF files
+    files_list = [file for file in files_in_directory if file.endswith(ext)]
+    return files_list
+
+
+def list_pdf_files(directory):
+    # List all files in the directory
+    files_in_directory = os.listdir(directory)
+    # Filter the list to only include PDF files
+    pdf_files = [file for file in files_in_directory if file.endswith(".pdf")]
+    return pdf_files
+
+
+
+def tiktoken_len(text):
+    # evaluate how many tokens for the given text
+    tokens = tokenizer.encode(text, disallowed_special=())
+    return len(tokens)
+
+
+def get_chunks(docs, chunk_size=500, chunk_overlap=20, length_function=tiktoken_len):
+    # docs should be the output of `loader.load()`
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size,
+                                                   chunk_overlap=chunk_overlap,
+                                                   length_function=length_function,
+                                                   separators=["\n\n", "\n", " ", ""])
+    chunks = []
+    for idx, page in enumerate(tqdm(docs)):
+        source = page.metadata.get('source')
+        content = page.page_content
+        if len(content) > 50:
+            texts = text_splitter.split_text(content)
+            chunks.extend([str({'content': texts[i], 'chunk': i, 'source': os.path.basename(source)}) for i in
+                           range(len(texts))])
+    return chunks
+
+
+#######################################################################################################################
+# Create FAISS object
+#######################################################################################################################
+
+# ["text-embedding-ada-002", "distilbert-dot-tas_b-b256-msmarco"]
+
+def create_faiss_index_from_zip(path_to_zip_file, embeddings=None, pdf_loader=None,
+                                chunk_size=500, chunk_overlap=20,
+                                project_name="Very_Cool_Project_Name"):
+    # initialize the file structure
+    # structure: project_name
+    #               - source data
+    #               - embeddings
+    #               - faiss_index
+    if isinstance(embeddings, str):
+        import copy
+        embeddings_str = copy.deepcopy(embeddings)
+    else:
+        embeddings_str = "other-embedding-model"
+
+    if embeddings is None or embeddings == "text-embedding-ada-002":
+        embeddings = EMBEDDINGS_MAPPING["text-embedding-ada-002"]
+    elif isinstance(embeddings, str):
+        embeddings = EMBEDDINGS_MAPPING[embeddings]
+    else:
+        embeddings = EMBEDDINGS_MAPPING["text-embedding-ada-002"]
+    # STEP 1:
+    #   Create a folder f"{project_name}" in the current directory.
+    current_directory = os.getcwd()
+    if not os.path.exists(project_name):
+        os.makedirs(project_name)
+        project_path = os.path.join(current_directory, project_name)
+        source_data = os.path.join(project_path, "source_data")
+        embeddings_data = os.path.join(project_path, "embeddings")
+        index_data = os.path.join(project_path, "faiss_index")
+        os.makedirs(source_data)     #./project/source_data
+        os.makedirs(embeddings_data) #./project/embeddings
+        os.makedirs(index_data)      #./project/faiss_index
+    else:
+        raise ValueError(f"The project {project_name} exists.")
+    with zipfile.ZipFile(path_to_zip_file, 'r') as zip_ref:
+        # extract everything to "source_data"
+        zip_ref.extractall(source_data)
+
+
+    db_meta = {"project_name": project_name,
+               "pdf_loader": pdf_loader.__name__, "chunk_size": chunk_size,
+               "chunk_overlap": chunk_overlap,
+               "embedding_model": embeddings_str,
+               "files": os.listdir(source_data),
+               "source_path": source_data}
+    with open(os.path.join(project_path, "db_meta.json"), "w", encoding="utf-8") as f:
+        # save db_meta.json to folder
+        json.dump(db_meta, f)
+
+
+    all_docs = []
+    for ext in [".txt", ".tex", ".md", ".pdf"]:
+        if ext in [".txt", ".tex", ".md"]:
+            loader = DirectoryLoader(source_data, glob=f"**/*{ext}", loader_cls=TextLoader,
+                                     loader_kwargs={'autodetect_encoding': True})
+        elif ext in [".pdf"]:
+            loader = DirectoryLoader(source_data, glob=f"**/*{ext}", loader_cls=pdf_loader)
+        else:
+            continue
+        docs = loader.load()
+        all_docs = all_docs + docs
+
+    # split pdf files into chunks and evaluate its embeddings; save all results into embeddings
+    chunks = get_chunks(all_docs, chunk_size, chunk_overlap)
+    text_embeddings = embeddings.embed_documents(chunks)
+    text_embedding_pairs = list(zip(chunks, text_embeddings))
+    embeddings_save_to = os.path.join(embeddings_data, 'text_embedding_pairs.pickle')
+    with open(embeddings_save_to, 'wb') as handle:
+        pickle.dump(text_embedding_pairs, handle, protocol=pickle.HIGHEST_PROTOCOL)
+    db = FAISS.from_embeddings(text_embedding_pairs, embeddings)
+
+    db.save_local(index_data)
+    print(db_meta)
+    print("Success!")
+    return db, project_name, db_meta
+
+
+def find_file(file_name, directory):
+    for root, dirs, files in os.walk(directory):
+        if file_name in files:
+            return os.path.join(root, file_name)
+    return None  # If the file was not found
+
+def find_file_dir(file_name, directory):
+    for root, dirs, files in os.walk(directory):
+        if file_name in files:
+            return root  # return the directory instead of the full path
+    return None  # If the file was not found
+
+
+def load_faiss_index_from_zip(path_to_zip_file):
+    # Extract the zip file. Read the db_meta
+    # base_name = os.path.basename(path_to_zip_file)
+    path_to_extract = os.path.join(os.getcwd())
+    with zipfile.ZipFile(path_to_zip_file, 'r') as zip_ref:
+        zip_ref.extractall(path_to_extract)
+
+    db_meta_json = find_file("db_meta.json" , path_to_extract)
+    if db_meta_json is not None:
+        with open(db_meta_json, "r", encoding="utf-8") as f:
+            db_meta_dict = json.load(f)
+    else:
+        raise ValueError("Cannot find `db_meta.json` in the .zip file. ")
+
+    try:
+        embeddings = EMBEDDINGS_MAPPING[db_meta_dict["embedding_model"]]
+    except:
+        from langchain.embeddings.openai import OpenAIEmbeddings
+        embeddings = OpenAIEmbeddings(model="text-embedding-ada-002")
+
+    # locate index.faiss
+    index_path = find_file_dir("index.faiss", path_to_extract)
+    if index_path is not None:
+        db = FAISS.load_local(index_path, embeddings)
+        return db
+    else:
+        raise ValueError("Failed to find `index.faiss` in the .zip file.")
+
+
+if __name__ == "__main__":
+    from langchain.document_loaders import PyPDFLoader
+    from langchain.embeddings.openai import OpenAIEmbeddings
+    from langchain.embeddings import HuggingFaceEmbeddings
+
+    model_name = "sebastian-hofstaetter/distilbert-dot-tas_b-b256-msmarco"
+    model_kwargs = {'device': 'cpu'}
+    encode_kwargs = {'normalize_embeddings': False}
+    embeddings = HuggingFaceEmbeddings(
+        model_name=model_name,
+        model_kwargs=model_kwargs,
+        encode_kwargs=encode_kwargs)
+    create_faiss_index_from_zip(path_to_zip_file="document.zip", pdf_loader=PyPDFLoader, embeddings=embeddings)

knowledge/img_handler.py

@@ -0,0 +1,44 @@
+import torch
+from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor
+from PIL import Image
+
+if torch.cuda.is_available():
+    device = "cuda"
+else:
+    device = "cpu"
+
+model_deplot = Pix2StructForConditionalGeneration.from_pretrained("google/deplot", torch_dtype=torch.bfloat16)
+if device == "cuda":
+    model_deplot = model_deplot.to(0)
+processor_deplot = Pix2StructProcessor.from_pretrained("google/deplot")
+
+
+
+def add_markup(table):
+    try:
+        parts = [p.strip() for p in table.splitlines(keepends=False)]
+        if parts[0].startswith('TITLE'):
+            result = f"Title: {parts[0].split(' | ')[1].strip()}\n"
+            rows = parts[1:]
+        else:
+            result = ''
+            rows = parts
+        prefixes = ['Header: '] + [f'Row {i+1}: ' for i in range(len(rows) - 1)]
+        return result + '\n'.join(prefix + row for prefix, row in zip(prefixes, rows))
+    except:
+        # just use the raw table if parsing fails
+        return table
+
+def process_image(image):
+    inputs = processor_deplot(images=image, text="Generate the underlying data table for the figure below:",
+                              return_tensors="pt").to(torch.bfloat16)
+    if device == "cuda":
+        inputs = inputs.to(0)
+    predictions = model_deplot.generate(**inputs, max_new_tokens=512)
+    table = processor_deplot.decode(predictions[0], skip_special_tokens=True).replace("<0x0A>", "\n")
+    return table
+
+
+if __name__ == "__main__":
+    im = Image.open(r"meat-image.png")
+    process_image(im)

llms/chatbot.py

@@ -0,0 +1,37 @@
+import openai
+import copy
+
+class OpenAIChatBot:
+    def __init__(self, model="gpt-3.5-turbo"):
+        self.system = "You are Q&A bot. A highly intelligent system that answers user questions based on the information provided by the user's local database. " \
+                      "User's question will include some references information above his question." \
+                      "You need to answer user's question based on the provided references and inform the user what is the source of that reference. " \
+                      "If you cannot find answer in the provided references, you still need to answer user's question but you also need to notice the user that your response is not based on the provided references."
+        self.model = model
+        self.message = [{"role": "system", "content": self.system}]
+        self.raw_message = [{"role": "system", "content": self.system}]
+
+    def load_message(self, message, role, original_message=None):
+        if original_message is None:
+            original_message = message
+        msg = {"role": role, "content": message}
+        self.message.append(msg)
+        msg = {"role": role, "content": original_message}
+        self.raw_message.append(msg)
+
+    def load_chat(self, chat):
+        msg = {"role": "user", "content": chat[0]}
+        self.message.append(msg)
+        msg = {"role": "assistant", "content": chat[1]}
+        self.raw_message.append(msg)
+
+
+    def __call__(self, message, original_message = None):
+        self.load_message(message, "user", original_message)
+        augmented_message = copy.deepcopy(self.message)
+        completion = openai.ChatCompletion.create(
+            model=self.model,
+            messages=augmented_message
+        )
+        assistant_message = completion.choices[0].message
+        return assistant_message["content"]

llms/embeddings.py

@@ -0,0 +1,30 @@
+from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.embeddings import HuggingFaceEmbeddings
+
+model_name = "sebastian-hofstaetter/distilbert-dot-tas_b-b256-msmarco"
+model_kwargs = {'device': 'cpu'}
+encode_kwargs = {'normalize_embeddings': False}
+hf_embeddings_1 = HuggingFaceEmbeddings(
+    model_name=model_name,
+    model_kwargs=model_kwargs,
+    encode_kwargs=encode_kwargs)
+
+openai_embedding = OpenAIEmbeddings(model="text-embedding-ada-002")
+
+
+model_name = "GanymedeNil/text2vec-large-chinese"
+hf_embeddings_2 = HuggingFaceEmbeddings(
+    model_name=model_name,
+    model_kwargs=model_kwargs,
+    encode_kwargs=encode_kwargs)
+
+
+EMBEDDINGS_MAPPING = {"text-embedding-ada-002": openai_embedding,
+                      "distilbert-dot-tas_b-b256-msmarco": hf_embeddings_1,
+                      "text2vec-large-chinese": hf_embeddings_2}
+
+def main():
+    pass
+
+if __name__ == "__main__":
+    main()

llms/tools.py

@@ -0,0 +1,45 @@
+import openai
+
+
+class BaseTool:
+    def __init__(self, model="gpt-3.5-turbo"):
+        self.system = ""
+        self.model = model
+        self.message = [{"role": "system", "content": self.system}]
+
+    def __call__(self, message):
+        user_message = {"role": "user", "content": message}
+        messages = self.message + [user_message]
+        completion = openai.ChatCompletion.create(
+            model=self.model,
+            messages=messages
+        )
+        assistant_message = completion.choices[0].message
+        return assistant_message["content"].replace("\n", " ")
+
+
+
+
+class PreprocessingBot(BaseTool):
+    def __init__(self, model="gpt-3.5-turbo"):
+        super().__init__(model)
+        self.system = r"""You are an AI assistant for raw data pre-processing. The user will input multiple raw references which may include unicode characters or ASCII code such as '\u001e'. Your task it to make it more readable by doing:
+                      - Change all unicode characters or ASCII code such as '\u001e' to LaTeX format and put them in formula environment $...$ or $$...$$. 
+                      - Re-write formulas or mathematical notations to LaTeX format in formula environment $...$ or $$...$$.
+                      - Remove meaningless contents. 
+                      - Response in the following format: {pdf-name-1: main contents from pdf-name-1, pdf-name-2: main contents from pdf-name-2, ...}. 
+                      """
+        self.message = [{"role": "system", "content": self.system}]
+
+class ToolBot(BaseTool):
+    def __init__(self, model="gpt-3.5-turbo"):
+        super().__init__(model)
+        self.system = r"""You need to pretend a Python function. You receive a string that is the user's question to a QA bot. You need to analyze the user's goal and decide if the QA bot needs to use the search engine to generate the response to the user. 
+        Response 1 if you think the QA bot needs to use the search engine to user's input and response 0 if the QA bot doesn't need that.
+                              """
+        self.message = [{"role": "system", "content": self.system}]
+
+if __name__ == "__main__":
+    bot = ToolBot()
+    rsp = bot("Hello!")
+    print(rsp)

utils.py

@@ -0,0 +1,12 @@
+import os
+import shutil
+
+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