用于与多个数据库聊天的智能 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 文档处理为可用于检索的矢量嵌入,并持久化存储。
  • 提供了灵活性,允许根据不同任务加载不同的文档和配置。
  • 适合用于构建文档检索或问答系统等应用场景。
相关推荐
user-猴子36 分钟前
2026国产AI智能体七强横评:WorkBuddy、AiPy、悟空、Kimi Work、TRAE Work、QoderWork、百度搭子怎么选?
人工智能·百度·dubbo
名字还没想好☜8 小时前
Python f-string 进阶:数字格式化、对齐填充、调试 = 号与嵌套表达式
开发语言·数据库·python·字符串格式化·f-string
IvanCodes8 小时前
RAG 实战教程(一):RAG 工作原理与完整流程——分片、索引、召回、重排和生成
人工智能·后端·agent
ltl8 小时前
Serverless 数据库弹性理论:Neon 与 Aurora Serverless v2
数据库
今天AI了吗9 小时前
合规场景下的 AI 推理可解释性:Attention 可视化与推理路径追踪的工程实践
人工智能
周末程序猿9 小时前
浅析大模型推理十二篇之KV Cache
人工智能
鱼樱前端9 小时前
用 AI 做内容变收入
前端·人工智能·ai编程
Dawson Zhu10 小时前
基于Palantir Foundry构建半导体制造AI友好型数据中台:从良率分析场景谈起
人工智能·语言模型·架构·制造·agi
fthux10 小时前
装闭 RenoPit 源码解析(04):装修图纸和合同文件上传处理流程
人工智能·ai·开源·github·open source·renopit