先看完整初级代码,再看我一块一块说明,每一块都可以扩展
import os
import torch
from langchain.text_splitter import MarkdownTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceBgeEmbeddings
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import LLMChainExtractor
from modelscope import AutoModelForCausalLM, AutoTokenizer
def pretty_print_docs(docs):
print(f"\n{'-' * 100}\n".join([f"Document {i+1}:\n\n" + d.page_content for i, d in enumerate(docs)]))
def load_data(path):
'''
加载数据
:param path: md文件路径
:return:
'''
with open(path,'r',encoding='utf-8') as files:
markdown_text = files.read()
markdown_splitter = MarkdownTextSplitter(chunk_size=500, chunk_overlap=50)
docs = markdown_splitter.create_documents([markdown_text]) # 分割
for i,m in enumerate(docs):
print(i,m)
'''
embedding 数据 无openai key
embeddings = OpenAIEmbeddings()
docsearch = Chroma.from_texts(docs,embeddings)
query = "什么是肤质?"
doc = docsearch.similarity_search(query)
print(doc[0].page_content)
'''
model_name = r"bge-large-zh"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': True}
model = HuggingFaceBgeEmbeddings(
model_name=model_name,
model_kwargs=model_kwargs,
encode_kwargs=encode_kwargs
)
'''
向量数据库,使用chroma
'''
# 向量数据库保存位置
persist_directory = 'chroma/'
# 检查持久化目录是否存在
if os.path.exists(persist_directory) and len(os.listdir(persist_directory)) > 0: # 检查目录是否为空
print("Chroma数据库已存在,无需重复创建.")
vectordb = Chroma(persist_directory=persist_directory, embedding_function=model)
else:
# 创建向量数据库
vectordb = Chroma.from_documents(
documents=docs,
embedding=model,
persist_directory=persist_directory
)
vectordb.persist() # 向量数据库的持久化
# 查看向量数据库中的文档数量
print("向量数据库中的文档数量",vectordb._collection.count())
'''
向原有数据库添加内容
#docsearch.add_texts(["Ankush went to Princeton"]) # 添加文本
#docsearch = Chroma.from_texts(docs, model)
'''
'''
从向量数据库中检索与问题相关的数据
'''
query = "什么是肤质?"
# doc = vectordb.similarity_search(query,k=3) # 寻常检索
doc = vectordb.max_marginal_relevance_search(query, k=3) # 使用mmr检索,可过滤相似度高的文档
# 上下文压缩检索
compressor = LLMChainExtractor.from_llm(llm)
compression_retriever = ContextualCompressionRetriever(
base_compressor=compressor,
base_retriever=vectordb.as_retriever()
)
question = "什么是肤质?"
compressed_docs = compression_retriever.get_relevant_documents(question)
pretty_print_docs(compressed_docs)
# 打印文档数量
print("文档数量 ",len(doc))
print("第一个文档 ",doc[0].page_content)
if __name__ == '__main__':
load_data('beautiful bible.md')
def pretty_print_docs(docs):
print(f"\n{'-' * 100}\n".join([f"Document {i+1}:\n\n" + d.page_content for i, d in enumerate(docs)]))
这个代码用来打印检索出的文档信息,类型基本是列表
接下来是load data用来加载、分割、检索数据
with open(path,'r',encoding='utf-8') as files:
markdown_text = files.read()
markdown_splitter = MarkdownTextSplitter(chunk_size=500, chunk_overlap=50)
docs = markdown_splitter.create_documents([markdown_text]) # 分割
这个代码是open我的markdown文档,然后根据500大小,50重复分割,具体见langchain技术文档Markdown文本分割器# – LangChain中文网
model_name = r"bge-large-zh"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': True}
model = HuggingFaceBgeEmbeddings(
model_name=model_name,
model_kwargs=model_kwargs,
encode_kwargs=encode_kwargs
)
这个代码是加载embedding模型bge-large-zh,如果有网络困难,可以手动去bge-large-zh · 模型库 (modelscope.cn)这里下载后放在代码同级,里面也有很多方式加载模型,我选择langchain方法。解释一下为什么要在"bge-large-zh"前面加r,因为这个模型文件与代码同级,但是我不想用绝对路径,用相对路径的话又会自动去huggingface下载模型,所以我加r就可以识别到本地的模型了
# 向量数据库保存位置
persist_directory = 'chroma/'
# 检查持久化目录是否存在
if os.path.exists(persist_directory) and len(os.listdir(persist_directory)) > 0: # 检查目录是否为空
print("Chroma数据库已存在,无需重复创建.")
vectordb = Chroma(persist_directory=persist_directory, embedding_function=model)
else:
# 创建向量数据库
vectordb = Chroma.from_documents(
documents=docs,
embedding=model,
persist_directory=persist_directory
)
vectordb.persist() # 向量数据库的持久化
# 查看向量数据库中的文档数量
print("向量数据库中的文档数量",vectordb._collection.count())
具体实现见Chroma# – LangChain中文网,需要判断persist_directory文件是否存在,不要重复加载了,用Chroma.from_documents()创建向量数据库,docs是文档,且要persist向量持久化。如果存在了,则可以直接用Chroma()加载存在的数据库
query = "什么是肤质?"
# doc = vectordb.similarity_search(query,k=3) # 寻常检索
doc = vectordb.max_marginal_relevance_search(query, k=3) # 使用mmr检索,可过滤相似度高的文档
正常检索和用mmr检索,mmr检索可以少点重复内容,返回k个答案。以上代码连起来可以运行,不需要加上下文压缩检索。
# 上下文压缩检索
compressor = LLMChainExtractor.from_llm(llm)
compression_retriever = ContextualCompressionRetriever(
base_compressor=compressor,
base_retriever=vectordb.as_retriever()
)
question = "什么是肤质?"
compressed_docs = compression_retriever.get_relevant_documents(question)
pretty_print_docs(compressed_docs)
检索方法使用上下文压缩检索,具体见将字符串压缩器和文档转换器连在一起# – LangChain中文网
大致有原始的压缩和llm的压缩,这样输出可以少点字数,否则只会输出原文。
其中本代码中compressor = LLMChainExtractor.from_llm(llm)是可以选择llm=openai(),但是作者想要用本地自定义模型的api,但是作者目前没写出来,所以就搁置一下。先用ollama的见Llama-cpp# – LangChain中文网
llm = Ollama(base_url="",
model="ggml-model-q4_0.gguf",
callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]))
使用ollama的gguf的模型,由于官网没有qwen的gguf,具体量化看我下一篇文章。但是在windows跑不起来,base_url也需要一个代码,后续补充。
以上所有内容都可以在界面中选择不同的类型,如知识库增加删除、模型选择等。后续再搞