Update app.py

app.py

@@ -14,22 +14,37 @@ history_max_len = 500  # 机器人记忆的最大长度
 all_max_len = 2000  # 输入的最大长度
 
 
-def get_text_emb(open_ai_key, text):
-    openai.api_key = open_ai_key
+def get_text_emb(open_ai_key, text):  # 文本向量化
+    openai.api_key = open_ai_key  # 设置openai的key
     response = openai.Embedding.create(
         input=text,
         model="text-embedding-ada-002"
-    )
-    return response['data'][0]['embedding']
+    )  # 调用openai的api
+    return response['data'][0]['embedding']  # 返回向量
 
 
 def doc_index_self(open_ai_key, doc):  # 文档向量化
     texts = doc.split('\n')  # 按行切分
-    emb_list = []
-    for text in texts:
-        emb_list.append(get_text_emb(open_ai_key, text))
+    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), 1
+        value="""操作说明 step 3：建立索引（by self）成功！ 🙋 可以开始对话啦~"""), gr.Chatbot.update(visible=True), 1
+
+
+def doc_index_llama(open_ai_key, txt):  # 建立索引
+    # 根据时间戳新建目录，保存txt文件
+    path = str(time.time())
+    import os
+    os.mkdir(path)
+    with open(path + '/doc.txt', mode='w', encoding='utf-8') as f:
+        f.write(txt)
+    openai.api_key = open_ai_key  # 设置OpenAI API Key
+    documents = SimpleDirectoryReader(path).load_data()  # 读取文档
+    index = GPTVectorStoreIndex.from_documents(documents)  # 建立索引
+    query_engine = index.as_query_engine()  # 建立查询引擎
+    return query_engine, gr.Textbox.update(visible=True), gr.Button.update(visible=True), gr.Markdown.update(
+        value="""操作说明 step 3：建立索引（by llama_index）成功！ 🙋 可以开始对话啦~"""), gr.Chatbot.update(visible=True), 0
 
 
 def get_response_by_self(open_ai_key, msg, bot, doc_text_list, doc_embeddings):  # 获取机器人回复
@@ -44,8 +59,8 @@ def get_response_by_self(open_ai_key, msg, bot, doc_text_list, doc_embeddings):
     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))
+    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))  # 计算相似度
@@ -111,7 +126,6 @@ def get_response_by_llama_index(open_ai_key, msg, bot, query_engine):  # 获取
         query_str += "机器人：" + his[1] + "\n"  # 加入机器人的历史记录
     query_str += "用户：" + msg + "\n"  # 加入用户的当前输入
     qa_template = Prompt(template)  # 将模板转换成Prompt对象
-    query_engine = query_engine.as_query_engine(text_qa_template=qa_template)  # 建立查询引擎
     res = query_engine.query(msg)  # 获取回答
     print(res)  # 显示回答
     bot.append([msg, res])  # 加入历史记录
@@ -119,9 +133,9 @@ def get_response_by_llama_index(open_ai_key, msg, bot, query_engine):  # 获取
 
 
 def get_response(open_ai_key, msg, bot, doc_text_list, doc_embeddings, query_engine, index_type):  # 获取机器人回复
-    if index_type == 1:
+    if index_type == 1:  # 如果是使用自己的索引
         return get_response_by_self(open_ai_key, msg, bot, doc_text_list, doc_embeddings)
-    else:
+    else:  # 如果是使用llama_index索引
         return get_response_by_llama_index(open_ai_key, msg, bot, query_engine)
 
 
@@ -166,21 +180,6 @@ def up_file(files):  # 上传文件
         value="操作说明 step 2：确认PDF解析结果（可修正），点击“建立索引”，随后进行对话")
 
 
-def doc_index_llama(open_ai_key, txt):  # 建立索引
-    # 根据时间戳新建目录，保存txt文件
-    path = str(time.time())
-    import os
-    os.mkdir(path)
-    with open(path + '/doc.txt', mode='w', encoding='utf-8') as f:
-        f.write(txt)
-    openai.api_key = open_ai_key  # 设置OpenAI API Key
-    documents = SimpleDirectoryReader(path).load_data()  # 读取文档
-    index = GPTVectorStoreIndex.from_documents(documents)  # 建立索引
-    query_engine = index.as_query_engine()  # 建立查询引擎
-    return query_engine, gr.Textbox.update(visible=True), gr.Button.update(visible=True), gr.Markdown.update(
-        value="""操作说明 step 3：PDF解析提交成功！ 🙋 可以开始对话啦~"""), gr.Chatbot.update(visible=True), 0
-
-
 with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():