fix/ format and modify __get_index_file sequence

app.py

@@ -1,35 +1,14 @@
-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
 from utils.chatbot import Chatbot
 from utils.utils import *
-    
-def create_chatbot():
-    bot = Chatbot()
-    return bot
 
+# start of gradio interface
 with gr.Blocks() as demo:
-    history = gr.State([])
-    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')
+        gr.HTML("Junyi Academy Chatbot")
     with gr.Row(equal_height=True):
         with gr.Column(scale=5):
             with gr.Row():
@@ -48,48 +27,55 @@ with gr.Blocks() as demo:
                 with gr.Column(min_width=70, scale=1):
                     submit_btn = gr.Button("傳送")
 
-                bot_args = dict(
-                    fn=bot,
-                    inputs=user_chatbot,
-                    outputs=chatbot,
-                )
-
-                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_args).then(lambda: gr.update(interactive=True), None, [user_input], queue=False)
-                
-                
-
     with gr.Row():
-        index_file = gr.File(file_count="multiple", file_types=["pdf"], label="Upload PDF file")
-    
+        index_file = gr.File(
+            file_count="multiple", file_types=["pdf"], label="Upload PDF file"
+        )
+
     with gr.Row():
-        instruction = gr.Markdown("""
+        instruction = gr.Markdown(
+            """
             ## 使用說明
             1. 上傳一個或多個 PDF 檔案，系統將自動進行摘要、翻譯等處理後建立知識庫
             2. 在上方輸入欄輸入問題，系統將自動回覆
             3. 可以根據下方的摘要內容來提問
             4. 每次對話會根據第一個問題的內容來檢索所有文件，並挑選最能回答問題的文件來回覆
-            5. 要切換檢索的文件，請點選「清除對話記錄」按鈕後再重新提問
-        """)
+            5. 要切換檢索的文件，請點選「清除」按鈕後再重新提問
+        """
+        )
 
     with gr.Row():
-        describe = gr.Markdown('', visible=True)
+        describe = gr.Markdown("", visible=True)
+
+    # end of gradio interface
 
-        
+    # start of workflow controller
+
+    # defining workflow of user bot interaction
+    bot_args = dict(
+        fn=bot,
+        inputs=user_chatbot,
+        outputs=chatbot,
+    )
+
+    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_args,
+    ).then(
+        **bot_args
+    ).then(lambda: gr.update(interactive=True), None, [user_input], queue=False)
+
+    # defining workflow of clear state
     clear_state_args = dict(
         fn=clear_state,
         inputs=user_chatbot,
@@ -98,6 +84,7 @@ with gr.Blocks() as demo:
 
     clear_btn.click(**clear_state_args)
 
+    # defining workflow of building knowledge base
     send_system_nofification_args = dict(
         fn=send_system_nofification,
         inputs=user_chatbot,
@@ -116,12 +103,13 @@ with gr.Blocks() as demo:
         outputs=[describe],
     )
 
-    index_file.upload(**send_system_nofification_args) \
-                  .then(lambda: gr.update(interactive=True), None, None, queue=False) \
-                  .then(**bulid_knowledge_base_args) \
-                  .then(**send_system_nofification_args) \
-                  .then(lambda: gr.update(interactive=True), None, None, queue=False) \
-                  .then(**change_md_args)
-    
+    index_file.upload(**send_system_nofification_args).then(
+        lambda: gr.update(interactive=True), None, None, queue=False
+    ).then(**bulid_knowledge_base_args).then(**send_system_nofification_args).then(
+        lambda: gr.update(interactive=True), None, None, queue=False
+    ).then(
+        **change_md_args
+    )
+
 if __name__ == "__main__":
     demo.launch()

utils/chatbot.py

@@ -9,57 +9,59 @@ from openai.embeddings_utils import distances_from_embeddings
 from .work_flow_controller import WorkFlowController
 from .gpt_processor import QuestionAnswerer
 
-class Chatbot():
+
+class Chatbot:
     def __init__(self) -> None:
         self.history = []
-        self.upload_state = 'waiting'
-        
+        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:
+        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)
+        knowledge_base["page_embedding"] = (
+            knowledge_base["page_embedding"].apply(eval).apply(np.array)
+        )
 
         self.knowledge_base = knowledge_base
