Update app.py

app.py

@@ -1,64 +1,41 @@
 import requests
 import json
 import gradio as gr
-# from concurrent.futures import ThreadPoolExecutor
 import pdfplumber
 import pandas as pd
 import time
 from cnocr import CnOcr
 from sentence_transformers import SentenceTransformer, models, util
-word_embedding_model = models.Transformer('uer/sbert-base-chinese-nli', do_lower_case=True)
-pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), pooling_mode='cls')
-embedder = SentenceTransformer(modules=[word_embedding_model, pooling_model])
-ocr = CnOcr()
-# chat_url = 'https://souljoy-my-api.hf.space/sale'
-chat_url = 'https://souljoy-my-api.hf.space/chatpdf'
+
+word_embedding_model = models.Transformer('uer/sbert-base-chinese-nli', do_lower_case=True)  # BERT模型
+pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), pooling_mode='cls')  # 取cls向量作为句向量
+embedder = SentenceTransformer(modules=[word_embedding_model, pooling_model])  # 定义模型
+ocr = CnOcr()  # 初始化ocr模型
+chat_url = 'https://souljoy-my-api.hf.space/chatgpt'  # 你的url
 headers = {
     'Content-Type': 'application/json',
-}
-# thread_pool_executor = ThreadPoolExecutor(max_workers=4)
-history_max_len = 500
-all_max_len = 3000
-
-
-def get_emb(text):
-    emb_url = 'https://souljoy-my-api.hf.space/embeddings'
-    data = {"content": text}
-    try:
-        result = requests.post(url=emb_url,
-                               data=json.dumps(data),
-                               headers=headers
-                               )
-        return result.json()['data'][0]['embedding']
-    except Exception as e:
-        print('data', data, 'result json', result.json())
+}  # 你的headers
+history_max_len = 500  # 机器人记忆的最大长度
+all_max_len = 3000  # 输入的最大长度
 
 
-def doc_emb(doc: str):
-    texts = doc.split('\n')
-    # futures = []
-    emb_list = embedder.encode(texts)
-    # for text in texts:
-    #     futures.append(thread_pool_executor.submit(get_emb, text))
-    # for f in futures:
-    #     emb_list.append(f.result())
+def doc_emb(doc):  # 文档向量化
+    texts = doc.split('\n')  # 按行切分
+    emb_list = embedder.encode(texts)  # 句向量化
     print('\n'.join(texts))
     return texts, emb_list, gr.Textbox.update(visible=True), gr.Button.update(visible=True), gr.Markdown.update(
         value="""操作说明 step 3：PDF解析提交成功！ 🙋 可以开始对话啦~"""), gr.Chatbot.update(visible=True)
 
 
-def get_response(msg, bot, doc_text_list, doc_embeddings):
-    # future = thread_pool_executor.submit(get_emb, msg)
+def get_response(open_ai_key, msg, bot, doc_text_list, doc_embeddings):
     now_len = len(msg)
-    req_json = {'question': msg}
     his_bg = -1
     for i in range(len(bot) - 1, -1, -1):
         if now_len + len(bot[i][0]) + len(bot[i][1]) > history_max_len:
             break
         now_len += len(bot[i][0]) + len(bot[i][1])
         his_bg = i
-    req_json['history'] = [] if his_bg == -1 else bot[his_bg:]
-    # query_embedding = future.result()
+    history = [] if his_bg == -1 else bot[his_bg:]
     query_embedding = embedder.encode([msg])
     cos_scores = util.cos_sim(query_embedding, doc_embeddings)[0]
     score_index = [[score, index] for score, index in zip(cos_scores, [i for i in range(len(cos_scores))])]
@@ -72,24 +49,33 @@ def get_response(msg, bot, doc_text_list, doc_embeddings):
         index_set.add(s_i[1])
         now_len += len(doc)
         # 可能段落截断错误，所以把上下段也加入进来
-        if s_i[1] > 0 and s_i[1] -1 not in index_set:
-            doc = doc_text_list[s_i[1]-1]
+        if s_i[1] > 0 and s_i[1] - 1 not in index_set:
+            doc = doc_text_list[s_i[1] - 1]
             if now_len + len(doc) > all_max_len:
                 break
