用于与多个数据库聊天的智能 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 文档处理为可用于检索的矢量嵌入,并持久化存储。
  • 提供了灵活性,允许根据不同任务加载不同的文档和配置。
  • 适合用于构建文档检索或问答系统等应用场景。
相关推荐
m0_748254663 分钟前
MySQL和SQL server的区别
数据库·mysql
秦南北5 分钟前
国内领先的宠物类电商代运营公司品融电商
大数据·人工智能·电商
补三补四5 分钟前
Yashan DB 实例
数据库·oracle·dba
JobsandCzj11 分钟前
PDF 分割工具
javascript·小程序·pdf
沈健_算法小生17 分钟前
基于DeepSeek实现PDF嵌入SVG图片无损放大
pdf
龚大龙19 分钟前
机器学习(李宏毅)——Domain Adaptation
人工智能·机器学习
源码姑娘27 分钟前
基于DeepSeek的智慧医药系统(源码+部署教程)
java·人工智能·程序人生·毕业设计·springboot·健康医疗·课程设计
椰椰椰耶36 分钟前
【redis】全局命令set、get、keys
数据库·redis·缓存
月落星还在43 分钟前
Redis 内存淘汰策略深度解析
数据库·redis·缓存
左灯右行的爱情1 小时前
Redis- 切片集群
数据库·redis·缓存