uP

app.py

@@ -6,7 +6,8 @@ from inference import ModelFactory
 from face import FaceAlgo
 from sam_everything import SamAnything
 from ocr import do_ocr
-
+from retriever import knowledgeBase
+import time
 
 components = {}
 
@@ -125,10 +126,55 @@ def create_ui():
             with gr.Row():
                 with gr.Group():
                     components["ocr_json_output"] = gr.JSON(label="推理结果")
+        with gr.Tab("知识库"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    with gr.Group():
+                        components["db_view"] = gr.Dataframe(
+                                                    headers=["列表"],
+                                                    datatype=["str"],
+                                                    row_count=8,
+                                                    col_count=(1, "fixed"),
+                                                    interactive=False
+                        )
+                with gr.Column(scale=2):
+                        with gr.Group():
+                            components["db_name"] = gr.Textbox(label="名称", info="请输入库名称", lines=1, value="")
+                            components["file_upload"] = gr.File(elem_id='file_upload',file_count='multiple',label='文档上传', file_types=[".pdf", ".doc", '.docx', '.json', '.csv'])
+                            components["db_submit_btn"] = gr.Button(value="提交")
+            with gr.Row():
+                with gr.Column(scale=2):
+                    components["db_input"] = gr.Textbox(label="关键词", lines=1, value="")
+                    
+                with gr.Column(scale=1):
+                    components["db_test_select"] = gr.Dropdown(
+                                        choices=knowledgeBase.get_bases(),value=None,multiselect=True, label="知识库选择"
+                    )
+                    components["dbtest_submit_btn"] = gr.Button(value="检索")
+            with gr.Row():
+                with gr.Group():
+                    components["db_search_result"] = gr.JSON(label="检索结果")
 
+        with gr.Tab("问答"):
+            with gr.Row():
+                with gr.Column():
+                    with gr.Group():
+                        components["chatbot"] = gr.Chatbot(
+                                            [(None,"What can I help you?")],
+                                            elem_id="chatbot",
+                                            bubble_full_width=False,
+                                            height=600
+                            )
+                        components["chat_input"] = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter message or upload file...", show_label=False)
+                        components["db_select"] = gr.CheckboxGroup(choices=knowledgeBase.get_bases(),value=None,label="知识库", info="可选择1个或多个知识库")
         create_event_handlers()
+        demo.load(init,None,gradio("db_view"))
     return demo
 
+def init():
+    # db_list = knowledgeBase.get_bases()
+    db_df_list = knowledgeBase.get_df_bases()
+    return db_df_list
 
 def create_event_handlers():
     params["algo_type"] = gr.State("全景分割")
@@ -172,6 +218,24 @@ def create_event_handlers():
         do_ocr,gradio('ocr_type','ocr_input'),gradio("ocr_output","ocr_json_output")
     )
 
+    components["db_submit_btn"].click(
+        file_handler,gradio('file_upload','db_name'),gradio("db_view",'db_select',"db_test_select")
+    )
+
+    components["chat_input"].submit(
+        do_llm_request, gradio("chatbot", "chat_input"), gradio("chatbot", "chat_input")
+    ).then(
+        do_llm_response, gradio("chatbot","db_select"), gradio("chatbot"), api_name="bot_response"
+    ).then(
+        lambda: gr.MultimodalTextbox(interactive=True), None, gradio('chat_input')
+    )
+
+    # components["chatbot"].like(print_like_dislike, None, None)
+
+    components['dbtest_submit_btn'].click(
+        do_search, gradio('db_test_select','db_input'), gradio('db_search_result')
+    )
+
 def do_refernce(algo_type,input_image):
 # def do_refernce():
     print("input image",input_image)
@@ -261,6 +325,93 @@ def point_to_mask(pil_image):
     points_array_reshaped = points_array.reshape(-1, 2)
     return points_array_reshaped
 