-            index_set.add(s_i[1]-1)
+            index_set.add(s_i[1] - 1)
             now_len += len(doc)
         if s_i[1] + 1 < len(doc_text_list) and s_i[1] + 1 not in index_set:
-            doc = doc_text_list[s_i[1]+1]
+            doc = doc_text_list[s_i[1] + 1]
             if now_len + len(doc) > all_max_len:
                 break
-            index_set.add(s_i[1]+1)
+            index_set.add(s_i[1] + 1)
             now_len += len(doc)
 
     index_list = list(index_set)
     index_list.sort()
     for i in index_list:
         sub_doc_list.append(doc_text_list[i])
-    req_json['doc'] = '' if len(sub_doc_list) == 0 else '\n'.join(sub_doc_list)
+    document = '' if len(sub_doc_list) == 0 else '\n'.join(sub_doc_list)
+    messages = [{
+        "role": "system",
+        "content": "你是一个有用的助手，可以使用文章内容准确地回答问题。使用提供的文章来生成你的答案，但避免逐字复制文章。尽可能使用自己的话。准确、有用、简洁、清晰。"
+    }, {"role": "system", "content": "文章内容：\n" + document}]
+    for his in history:
+        messages.append({"role": "user", "content": his[0]})
+        messages.append({"role": "assistant", "content": his[1]})
+    messages.append({"role": "user", "content": msg})
+    req_json = {'messages': messages, 'key': open_ai_key, 'model': "gpt-3.5-turbo"}
     data = {"content": json.dumps(req_json)}
     print('data:\n', req_json)
     result = requests.post(url=chat_url,
@@ -146,21 +132,23 @@ def up_file(files):
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
-            file = gr.File(file_types=['.pdf'], label='点击上传PDF，进行解析(支持多文档、表格、OCR)', file_count='multiple')
-            doc_bu = gr.Button(value='提交解析结果', visible=False)
-            txt = gr.Textbox(label='PDF解析结果', visible=False)
-            doc_text_state = gr.State([])
-            doc_emb_state = gr.State([])
+            open_ai_key = gr.Textbox(label='OpenAI API Key', placeholder='输入你的OpenAI API Key')  # 你的OpenAI API Key
+            file = gr.File(file_types=['.pdf'], label='点击上传PDF，进行解析(支持多文档、表格、OCR)',
+                           file_count='multiple')  # 支持多文档、表格、OCR
+            doc_bu = gr.Button(value='开始PDF解析', visible=False)  # 开始PDF解析
+            txt = gr.Textbox(label='PDF解析结果', visible=False)  # PDF解析结果
+            doc_text_state = gr.State([])  # 存储PDF解析结果
+            doc_emb_state = gr.State([])  # 存储PDF解析结果的embedding
         with gr.Column():
-            md = gr.Markdown("""操作说明 step 1：点击左侧区域，上传PDF，进行解析""")
-            chat_bot = gr.Chatbot(visible=False)
-            msg_txt = gr.Textbox(label='消息框', placeholder='输入消息，点击发送', visible=False)
-            chat_bu = gr.Button(value='发送', visible=False)
+            md = gr.Markdown("""操作说明 step 1：点击左侧区域，上传PDF，进行解析""")  # 操作说明
+            chat_bot = gr.Chatbot(visible=False)  # 聊天机器人
+            msg_txt = gr.Textbox(label='消息框', placeholder='输入消息，点击发送', visible=False)  # 消息框
+            with gr.Row():
+                chat_bu = gr.Button(value='发送', visible=False)
 
     file.change(up_file, [file], [txt, doc_bu, md])
     doc_bu.click(doc_emb, [txt], [doc_text_state, doc_emb_state, msg_txt, chat_bu, md, chat_bot])
-    chat_bu.click(get_response, [msg_txt, chat_bot, doc_text_state, doc_emb_state], [chat_bot])
+    chat_bu.click(get_response, [open_ai_key, msg_txt, chat_bot, doc_text_state, doc_emb_state], [chat_bot])
 
 if __name__ == "__main__":
     demo.queue().launch()
-    # demo.queue().launch(share=False, server_name='172.22.2.54', server_port=9191)