refactor/ enable mutiple usage

app.py

@@ -11,126 +11,25 @@ from openai.embeddings_utils import distances_from_embeddings
 
 from utils.gpt_processor import QuestionAnswerer
 from utils.work_flow_controller import WorkFlowController
-
-qa_processor = QuestionAnswerer()
-CSV_FILE_PATHS = ''
-JSON_FILE_PATHS = ''
-KNOWLEDGE_BASE = None
-CONTEXT = None
-CONTEXT_PAGE_NUM = None
-CONTEXT_FILE_NAME = None
-
-def build_knowledge_base(files):
-    global CSV_FILE_PATHS
-    global JSON_FILE_PATHS
-    global KNOWLEDGE_BASE
-
-    work_flow_controller = WorkFlowController(files)
-    CSV_FILE_PATHS = work_flow_controller.csv_result_path
-    JSON_FILE_PATHS = work_flow_controller.result_path
-    with open(CSV_FILE_PATHS, 'r', encoding='UTF-8') as fp:
-        knowledge_base = pd.read_csv(fp)
-    knowledge_base['page_embedding'] = knowledge_base['page_embedding'].apply(eval).apply(np.array)
-    KNOWLEDGE_BASE = knowledge_base
-
-def construct_summary():
-    with open(JSON_FILE_PATHS, 'r', encoding='UTF-8') as fp:
-        knowledge_base = json.load(fp)
-
-    context = """"""
-    for key in knowledge_base.keys():
-        file_name = knowledge_base[key]['file_name']
-        total_page = knowledge_base[key]['total_pages']
-        summary = knowledge_base[key]['summarized_content']
-        file_context = f"""
-            ### 文件摘要
-            {file_name}  (共 {total_page} 頁)<br><br>
-            {summary}<br><br>
-        """
-        context += file_context
-    return context
-
-def change_md():
-    content = construct_summary()
-    return gr.Markdown.update(content, visible=True)
-
-def user(message, history):
-    return "", history + [[message, None]]
-
-def system_notification(action):
-    if action == 'upload':
-        return [['已上傳文件', '文件處理中（摘要、翻譯等），結束後將自動回覆']]
-    else:
-        return [['已上傳文件', '文件處理完成，請開始提問']]
-    
-def get_index_file(user_message):
-    global KNOWLEDGE_BASE
-    global CONTEXT
-    global CONTEXT_PAGE_NUM
-    global CONTEXT_FILE_NAME
-
-    user_message_embedding = openai.Embedding.create(input=user_message, engine='text-embedding-ada-002')['data'][0]['embedding']
-    KNOWLEDGE_BASE['distance'] = distances_from_embeddings(user_message_embedding, KNOWLEDGE_BASE['page_embedding'].values, distance_metric='cosine')
-    KNOWLEDGE_BASE = KNOWLEDGE_BASE.sort_values(by='distance', ascending=True).head(1)
-    if KNOWLEDGE_BASE['distance'].values[0] > 0.2:
-        CONTEXT = None
-    else:
-
-        CONTEXT = KNOWLEDGE_BASE['page_content'].values[0]
-        CONTEXT_PAGE_NUM = KNOWLEDGE_BASE['page_num'].values[0]
-        CONTEXT_FILE_NAME = KNOWLEDGE_BASE['file_name'].values[0]
-
-def bot(history):
-    user_message = history[-1][0]
-    global CONTEXT
-    print(f'user_message: {user_message}')
-
-    if KNOWLEDGE_BASE is None:
-        response = [
-            [user_message, "請先上傳文件"],
-        ]
-        history = response
-        return history
-    elif CONTEXT is None:
-        get_index_file(user_message)
-        print(f'CONTEXT: {CONTEXT}')
-        if CONTEXT is None:
-            response = [
-                [user_message, "無法找到相關文件，請重新提問"],
-            ]
-            history = response
-            return history
-    else:
-        pass
-
-    if CONTEXT is not None:
-        bot_message = qa_processor.answer_question(CONTEXT, CONTEXT_PAGE_NUM, CONTEXT_FILE_NAME, history)
-        print(f'bot_message: {bot_message}')
-        response = [
-            [user_message, bot_message],
-        ]
-        history[-1] = response[0]
-        return history
+from utils.chatbot import Chatbot
+from utils.utils import *
     
-def clear_state():
-    global KNOWLEDGE_BASE
-    global CONTEXT
-    global CONTEXT_PAGE_NUM
-    global CONTEXT_FILE_NAME
-
-    CONTEXT = None
-    CONTEXT_PAGE_NUM = None
-    CONTEXT_FILE_NAME = None
-    KNOWLEDGE_BASE = None
+def create_chatbot():
+    bot = Chatbot()
+    return bot
 
 with gr.Blocks() as demo:
     history = gr.State([])
-    upload_state = gr.State("upload")
-    finished = gr.State("finished")
     user_question = gr.State("")
+    chatbot_utils = Chatbot()
+
+    user_chatbot = gr.State(Chatbot())
+    
+    upload_state = gr.State("wating")
+    finished = gr.State("finished")
+
     with gr.Row():
         gr.HTML('Junyi Academy Chatbot')
