add upload feature and optimize user experience

app.py

@@ -1,100 +1,130 @@
 import json
 import time
 import random
+import os
 
+import openai
 import gradio as gr
 import pandas as pd
+import numpy as np
+from openai.embeddings_utils import distances_from_embeddings
 
 from utils.gpt_processor import QuestionAnswerer
+from utils.work_flow_controller import WorkFlowController
 
 qa_processor = QuestionAnswerer()
-current_file = None
-context = None
+CSV_FILE_PATHS = ''
+JSON_FILE_PATHS = ''
+KNOWLEDGE_BASE = None
+CONTEXT = None
+CONTEXT_PAGE_NUM = None
+CONTEXT_FILE_NAME = None
 
-with open("final_result.json", 'r', encoding='UTF-8') as fp:
-    db = json.load(fp)
+def build_knowledge_base(files):
+    global CSV_FILE_PATHS
+    global JSON_FILE_PATHS
+    global KNOWLEDGE_BASE
 
-def read_examples():
-    df = pd.read_csv(r'examples.csv')
-    return [f"{keyword}" for keyword in df['word'].tolist()]
+    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 gr.update(value="", interactive=False), history + [[message, None]]
     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 current_file
-    global context
-    #check if user input has "我想了解"
-    if "我想了解" in user_message:
-        # get keyword from "「」"
-        keyword = user_message.split("「")[1].split("」")[0]
-        # check if keyword is in db
-        file_list = []
-        for key in db.keys():
-            if keyword in db[key]['keywords']:
-                file_list.append(key)
-        if len(file_list) == 0:
-            response = [
-                [user_message, "Sorry, I can't find any documents about this topic. Please try again."],
-            ]
-        else:
-            bot_message = "以下是我所找到的文件："
-            for file in file_list:
-                bot_message += "\n" + file
-            bot_message += "\n\n" + "請複製貼上想要了解的文件，我會給你該文件的摘要"
-            response = [
-                [user_message, bot_message],
-            ]
+    global CONTEXT
+    print(f'user_message: {user_message}')
+
+    if KNOWLEDGE_BASE is None:
+        response = [
+            [user_message, "請先上傳文件"],
+        ]
         history = response
-        # history[-1][1] = ""
-        # for character in bot_message:
-        #     history[-1][1] += character
-        #     time.sleep(random.uniform(0.01, 0.05))
-        #     yield history
         return history
-    
-    # check if user input has a pdf file name
-    if ".pdf" in user_message or ".docx" in user_message:
-        current_file = user_message
-        context = db[current_file]['file_full_content']
-        # check if file name is in db
-        if user_message in db.keys():
-            bot_message = f"文件 {user_message} 的摘要如下："
-            bot_message += "\n\n" + db[user_message]['summarized_content']
-            bot_message += "\n\n" + "可以透過詢問來了解更多這個文件的內容"
-            response = [
-                [user_message, bot_message],
-            ]
-        else:
+    elif CONTEXT is None:
+        get_index_file(user_message)
+        print(f'CONTEXT: {CONTEXT}')
+        if CONTEXT is None:
             response = [
-                [user_message, "Sorry, I can't find this file. Please try again."],
+                [user_message, "無法找到相關文件，請重新提問"],
             ]
