从零学习大模型——使用GLM-4-9B-Chat + BGE-M3 + langchain + chroma建立的本地RAG应用（四）——使用BGE-M3与向量数据库chroma进行数据向量化与持久化。

在设置好embedding模型后，我们就可以对数据集进行embeddings处理了。注意在处理之前要先对数据进行清洗。方便进行后续分割的过程。

LangChain提供了多种格式的数据集处理方式。本文以pdf为例：

在/root/autodl-tmp路径下新建datasetPDF.py文件

from langchain_community.document_loaders import PyMuPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from testEmbeddings import BGEM3Embeddings
from transformers import AutoTokenizer, AutoModel
from langchain_community.vectorstores import Chroma
import re
# 创建一个 PyMuPDFLoader Class 实例，输入为待加载的 pdf 文档路径
loader = PyMuPDFLoader("/root/autodl-fs/dataset/LLM-v1.0.0.pdf")
# 调用 PyMuPDFLoader Class 的函数 load 对 pdf 文件进行加载
pdf_pages = loader.load()

# 进行文档格式清洗
for pdf_page in pdf_pages:
    pattern = re.compile(r'[^\u4e00-\u9fff](\n)[^\u4e00-\u9fff]', re.DOTALL)
    pdf_page.page_content = re.sub(pattern, lambda match: match.group(0).replace('\n', ''), pdf_page.page_content)
    pdf_page.page_content = pdf_page.page_content.replace('•', '')
    pdf_page.page_content = pdf_page.page_content.replace(' ', '')
    # pdf_page.page_content = pdf_page.page_content.replace('\n', '')
    # print(pdf_page.page_content)

# 知识库中单段文本长度
CHUNK_SIZE = 500

# 知识库中相邻文本重合长度
OVERLAP_SIZE = 50

# 使用递归字符文本分割器
# RecursiveCharacterTextSplitter 递归字符文本分割
# RecursiveCharacterTextSplitter 将按不同的字符递归地分割(按照这个优先级["\n\n", "\n", " ", ""])，
#     这样就能尽量把所有和语义相关的内容尽可能长时间地保留在同一位置
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=CHUNK_SIZE,
    chunk_overlap=OVERLAP_SIZE,
    separators=["\n\n", "\n", " ", "", "。", "，"]
)

split_docs = text_splitter.split_documents(pdf_pages)

# print(split_docs[150].page_content)


# 定义本地embedding模型
model_name = "bge-m3"  
save_directory = "/root/autodl-tmp/bge-m3"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)
embedding = BGEM3Embeddings(model, tokenizer)

persist_directory='/root/autodl-tmp/vectorDatabase/chroma'

vectordb = Chroma.from_documents(
    documents=split_docs,
    embedding=embedding,
    # 允许我们将persist_directory目录保存到磁盘上
    persist_directory=persist_directory  
)

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

注意在执行之前要先在/root/autodl-tmp路径下新建好vectorDatabase文件夹，再在该目录下新建chroma文件夹用于存放持久化的向量数据。

如果最后输出的数量不为空代表成功进行向量化。

新建一个测试文件testquerySearch.py用于测试相似度检索。

from langchain_community.document_loaders import PyMuPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from testEmbeddings import BGEM3Embeddings
from transformers import AutoTokenizer, AutoModel
from langchain_community.vectorstores import Chroma
import re

model_name = "bge-m3" 
save_directory = "/root/autodl-tmp/bge-m3"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)
embedding = BGEM3Embeddings(model, tokenizer)

persist_directory='/root/autodl-tmp/vectorDatabase/chroma'

vectordb_load = Chroma(
    persist_directory=persist_directory,
    embedding_function=embedding
)
# 相似度检索，k代表返回多少条结果
sim_docs_similarity_search=vectordb_load.similarity_search("请于此处输入与数据集相关的语句",k=5)
# MMR检索
# sim_docs_MMR_search=vectordb_load.max_marginal_relevance_search("请于此处输入与数据集相关的语句",k=3)

for i, sim_doc in enumerate(sim_docs_similarity_search):
    print(f"通过相似度检索到的第{i}个内容: \n{sim_doc.page_content}", end="\n--------------\n")
# for i, sim_doc in enumerate(sim_docs_MMR_search):
#     print(f"通过MMR检索到的第{i}个内容: \n{sim_doc.page_content}", end="\n--------------\n")

 如果运行正常应该返回内容意思相近的数据集片段，如果结果不准确也无所谓，因为当前返回的数据集片段并不是最终结果，后续LLM会结合搜索结果进行回答，此时准确性会提升很多。