-        #status_display = gr.Markdown("Success", elem_id="status_display")
     with gr.Row(equal_height=True):
         with gr.Column(scale=5):
             with gr.Row():
@@ -143,53 +42,86 @@ with gr.Blocks() as demo:
                         placeholder="Enter text",
                         container=False,
                     )
-                # with gr.Column(min_width=70, scale=1):
-                #     submit_btn = gr.Button("Send")
+
                 with gr.Column(min_width=70, scale=1):
                     clear_btn = gr.Button("清除")
                 with gr.Column(min_width=70, scale=1):
                     submit_btn = gr.Button("傳送")
 
-                response = user_input.submit(user, 
-                                  [user_input, chatbot], 
-                                  [user_input, chatbot], 
-                                  queue=False,
-                                  ).then(bot, chatbot, chatbot)
-                response.then(lambda: gr.update(interactive=True), None, [user_input], queue=False)
+                bot_args = dict(
+                    fn=bot,
+                    inputs=user_chatbot,
+                    outputs=chatbot,
+                )
 
-                clear_btn.click(lambda: None, None, chatbot, queue=False)
+                user_args = dict(
+                    fn=user,
+                    inputs=[user_chatbot, user_input],
+                    outputs=[user_input, chatbot],
+                    queue=False,
+                )
+
+                response = user_input.submit(**user_args).then(**bot_args)
+                
+                response.then(lambda: gr.update(interactive=True), None, [user_input], queue=False)
 
                 submit_btn.click(user,
                                 [user_input, chatbot], 
                                 [user_input, chatbot],
                                 chatbot,
-                                queue=False).then(bot, chatbot, chatbot).then(lambda: gr.update(interactive=True), None, [user_input], queue=False)
+                                queue=False).then(**bot_args).then(lambda: gr.update(interactive=True), None, [user_input], queue=False)
+                
                 
-                clear_btn.click(clear_state, None, None, queue=False)
 
     with gr.Row():
         index_file = gr.File(file_count="multiple", file_types=["pdf"], label="Upload PDF file")
     
     with gr.Row():
         instruction = gr.Markdown("""
-        ## 使用說明
-        1. 上傳一個或多個 PDF 檔案，系統將自動進行摘要、翻譯等處理後建立知識庫
-        2. 在上方輸入欄輸入問題，系統將自動回覆
-        3. 可以根據下方的摘要內容來提問
-        4. 每次對話會根據第一個問題的內容來檢索所有文件，並挑選最能回答問題的文件來回覆
-        5. 要切換檢索的文件，請點選「清除」按鈕後再重新提問
+            ## 使用說明
+            1. 上傳一個或多個 PDF 檔案，系統將自動進行摘要、翻譯等處理後建立知識庫
+            2. 在上方輸入欄輸入問題，系統將自動回覆
+            3. 可以根據下方的摘要內容來提問
+            4. 每次對話會根據第一個問題的內容來檢索所有文件，並挑選最能回答問題的文件來回覆
+            5. 要切換檢索的文件，請點選「清除對話記錄」按鈕後再重新提問
         """)
 
     with gr.Row():
         describe = gr.Markdown('', visible=True)
 
-        index_file.upload(system_notification, [upload_state], chatbot) \
+        
+    clear_state_args = dict(
+        fn=clear_state,
+        inputs=user_chatbot,
+        outputs=None,
+    )
+
+    clear_btn.click(**clear_state_args)
+
+    send_system_nofification_args = dict(
+        fn=send_system_nofification,
+        inputs=user_chatbot,
+        outputs=chatbot,
+    )
+
+    bulid_knowledge_base_args = dict(
+        fn=build_knowledge_base,
+        inputs=[user_chatbot, index_file],
+        outputs=None,
+    )
+
+    change_md_args = dict(
+        fn=change_md,
+        inputs=[user_chatbot],
+        outputs=[describe],
+    )
+
+    index_file.upload(**send_system_nofification_args) \
                   .then(lambda: gr.update(interactive=True), None, None, queue=False) \
-                  .then(build_knowledge_base, [index_file]) \
-                  .then(system_notification, [finished], chatbot) \
+                  .then(**bulid_knowledge_base_args) \
+                  .then(**send_system_nofification_args) \
                   .then(lambda: gr.update(interactive=True), None, None, queue=False) \
-                  .then(change_md, None, describe)
-                  
+                  .then(**change_md_args)
     
 if __name__ == "__main__":
     demo.launch()

utils/chatbot.py

