Update core/chatbot/retrieval_chatbot.py

core/chatbot/retrieval_chatbot.py

@@ -6,31 +6,9 @@ from models import BaseModel, GPT4Model
 from prompts import DecomposePrompt, QAPrompt, SummaryPrompt, ReferencePrompt
 import ast
 from utils.image_encoder import encode_image
+import asyncio
+import time
 
-# QA_PROMPT = "\
-# You are a Question-Answering Chatbot. \
-# Given some references and a question, please answer the question according to the references. \
-# If you find the references insufficient, you can answer the question according to your own knowledge. \
-# ONLY output the answer. \
-# "
-# QUESTION_PROMPT = "\
-# You are a Question Refiner. \
-# Given a question, you need to break it down to several subquestions and output a list of string: [\"<subquestion1>\", \"<subquestion2>\", ...]. \
-# MAKE SURE there are no vague concepts in each subquestion that require reference to other subquestions, such as determiners, pronominal and so on. \
-# If the question cannot be broken down, you need to rephrase it in 3 ways and output a list of string: [\"<rephrase1>\", \"<rephrase2>\", \"<rephrase3>\"]. \
-# ONLY output the list of subquestions or rephrases. \
-# "
-# SUMMARY_PROMPT = "\
-# You are a Summary Refiner. \
-# Given a question and several answers to it, you need to organize and summarize the answers to form one coherent answer to the question. \
-# ONLY output the summarized answer. \
-# "
-# REFERENCE_PROMPT = "\
-# You are a Reference Refiner. \
-# Given paragraphs extract from a paper, you need to remove the unnecessary and messy symbols to make it more readable. \
-# But keep the original expression and sentences as much as possible. \
-# ONLY output the refined paragraphs. \
-# "
 class RetrievalChatbot(BaseChatbot):
     def __init__(self, 
                  model: BaseModel = None, 
@@ -56,57 +34,64 @@ class RetrievalChatbot(BaseChatbot):
         self.summarizer = summarizer if summarizer \
                                      else SimpleRefiner(model=GPT4Model(), sys_prompt=SummaryPrompt.content)
     
-    def response(self, message: str, image_path=None, return_logs=False) -> str:
+    async def response(self, message: str, image_paths=None, return_logs=False) -> str:
+        time1 = time.time()
         print("Query: {message}".format(message=message))
-        question = self.decomposer.refine(message, None, image_path)
+        question = self.decomposer.refine(message, None, image_paths)
         print(question)
         question = question.replace('"', "'").replace("', '", "','").lstrip("['").rstrip("']")
         sub_questions = question.split("','")
         print("Decomposed your query into subquestions: {sub_questions}".format(sub_questions=sub_questions))
-        references = ""
+        tasks = []
+        time2 = time.time()
         for sub_question in sub_questions:
             print("="*20)
-            print(f"Subquestion: {sub_question}")
-            
+            print(f"Subquestion: {sub_question}")            
             print(f"Retrieving pdf papers for references...\n")
-            sub_retrieve_reference = references
-            sub_retrieve = self.retriever.retrieve(sub_question)
-            for ref in sub_retrieve:
-                sub_retrieve_reference += "Related research: {ref}\n".format(ref=ref)
-            # context = self.memory.messages + [{"role": "user", "content": "References: {references}\nQuestion: {question}".format(references=reference, question=sub_question)}]
-            # sub_answer = self.model.respond(context)
-            sub_answerer_context = "Sub Question References: {sub_retrieve_reference}\nQuestion: {question}\n".format(sub_retrieve_reference=sub_retrieve_reference, question=sub_question)
-
-            print(sub_answerer_context)
-            print(self.memory)
-            print(image_path)
-            sub_answer = self.answerer.refine(sub_answerer_context, self.memory, image_path)
-            
-            print(f"Subanswer: {sub_answer}")
-            
-            references += "Subquestion: {sub_question}\nSubanswer: {sub_answer}\n\n\n".format(sub_question=sub_question, sub_answer=sub_answer)
-
+            task = asyncio.create_task(self.subquestion_answerer(sub_question, image_paths))
+            tasks.append(task)
+        results = await asyncio.gather(*tasks)
+        references = "".join(results)
+        time3 = time.time()
+        print("Sub references are ",references)
         refs = self.retriever.retrieve(message)
         for ref in refs:
             references += "Related research for the user query: {ref}\n".format(ref=ref)
 
         summarizer_context = "Question References: {references}\nQuestion: {message}\n".format(references=references, message=message)
-        answer = self.summarizer.refine(summarizer_context, None, image_path)
-        
+        answer = self.summarizer.refine(summarizer_context, None, image_paths)
+        time4 = time.time()
+
         #todo 记忆管理
-        if image_path is None:
+        if image_paths is None:
             self.memory.append([{"role": "user", "content": [
                             {"type": "text", "text": f"{message}"},
                         ]}, {"role": "assistant", "content": answer}])
         else:
-            self.memory.append([{"role": "user", "content": [
-                            {"type": "text", "text": f"{message}"},
-                            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(image_path)}"}},
-                        ]}, {"role": "assistant", "content": answer}])
+            if not isinstance(image_paths, list):
+                image_paths = [image_paths]
+            memory_user = [{"type": "text", "text": f"{message}"},]
+            for image_path in image_paths:
+                memory_user.append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(image_path.name)}"}},)
+            self.memory.append([{"role": "user", "content": memory_user}, {"role": "assistant", "content": answer}])
         print("="*20)
         print(f"Final answer: {answer}".format(answer=answer))
+
+        print(f"Decompose: {time2-time1}")
+        print(f"Answer Subquestions: {time3-time2}")
+        print(f"Summarize: {time4-time3}")
         
         if return_logs:
             return answer, references
         else:
-            return answer
+            return answer
+        
+    async def subquestion_answerer(self, sub_question: str, image_paths=None, return_logs=False) -> str:
+        sub_retrieve_reference=""
+        sub_retrieve = self.retriever.retrieve(sub_question)
+        for ref in sub_retrieve:
+            sub_retrieve_reference += "Related research: {ref}\n".format(ref=ref)
+        sub_answerer_context = "Sub Question References: {sub_retrieve_reference}\nQuestion: {question}\n".format(sub_retrieve_reference=sub_retrieve_reference, question=sub_question)
+        sub_answer = self.answerer.refine(sub_answerer_context, self.memory, image_paths)
+        print(f"Subanswer: {sub_answer}")
+        return "Subquestion: {sub_question}\nSubanswer: {sub_answer}\n\n\n".format(sub_question=sub_question, sub_answer=sub_answer)