langchain知识库分割、加载、检索

先看完整初级代码,再看我一块一块说明,每一块都可以扩展

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,但是作者目前没写出来,所以就搁置一下。

以上所有内容都可以在界面中选择不同的类型,如知识库增加删除、模型选择等。后续再搞