用于与多个数据库聊天的智能 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 文档处理为可用于检索的矢量嵌入,并持久化存储。
  • 提供了灵活性,允许根据不同任务加载不同的文档和配置。
  • 适合用于构建文档检索或问答系统等应用场景。
相关推荐
酒酿小圆子~39 分钟前
NLP中常见的分词算法(BPE、WordPiece、Unigram、SentencePiece)
人工智能·算法·自然语言处理
新加坡内哥谈技术2 小时前
Virgo:增强慢思考推理能力的多模态大语言模型
人工智能·语言模型·自然语言处理
martian6652 小时前
深入详解人工智能计算机视觉之图像生成与增强:生成对抗网络(GAN)
人工智能·计算机视觉
qq_273900232 小时前
pytorch torch.isclose函数介绍
人工智能·pytorch·python
说私域2 小时前
阿里巴巴新零售模式下的创新实践:结合开源AI智能名片2+1链动模式S2B2C商城小程序的应用探索
人工智能·开源·零售
致Great2 小时前
《你的RAG出错了?快来Get这份改进秘籍》
人工智能·llm·nlp
MiniFlyZt3 小时前
省市区三级联动(后端)
数据库·spring boot
这我可不懂3 小时前
2025低代码与人工智能AI新篇
人工智能·低代码
XianxinMao3 小时前
企业通过私有安全端点访问大型语言模型的益处
人工智能·安全·语言模型
itwangyang5203 小时前
AIDD-人工智能药物设计-可扩展!更快!更便宜!大规模基因组数据存储新结构
人工智能