rag知识库问答
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

495 lines
15 KiB

# 这是一个示例 Python 脚本。
# 按 Shift+F10 执行或将其替换为您的代码。
# 按 双击 Shift 在所有地方搜索类、文件、工具窗口、操作和设置。
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
import faiss
import numpy as np
from tqdm import tqdm
from sentence_transformers import SentenceTransformer
import requests
import time
from flask import Flask, jsonify, Response, request
from openai import OpenAI
from flask_cors import CORS
import pandas as pd
import concurrent.futures
import json
import torch
import uuid
# flask配置
app = Flask(__name__)
CORS(app)
app.config["JSON_AS_ASCII"] = False
openai_api_key = "token-abc123"
openai_api_base = "http://127.0.0.1:12011/v1"
client = OpenAI(
api_key=openai_api_key,
base_url=openai_api_base,
)
# 模型配置
models = client.models.list()
model = models.data[0].id
# model = "1"
model_encode = SentenceTransformer('/home/majiahui/project/models-llm/bge-large-zh-v1.5')
# 提示配置
propmt_connect = '''我是一名中医,你是一个中医的医生的助理,我的患者有一个症状,症状如下:
{}
根据这些症状,我通过查找资料,{}
请根据上面的这些资料和方子,并根据每篇文章的转发数确定文章的重要程度,转发数越高的文章,最终答案的参考度越高,反之越低。根据患者的症状和上面的文章的资料的重要程度以及文章和症状的匹配程度,帮我开出正确的药方和治疗方案'''
propmt_connect_ziliao = '''在“{}”资料中,有如下相关内容:
{}'''
def dialog_line_parse(text):
"""
将数据输入模型进行分析并输出结果
:param url: 模型url
:param text: 进入模型的数据
:return: 模型返回结果
"""
url_predict = "http://118.178.228.101:12004/predict"
response = requests.post(
url_predict,
json=text,
timeout=100000
)
if response.status_code == 200:
return response.json()
else:
# logger.error(
# "【{}】 Failed to get a proper response from remote "
# "server. Status Code: {}. Response: {}"
# "".format(url, response.status_code, response.text)
# )
print("{}】 Failed to get a proper response from remote "
"server. Status Code: {}. Response: {}"
"".format(url_predict, response.status_code, response.text))
return {}
# ['choices'][0]['message']['content']
#
# text = text['messages'][0]['content']
# return_text = {
# 'code': 200,
# 'id': "1",
# 'object': 0,
# 'created': 0,
# 'model': 0,
# 'choices': [
# {
# 'index': 0,
# 'message': {
# 'role': 'assistant',
# 'content': text
# },
# 'logprobs': None,
# 'finish_reason': 'stop'
# }
# ],
# 'usage': 0,
# 'system_fingerprint': 0
# }
# return return_text
def shengcehng_array(data):
'''
模型生成向量
:param data:
:return:
'''
embs = model_encode.encode(data, normalize_embeddings=True)
return embs
def Building_vector_database(title, df):
'''
次函数暂时弃用
:param title:
:param df:
:return:
'''
# 加载需要处理的数据(有效且未向量化)
to_process = df[(df["有效"] == True) & (df["已向量化"] == False)]
if len(to_process) == 0:
print("无新增数据需要向量化")
return
# 生成向量数组
new_vectors = shengcehng_array(to_process["总结"].tolist()) # 假设这是你的向量生成函数
# 加载现有向量库和索引
vector_path = f"data_np/{title}.npy"
index_path = f"data_np/{title}_index.json"
vectors = np.load(vector_path) if os.path.exists(vector_path) else np.empty((0, 1024))
index_data = {}
if os.path.exists(index_path):
with open(index_path, "r") as f:
index_data = json.load(f)
# 更新索引和向量库
start_idx = len(vectors)
vectors = np.vstack([vectors, new_vectors])
for i, (_, row) in enumerate(to_process.iterrows()):
index_data[row["ID"]] = {
"row": start_idx + i,
"valid": True
}
# 保存数据
np.save(vector_path, vectors)
with open(index_path, "w") as f:
json.dump(index_data, f)
# 标记已向量化
df.loc[to_process.index, "已向量化"] = True
df.to_csv(f"data_file_res/{title}.csv", sep="\t", index=False)
def delete_data(title, new_id):
'''
假删除,只是标记有效无效
:param title:
:param new_id:
:return:
'''
new_id = str(new_id)
# 更新CSV标记
csv_path = f"data_file_res/{title}.csv"
df = pd.read_csv(csv_path, sep="\t")
# df.loc[df["ID"] == new_id, "有效"] = False
df.loc[df['ID'] == new_id, "有效"] = False
df.to_csv(csv_path, sep="\t", index=False)
return "删除完成"
def check_file_exists(file_path):
"""
检查文件是否存在
参数:
file_path (str): 要检查的文件路径
返回:
bool: 文件存在返回True,否则返回False
"""
return os.path.isfile(file_path)
def ulit_request_file(sentence, title, zongjie):
'''
上传文件,生成固定内容,"ID", "正文", "总结", "有效", "向量"
:param sentence:
:param title:
:param zongjie:
:return:
'''
file_name_res_save = f"data_file_res/{title}.csv"
# 初始化或读取CSV文件,如果存在文件,读取文件,并添加行,
# 如果不存在文件,新建DataFrame
if os.path.exists(file_name_res_save):
df = pd.read_csv(file_name_res_save, sep="\t")
# 检查是否已存在相同正文
if sentence in df["正文"].values:
if zongjie == None:
return "正文已存在,跳过处理"
else:
result = df[df['正文'] == sentence]
id_ = result['ID'].values[0]
print(id_)
return ulit_request_file_zongjie(id_, sentence, zongjie, title)
else:
df = pd.DataFrame(columns=["ID", "正文", "总结", "有效", "向量"])
# 添加新数据(生成唯一ID)
if zongjie == None:
id_ = str(uuid.uuid1())
new_row = {
"ID": id_,
"正文": sentence,
"总结": None,
"有效": True,
"向量": None
}
df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
# 需要根据不同的项目修改提示,目的是精简内容,为了方便匹配
data_dan = {
"model": "gpt-4-turbo",
"messages": [{
"role": "user",
"content": f"{sentence}\n以上这条中可能包含了一些病情或者症状,请帮我归纳这条中所对应的病情或者症状是哪些,总结出来,不需要很长,简单归纳即可,直接输出症状或者病情,可以包含一些形容词来辅助描述,不需要有辅助词汇"
}],
"top_p": 0.9,
"temperature": 0.3
}
results = dialog_line_parse(data_dan)
summary = results['choices'][0]['message']['content']
# 这是你的向量生成函数,来生成总结的词汇的向量
new_vectors = shengcehng_array([summary])
df.loc[df['ID'] == id_, '总结'] = summary
df.loc[df['ID'] == id_, '向量'] = str(new_vectors[0].tolist())
else:
id_ = str(uuid.uuid1())
new_row = {
"ID": id_ ,
"正文": sentence,
"总结": zongjie,
"有效": True,
"向量": None
}
df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
new_vectors = shengcehng_array([zongjie]) # 假设这是你的向量生成函数
df.loc[df['ID'] == id_, '总结'] = zongjie
df.loc[df['ID'] == id_, '向量'] = str(new_vectors[0].tolist())
# 保存更新后的CSV
df.to_csv(file_name_res_save, sep="\t", index=False)
return "上传完成"
def ulit_request_file_zongjie(new_id, sentence, zongjie, title):
new_id = str(new_id)
print(new_id)
print(type(new_id))
file_name_res_save = f"data_file_res/{title}.csv"
# 初始化或读取CSV文件
df = pd.read_csv(file_name_res_save, sep="\t")
df.loc[df['ID'] == new_id, '正文'] = sentence
if zongjie == None:
pass
else:
df.loc[df['ID'] == new_id, '总结'] = zongjie
new_vectors = shengcehng_array([zongjie]) # 假设这是你的向量生成函数
df.loc[df['ID'] == new_id, '向量'] = str(new_vectors[0].tolist())
# 保存更新后的CSV
df.to_csv(file_name_res_save, sep="\t", index=False)
return "修改完成"
def ulit_request_file_check(title):
file_name_res_save = f"data_file_res/{title}.csv"
# 初始化或读取CSV文件
# 初始化或读取CSV文件
if os.path.exists(file_name_res_save):
df = pd.read_csv(file_name_res_save, sep="\t").values.tolist()
data_new = []
for i in df:
if i[3] == True:
data_new.append([i[0], i[1], i[2]])
return data_new
else:
return "无可展示文件"
def ulit_request_file_check_dan(new_id, title):
new_id = str(new_id)
file_name_res_save = f"data_file_res/{title}.csv"
# 初始化或读取CSV文件
# 初始化或读取CSV文件
if os.path.exists(file_name_res_save):
df = pd.read_csv(file_name_res_save, sep="\t")
zhengwen = df.loc[df['ID'] == new_id, '正文'].values
zongjie = df.loc[df['ID'] == new_id, '总结'].values
# 输出结果
if len(zhengwen) > 0:
if df.loc[df['ID'] == new_id, '有效'].values == True:
return [new_id, zhengwen[0], zongjie[0]]
else:
return "未找到对应的ID"
else:
return "未找到对应的ID"
else:
return "无可展示文件"
def main(question, title, top):
db_dict = {
"1": "yetianshi"
}
'''
定义文件路径
'''
'''
加载文件
'''
'''
文本分割
'''
'''
构建向量数据库
1. 正常匹配
2. 把文本使用大模型生成一个问题之后再进行匹配
'''
'''
根据提问匹配上下文
'''
d = 1024
db_type_list = title.split(",")
paper_list_str = ""
for title_dan in db_type_list:
embs = shengcehng_array([question])
index = faiss.IndexFlatIP(d) # buid the index
# 查找向量
# vector_path = f"data_np/{title_dan}.npy"
# vectors = np.load(vector_path)
data_str = pd.read_csv(f"data_file_res/{title_dan}.csv", sep="\t", encoding="utf-8").values.tolist()
data_str_valid = []
for i in data_str:
if i[3] == True:
data_str_valid.append(i)
data_str_vectors_list = []
for i in data_str_valid:
data_str_vectors_list.append(eval(i[-1]))
vectors = np.array(data_str_vectors_list)
index.add(vectors)
D, I = index.search(embs, int(top))
print(I)
reference_list = []
for i,j in zip(I[0], D[0]):
reference_list.append([data_str_valid[i], j])
for i,j in enumerate(reference_list):
paper_list_str += "{}\n{},此篇文章跟问题的相关度为{}%\n".format(str(i+1), j[0][1], j[1])
'''
构造prompt
'''
print("paper_list_str", paper_list_str)
propmt_connect_ziliao_input = []
for i in db_type_list:
propmt_connect_ziliao_input.append(propmt_connect_ziliao.format(i, paper_list_str))
propmt_connect_ziliao_input_str = "".join(propmt_connect_ziliao_input)
propmt_connect_input = propmt_connect.format(question, propmt_connect_ziliao_input_str)
'''
生成回答
'''
return model_generate_stream(propmt_connect_input)
def model_generate_stream(prompt):
messages = [
{"role": "user", "content": prompt}
]
stream = client.chat.completions.create(model=model,
messages=messages,
stream=True)
printed_reasoning_content = False
printed_content = False
for chunk in stream:
reasoning_content = None
content = None
# Check the content is reasoning_content or content
if hasattr(chunk.choices[0].delta, "reasoning_content"):
reasoning_content = chunk.choices[0].delta.reasoning_content
elif hasattr(chunk.choices[0].delta, "content"):
content = chunk.choices[0].delta.content
if reasoning_content is not None:
if not printed_reasoning_content:
printed_reasoning_content = True
print("reasoning_content:", end="", flush=True)
print(reasoning_content, end="", flush=True)
elif content is not None:
if not printed_content:
printed_content = True
print("\ncontent:", end="", flush=True)
# Extract and print the content
# print(content, end="", flush=True)
print(content, end="")
yield content
@app.route("/upload_file_check", methods=["POST"])
def upload_file_check():
print(request.remote_addr)
sentence = request.form.get('sentence')
title = request.form.get("title")
new_id = request.form.get("id")
zongjie = request.form.get("zongjie")
state = request.form.get("state")
'''
{
"1": "csv",
"2": "xlsx",
"3": "txt",
"4": "pdf"
}
'''
# 增
state_res = ""
if state == "1":
state_res = ulit_request_file(sentence, title, zongjie)
# 删
elif state == "2":
state_res = delete_data(title, new_id)
# 改
elif state == "3":
state_res = ulit_request_file_zongjie(new_id, sentence, zongjie,title)
# 查
elif state == "4":
state_res = ulit_request_file_check(title)
# 通过uuid查单条数据
elif state == "5":
ulit_request_file_check_dan(new_id, title)
state_res = ""
return_json = {
"code": 200,
"info": state_res
}
return jsonify(return_json) # 返回结果
@app.route("/search", methods=["POST"])
def search():
print(request.remote_addr)
texts = request.json["texts"]
title = request.json["title"]
top = request.json["top"]
response = main(texts, title, top)
return Response(response, mimetype='text/plain; charset=utf-8') # 返回结果
if __name__ == "__main__":
app.run(host="0.0.0.0", port=27000, threaded=True, debug=False)