Skip to main content

pg_embedding

pg_embedding是一个开源的用于Postgres的向量相似性搜索工具，它使用Hierarchical Navigable Small Worlds进行近似最近邻搜索。

它支持以下功能：

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

本笔记本展示了如何使用Postgres向量数据库（PGEmbedding）。

PGEmbedding集成会为您创建pg_embedding扩展，但您需要运行以下Postgres查询来添加它：

CREATE EXTENSION embedding;
# 安装必要的包
pip install openai
pip install psycopg2-binary
pip install tiktoken

将OpenAI API密钥添加到环境变量中以使用OpenAIEmbeddings。

import os
import getpass

os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")
    OpenAI API Key:········
## 加载环境变量
from typing import List, Tuple
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import PGEmbedding
from langchain.document_loaders import TextLoader
from langchain.docstore.document import Document
os.environ["DATABASE_URL"] = getpass.getpass("Database Url:")
    Database Url:········
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)

在Postgres中使用向量存储​

上传向量存储到PG​

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

创建HNSW索引​

默认情况下，该扩展执行顺序扫描搜索，具有100%的召回率。您可以考虑创建一个HNSW索引，用于加速similarity_search_with_score的执行时间，以进行近似最近邻（ANN）搜索。要在向量列上创建HNSW索引，请使用create_hnsw_index函数：

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

上述语句中使用的HNSW索引选项包括：

  • maxelements：定义索引的最大元素数。这是一个必需的参数。上面的示例中的值为3。实际应用中，该值会更大，例如1000000。"元素"是指数据集中的数据点（向量），它在HNSW图中表示为一个节点。通常，您会将此选项设置为能够容纳数据集中行数的值。

  • dims：定义向量数据中的维数数。这是一个必需的参数。上面的示例中使用了一个较小的值。如果您存储的是使用OpenAI的text-embedding-ada-002模型生成的数据，该模型支持1536个维度，您可以定义一个值为1536的维数。

  • m：定义在图构建过程中为每个节点创建的双向链接（也称为"边"）的最大数量。 还支持以下其他索引选项：

  • efConstruction：定义在索引构建过程中考虑的最近邻数。默认值为32。

  • efsearch：定义在索引搜索过程中考虑的最近邻数。默认值为32。 有关如何配置这些选项以影响HNSW算法的信息，请参阅Tuning the HNSW algorithm。

从PG检索向量存储​

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

retriever = store.as_retriever()
retriever
    VectorStoreRetriever(vectorstore=<langchain.vectorstores.pghnsw.HNSWVectoreStore object at 0x121d3c8b0>, search_type='similarity', search_kwargs={})
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: List[Tuple[Document, float]] = db1.similarity_search_with_score(query)
for doc, score in docs_with_score:
    print("-" * 80)
    print("Score: ", score)
    print(doc.page_content)
    print("-" * 80)