feat: add summerizer map-reduce

utils/chatbot_diff.py

@@ -0,0 +1,249 @@
+import io
+import os
+import json
+import logging
+import secrets
+
+import gradio as gr
+import numpy as np
+import openai
+import pandas as pd
+from google.oauth2.service_account import Credentials
+from googleapiclient.discovery import build
+from googleapiclient.http import MediaIoBaseDownload, MediaFileUpload
+from openai.embeddings_utils import distances_from_embeddings
+
+from .gpt_processor import QuestionAnswerer
+from .work_flow_controller import WorkFlowController
+
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+openai.api_key = OPENAI_API_KEY
+
+
+class Chatbot:
+    def __init__(self):
+        self.history = []
+        self.upload_state = "waiting"
+        self.uid = self.__generate_uid()
+
+        self.g_drive_service = self.__init_drive_service()
+        self.knowledge_base = None
+        self.context = None
+        self.context_page_num = None
+        self.context_file_name = None
+
+    def build_knowledge_base(self, files, upload_mode="once"):
+        work_flow_controller = WorkFlowController(files, self.uid)
+        self.csv_result_path = work_flow_controller.csv_result_path
+        self.json_result_path = work_flow_controller.json_result_path
+
+        if upload_mode == "Upload to Database":
+            self.__get_db_knowledge_base()
+        else:
+            self.__get_local_knowledge_base()
+
+    def __get_db_knowledge_base(self):
+        filename = "knowledge_base.csv"
+        db = self.__read_db(self.g_drive_service)
+        cur_content = pd.read_csv(self.csv_result_path)
+        for _ in range(10):
+            try:
+                self.__write_into_db(self.g_drive_service, db, cur_content)
+                break
+            except Exception as e:
+                logging.error(e)
+                logging.error("Failed to upload to database, retrying...")
+                continue
+        self.knowledge_base = db
+        self.upload_state = "done"
+
+    def __get_local_knowledge_base(self):
+        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 __write_into_db(self, service, db: pd.DataFrame, cur_content: pd.DataFrame):
+        db = pd.concat([db, cur_content], ignore_index=True)
+        db.to_csv(f"{self.uid}_knowledge_base.csv", index=False)
+        media = MediaFileUpload(f"{self.uid}_knowledge_base.csv", resumable=True)
+        request = (
+            service.files()
+            .update(fileId="1m3ozrphHP221hhdCFMFX9-10nzSDfNyW", media_body=media)
+            .execute()
+        )
+
+    def __init_drive_service(self):
+        SCOPES = ["https://www.googleapis.com/auth/drive"]
+        SERVICE_ACCOUNT_INFO = os.getenv("CREDENTIALS")
+        service_account_info_dict = json.loads(SERVICE_ACCOUNT_INFO)
+
+        creds = Credentials.from_service_account_info(
+            service_account_info_dict, scopes=SCOPES
+        )
+
+        return build("drive", "v3", credentials=creds)
+
+    def __read_db(self, service):
+        request = service.files().get_media(fileId="1m3ozrphHP221hhdCFMFX9-10nzSDfNyW")
+        fh = io.BytesIO()
+        downloader = MediaIoBaseDownload(fh, request)
+
+        done = False
+        while done is False:
+            status, done = downloader.next_chunk()
+            print(f"Download {int(status.progress() * 100)}%.")
+
+        fh.seek(0)
+
+        return pd.read_csv(fh)
+
+    def __read_file(self, service, filename) -> pd.DataFrame:
+        query = f"name='{filename}'"
+        results = service.files().list(q=query).execute()
+        files = results.get("files", [])
+
+        file_id = files[0]["id"]
+
+        request = service.files().get_media(fileId=file_id)
+        fh = io.BytesIO()
+        downloader = MediaIoBaseDownload(fh, request)
+
+        done = False
+        while done is False:
+            status, done = downloader.next_chunk()
+            print(f"Download {int(status.progress() * 100)}%.")
+
+        fh.seek(0)
+
+        return pd.read_csv(fh)
+
+    def __upload_file(self, service):
+        results = service.files().list(pageSize=10).execute()
+        items = results.get("files", [])
+        if not items:
+            print("No files found.")
+        else:
+            print("Files:")
+            for item in items:
+                print(f"{item['name']} ({item['id']})")
+
+        media = MediaFileUpload(self.csv_result_path, resumable=True)
+        filename_prefix = "ex_bot_database_"
+        filename = filename_prefix + self.uid + ".csv"
+        request = (
+            service.files()
+            .create(
+                media_body=media,
+                body={
+                    "name": filename,
+                    "parents": [
+                        "1Lp21EZlVlqL-c27VQBC6wTbUC1YpKMsG"
+                    ],
+                },
+            )
+            .execute()
+        )
+
+    def clear_state(self):
+        self.context = None
+        self.context_page_num = None
+        self.context_file_name = None
+        self.knowledge_base = None
+        self.upload_state = "waiting"
+        self.history = []
+
+    def send_system_notification(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
+
+        else:
+            self.__get_index_file(user_message)
+            if self.context is None:
+                response = [
+                    [user_message, "無法找到相關文件，請重新提問"],
+                ]
+                self.history = response
+                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
+        )
+
+        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]
+
+    def __generate_uid(self):
+        return secrets.token_hex(8)

