Update app.py

app.py

@@ -5,24 +5,35 @@ import pdfplumber
 import pandas as pd
 import time
 from cnocr import CnOcr
-from sentence_transformers import SentenceTransformer, models, util
+import numpy as np
 
-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',
-}  # 你的headers
 history_max_len = 500  # 机器人记忆的最大长度
 all_max_len = 3000  # 输入的最大长度
 
 
-def doc_emb(doc):  # 文档向量化
+def get_text_emb(open_ai_key, text):
+    url = 'https://api.openai.com/v1/embeddings'
+    headers = {
+        'Content-Type': 'application/json',
+        'Authorization': 'Bearer ' + open_ai_key
+    }
+    data = {
+        "model": "text-embedding-ada-002",
+        "input": text
+    }
+    result = requests.post(url=url,
+                           data=json.dumps(data),
+                           headers=headers
+                           )
+    return result.json()['data'][0]['embedding']
+
+
+def doc_index_self(open_ai_key, doc):  # 文档向量化
     texts = doc.split('\n')  # 按行切分
-    emb_list = embedder.encode(texts)  # 句向量化
-    print('\n'.join(texts))
+    emb_list = []
+    for text in texts:
+        emb_list.append(get_text_emb(open_ai_key, text))
     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)
 
@@ -36,9 +47,17 @@ def get_response(open_ai_key, msg, bot, doc_text_list, doc_embeddings):  # 获
         now_len += len(bot[i][0]) + len(bot[i][1])  # 更新当前长度
         his_bg = i  # 更新历史记录的起始位置
     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))])]  # 相似度和索引对应
+    query_embedding = get_text_emb(open_ai_key, msg)  # 获取输入的向量
+    cos_scores = []  # 用于存储相似度
+
+    def cos_sim(a, b):
+        return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
+
+    for doc_embedding in doc_embeddings:  # 遍历文档向量
+        cos_scores.append(cos_sim(query_embedding, doc_embedding))  # 计算相似度
+    score_index = []  # 用于存储相似度和索引对应
+    for i in range(len(cos_scores)):  # 遍历相似度
+        score_index.append((cos_scores[i], i))  # 加入相似度和索引对应
     score_index.sort(key=lambda x: x[0], reverse=True)  # 按相似度排序
     print('score_index:\n', score_index)
     index_set, sub_doc_list = set(), []  # 用于存储最终的索引和文档
@@ -75,14 +94,24 @@ def get_response(open_ai_key, msg, bot, doc_text_list, doc_embeddings):  # 获
         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"}  # 请求json
-    data = {"content": json.dumps(req_json)}  # 请求data
-    print('data:\n', req_json)
-    result = requests.post(url=chat_url,
+
+    url = 'https://api.openai.com/v1/chat/completions'
+
+    data = {
+        "model": "gpt-3.5-turbo",
+        "messages": messages
+    }
+    print("data = \n", data)
+
+    headers = {
+        'Content-Type': 'application/json',
+        'Authorization': 'Bearer ' + open_ai_key
+    }
+    result = requests.post(url=url,
                            data=json.dumps(data),
                            headers=headers
-                           )  # 请求
-    res = result.json()['content']  # 获取回复
+                           )
+    res = str(result.json()['choices'][0]['message']['content']).strip()
     bot.append([msg, res])  # 加入历史记录
     return bot[max(0, len(bot) - 3):]  # 返回最近3轮的历史记录
 
@@ -124,7 +153,7 @@ def up_file(files):  # 上传文件
     print(doc_text_list)
     return gr.Textbox.update(value='\n'.join(doc_text_list), visible=True), gr.Button.update(
         visible=True), gr.Markdown.update(
-        value="操作说明 step 2：确认PDF解析结果（可修正），点击“提交解析结果”，随后进行对话")
+        value="操作说明 step 2：确认PDF解析结果（可修正），点击“建立索引”，随后进行对话")
 
 
 with gr.Blocks() as demo:
@@ -134,7 +163,8 @@ with gr.Blocks() as demo:
             file = gr.File(file_types=['.pdf'], label='点击上传PDF，进行解析(支持多文档、表格、OCR)',
                            file_count='multiple')  # 支持多文档、表格、OCR
             txt = gr.Textbox(label='PDF解析结果', visible=False)  # PDF解析结果
-            doc_bu = gr.Button(value='提交解析结果', visible=False)  # 提交解析结果
+            index_self_bu = gr.Button(value='建立索引（by self）', visible=False)  #
+            index_llama_bu = gr.Button(value='建立索引（by llama_index）', visible=False)  #
             doc_text_state = gr.State([])  # 存储PDF解析结果
             doc_emb_state = gr.State([])  # 存储PDF解析结果的embedding
         with gr.Column():
@@ -144,8 +174,9 @@ with gr.Blocks() as demo:
             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])  # 提交解析结果
+    file.change(up_file, [file], [txt, index_self_bu, md])  # 上传文件
+    index_self_bu.click(doc_index_self, [txt],
+                        [doc_text_state, doc_emb_state, msg_txt, chat_bu, md, 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__":