Postgres Embedding
Postgres Embedding 是一個開源向量相似度搜尋工具，適用於
Postgres，它使用階層式可導航小世界 (HNSW)來進行近似最近鄰搜尋。
它支援
- 使用 HNSW 的精確和近似最近鄰搜尋
- L2 距離
本筆記本展示如何使用 Postgres 向量資料庫 (PGEmbedding)。
PGEmbedding 整合為您建立 pg_embedding 擴充功能，但您需要執行以下 Postgres 查詢來新增它
CREATE EXTENSION embedding;
# Pip install necessary package
%pip install --upgrade --quiet langchain-openai langchain-community
%pip install --upgrade --quiet psycopg2-binary
%pip install --upgrade --quiet tiktoken
將 OpenAI API 金鑰新增至環境變數以使用 OpenAIEmbeddings。
import getpass
import os
if "OPENAI_API_KEY" not in os.environ:
os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")
OpenAI API Key:········
## Loading Environment Variables
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
if "DATABASE_URL" not in os.environ:
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 索引以進行近似最近鄰 (ANN) 搜尋，以加快 similarity_search_with_score 的執行時間。若要在您的向量欄位上建立 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_community.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)