262 lines
9.9 KiB
Python
262 lines
9.9 KiB
Python
import json
|
||
import pickle
|
||
from loguru import logger
|
||
from sentence_transformers import SentenceTransformer
|
||
|
||
from config.config import embedding_path, doc_dir, qa_dir, knowledge_pkl_path, data_dir, base_dir, vector_db_dir
|
||
import os
|
||
import faiss
|
||
import platform
|
||
from langchain_community.document_loaders import DirectoryLoader, TextLoader, JSONLoader
|
||
from langchain_text_splitters import CharacterTextSplitter, RecursiveCharacterTextSplitter
|
||
from BCEmbedding import EmbeddingModel, RerankerModel
|
||
from util.pipeline import EmoLLMRAG
|
||
import pickle
|
||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||
import torch
|
||
import streamlit as st
|
||
from openxlab.model import download
|
||
|
||
|
||
'''
|
||
1）根据QA对/TXT 文本生成 embedding
|
||
2）调用 langchain FAISS 接口构建 vector DB
|
||
3）存储到 openxlab.dataset 中，方便后续调用
|
||
4）提供 embedding 的接口函数，方便后续调用
|
||
5）提供 rerank 的接口函数，方便后续调用
|
||
'''
|
||
|
||
"""
|
||
加载向量模型
|
||
"""
|
||
def load_embedding_model():
|
||
logger.info('Loading embedding model...')
|
||
# model = EmbeddingModel(model_name_or_path="huggingface/bce-embedding-base_v1")
|
||
model = EmbeddingModel(model_name_or_path="maidalun1020/bce-embedding-base_v1")
|
||
logger.info('Embedding model loaded.')
|
||
return model
|
||
|
||
def load_rerank_model():
|
||
logger.info('Loading rerank_model...')
|
||
model = RerankerModel(model_name_or_path="maidalun1020/bce-reranker-base_v1")
|
||
# model = RerankerModel(model_name_or_path="huggingface/bce-reranker-base_v1")
|
||
logger.info('Rerank model loaded.')
|
||
return model
|
||
|
||
|
||
def split_document(data_path, chunk_size=1000, chunk_overlap=100):
|
||
# text_spliter = CharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
||
text_spliter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
||
split_docs = []
|
||
logger.info(f'Loading txt files from {data_path}')
|
||
if os.path.isdir(data_path):
|
||
# 如果是文件夹，则遍历读取
|
||
for root, dirs, files in os.walk(data_path):
|
||
for file in files:
|
||
if file.endswith('.txt'):
|
||
file_path = os.path.join(root, file)
|
||
# logger.info(f'splitting file {file_path}')
|
||
text_loader = TextLoader(file_path, encoding='utf-8')
|
||
text = text_loader.load()
|
||
|
||
splits = text_spliter.split_documents(text)
|
||
# logger.info(f"splits type {type(splits[0])}")
|
||
# logger.info(f'splits size {len(splits)}')
|
||
split_docs += splits
|
||
elif data_path.endswith('.txt'):
|
||
file_path = os.path.join(root, data_path)
|
||
# logger.info(f'splitting file {file_path}')
|
||
text_loader = TextLoader(file_path, encoding='utf-8')
|
||
text = text_loader.load()
|
||
splits = text_spliter.split_documents(text)
|
||
# logger.info(f"splits type {type(splits[0])}")
|
||
# logger.info(f'splits size {len(splits)}')
|
||
split_docs = splits
|
||
logger.info(f'split_docs size {len(split_docs)}')
|
||
return split_docs
|
||
|
||
|
||
##TODO 1、读取system prompt 2、限制序列长度
|
||
def split_conversation(path):
|
||
'''
|
||
data format:
|
||
[
|
||
{
|
||
"conversation": [
|
||
{
|
||
"input": Q1
|
||
"output": A1
|
||
},
|
||
{
|
||
"input": Q2
|
||
"output": A2
|
||
},
|
||
]
|
||
},
|
||
]
|
||
'''
|
||
qa_pairs = []
|
||
logger.info(f'Loading json files from {path}')
|
||
if os.path.isfile(path):
|
||
with open(path, 'r', encoding='utf-8') as file:
|
||
data = json.load(file)
|
||
for conversation in data:
|
||
for dialog in conversation['conversation']:
|
||
# input_text = dialog['input']
|
||
# output_text = dialog['output']
|
||
# if len(input_text) > max_length or len(output_text) > max_length:
|
||
# continue
|
||
qa_pairs.append(dialog)
|
||
elif os.path.isdir(path):
|
||
# 如果是文件夹，则遍历读取
|
||
for root, dirs, files in os.walk(path):
|
||
for file in files:
|
||
if file.endswith('.json'):
|
||
file_path = os.path.join(root, file)
|
||
logger.info(f'splitting file {file_path}')
|
||
with open(file_path, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
for conversation in data:
|
||
for dialog in conversation['conversation']:
|
||
qa_pairs.append(dialog)
|
||
return qa_pairs
|
||
|
||
|
||
|
||
# 加载本地索引
|
||
def load_index_and_knowledge():
|
||
current_os = platform.system()
|
||
split_doc = []
|
||
split_qa = []
|
||
#读取知识库
|
||
if not os.path.exists(knowledge_pkl_path):
|
||
split_doc = split_document(doc_dir)
|
||
split_qa = split_conversation(qa_dir)
|
||
# logger.info(f'split_qa size:{len(split_qa)}')
|
||
# logger.info(f'type of split_qa:{type(split_qa[0])}')
|
||
# logger.info(f'split_doc size:{len(split_doc)}')
|
||
# logger.info(f'type of doc:{type(split_doc[0])}')
|
||
knowledge_chunks = split_doc + split_qa
|
||
with open(knowledge_pkl_path, 'wb') as file:
|
||
pickle.dump(knowledge_chunks, file)
|
||
else:
|
||
with open(knowledge_pkl_path , 'rb') as f:
|
||
knowledge_chunks = pickle.load(f)
|
||
|
||
#读取vector DB
|
||
if not os.path.exists(vector_db_dir):
|
||
logger.info(f'Creating index...')
|
||
emb_model = load_embedding_model()
|
||
if not split_doc:
|
||
split_doc = split_document(doc_dir)
|
||
if not split_qa:
|
||
split_qa = split_conversation(qa_dir)
|
||
# 创建索引,windows不支持faiss-gpu
|
||
if current_os == 'Linux':
|
||
index = create_index_gpu(split_doc, split_qa, emb_model, vector_db_dir)
|
||
else:
|
||
index = create_index_cpu(split_doc, split_qa, emb_model, vector_db_dir)
|
||
else:
|
||
if current_os == 'Linux':
|
||
res = faiss.StandardGpuResources()
|
||
index = faiss.index_cpu_to_gpu(res, 0, index, vector_db_dir)
|
||
else:
|
||
index = faiss.read_index(vector_db_dir)
|
||
|
||
return index, knowledge_chunks
|
||
|
||
|
||
def create_index_cpu(split_doc, split_qa, emb_model, knowledge_pkl_path, dimension = 768, question_only=False):
|
||
# 假设BCE嵌入的维度是768，根据你选择的模型可能不同
|
||
faiss_index_cpu = faiss.IndexFlatIP(dimension) # 创建一个使用内积的FAISS索引
|
||
# 将问答对转换为向量并添加到FAISS索引中
|
||
for doc in split_doc:
|
||
# type_of_docs = type(split_doc)
|
||
text = f"{doc.page_content}"
|
||
vector = emb_model.encode([text])
|
||
faiss_index_cpu.add(vector)
|
||
for qa in split_qa:
|
||
#仅对Q对进行编码
|
||
text = f"{qa['input']}"
|
||
vector = emb_model.encode([text])
|
||
faiss_index_cpu.add(vector)
|
||
faiss.write_index(faiss_index_cpu, knowledge_pkl_path)
|
||
return faiss_index_cpu
|
||
|
||
def create_index_gpu(split_doc, split_qa, emb_model, knowledge_pkl_path, dimension = 768, question_only=False):
|
||
res = faiss.StandardGpuResources()
|
||
index = faiss.IndexFlatIP(dimension)
|
||
faiss_index_gpu = faiss.index_cpu_to_gpu(res, 0, index)
|
||
for doc in split_doc:
|
||
# type_of_docs = type(split_doc)
|
||
text = f"{doc.page_content}"
|
||
vector = emb_model.encode([text])
|
||
faiss_index_gpu.add(vector)
|
||
for qa in split_qa:
|
||
#仅对Q对进行编码
|
||
text = f"{qa['input']}"
|
||
vector = emb_model.encode([text])
|
||
faiss_index_gpu.add(vector)
|
||
faiss.write_index(faiss_index_gpu, knowledge_pkl_path)
|
||
return faiss_index_gpu
|
||
|
||
|
||
|
||
# 根据query搜索相似文本
|
||
def find_top_k(query, faiss_index, k=5):
|
||
emb_model = load_embedding_model()
|
||
emb_query = emb_model.encode([query])
|
||
distances, indices = faiss_index.search(emb_query, k)
|
||
return distances, indices
|
||
|
||
def rerank(query, indices, knowledge_chunks):
|
||
passages = []
|
||
for index in indices[0]:
|
||
content = knowledge_chunks[index]
|
||
'''
|
||
txt: 'langchain_core.documents.base.Document'
|
||
json: dict
|
||
'''
|
||
# logger.info(f'retrieved content:{content}')
|
||
# logger.info(f'type of content:{type(content)}')
|
||
if type(content) == dict:
|
||
content = content["input"] + '\n' + content["output"]
|
||
else:
|
||
content = content.page_content
|
||
passages.append(content)
|
||
|
||
model = load_rerank_model()
|
||
rerank_results = model.rerank(query, passages)
|
||
return rerank_results
|
||
|
||
@st.cache_resource
|
||
def load_model():
|
||
model = (
|
||
AutoModelForCausalLM.from_pretrained("model", trust_remote_code=True)
|
||
.to(torch.bfloat16)
|
||
.cuda()
|
||
)
|
||
tokenizer = AutoTokenizer.from_pretrained("model", trust_remote_code=True)
|
||
return model, tokenizer
|
||
|
||
if __name__ == "__main__":
|
||
logger.info(data_dir)
|
||
if not os.path.exists(data_dir):
|
||
os.mkdir(data_dir)
|
||
faiss_index, knowledge_chunks = load_index_and_knowledge()
|
||
# 按照query进行查询
|
||
# query = "她要阻挠姐姐的婚姻，即使她自己的尸体在房门跟前"
|
||
# query = "肯定的。我最近睡眠很差，总是做噩梦。而且我吃得也不好，体重一直在下降"
|
||
# query = "序言 （一） 变态心理学是心理学本科生的必修课程之一，教材更新的问题一直在困扰着我们。"
|
||
query = "心理咨询师，我觉得我的胸闷症状越来越严重了，这让我很害怕"
|
||
distances, indices = find_top_k(query, faiss_index, 5)
|
||
logger.info(f'distances==={distances}')
|
||
logger.info(f'indices==={indices}')
|
||
|
||
|
||
# rerank无法返回id，先实现按整个问答对排序
|
||
rerank_results = rerank(query, indices, knowledge_chunks)
|
||
for passage, score in zip(rerank_results['rerank_passages'], rerank_results['rerank_scores']):
|
||
print(str(score)+'\n')
|
||
print(passage+'\n')
|
||
|