utils/gpt_processor.py

@@ -24,38 +24,30 @@ class GPTAgent:
         response = self.agent.complete(messages=messages)
         return response.choices[0].message["content"]
 
-    def split_into_many(self, text) -> List[str]:
+    def split_into_many(text):
         tokenizer = tiktoken.get_encoding("cl100k_base")
-        # Split the text into sentences
-        sentences = text.split("。")
 
-        # Get the number of tokens for each sentence
+        sentences = text.split("。")
         n_tokens = [len(tokenizer.encode(" " + sentence)) for sentence in sentences]
 
         chunks = []
         tokens_so_far = 0
         chunk = []
 
-        # Loop through the sentences and tokens joined together in a tuple
         for sentence, token in zip(sentences, n_tokens):
-            # If the number of tokens so far plus the number of tokens in the current sentence is greater
-            # than the max number of tokens, then add the chunk to the list of chunks and reset
-            # the chunk and tokens so far
-            if tokens_so_far + token > self.split_max_tokens:
+
+            if tokens_so_far + token > 500:
                 chunks.append("。".join(chunk) + "。")
                 chunk = []
                 tokens_so_far = 0
 
-            # If the number of tokens in the current sentence is greater than the max number of
-            # tokens, go to the next sentence
-            if token > self.split_max_tokens:
-                continue
-
-            # Otherwise, add the sentence to the chunk and add the number of tokens to the total
+            if token > 500:
+                    continue
             chunk.append(sentence)
             tokens_so_far += token + 1
 
-        # if the length of the text is less than the max number of tokens, then return the text
+        chunks.append("。".join(chunk) + "。")
+
         return [text] if len(chunks) == 0 else chunks
 
     def preprocess(self, text):
@@ -202,10 +194,59 @@ class Summarizer(GPTAgent):
         system_prompt = """
             請幫我總結以下的文章。  
         """
-        messages = [
-            {"role": "system", "content": f"{system_prompt}"},
-            {"role": "user", "content": text},
-        ]
+
+        text_chunks = self.split_into_many(text)
+        if len(text_chunks) > 1:
+            concated_summary = ""
+            for i in range(len(text_chunks)):
+                text_chunk = text[i].replace("\n", " ").replace("\r", "")
+                messages = [
+                    {"role": "system", "content": f"{system_prompt}"},
+                    {"role": "user", "content": text_chunk},
+                ]
+                try:
+                    response = openai.ChatCompletion.create(
+                        model=self.model,
+                        messages=messages,
+                        temperature=self.temperature,
+                        max_tokens=self.max_tokens,
+                        frequency_penalty=self.frequency_penalty,
+                        presence_penalty=self.presence_penalty,
+                    )
+                except Exception as e:
+                    logging.error(e)
+                    logging.error("Failed to summarize text_chunk")
+                chinese_converter = OpenCC("s2tw")
+                concated_summary += chinese_converter.convert(
+                    response["choices"][0]["message"]["content"].strip()
+                )
+
+            # summarize concated_summary
+            messages = [
+                {"role": "system", "content": f"{system_prompt}"},
+                {"role": "user", "content": concated_summary},
+            ]
+            try:
+                response = openai.ChatCompletion.create(
+                    model=self.model,
+                    messages=messages,
+                    temperature=self.temperature,
+                    max_tokens=self.max_tokens,
+                    frequency_penalty=self.frequency_penalty,
+                    presence_penalty=self.presence_penalty,
+                )
+            except Exception as e:
+                logging.error(e)
+                logging.error("Failed to summarize concated_summary")
+            chinese_converter = OpenCC("s2tw")
+            return chinese_converter.convert(
+                response["choices"][0]["message"]["content"].strip()
+            )
+        else:
+            messages = [
+                {"role": "system", "content": f"{system_prompt}"},
+                {"role": "user", "content": text},
+            ]
         try:
             response = openai.ChatCompletion.create(
                 model=self.model,
@@ -224,7 +265,7 @@ class Summarizer(GPTAgent):
             response["choices"][0]["message"]["content"]
         )
 
-        return re.sub(r"\n+", "<br>", response)
+        return response
 
 
 class QuestionAnswerer(GPTAgent):