用于与多个数据库聊天的智能 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 文档处理为可用于检索的矢量嵌入,并持久化存储。
  • 提供了灵活性,允许根据不同任务加载不同的文档和配置。
  • 适合用于构建文档检索或问答系统等应用场景。
相关推荐
weiwuxian4 分钟前
揭开智能体的神秘面纱:原来你不是"超级AI"!
人工智能
Codebee5 分钟前
“自举开发“范式:OneCode如何用低代码重构自身工具链
java·人工智能·架构
说私域16 分钟前
基于开源AI智能名片链动2+1模式的S2B2C商城小程序:门店私域流量与视频号直播融合的生态创新研究
人工智能·小程序·开源
Ronin-Lotus19 分钟前
深度学习篇---Yolov系列
人工智能·深度学习
阿蒙Amon29 分钟前
C#读写文件:多种方式详解
开发语言·数据库·c#
静心问道1 小时前
GoT:超越思维链:语言模型中的有效思维图推理
人工智能·计算机视觉·语言模型
东窗西篱梦1 小时前
Redis集群部署指南:高可用与分布式实践
数据库·redis·分布式
aneasystone本尊1 小时前
学习 Claude Code 的工具使用(三)
人工智能
szxinmai主板定制专家1 小时前
【精密测量】基于ARM+FPGA的多路光栅信号采集方案
服务器·arm开发·人工智能·嵌入式硬件·fpga开发
T__TIII1 小时前
Dify 自定义插件
人工智能·github