-        history[-1] = response[0]
-        # history[-1][1] = ""
-        # for character in bot_message:
-        #     history[-1][1] += character
-        #     time.sleep(random.uniform(0.01, 0.05))
-        #     yield history
-        return history
-    if context is None:
-        response = [
-            [user_message, "請輸入一個文件名稱或是點選下方的範例"],
-        ]
-        history[-1] = response[0]
-        return history
-    
-    if context is not None:
-        bot_message = qa_processor.answer_question(context, 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
+    
+def clear_state():
+    global CONTEXT
+    global CONTEXT_PAGE_NUM
+    global CONTEXT_FILE_NAME
+
+    CONTEXT = None
+    CONTEXT_PAGE_NUM = None
+    CONTEXT_FILE_NAME = None
 
 with gr.Blocks() as demo:
     history = gr.State([])
+    upload_state = gr.State("upload")
+    finished = gr.State("finished")
     user_question = gr.State("")
     with gr.Row():
         gr.HTML('Junyi Academy Chatbot')
@@ -114,7 +144,9 @@ with gr.Blocks() as demo:
                 # 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("Clear")
+                    clear_btn = gr.Button("清除")
+                with gr.Column(min_width=70, scale=1):
+                    submit_btn = gr.Button("傳送")
 
                 response = user_input.submit(user, 
                                   [user_input, chatbot], 
@@ -122,11 +154,40 @@ with gr.Blocks() as demo:
                                   queue=False,
                                   ).then(bot, chatbot, chatbot)
                 response.then(lambda: gr.update(interactive=True), None, [user_input], queue=False)
+
                 clear_btn.click(lambda: None, None, chatbot, 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)
+                
+                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. 要切換檢索的文件，請點選「清除對話記錄」按鈕後再重新提問
+        """)
+
+    with gr.Row():
+        describe = gr.Markdown('', visible=True)
 
-    examples = gr.Examples(examples=read_examples(),
-                           inputs=[user_input])
+        index_file.upload(system_notification, [upload_state], chatbot) \
+                  .then(lambda: gr.update(interactive=True), None, None, queue=False) \
+                  .then(build_knowledge_base, [index_file]) \
+                  .then(system_notification, [finished], chatbot) \
+                  .then(lambda: gr.update(interactive=True), None, None, queue=False) \
+                  .then(change_md, None, describe)
+                  
     
 if __name__ == "__main__":
     demo.launch()

examples.csv

@@ -1,8 +0,0 @@
-word,count
-_KTX CARES.Non-Negotiables.docx,0
-🄣 3.5小時 Getter Better Faster Rubric.pdf,0
-02 - IP.Internalization and Planning Mastery Rubric.docx,0
-KTX Houst First 21 Days Rubric 2019 Final-2 (2).docx,0
-KTX Rubric for Equity _ Excellence.pdf,0
-Leading and Coaching through States of Being 1 pager and integration guides.pdf,0
-SEAMS Tool.FINAL (1).pdf,0

utils/pdf_processor.py

@@ -3,7 +3,6 @@ import unicodedata
 import re
 import logging
 
-from datamodel.data_model import PDFRawData
 from .gpt_processor import Translator
 
 class PDFProcessor:
@@ -15,6 +14,7 @@ class PDFProcessor:
             'total_pages': 0,
             'file_content': {},
             'file_full_content': '',
+            'is_chinese': '',
         }
         self.__build_info()
 
@@ -31,14 +31,12 @@ class PDFProcessor:
                     text = re.sub(' +', ' ', text)
                     self.file_info['is_chinese'] = self.__is_chinese(text)
 
-                    temp = {}
+                    page_info = {}
                     logging.info(f"Processing page {i + 1}...")
-                    temp['page_num'] = i + 1
-                    tranlator = Translator()
-                    temp['page_content'] = tranlator.translate_to_chinese(text) if not self.file_info['is_chinese'] else text
-                    self.file_info['file_content'][i + 1] = temp
-                    self.file_info['file_full_content'] = self.file_info['file_full_content'] + temp['page_content']
-
+                    page_info['page_num'] = i + 1
+                    page_info['page_content'] = text
+                    self.file_info['file_content'][i + 1] = page_info
+                    self.file_info['file_full_content'] = self.file_info['file_full_content'] + page_info['page_content']
         except FileNotFoundError:
             print(f"File not found: {self.file_path}")
         except Exception as e:

utils/work_flow_controller.py

@@ -1,31 +1,132 @@
+import os
 import json
+import logging
+import hashlib
 
+import pandas as pd
+
+from .gpt_processor import (EmbeddingGenerator, KeywordsGenerator, Summarizer,
+                            TopicsGenerator, Translator)
 from .pdf_processor import PDFProcessor
-from .gpt_processor import Translator, EmbeddingGenerator, KeywordsGenerator, TopicsGenerator
 
 processors = {
     'pdf': PDFProcessor,
 }
 
 class WorkFlowController():
-    def __init__(self, file_path: str, file_name: str) -> None:
-        # get file raw content
-        self.file_name = file_name
-        file_format = file_path.split('.')[-1]
-        self.file_processor = processors[file_format]
-        self.file_info = self.file_processor(file_path).file_info
-
-    def process_file(self):
+    def __init__(self, file_src) -> None:
+        # check if the file_path is list
+        # self.file_paths = self.__get_file_name(file_src)
+        self.file_paths = [x.name for x in file_src]
+
+        print(self.file_paths)
+
+        self.files_info = {}
+
+        for file_path in self.file_paths:
+            file_name = file_path.split('/')[-1]
+            file_format = file_path.split('.')[-1]
+            self.file_processor = processors[file_format]
+            file = self.file_processor(file_path).file_info
+            file = self.__process_file(file)
+            self.files_info[file_name] = file
+
+        self.__dump_to_json()
+        self.__dump_to_csv()
+
+
+    def __get_summary(self, file: dict):
+        # get summary from file content
+        
+        summarizer = Summarizer()
+        file['summarized_content'] = summarizer.summarize(file['file_full_content'])
+        return file
+
+    def __get_keywords(self, file: dict):
+        # get keywords from file content
+        keywords_generator = KeywordsGenerator()
+        file['keywords'] = keywords_generator.extract_keywords(file['file_full_content'])
+        return file
+
+    def __get_topics(self, file: dict):
+        # get topics from file content
+        topics_generator = TopicsGenerator()
+        file['topics'] = topics_generator.extract_topics(file['file_full_content'])
+        return file
+
+    def __get_embedding(self, file):
+        # get embedding from file content
+        # return embedding
+        embedding_generator = EmbeddingGenerator()
+
+        for i, _ in enumerate(file['file_content']):
+            # use i+1 to meet the index of file_content
+            file['file_content'][i+1]['page_embedding'] = embedding_generator.get_embedding(file['file_content'][i+1]['page_content'])
+        return file
+        
+
+    def __translate_to_chinese(self, file: dict):
+        # translate file content to chinese
+        translator = Translator()
+        # reset the file full content
+        file['file_full_content'] = ''
+
+        for i, _ in enumerate(file['file_content']):
+            # use i+1 to meet the index of file_content
+            file['file_content'][i+1]['page_content'] = translator.translate_to_chinese(file['file_content'][i+1]['page_content'])
+            file['file_full_content'] = file['file_full_content'] + file['file_content'][i+1]['page_content']
+        return file
+        
+    def __process_file(self, file: dict):
         # process file content
         # return processed data
-        if not self.file_info['is_chinese']:
-            translator = Translator()
-            self.file_info[1]['file_content'] = translator.translate_to_chinese(self.file_info[1]['file_content'])
+        if not file['is_chinese']:
+            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')
+            json.dump(self.files_info, f, indent=4, ensure_ascii=False)
+
+    def __construct_knowledge_base_dataframe(self):
+
+        rows = []
+        for file_path, content in self.files_info.items():
+            file_full_content = content["file_full_content"]
+            for page_num, page_details in content["file_content"].items():
+                row = {
+                    "file_name": content["file_name"],
+                    "page_num": page_details["page_num"],
+                    "page_content": page_details["page_content"],
+                    "page_embedding": page_details["page_embedding"],
+                    "file_full_content": file_full_content,
+                }
+                rows.append(row)
+
+        columns = ["file_name", "page_num", "page_content", "page_embedding", "file_full_content"]
+        df = pd.DataFrame(rows, columns=columns)
+        return df
 
-    # save file_info data to json file
-    def dump_to_json(self) -> None:
-        with open(f'{self.file_name}.json', 'w', encoding='utf-8') as f:
-            json.dump(self.file_info, f, indent=4, ensure_ascii=False)
+    def __dump_to_csv(self):
+        df = self.__construct_knowledge_base_dataframe()
+        df.to_csv(os.path.join(os.getcwd(), 'knowledge_base.csv'), index=False)
+        print("Dumping to csv, the path is: " + os.path.join(os.getcwd(), 'knowledge_base.csv'))
+        self.csv_result_path = os.path.join(os.getcwd(), 'knowledge_base.csv')
 
+    def __get_file_name(self, file_src):
+        file_paths = [x.name for x in file_src]
+        file_paths.sort(key=lambda x: os.path.basename(x))
 
+        md5_hash = hashlib.md5()
+        for file_path in file_paths:
+            with open(file_path, "rb") as f:
+                while chunk := f.read(8192):
+                    md5_hash.update(chunk)
 
+        return md5_hash.hexdigest()