从 PDF 文档生成矢量数据库 (VectorDB),然后存储文本的嵌入向量。它结合了多个工具和库(如 langchain
, PyPDFLoader
, 和 Chroma
)来处理文档,并将其转换为嵌入形式,以便用于基于向量的检索任务。以下是代码逻辑的详细解释:
首先定义一个 PrepareVectorDB
类,
类与方法
PrepareVectorDB
类
-
作用:
- 加载 PDF 文档。
- 将文档分块并生成文本嵌入。
- 将生成的嵌入存储在一个持久化的向量数据库目录中。
-
构造函数 (
__init__
) 参数:doc_dir
:文档所在的目录路径。chunk_size
:分块时每块的最大字符数。chunk_overlap
:分块时块之间的字符重叠数。embedding_model
:生成嵌入向量的模型。vectordb_dir
:向量数据库存储目录。collection_name
:向量数据库中的集合名称。
path_maker
方法
- 作用:生成文件的完整路径。
- 实现 :
- 使用
os.path.join
将目录路径和文件名组合为完整路径。
- 使用
run
方法
-
作用:
- 检查
vectordb_dir
是否存在:- 如果不存在:创建目录,加载文档,分块生成嵌入,并存储到向量数据库中。
- 如果存在:跳过向量数据库的生成过程,提示目录已经存在。
- 检查
-
步骤:
- 创建目录:检查目录是否存在,若不存在则创建。
- 加载文档 :使用
PyPDFLoader
加载目录下所有 PDF 文档并提取其内容。 - 分块 :使用
RecursiveCharacterTextSplitter
将文档分块,支持分块大小及重叠。 - 生成向量嵌入 :通过
OpenAIEmbeddings
模型生成文本块的嵌入。 - 存储向量数据库 :通过
Chroma
保存嵌入向量到本地持久化数据库。 - 打印状态:提示数据库是否已生成,以及向量数量。
主要逻辑
-
加载环境变量:
pythonload_dotenv() os.environ['OPENAI_API_KEY'] = os.getenv("OPENAI_API_KEY")
- 使用
.env
文件中的OPENAI_API_KEY
设置 OpenAI API 密钥。
- 使用
-
加载配置文件:
pythonwith open(here("configs/tools_config.yml")) as cfg: app_config = yaml.load(cfg, Loader=yaml.FullLoader)
- 从
tools_config.yml
文件加载配置参数。 - 配置文件包含了不同任务(如
swiss_airline_policy_rag
和stories_rag
)的相关参数。
- 从
-
执行向量数据库生成:
- 读取配置中
swiss_airline_policy_rag
和stories_rag
的参数。 - 依次创建对应的
PrepareVectorDB
实例,调用run
方法生成和保存向量数据库。
- 读取配置中
代码运行示例
输入
假设tools_config.yml文件配置文件包含以下内容:
yaml
swiss_airline_policy_rag:
chunk_size: 500
chunk_overlap: 50
embedding_model: "openai"
vectordb: "swiss_airline_vectordb"
collection_name: "swiss_policies"
unstructured_docs: "data/swiss_policies"
stories_rag:
chunk_size: 400
chunk_overlap: 40
embedding_model: "openai"
vectordb: "stories_vectordb"
collection_name: "stories"
unstructured_docs: "data/stories"
运行后:
- 如果目录
swiss_airline_vectordb
不存在:- 加载
data/swiss_policies
下的 PDF 文件。 - 分块并生成嵌入。
- 保存嵌入到
swiss_airline_vectordb
。
- 加载
- 同理处理
stories_vectordb
。
代码运行结果:
代码依赖
- 主要依赖库 :
os
,yaml
: 用于文件和配置管理。dotenv
: 加载环境变量。langchain_chroma
: 用于持久化向量数据库。langchain_community.document_loaders
: 加载文档。langchain_openai
: 嵌入向量生成。pyprojroot
: 管理项目路径。
- 配置文件 :需要一个
tools_config.yml
文件定义参数。
总结
- 这段代码是一个自动化流程,用于将 PDF 文档处理为可用于检索的矢量嵌入,并持久化存储。
- 提供了灵活性,允许根据不同任务加载不同的文档和配置。
- 适合用于构建文档检索或问答系统等应用场景。