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"\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也需要一个代码，后续补充。

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