ZenML项目实战：在向量数据库中高效存储嵌入向量-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_01050/article/details/148524542

ZenML项目实战：在向量数据库中高效存储嵌入向量

zenml 项目地址: https://gitcode.com/gh_mirrors/zen/zenml

引言

在现代机器学习应用中，特别是涉及自然语言处理(NLP)的场景中，嵌入向量(Embeddings)扮演着至关重要的角色。本文将深入探讨如何在ZenML项目中，将生成的嵌入向量高效地存储在PostgreSQL向量数据库中，为后续的检索增强生成(RAG)流程奠定基础。

为什么需要向量数据库

嵌入向量生成过程虽然在有GPU的机器上速度较快，但每次查询都重新生成显然不是高效的做法。向量数据库提供了以下关键优势：

高效检索：专门为高维向量相似性搜索优化
持久存储：避免重复计算，节省计算资源
可扩展性：轻松处理大规模向量数据集

PostgreSQL作为向量数据库的选择

在ZenML项目中，我们选择PostgreSQL作为向量数据库解决方案，原因在于：

成熟稳定：PostgreSQL是经过时间考验的关系型数据库
扩展性强：通过pgvector扩展支持向量操作
生态完善：丰富的工具链和社区支持

实现步骤详解

1. 数据库连接与初始化

conn = get_db_conn()
with conn.cursor() as cur:
    # 安装pgvector扩展
    cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
    conn.commit()

首先建立数据库连接，并确保pgvector扩展已安装，这是PostgreSQL支持向量操作的基础。

2. 创建嵌入向量表

table_create_command = f"""
CREATE TABLE IF NOT EXISTS embeddings (
    id SERIAL PRIMARY KEY,
    content TEXT,
    token_count INTEGER,
    embedding VECTOR({EMBEDDING_DIMENSIONALITY}),
    filename TEXT,
    parent_section TEXT,
    url TEXT
);
"""

我们设计了包含以下字段的表结构：

content：存储原始文本内容
embedding：存储向量数据，维度由EMBEDDING_DIMENSIONALITY决定
其他元数据字段便于后续检索和调试

3. 数据插入策略

# 检查内容是否已存在
cur.execute("SELECT COUNT(*) FROM embeddings WHERE content = %s", (content,))
count = cur.fetchone()[0]
if count == 0:
    # 执行插入操作

采用"不存在才插入"的策略，避免数据重复，这在增量更新场景中特别有用。

4. 索引优化

# 根据记录数计算最佳列表数
num_lists = max(num_records / 1000, 10)
if num_records > 1000000:
    num_lists = math.sqrt(num_records)

# 创建IVFFlat索引
cur.execute(f"CREATE INDEX IF NOT EXISTS embeddings_idx ON embeddings USING ivfflat (embedding vector_cosine_ops) WITH (lists = {num_lists});")

IVFFlat(倒排文件索引)是PostgreSQL中高效的向量索引方法，我们：