-        self.upload_state = 'done'
+        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.upload_state = "waiting"
         self.history = []
 
     def send_system_nofification(self):
-        if self.upload_state == 'waiting':
-            conversation = [['已上傳文件', '文件處理中（摘要、翻譯等），結束後將自動回覆']]
+        if self.upload_state == "waiting":
+            conversation = [["已上傳文件", "文件處理中（摘要、翻譯等），結束後將自動回覆"]]
             return conversation
-        elif self.upload_state == 'done':
-            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:
+        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_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>
@@ -67,14 +69,14 @@ class Chatbot():
             """
             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}')
+        print(f"user_message: {user_message}")
 
         if self.knowledge_base is None:
             response = [
@@ -82,41 +84,47 @@ class Chatbot():
             ]
             self.history = response
             return self.history
-        elif self.context is None:
+
+        else:
             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
-    
+            else:
+                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)
+        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
+        )
 
-        if self.knowledge_base['distance'].values[0] > 0.2:
+        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]
+            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/docx_processor.py

@@ -6,13 +6,14 @@ import docx2txt
 
 from gpt_processor import Translator
 
+
 class DOCXProcessor:
     def __init__(self, file_path: str) -> None:
         self.file_path = file_path
         self.file_info = {
-            'file_name': self.file_path.split('/')[-1],
-            'file_format': 'DOCX',
-            'file_full_content': '',
+            "file_name": self.file_path.split("/")[-1],
+            "file_format": "DOCX",
+            "file_full_content": "",
         }
         self.__build_info()
 
@@ -20,21 +21,24 @@ class DOCXProcessor:
         try:
             text = docx2txt.process(self.file_path)
             text = unicodedata.normalize("NFKD", text)
-            text = text.replace('\n', ' ').replace('\r', '')
-            text = re.sub(' +', ' ', text)
-            self.file_info['is_chinese'] = self.__is_chinese(text)
+            text = text.replace("\n", " ").replace("\r", "")
+            text = re.sub(" +", " ", text)
+            self.file_info["is_chinese"] = self.__is_chinese(text)
 
             tranlator = Translator()
-            self.file_info['file_full_content'] = tranlator.translate_to_chinese(text) if not self.file_info['is_chinese'] else text
+            self.file_info["file_full_content"] = (
+                tranlator.translate_to_chinese(text)
+                if not self.file_info["is_chinese"]
+                else text
+            )
 
-                
         except FileNotFoundError:
             print(f"File not found: {self.file_path}")
         except Exception as e:
             print(f"An error occurred: {str(e)}")
-        
+
     def __is_chinese(self, text: str) -> bool:
         for char in text:
-            if char >= '\u4e00' and char <= '\u9fff':
+            if char >= "\u4e00" and char <= "\u9fff":
                 return True
-        return False
+        return False

utils/pdf_processor.py

@@ -5,45 +5,48 @@ import logging
 
 from .gpt_processor import Translator
 
+
 class PDFProcessor:
     def __init__(self, file_path: str) -> None:
         self.file_path = file_path
         self.file_info = {
-            'file_name': self.file_path.split('/')[-1],
-            'file_format': 'PDF',
-            'total_pages': 0,
-            'file_content': {},
-            'file_full_content': '',
-            'is_chinese': '',
+            "file_name": self.file_path.split("/")[-1],
+            "file_format": "PDF",
+            "total_pages": 0,
+            "file_content": {},
+            "file_full_content": "",
+            "is_chinese": "",
         }
         self.__build_info()
 
     def __build_info(self) -> None:
         try:
-            with open(self.file_path, 'rb') as pdf_file:
+            with open(self.file_path, "rb") as pdf_file:
                 pdf_reader = PyPDF2.PdfReader(pdf_file)
                 pages = len(pdf_reader.pages)
-                self.file_info['total_pages'] = pages
+                self.file_info["total_pages"] = pages
                 for i, page in enumerate(pdf_reader.pages):
                     text = page.extract_text()
                     text = unicodedata.normalize("NFKD", text)
-                    text = text.replace('\n', ' ').replace('\r', '')
-                    text = re.sub(' +', ' ', text)
-                    self.file_info['is_chinese'] = self.__is_chinese(text)
+                    text = text.replace("\n", " ").replace("\r", "")
+                    text = re.sub(" +", " ", text)
+                    self.file_info["is_chinese"] = self.__is_chinese(text)
 
                     page_info = {}
                     logging.info(f"Processing page {i + 1}...")
-                    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']
+                    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:
             print(f"An error occurred: {str(e)}")
-        
+
     def __is_chinese(self, text: str) -> bool:
         for char in text:
-            if char >= '\u4e00' and char <= '\u9fff':
+            if char >= "\u4e00" and char <= "\u9fff":
                 return True
-        return False
+        return False

utils/utils.py

@@ -1,21 +1,26 @@
-
 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)
+    return chatbot.bot(*args)

utils/work_flow_controller.py

@@ -5,15 +5,21 @@ import hashlib
 
 import pandas as pd
 
-from .gpt_processor import (EmbeddingGenerator, KeywordsGenerator, Summarizer,
-                            TopicsGenerator, Translator)
+from .gpt_processor import (
+    EmbeddingGenerator,
+    KeywordsGenerator,
+    Summarizer,
+    TopicsGenerator,
+    Translator,
+)
 from .pdf_processor import PDFProcessor
 
 processors = {
-    'pdf': PDFProcessor,
+    "pdf": PDFProcessor,
 }
 
-class WorkFlowController():
+
+class WorkFlowController:
     def __init__(self, file_src) -> None:
         # check if the file_path is list
         # self.file_paths = self.__get_file_name(file_src)
@@ -24,8 +30,8 @@ class WorkFlowController():
         self.files_info = {}
 
         for file_path in self.file_paths:
-            file_name = file_path.split('/')[-1]
-            file_format = file_path.split('.')[-1]
+            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)
@@ -34,24 +40,25 @@ class WorkFlowController():
         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'])
+        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'])
+        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'])
+        file["topics"] = topics_generator.extract_topics(file["file_full_content"])
         return file
 
     def __get_embedding(self, file):
@@ -59,41 +66,54 @@ class WorkFlowController():
         # return embedding
         embedding_generator = EmbeddingGenerator()
 
-        for i, _ in enumerate(file['file_content']):
+        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'])
+            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'] = ''
+        file["file_full_content"] = ""
 
-        for i, _ in enumerate(file['file_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']
+            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 file['is_chinese']:
+        if not file["is_chinese"]:
             file = self.__translate_to_chinese(file)
         file = self.__get_embedding(file)
         file = self.__get_summary(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.json_result_path = os.path.join(os.getcwd(), 'knowledge_base.json')
+        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.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):
-
         rows = []
         for file_path, content in self.files_info.items():
             file_full_content = content["file_full_content"]
@@ -107,15 +127,24 @@ class WorkFlowController():
                 }
                 rows.append(row)
 
-        columns = ["file_name", "page_num", "page_content", "page_embedding", "file_full_content"]
+        columns = [
+            "file_name",
+            "page_num",
+            "page_content",
+            "page_embedding",
+            "file_full_content",
+        ]
         df = pd.DataFrame(rows, columns=columns)
         return df
 
     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')
+        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]
@@ -127,4 +156,4 @@ class WorkFlowController():
                 while chunk := f.read(8192):
                     md5_hash.update(chunk)
 
-        return md5_hash.hexdigest()
+        return md5_hash.hexdigest()