Upload 86 files

app.py

@@ -0,0 +1,58 @@
+import core
+import openai
+import models
+import time
+import gradio as gr
+import os
+import asyncio
+import time
+
+api_key = os.environ["OPENAI_API_KEY"]
+api_base = os.environ["OPENAI_API_BASE"]
+
+# dddd
+# def embed(texts: list):
+#         return openai.Embedding.create(input=texts, model="text-embedding-ada-002")["data"]["embedding"]
+
+def chatbot_initialize():
+    retriever = core.retriever.ChromaRetriever(pdf_dir="", 
+                                               collection_name="pdfs_1000", 
+                                               split_args={"size": 2048, "overlap": 10}, #embedding_model="text-embedding-ada-002"
+                                               embed_model=models.BiomedModel()
+                                               )
+    Chatbot = core.chatbot.RetrievalChatbot(retriever=retriever)
+    return Chatbot
+
+async def respond(query, chat_history, img_path_list, chat_history_string):
+    time1 = time.time()
+    global Chatbot
+    result = await Chatbot.response(query, image_paths=img_path_list)
+    response = result["answer"]
+    logs = result["logs"]
+    titles_set = result["titles"]
+    titles = "\n".join(list(titles_set))
+    chat_history.append((query, response))
+    if img_path_list is None:
+        chat_history_string += "Query: " + query + "\nImage: None" + "\nResponse: " + response + "\n\n\n"
+    else:
+        chat_history_string += "Query: " + query + "\nImages: " + "\n".join([path.name for path in img_path_list]) + "\nResponse: " + response + "\n\n\n"
+    time2 = time.time()
+    print(f"Total: {time2-time1}")
+    return "", chat_history, chat_history_string
+
+if __name__ == "__main__":
+    global Chatbot
+    Chatbot=chatbot_initialize()
+
+    with gr.Blocks() as demo:
+        with gr.Row():
+            with gr.Column(scale=2):
+                chatbot = gr.Chatbot()
+                msg = gr.Textbox(label="Query", show_label=True)
+                imgs = gr.File(file_count='multiple', file_types=['image'], type="filepath", label='Upload Images')
+                clear = gr.ClearButton([msg, chatbot])
+            with gr.Column(scale=1):
+                # titles = gr.Textbox(label="Referenced Article Titles", show_label=True, show_copy_button=True, interactive=False)
+                history = gr.Textbox(label="Copy Chat History", show_label=True, show_copy_button=True, interactive=False, max_lines=5)
+            msg.submit(respond, inputs=[msg, chatbot, imgs, history], outputs=[msg, chatbot, history])
+    demo.queue().launch()

core/__init__.py

@@ -0,0 +1,9 @@
+from .chain import *
+from .chatbot import *
+from .memory import *
+from .planner import *
+from .refiner import *
+from .retriever import *
+
+from models import *
+from prompts import *

core/chain/__init__.py

@@ -0,0 +1 @@
+from .base_chain import BaseChain

core/chain/base_chain.py

@@ -0,0 +1,10 @@
+
+class BaseChain:
+    def __init__(self, chain: list):
+        raise NotImplementedError
+    
+    def append(self, item: str):
+        raise NotImplementedError
+    
+    def execute(self):
+        raise NotImplementedError

core/chain/simple_chain.py

@@ -0,0 +1,19 @@
+from base_chain import BaseChain
+
+# class SimpleChain(BaseChain):
+#     def __init__(self, chain: list[str]):
+#         self.chain = chain if chain else []
+    
+#     def append(self, item: str):
+#         self.chain.append(item)
+    
+#     def execute(self):
+#         # raise NotImplementedError
+#         for item in self.chain:
+#             pass
+#             #todo: execute item
+#             # item --> result
+#             item.execute(param=param)
+#             # result --> next item
+            
+#         return result

core/chatbot/__init__.py

@@ -0,0 +1,2 @@
+from .base_chatbot import BaseChatbot
+from .retrieval_chatbot import RetrievalChatbot

core/chatbot/base_chatbot.py