@@ -0,0 +1,122 @@
+import json
+
+import openai
+import pandas as pd
+import numpy as np
+import gradio as gr
+from openai.embeddings_utils import distances_from_embeddings
+
+from .work_flow_controller import WorkFlowController
+from .gpt_processor import QuestionAnswerer
+
+class Chatbot():
+    def __init__(self) -> None:
+        self.history = []
+        self.upload_state = 'waiting'
+        
+        self.knowledge_base = None
+        self.context = None
+        self.context_page_num = None
+        self.context_file_name = None
+        
+
+    def build_knowledge_base(self, files):
+        work_flow_controller = WorkFlowController(files)
+        self.csv_result_path = work_flow_controller.csv_result_path
+        self.json_result_path = work_flow_controller.json_result_path
+
+        with open(self.csv_result_path, 'r', encoding='UTF-8') as fp:
+            knowledge_base = pd.read_csv(fp)
+        knowledge_base['page_embedding'] = knowledge_base['page_embedding'].apply(eval).apply(np.array)
+
+        self.knowledge_base = knowledge_base
+        self.upload_state = 'done'
+
+    def clear_state(self):
+        self.context = None
+        self.context_page_num = None
+        self.context_file_name = None
+        self.upload_state = 'waiting'
+        self.history = []
+
+    def send_system_nofification(self):
+        if self.upload_state == 'waiting':
+            conversation = [['已上傳文件', '文件處理中（摘要、翻譯等），結束後將自動回覆']]
+            return conversation
+        elif self.upload_state == 'done':
+            conversation = [['已上傳文件', '文件處理完成，請開始提問']]
+            return conversation
+        
+    def change_md(self):
+        content = self.__construct_summary()
+        return gr.Markdown.update(content, visible=True)
+            
+    def __construct_summary(self):
+        with open(self.json_result_path, 'r', encoding='UTF-8') as fp:
+            knowledge_base = json.load(fp)
+
+        context = """"""
+        for key in knowledge_base.keys():
+            file_name = knowledge_base[key]['file_name']
+            total_page = knowledge_base[key]['total_pages']
+            summary = knowledge_base[key]['summarized_content']
+            file_context = f"""
+                ### 文件摘要
+                {file_name}  (共 {total_page} 頁)<br><br>
+                {summary}<br><br>
+            """
+            context += file_context
+        return context
+    
+    def user(self, message):
+        self.history += [[message, None]]
+        return "", self.history
+    
+    def bot(self):
+        user_message = self.history[-1][0]
+        print(f'user_message: {user_message}')
+
+        if self.knowledge_base is None:
+            response = [
+                [user_message, "請先上傳文件"],
+            ]
+            self.history = response
+            return self.history
+        elif self.context is None:
+            self.__get_index_file(user_message)
+            print(f'CONTEXT: {self.context}')
+            if self.context is None:
+                response = [
+                    [user_message, "無法找到相關文件，請重新提問"],
+                ]
+                self.history = response
+                return self.history
+        else:
+            pass
+
+        if self.context is not None:
+            qa_processor = QuestionAnswerer()
+            bot_message = qa_processor.answer_question(
+                self.context,
+                self.context_page_num,
+                self.context_file_name,
+                self.history
+            )
+            print(f'bot_message: {bot_message}')
+            response = [
+                [user_message, bot_message],
+            ]
+            self.history[-1] = response[0]
+        return self.history
+    
+    def __get_index_file(self, user_message):
+        user_message_embedding = openai.Embedding.create(input=user_message, engine='text-embedding-ada-002')['data'][0]['embedding']
+        self.knowledge_base['distance'] = distances_from_embeddings(user_message_embedding, self.knowledge_base['page_embedding'].values, distance_metric='cosine')
+        self.knowledge_base = self.knowledge_base.sort_values(by='distance', ascending=True).head(1)
+
+        if self.knowledge_base['distance'].values[0] > 0.2:
+            self.context = None
+        else:
+            self.context = self.knowledge_base['page_content'].values[0]
+            self.context_page_num = self.knowledge_base['page_num'].values[0]
+            self.context_file_name = self.knowledge_base['file_name'].values[0]

utils/utils.py

@@ -0,0 +1,21 @@
+
+def clear_state(chatbot, *args):
+    return chatbot.clear_state(*args)
+
+def send_system_nofification(chatbot, *args):
+    return chatbot.send_system_nofification(*args)
+
+def build_knowledge_base(chatbot, *args):
+    return chatbot.build_knowledge_base(*args)
+
+def change_md(chatbot, *args):
+    return chatbot.change_md(*args)
+
+def get_index_file(chatbot, *args):
+    return chatbot.get_index_file(*args)
+
+def user(chatbot, *args):
+    return chatbot.user(*args)
+
+def bot(chatbot, *args):
+    return chatbot.bot(*args)

utils/work_flow_controller.py

@@ -84,14 +84,12 @@ class WorkFlowController():
             file = self.__translate_to_chinese(file)
         file = self.__get_embedding(file)
         file = self.__get_summary(file)
-        # file = self.__get_keywords(file)
-        # file = self.__get_topics(file)
         return file
 
     def __dump_to_json(self):
         with open(os.path.join(os.getcwd(), 'knowledge_base.json'), 'w', encoding='utf-8') as f:
             print("Dumping to json, the path is: " + os.path.join(os.getcwd(), 'knowledge_base.json'))
-            self.result_path = os.path.join(os.getcwd(), 'knowledge_base.json')
+            self.json_result_path = os.path.join(os.getcwd(), 'knowledge_base.json')
             json.dump(self.files_info, f, indent=4, ensure_ascii=False)
 
     def __construct_knowledge_base_dataframe(self):