[探索Postgres Embedding:如何在Postgres中进行高效向量相似性搜索]

117 阅读3分钟

探索Postgres Embedding:如何在Postgres中进行高效向量相似性搜索

引言

随着大数据时代的到来,向量相似性搜索在诸如推荐系统、自然语言处理和计算机视觉等领域变得越来越重要。Postgres Embedding作为一种面向Postgres数据库的开源解决方案,使用了Hierarchical Navigable Small Worlds (HNSW)算法,提供了高效的近似最近邻搜索功能。本篇文章旨在为读者介绍如何在Postgres数据库中使用Postgres Embedding实现向量存储与搜索。

主要内容

1. 向量存储与搜索概述

Postgres Embedding利用HNSW算法来实现高效的向量相似性搜索。支持精确和近似最近邻搜索,以及L2距离度量。HNSW算法因其优越的性能和高检索精度,被广泛应用于处理大规模数据的相似性搜索。

2. 安装和设置

要开始使用Postgres Embedding,我们需要安装必要的软件包并设置环境变量。

# 安装必要的Python包
%pip install --upgrade --quiet langchain-openai langchain-community
%pip install --upgrade --quiet psycopg2-binary
%pip install --upgrade --quiet tiktoken

还需要在环境变量中添加OpenAI API密钥,以便使用OpenAI提供的嵌入功能。

import getpass
import os

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

3. 在Postgres中创建扩展

在使用向量存储功能前,需要首先在Postgres中添加相关扩展。

CREATE EXTENSION embedding;

4. 使用代码实现向量存储和搜索

下面的Python示例展示了如何使用langchain库加载文本数据并执行相似性搜索:

from typing import List, Tuple
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import PGEmbedding
from langchain_core.documents import Document
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)

# 创建嵌入并存储到Postgres
embeddings = OpenAIEmbeddings()
connection_string = os.environ.get("DATABASE_URL")  # 使用API代理服务提高访问稳定性
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: List[Tuple[Document, float]] = db.similarity_search_with_score(query)

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

5. 创建HNSW索引

创建HNSW索引可以显著提高搜索效率,以下代码展示了如何创建一个HNSW索引:

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

6. 检索向量存储

可以通过向量检索器来方便地进行搜索操作:

store = PGEmbedding(
    connection_string=connection_string,
    embedding_function=embeddings,
    collection_name=collection_name,
)

retriever = store.as_retriever()

常见问题和解决方案

  • 网络访问问题:由于某些地区的网络限制,API访问可能不稳定,建议使用API代理服务。
  • 内存限制:处理大型数据时,确保Postgres服务器有足够的内存来处理向量计算。

总结和进一步学习资源

Postgres Embedding为向量相似性搜索提供了一种高效且易于扩展的解决方案。可以结合OpenAI的嵌入技术,以满足不同领域的需求。对于想要深入了解HNSW算法和Postgres Embedding的读者,建议参考以下资源:

参考资料

  1. Langchain官方文档
  2. Postgres Embedding介绍
  3. HNSW论文

如果这篇文章对你有帮助,欢迎点赞并关注我的博客。您的支持是我持续创作的动力!

---END---