+
+def print_like_dislike(x: gr.LikeData):
+    print(x.index, x.value, x.liked)
+
+def do_llm_request(history, message):
+    for x in message["files"]:
+        history.append(((x,), None))
+    if message["text"] is not None:
+        history.append((message["text"], None))
+    return history, gr.MultimodalTextbox(value=None, interactive=False)
+
+def do_llm_response(history,selected_dbs):
+    user_input = history[-1][0]
+    prompt = ""
+    quote = ""
+    print("----------",selected_dbs)
+    if selected_dbs is not None and len(selected_dbs) != 0:
+        knowledge = knowledgeBase.retrieve_documents(selected_dbs,user_input)
+        print("do_llm_response context:",knowledge)
+        prompt = f'''
+背景1：{knowledge[0]["content"]}
+背景2：{knowledge[1]["content"]}
+背景3：{knowledge[2]["content"]}
+基于以上事实回答问题：{user_input}
+'''
+        print("do_llm_response prompt:",prompt)
+        quote = f'''
+> 文档：{knowledge[0]["meta"]["source"]}，页码：{knowledge[0]["meta"]["page"]}
+> 文档：{knowledge[1]["meta"]["source"]}，页码：{knowledge[1]["meta"]["page"]}
+> 文档：{knowledge[2]["meta"]["source"]}，页码：{knowledge[2]["meta"]["page"]}
+'''
+    else:
+        prompt = user_input
+
+    response = llm(prompt)
+    history[-1][1] = ""
+    response = response.removeprefix(prompt)
+    response += quote
+    for character in response:
+        history[-1][1] += character
+        time.sleep(0.01)
+        yield history
+
+def llm(input):
+    import requests
+    API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.3"
+    headers = {"Authorization": "Bearer "}
+
+    def query(payload):
+        response = requests.post(API_URL, headers=headers, json=payload)
+        return response.json()
+        
+    output = query({
+        "inputs": input,
+    })
+    print(output)
+    if len(output) >0:
+        return output[0]['generated_text']
+    return ""
+
+
+
+def file_handler(file_objs,name):
+    import shutil
+    import os
+    
+    print("file_obj:",file_objs)
+    
+    os.makedirs(os.path.dirname("./files/input/"), exist_ok=True)
+
+    for idx, file in enumerate(file_objs):
+        print(file)
+        file_path = "./files/input/" +  os.path.basename(file.name)
+        if not os.path.exists(file_path):
+            shutil.move(file.name,"./files/input/")
+        
+        knowledgeBase.add_documents_to_kb(name,[file_path])
+
+    dbs = knowledgeBase.get_bases()
+    dfs = knowledgeBase.get_df_bases()
+    return dfs,gr.CheckboxGroup(dbs,label="知识库", info="可选择1个或多个知识库"),gr.Dropdown(dbs,multiselect=True, label="知识库选择")
+
+def do_search(selected_dbs,user_input):
+    print("do_search:",selected_dbs,user_input)
+    context = knowledgeBase.retrieve_documents(selected_dbs,user_input)
+    return context
+
 if __name__ == "__main__":
     demo = create_ui()
     # demo.launch(server_name="10.151.124.137")

embedding.py

