说明

本文使用openai提供的embedding模型作为框架基础模型，知识库的搭建目的就是为了让大模型减少幻觉出现，实现起来也很简单，假如你要做一个大模型的客服问答系统，那么就把历史客服问答数据整理好，先做数据处理，在做数据向量化，最后保存到向量库中就可以了，下面文章中只是一个简单工作流程，只能用来参考，希望对大家有所帮助！

流程

上传知识库的文档不限于txt,pdf,markdown等数据格式，不同的数据格式用不同的方法来处理，文章内仅使用pdf文件做测试

1.数据加载

def load_data():
    from langchain.document_loaders.pdf import PyMuPDFLoader
    # 本地pdf文档路径
    loader = PyMuPDFLoader("./knowledge_db/pumkin_book/pumpkin_book.pdf")
    pdf_pages = loader.load()
    print(f"载入后的变量类型为：{type(pdf_pages)}，", f"该 PDF 一共包含 {len(pdf_pages)} 页")
    pdf_page = pdf_pages[1]
    page_content = pdf_page.page_content
    print(f"每一个元素的类型：{type(pdf_page)}.",
          f"该文档的描述性数据：{pdf_page.metadata}",
          f"查看该文档的内容:\n{pdf_page.page_content}",
          sep="\n------\n")
    return page_content,pdf_pages

2.数据清洗

def clean_data(pdf_content):
    # 匹配非中文字符和换行符
    pattern = re.compile(r'[^\u4e00-\u9fff](\n)[^\u4e00-\u9fff]', re.DOTALL)
    # 将匹配到的换行符替换为空字符串
    new_pdf_content = re.sub(pattern, lambda match: match.group(0).replace('\n', ''), pdf_content)
    # 去除。和空格符号
    new_pdf_content = new_pdf_content.replace('。', '').replace(' ', '')

    return new_pdf_content

3.数据切分

def split_data(pdf_pages,new_pdf_content):
    '''
    * RecursiveCharacterTextSplitter 递归字符文本分割
    RecursiveCharacterTextSplitter 将按不同的字符递归地分割(按照这个优先级["\n\n", "\n", " ", ""])，
        这样就能尽量把所有和语义相关的内容尽可能长时间地保留在同一位置
    RecursiveCharacterTextSplitter需要关注的是4个参数：

    * separators - 分隔符字符串数组
    * chunk_size - 每个文档的字符数量限制
    * chunk_overlap - 两份文档重叠区域的长度
    * length_function - 长度计算函数
    '''
    from langchain.text_splitter import RecursiveCharacterTextSplitter
    # 知识库中单段文本长度
    CHUNK_SIZE = 500
    # 知识库中相邻文本重合长度
    OVERLAP_SIZE = 50
    # 使用递归字符文本分割器
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=OVERLAP_SIZE
    )

    text_splitter.split_text(new_pdf_content[0:1000])
    split_docs = text_splitter.split_documents(pdf_pages)
    print(f"切分后的文件数量：{len(split_docs)}")
    print(f"切分后的字符数（可以用来大致评估 token 数）：{sum([len(doc.page_content) for doc in split_docs])}")
    return split_docs

4.获取向量

def gpt_config():
    import httpx
    # 使用httpx设置代理
    proxy = 'http://127.0.0.1:8080' # 修改为自己的代理地址
    proxies = {'http://': proxy, 'https://': proxy}
    http_client = httpx.Client(proxies=proxies, verify=True)
    return http_client

def get_vector(split_docs):
    # from langchain.embeddings import OpenAIEmbeddings
    from langchain_openai import OpenAIEmbeddings
    from langchain.vectorstores.chroma import Chroma
    from dotenv import load_dotenv, find_dotenv
	
	# 获取key
    _ = load_dotenv(find_dotenv()) # 可注释
    api_key = os.environ.get("OPENAI_API_KEY")
    http_client = gpt_config()
    # 官网有提供3个embedding模型，按需选择
    embedding = OpenAIEmbeddings(model="text-embedding-3-small",
                                 openai_api_key=api_key,
                                 http_client=http_client)
    # 保存路径
    persist_directory = './vector_db/chroma'
    vectordb = Chroma.from_documents(
        documents=split_docs[:20],  # 为了速度，只选择前 20 个切分的 doc 进行生成
        embedding=embedding,
        persist_directory=persist_directory  # 允许我们将persist_directory目录保存到磁盘上
    )

    return vectordb

5.向量库保存到本地

def save_vector(vectordb):
    vectordb.persist()
    print(f"向量库中存储的数量：{vectordb._collection.count()}")

6.向量搜索

def search_vector(vectordb):
    question = '什么是机器学习'
    # 余弦相似度搜索
    search_result = vectordb.similarity_search(question, k=2) # k表示返回的相似文档数量
    print(f"检索到的内容数：{len(search_result)}")
    for i, sim_doc in enumerate(search_result):
        print(f"检索到的第{i}个内容: \n{sim_doc.page_content[:200]}", end="\n--------------\n")

    # MMR搜索
    # 核心思想是在已经选择了一个相关性高的文档之后，再选择一个与已选文档相关性较低但是信息丰富的文档。这样可以在保持相关性的同时，增加内容的多样性，避免过于单一的结果。
    mmr_docs = vectordb.max_marginal_relevance_search(question, k=2)
    for i, sim_doc in enumerate(mmr_docs):
        print(f"MMR 检索到的第{i}个内容: \n{sim_doc.page_content[:200]}", end="\n--------------\n")

7.汇总调用

def main_task():
    # 加载数据
    pdf_content,pdf_pages = load_data()
    # 数据清洗
    new_pdf_content = clean_data(pdf_content)
    # 切分数据
    split_docs = split_data(pdf_pages,new_pdf_content)
    # 获取向量
    vectordb = get_vector(split_docs)
    # 将向量库内容保存到本地文件中
    # save_vector(vectordb)
    # 向量搜索
    search_vector(vectordb)