用于与多个数据库聊天的智能 SQL 代理问答和 RAG 系统(2) —— 从 PDF 文档生成矢量数据库 (VectorDB),然后存储文本的嵌入向量

从 PDF 文档生成矢量数据库 (VectorDB),然后存储文本的嵌入向量。它结合了多个工具和库(如 langchain, PyPDFLoader, 和 Chroma)来处理文档,并将其转换为嵌入形式,以便用于基于向量的检索任务。以下是代码逻辑的详细解释:


首先定义一个 PrepareVectorDB 类,

类与方法

PrepareVectorDB
  • 作用

    • 加载 PDF 文档。
    • 将文档分块并生成文本嵌入。
    • 将生成的嵌入存储在一个持久化的向量数据库目录中。
  • 构造函数 (__init__) 参数

    1. doc_dir:文档所在的目录路径。
    2. chunk_size:分块时每块的最大字符数。
    3. chunk_overlap:分块时块之间的字符重叠数。
    4. embedding_model:生成嵌入向量的模型。
    5. vectordb_dir:向量数据库存储目录。
    6. collection_name:向量数据库中的集合名称。
path_maker 方法
  • 作用:生成文件的完整路径。
  • 实现
    • 使用 os.path.join 将目录路径和文件名组合为完整路径。
run 方法
  • 作用

    • 检查 vectordb_dir 是否存在:
      • 如果不存在:创建目录,加载文档,分块生成嵌入,并存储到向量数据库中。
      • 如果存在:跳过向量数据库的生成过程,提示目录已经存在。
  • 步骤

    1. 创建目录:检查目录是否存在,若不存在则创建。
    2. 加载文档 :使用 PyPDFLoader 加载目录下所有 PDF 文档并提取其内容。
    3. 分块 :使用 RecursiveCharacterTextSplitter 将文档分块,支持分块大小及重叠。
    4. 生成向量嵌入 :通过 OpenAIEmbeddings 模型生成文本块的嵌入。
    5. 存储向量数据库 :通过 Chroma 保存嵌入向量到本地持久化数据库。
    6. 打印状态:提示数据库是否已生成,以及向量数量。

主要逻辑

  1. 加载环境变量

    python 复制代码
    load_dotenv()
    os.environ['OPENAI_API_KEY'] = os.getenv("OPENAI_API_KEY")
    • 使用 .env 文件中的 OPENAI_API_KEY 设置 OpenAI API 密钥。
  2. 加载配置文件

    python 复制代码
    with open(here("configs/tools_config.yml")) as cfg:
        app_config = yaml.load(cfg, Loader=yaml.FullLoader)
    • tools_config.yml 文件加载配置参数。
    • 配置文件包含了不同任务(如 swiss_airline_policy_ragstories_rag)的相关参数。
  3. 执行向量数据库生成

    • 读取配置中 swiss_airline_policy_ragstories_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"

运行后:

  1. 如果目录 swiss_airline_vectordb 不存在:
    • 加载 data/swiss_policies 下的 PDF 文件。
    • 分块并生成嵌入。
    • 保存嵌入到 swiss_airline_vectordb
  2. 同理处理 stories_vectordb

代码运行结果:


代码依赖

  1. 主要依赖库
    • os, yaml: 用于文件和配置管理。
    • dotenv: 加载环境变量。
    • langchain_chroma: 用于持久化向量数据库。
    • langchain_community.document_loaders: 加载文档。
    • langchain_openai: 嵌入向量生成。
    • pyprojroot: 管理项目路径。
  2. 配置文件 :需要一个 tools_config.yml 文件定义参数。

总结

  • 这段代码是一个自动化流程,用于将 PDF 文档处理为可用于检索的矢量嵌入,并持久化存储。
  • 提供了灵活性,允许根据不同任务加载不同的文档和配置。
  • 适合用于构建文档检索或问答系统等应用场景。
相关推荐
ZWZhangYu1 小时前
LangChain 构建向量数据库和检索器
数据库·langchain·easyui
千宇宙航1 小时前
闲庭信步使用SV搭建图像测试平台:第三十一课——基于神经网络的手写数字识别
图像处理·人工智能·深度学习·神经网络·计算机视觉·fpga开发
onceco1 小时前
领域LLM九讲——第5讲 为什么选择OpenManus而不是QwenAgent(附LLM免费api邀请码)
人工智能·python·深度学习·语言模型·自然语言处理·自动化
feifeigo1232 小时前
升级到MySQL 8.4,MySQL启动报错:io_setup() failed with EAGAIN
数据库·mysql·adb
火龙谷3 小时前
【nosql】有哪些非关系型数据库?
数据库·nosql
jndingxin4 小时前
OpenCV CUDA模块设备层-----高效地计算两个 uint 类型值的带权重平均值
人工智能·opencv·计算机视觉
焱焱枫4 小时前
Oracle获取执行计划之10046 技术详解
数据库·oracle
Sweet锦4 小时前
零基础保姆级本地化部署文心大模型4.5开源系列
人工智能·语言模型·文心一言
双力臂4045 小时前
MyBatis动态SQL进阶:复杂查询与性能优化实战
java·sql·性能优化·mybatis
hie988945 小时前
MATLAB锂离子电池伪二维(P2D)模型实现
人工智能·算法·matlab