@@ -0,0 +1,12 @@
+from models import BaseModel
+from ..memory import BaseMemory
+class BaseChatbot:
+    def __init__(self,
+                 model: BaseModel, 
+                 memory: BaseMemory
+        ) -> None:
+        self.model = model
+        self.memory = memory
+    
+    def respond(self, message: str) -> str:
+        raise NotImplementedError

core/chatbot/retrieval_chatbot.py

@@ -0,0 +1,69 @@
+from .base_chatbot import BaseChatbot
+from ..memory import BaseMemory, ChatMemory
+from ..retriever import BaseRetriever, ChromaRetriever, FaissRetriever
+from ..refiner import BaseRefiner, SimpleRefiner
+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
+
+class RetrievalChatbot(BaseChatbot):
+    def __init__(self, 
+                 model: BaseModel = None, 
+                 memory: BaseMemory = None,
+                 retriever: BaseRetriever = None,
+                 decomposer: BaseRefiner = None,
+                 answerer: BaseRefiner = None,
+                 summarizer: BaseRefiner = None,
+        ) -> None:
+        self.model = model if model \
+                           else GPT4Model()
+        self.memory = memory if memory \
+                             else ChatMemory(sys_prompt=SummaryPrompt.content)
+        self.retriever = retriever if retriever \
+                                   else ChromaRetriever(pdf_dir="papers_all", 
+                                                        collection_name="pdfs", 
+                                                        split_args={"size": 2048, "overlap": 10},
+                                                        embed_model=GPT4Model())
+        self.decomposer = decomposer if decomposer \
+                                     else SimpleRefiner(model=GPT4Model(), sys_prompt=DecomposePrompt.content)
+        self.answerer = answerer if answerer \
+                                 else SimpleRefiner(model=GPT4Model(), sys_prompt=QAPrompt.content)
+        self.summarizer = summarizer if summarizer \
+                                     else SimpleRefiner(model=GPT4Model(), sys_prompt=SummaryPrompt.content)
+    
+    async def response(self, message: str, image_paths=None, return_logs=False) -> str:
+        time1 = time.time()
+        print("Query: {message}".format(message=message))
+        retrieved_reference=""
+        time_s = time.time()
+        results = self.retriever.retrieve(message)
+        refs, titles = results
+        for ref in refs:
+            retrieved_reference += "Related research: {ref}\n".format(ref=ref)
+        answerer_context = "Sub Question References: {retrieved_reference}\nQuestion: {message}\n".format(retrieved_reference=retrieved_reference, message=message)
+        answer = self.answerer.refine(answerer_context, self.memory, image_paths)
+        time_e = time.time()
+
+        #todo 记忆管理
+        if image_paths is None:
+            self.memory.append([{"role": "user", "content": [
+                            {"type": "text", "text": f"{message}"},
+                        ]}, {"role": "assistant", "content": answer}])
+        else:
+            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))
+        
+        return {
+            "answer": answer,
+            "titles": set(titles),
+            "logs": ""
+        }

core/memory/__init__.py

@@ -0,0 +1,2 @@
+from .base_memory import BaseMemory
+from .chat_memory import ChatMemory

core/memory/base_memory.py

@@ -0,0 +1,18 @@
+class BaseMemory:
+    def __init__(self) -> None:
+        raise NotImplementedError
+    
+    def append(self, message: str) -> None:
+        raise NotImplementedError
+    
+    def pop(self) -> None:
+        raise NotImplementedError
+    
+    def clear(self) -> None:
+        raise NotImplementedError
+    
+    def load(self) -> None:
+        raise NotImplementedError
+    
+    def save(self) -> None:
+        raise NotImplementedError

core/memory/chat_memory.py

@@ -0,0 +1,22 @@
+from .base_memory import BaseMemory
+
+class ChatMemory(BaseMemory):
+    def __init__(self, sys_prompt = None) -> None:
+        self.sys_prompt = sys_prompt
+        self.messages = [{"role": "system", "content": sys_prompt}] if sys_prompt else []
+    
+    def append(self, message: list) -> None:
+        # assert
+        self.messages += message
+        
+    def pop(self) -> None:
+        self.messages.pop()
+        
+    def clear(self) -> None:
+        self.messages = [{"role": "system", "content": self.sys_prompt}]
+    
+    def load(self) -> None:
+        pass
+    
+    def save(self) -> None:
+        pass