@@ -0,0 +1,69 @@
+from transformers import AutoModel, AutoTokenizer
+from sklearn.preprocessing import normalize
+from typing import Any, List, Mapping, Optional,Union
+from langchain.callbacks.manager import (
+    CallbackManagerForLLMRun
+)
+from langchain_core.embeddings import Embeddings
+import torch
+
+class Embedding(Embeddings):
+
+    def __init__(self,**kwargs):
+        self.model=AutoModel.from_pretrained('BAAI/bge-small-zh-v1.5')
+        self.tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-small-zh-v1.5')
+        self.model.eval()
+       
+    @property
+    def _llm_type(self) -> str:
+        return "BAAI/bge-small-zh-v1.5"
+    
+    @property
+    def model_name(self) -> str:
+        return "embedding"
+    
+    def _call(
+        self,
+        prompt: List[str],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        encoded_input = self.tokenizer(prompt, padding=True, truncation=True, return_tensors='pt')
+
+        with torch.no_grad():
+            model_output = self.model(**encoded_input)
+            # Perform pooling. In this case, cls pooling.
+            sentence_embeddings = model_output[0][:, 0]
+            print(sentence_embeddings.shape)
+        # normalize embeddings
+        sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
+        return sentence_embeddings.numpy()
+
+    @property
+    def _identifying_params(self) -> Mapping[str, Any]:
+        """Get the identifying parameters."""
+        return {"model_path": self.model_path}
+    
+    def embed_documents(self, texts) -> List[List[float]]:
+        # Embed a list of documents
+        embeddings = []
+        print("embed_documents:",len(texts),type(texts))
+        embedding = self._call(texts)
+        for row in embedding:
+            embeddings.append(row)
+        # print("embed_documents: shape",embeddings.shape)
+        return embeddings
+    
+    def embed_query(self, text) -> List[float]:
+        # Embed a single query
+        embedding = self._call([text])
+        return embedding[0]
+    
+
+# if __name__ == '__main__':
+#     sd = Embedding()
+#     v1 = sd.embed_query("他是一个人")
+#     v2 = sd.embed_query("他是一个好人")
+#     v3 = sd.embed_documents(["她是一条狗","他是一个人"])
+#     print(v1 @ v2.T)

requirements.txt

@@ -14,4 +14,10 @@ cloudpickle==2.2.1
 segment_anything @ git+https://github.com/facebookresearch/segment-anything.git
 paddlepaddle==2.6.1
 paddleocr==2.7.3
-easyocr==1.7.1
+easyocr==1.7.1
+scikit-learn==1.5.0
+faiss-cpu==1.8.0
+pypdf==4.2.0
+langchain==0.2.5
+langchain-community==0.2.5
+transformers==4.32.1

retriever.py

@@ -0,0 +1,150 @@
+from langchain_community.vectorstores import FAISS
+from langchain_community.document_loaders import TextLoader, JSONLoader, PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.docstore.in_memory import InMemoryDocstore
+import faiss
+import os
+import glob
+from typing import Any,List,Dict
+from embedding import Embedding
+
+
+class KnowledgeBaseManager:
+    def __init__(self, base_path="./knowledge_bases", embedding_dim=512, batch_size=16):
+        self.base_path = base_path
+        self.embedding_dim = embedding_dim
+        self.batch_size = batch_size
+        self.embeddings = Embedding()
+        self.knowledge_bases: Dict[str, FAISS] = {}
+        os.makedirs(self.base_path, exist_ok=True)
+
+        faiss_files = glob.glob(os.path.join(base_path, '*.faiss'))
+        # 获取不带后缀的名称
+        file_names_without_extension = [os.path.splitext(os.path.basename(file))[0] for file in faiss_files]
+        for name in file_names_without_extension:
+            self.load_knowledge_base(name)
+
+
+    def create_knowledge_base(self, name: str):
+        index = faiss.IndexFlatL2(self.embedding_dim)
+        kb = FAISS(self.embeddings, index, InMemoryDocstore(), {})
+        if name in self.knowledge_bases:
+            print(f"Knowledge base '{name}' already exists.")
+            return
+        
+        self.knowledge_bases[name] = kb
+        self.save_knowledge_base(name)
+        print(f"Knowledge base '{name}' created.")
+
+    def delete_knowledge_base(self, name: str):
+        if name in self.knowledge_bases:
+            del self.knowledge_bases[name]
+            os.remove(os.path.join(self.base_path, f"{name}.faiss"))
+            print(f"Knowledge base '{name}' deleted.")
+        else:
+            print(f"Knowledge base '{name}' does not exist.")
+
+    def load_knowledge_base(self, name: str):
+        kb_path = os.path.join(self.base_path, f"{name}.faiss")
+        if os.path.exists(kb_path):
+            self.knowledge_bases[name] = FAISS.load_local(self.base_path, self.embeddings, name, allow_dangerous_deserialization=True)
+            print(f"Knowledge base '{name}' loaded.")
+        else:
+            print(f"Knowledge base '{name}' does not exist.")
+
+    def save_knowledge_base(self, name: str):
+        if name in self.knowledge_bases:
+            self.knowledge_bases[name].save_local(self.base_path, name)
+            print(f"Knowledge base '{name}' saved.")
+        else:
+            print(f"Knowledge base '{name}' does not exist.")
+
+    # Document(page_content = '渠道版', metadata = {
+	# 'source': './files/input/PS004.pdf',
+	# 'page': 0
+    # }), Document(page_content = '2/20.', metadata = {
+    #     'source': './files/input/PS004.pdf',
+    #     'page': 1
+    # })
+    def add_documents_to_kb(self, name: str, file_paths: List[str]):
+        if name not in self.knowledge_bases:
+            print(f"Knowledge base '{name}' does not exist.")
+            self.create_knowledge_base(name)
+        
+        kb = self.knowledge_bases[name]
+        documents = self.load_documents(file_paths)
+        print(f"Loaded {len(documents)} documents.")
+        print(documents)
+        pages = self.split_documents(documents)
+        print(f"Split documents into {len(pages)} pages.")
+        # print(pages)
+        
+        doc_ids = []
+        for i in range(0, len(pages), self.batch_size):
+            batch = pages[i:i+self.batch_size]
+            doc_ids.extend(kb.add_documents(batch))
+        
+        self.save_knowledge_base(name)
+        return doc_ids
+
+    def load_documents(self, file_paths: List[str]):
+        documents = []
+        for file_path in file_paths:
+            loader = self.get_loader(file_path)
+            documents.extend(loader.load())
+        return documents
+
+    def get_loader(self, file_path: str):
+        if file_path.endswith('.txt'):
+            return TextLoader(file_path)
+        elif file_path.endswith('.json'):
+            return JSONLoader(file_path)
+        elif file_path.endswith('.pdf'):
+            return PyPDFLoader(file_path)
+        else:
+            raise ValueError("Unsupported file format")
+
+    def split_documents(self, documents):
+        text_splitter = RecursiveCharacterTextSplitter(separators=[
+                                                    "\n\n",
+                                                    "\n",
+                                                    " ",
+                                                    ".",
+                                                    ",",
+                                                    "\u200b",  # Zero-width space
+                                                    "\uff0c",  # Fullwidth comma
+                                                    "\u3001",  # Ideographic comma
+                                                    "\uff0e",  # Fullwidth full stop
+                                                    "\u3002",  # Ideographic full stop
+                                                    "",
+                                                ],
+                                                chunk_size=512, chunk_overlap=0)
+        return text_splitter.split_documents(documents)
+
+    def retrieve_documents(self, names: List[str], query: str):
+        results = []
+        for name in names:
+            if name not in self.knowledge_bases:
+                print(f"Knowledge base '{name}' does not exist.")
+                continue
+            
+            retriever = self.knowledge_bases[name].as_retriever(
+                search_type="mmr",
+                search_kwargs={"score_threshold": 0.5, "k": 3}
+            )
+            docs = retriever.get_relevant_documents(query)
+            results.extend([{"name": name, "content": doc.page_content,"meta": doc.metadata} for doc in docs])
+            
+        
+        return results
+    
+    def get_bases(self):
+        data = self.knowledge_bases.keys()
+        return list(data)
+    
+    def get_df_bases(self):
+        import pandas as pd
+        data = self.knowledge_bases.keys()
+        return pd.DataFrame(list(data), columns=['列表'])
+
+knowledgeBase = KnowledgeBaseManager()