app.py

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)