先看完整初级代码,再看我一块一块说明,每一块都可以扩展
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" {'-' * 100} ".join([f"Document {i+1}: " + 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" {'-' * 100} ".join([f"Document {i+1}: " + 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,但是作者目前没写出来,所以就搁置一下。
以上所有内容都可以在界面中选择不同的类型,如知识库增加删除、模型选择等。后续再搞