finish rag

.gitignore

@@ -2,3 +2,6 @@
 __pycache__/
 *.bin
 .vscode/
+files/input/ir2023_ashare.pdf
+knowledge_bases/中国移动.faiss
+knowledge_bases/中国移动.pkl

app.py

@@ -133,23 +133,23 @@ def create_ui():
                         components["db_view"] = gr.Dataframe(
                                                     headers=["列表"],
                                                     datatype=["str"],
-                                                    row_count=8,
+                                                    row_count=2,
                                                     col_count=(1, "fixed"),
                                                     interactive=False
                         )
                 with gr.Column(scale=2):
-                        with gr.Group():
+                    with gr.Row():
+                        with gr.Column(scale=2):
                             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'])
+                        with gr.Column(scale=2):
                             components["db_submit_btn"] = gr.Button(value="提交")
+                    components["file_upload"] = gr.File(elem_id='file_upload',file_count='multiple',label='文档上传', file_types=[".pdf", ".doc", '.docx', '.json', '.csv'])
             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["db_test_select"] = gr.Dropdown(knowledgeBase.get_bases(),multiselect=True, label="知识库选择")
+                with gr.Column(scale=1):
                     components["dbtest_submit_btn"] = gr.Button(value="检索")
             with gr.Row():
                 with gr.Group():
@@ -157,16 +157,22 @@ def create_ui():
 
         with gr.Tab("问答"):
             with gr.Row():
-                with gr.Column():
+                with gr.Column(scale=1):
+                    with gr.Group():
+                        components["ak"] = gr.Textbox(label="appid")
+                        components["sk"] = gr.Textbox(label="secret")
+                        components["llm_client"] =gr.Radio(["Wenxin", "Tongyi","Huggingface"],value="Wenxin", label="llm")
+                        components["llm_setting_btn"] =  gr.Button(value="设置")
+                with gr.Column(scale=2):
                     with gr.Group():
                         components["chatbot"] = gr.Chatbot(
-                                            [(None,"What can I help you?")],
+                                            [(None,"你好，有什么需要帮助的？")],
                                             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个或多个知识库")
+                        components["db_select"] = gr.CheckboxGroup(knowledgeBase.get_bases(),label="知识库", info="可选择1个或多个知识库")
         create_event_handlers()
         demo.load(init,None,gradio("db_view"))
     return demo
@@ -236,6 +242,10 @@ def create_event_handlers():
         do_search, gradio('db_test_select','db_input'), gradio('db_search_result')
     )
 
+    components['llm_setting_btn'].click(
+        llm, gradio('ak','sk','llm_client'), None
+    )
+
 def do_refernce(algo_type,input_image):
 # def do_refernce():
     print("input image",input_image)
@@ -307,9 +317,6 @@ def do_sam_everything(im):
         
     return images
 
-
-    
-
 def point_to_mask(pil_image):
     # 遍历每个像素
     width, height = pil_image.size
@@ -337,11 +344,11 @@ def do_llm_request(history, message):
     return history, gr.MultimodalTextbox(value=None, interactive=False)
 
 def do_llm_response(history,selected_dbs):
+    print("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:
+    if len(selected_dbs) > 0:
         knowledge = knowledgeBase.retrieve_documents(selected_dbs,user_input)
         print("do_llm_response context:",knowledge)
         prompt = f'''
@@ -349,8 +356,8 @@ def do_llm_response(history,selected_dbs):
 背景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"]}
@@ -358,33 +365,41 @@ def do_llm_response(history,selected_dbs):
 '''
     else:
         prompt = user_input
-
-    response = llm(prompt)
+    
     history[-1][1] = ""
-    response = response.removeprefix(prompt)
-    response += quote
+    if llm_client is None:
+        gr.Warning("请先设置大模型")
+        response = "模型参数未设置"
+    else:
+        print("do_llm_response prompt:",prompt)
+        response = llm_client(prompt)
+        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 ""
+llm_client = None
+def llm(ak,sk,client):
+    global llm_client
+    import llm
+    llm.init_param(ak,sk)
+    if client == "Wenxin":
+        llm_client = llm.baidu_client
+    elif client == "Tongyi":
+        llm_client = llm.qwen_agent_app
+    elif client == "Huggingface":
+        llm_client = llm.hg_client
+    
+    if ak == "" and sk == "":
+        gr.Info("重置成功")
+    else:
+        gr.Info("设置成功")
 
+    return llm_client
 
 
 def file_handler(file_objs,name):

llm.py

@@ -0,0 +1,89 @@
+import requests
+import json
+from http import HTTPStatus
+from dashscope import Application
+
+ak = ""
+sk = ""
+
+def init_param(access_key,secret_key):
+    global ak, sk
+    ak = access_key
+    sk = secret_key
+
+
+def baidu_client(input):
+    global ak, sk   
+    if ak == "" or sk == "":
+        return ""
+    
+    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie-lite-8k?access_token=" + get_access_token()
+    
+    payload = json.dumps({
+        "temperature": 0.95,
+        "top_p": 0.7,
+        "penalty_score": 1,
+        "messages": [
+            {
+                "role": "user",
+                "content": input
+            }
+        ],
+        "system": ""
+    })
+    headers = {
+        'Content-Type': 'application/json'
+    }
+    
+    response = requests.request("POST", url, headers=headers, data=payload)
+    
+    print("baidu_client",response.text)
+    return response.json()["result"]
+    
+
+def get_access_token():
+    """
+    使用 AK，SK 生成鉴权签名（Access Token）
+    :return: access_token，或是None(如果错误)
+    """
+    url = "https://aip.baidubce.com/oauth/2.0/token"
+    params = {"grant_type": "client_credentials", "client_id": ak, "client_secret": sk}
+    return str(requests.post(url, params=params).json().get("access_token"))
+
+
+def qwen_agent_app(input):
+    global ak, sk   
+    if ak == "" or sk == "":
+        return ""
+    response = Application.call(app_id=ak,
+                                prompt=input,
+                                api_key=sk,
+                                )
+
+    if response.status_code != HTTPStatus.OK:
+        print('request_id=%s, code=%s, message=%s\n' % (response.request_id, response.status_code, response.message))
+        return ""
+    else:
+        print('request_id=%s\n output=%s\n usage=%s\n' % (response.request_id, response.output, response.usage))
+        return response.output["text"]
+    
+
+def hg_client(input):
+    global ak, sk   
+    if sk == "":
+        return ""
+    import requests
+    API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.3"
+    headers = {"Authorization": f"Bearer {sk}"}
+
+    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 ""

requirements.txt

@@ -20,4 +20,5 @@ faiss-cpu==1.8.0
 pypdf==4.2.0
 langchain==0.2.5
 langchain-community==0.2.5
-transformers==4.32.1
+transformers==4.32.1
+dashscope==1.20.0