探索Postgres Embedding:使用HNSW进行高效的向量相似性搜索

122 阅读3分钟
# 探索Postgres Embedding:使用HNSW进行高效的向量相似性搜索

在现代应用中,向量相似性搜索是实现高效信息检索和个性化服务的重要组成部分。Postgres Embedding作为一个开源解决方案,提供了在Postgres中使用Hierarchical Navigable Small Worlds (HNSW)算法进行近似最近邻搜索的能力。在这篇文章中,我们将介绍如何使用Postgres Embedding进行向量搜索,并提供完整代码示例以供参考。

## 主要内容

### 1. Postgres Embedding简介

Postgres Embedding是Postgres的一个扩展,它允许用户在数据库中进行向量化数据的存储和相似性搜索。该扩展支持使用HNSW进行近似最近邻搜索,结合L2距离,能够高效处理大规模数据集的相似性计算。

### 2. 安装和环境配置

首先,我们需要安装必要的软件包并配置环境变量。

```python
# 安装必要的软件包
%pip install --upgrade --quiet langchain-openai langchain-community
%pip install --upgrade --quiet psycopg2-binary
%pip install --upgrade --quiet tiktoken

# 配置OpenAI API Key
import getpass
import os

os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")
# 使用API代理服务提高访问稳定性
os.environ["DATABASE_URL"] = getpass.getpass("Database Url:")

3. 向量数据库的使用

使用PGEmbedding扩展可以很方便地将OpenAI生成的向量数据存储到Postgres中并进行相似性搜索。以下是一个基本的使用示例:

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)

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: 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)

4. 创建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);

5. 常见问题和解决方案

  • 连接问题:由于某些地区的网络限制,开发者在使用API时可能需要考虑使用API代理服务以提高访问稳定性。
  • 索引性能:如果在查询中发现性能问题,可以调整HNSW的参数以优化结果。

总结和进一步学习资源

Postgres Embedding提供了一种高效、灵活的向量相似性搜索解决方案。在实践中,可以根据具体应用场景调整HNSW参数,进一步提高检索性能。更多细节可以参考以下资源:

参考资料

  1. Postgres Embedding官方文档
  2. OpenAI Embeddings使用指南
  3. HNSW算法研究论文

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

---END---