core/memory/plan_memory.py

@@ -0,0 +1,38 @@
+from .base_memory import BaseMemory
+
+from dataclasses import dataclass
+
+@dataclass
+class Task:
+    def __init__(self, name: str, description: str):
+        self.name = name
+        self.description = description
+        
+class TaskChain:
+    def __init__(self, tasks: list):
+        self.tasks = tasks
+    def append(self, task: Task):
+        self.tasks.append(task)
+    def clear(self):
+        self.tasks = []
+    def __str__(self):
+        return "\n".join([f"{task.name}: {task.description}" for task in self.tasks])
+
+class PlanMemory(BaseMemory):
+    def __init__(self, initial_message, initial_task) -> None:
+        self.messages = initial_message if initial_message else []
+        self.tasks = TaskChain(initial_task) if initial_task else TaskChain([])
+    
+    def append(self, message: str) -> None:
+        self.messages.append(message)
+        #todo: parse message for tasks & add to task chain
+        self.tasks.append(Task("Task", message))
+    
+    def clear(self) -> None:
+        self.messages = []
+    
+    def load(self) -> None:
+        pass
+    
+    def save(self) -> None:
+        pass

core/planner/__init__.py

@@ -0,0 +1 @@
+from .base_planner import BasePlanner

core/planner/base_planner.py

@@ -0,0 +1,6 @@
+
+class BasePlanner:
+    def __init__(self):
+        raise NotImplementedError
+    def plan(self, message: str) -> list[str]:
+        raise NotImplementedError

core/refiner/__init__.py

@@ -0,0 +1,2 @@
+from .base_refiner import BaseRefiner
+from .simple_refiner import SimpleRefiner

core/refiner/base_refiner.py

@@ -0,0 +1,11 @@
+from models import BaseModel
+class BaseRefiner:
+    def __init__(self,
+                 sys_prompt: str, 
+                 model: BaseModel,
+        ) -> None:
+        self.sys_prompt = sys_prompt
+        self.model = model
+    
+    def refine(self, message: str) -> str:
+        raise NotImplementedError

core/refiner/simple_refiner.py

@@ -0,0 +1,50 @@
+from models import BaseModel
+from .base_refiner import BaseRefiner
+from utils.image_encoder import encode_image
+import asyncio
+
+class SimpleRefiner(BaseRefiner):
+    def __init__(self,
+                 sys_prompt: str, 
+                 model: BaseModel,
+        ) -> None:
+        BaseRefiner.__init__(self, sys_prompt=sys_prompt, model=model)
+    
+    async def refine_async(self, message: str, memory, image_paths=None) -> str:
+        if memory is None:
+            memory = []
+        else:
+            memory = memory.messages[1:]
+        
+        user_context = [{"role": "user", "content": [
+                        {"type": "text", "text": f"{message}"},]}]
+        if image_paths:
+            if not isinstance(image_paths, list):
+                image_paths = [image_paths]
+            for image_path in image_paths:
+                user_context[0]["content"].append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(image_path.name)}"}})
+        context = [{"role": "system", "content": self.sys_prompt}] + memory + user_context
+
+        respond_task = asyncio.create_task(self.model.respond_async(context))
+        await respond_task
+        response = respond_task.result()
+        return response
+    
+    def refine(self, message: str, memory, image_paths=None) -> str:
+        if memory is None:
+            memory = []
+        else:
+            memory = memory.messages[1:]
+        
+        user_context = [{"role": "user", "content": [
+                        {"type": "text", "text": f"{message}"},]}]
+        if image_paths:
+            if not isinstance(image_paths, list):
+                image_paths = [image_paths]
+            for image_path in image_paths:
+                user_context[0]["content"].append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(image_path.name)}"}})
+        context = [{"role": "system", "content": self.sys_prompt}] + memory + user_context
+        
+        response = self.model.respond(context)
+
+        return response

core/retriever/__init__.py

@@ -0,0 +1,3 @@
+from .base_retriever import BaseRetriever
+from .chroma_retriever import ChromaRetriever
+from .faiss_retriever import FaissRetriever