PgvectorEmbeddingRetriever

PgvectorEmbeddingRetriever

Pgvector Document Store와 호환되는 임베딩 기반 Retriever예요. 쿼리와 문서의 임베딩을 비교해 관련성이 가장 높은 문서를 가져오죠.

출처: 문서

본문

개요

PgvectorEmbeddingRetriever는 PgvectorDocumentStore와 호환되는 임베딩 기반 Retriever예요. 쿼리와 문서 임베딩을 비교한 결과에 따라 PgvectorDocumentStore에서 쿼리에 가장 관련 있는 문서를 가져와요.

파이프라인에서 PgvectorEmbeddingRetriever를 사용할 때는 쿼리와 문서 임베딩이 준비되어 있어야 해요. 인덱싱 파이프라인에 Document Embedder를, 쿼리 파이프라인에 Text Embedder를 추가하면 되죠.

query_embedding 외에도 PgvectorEmbeddingRetriever는 top_k(가져올 최대 문서 수)와 filters(검색 공간을 좁히기) 같은 선택적 파라미터를 받아요.

임베딩 검색에 영향을 주는 몇 가지 파라미터는 해당 PgvectorDocumentStore가 초기화될 때 정의해야 해요. 임베딩 차원, 벡터 함수, 그리고 검색 전략(정확 최근접 이웃 또는 HNSW)과 관련된 몇 가지가 여기에 해당하죠.

설치

pgvector가 있는 PostgreSQL 데이터베이스를 빠르게 구성하려면 Docker를 사용할 수 있어요:

docker run -d -p 5432:5432 -e POSTGRES_USER=postgres -e POSTGRES_PASSWORD=postgres -e POSTGRES_DB=postgres pgvector/pgvector:pg17

pgvector 설치에 대한 자세한 내용은 pgvector GitHub 저장소를 방문하세요.

Haystack에서 pgvector를 사용하려면 pgvector-haystack 통합 패키지를 설치하세요:

pip install pgvector-haystack

이 페이지의 예제는 sentence-transformers-haystack 패키지의 Sentence Transformers 임베더를 사용해요. 예제를 실행하려면 설치하세요:

pip install sentence-transformers-haystack

사용법

단독으로 사용하기

이 Retriever는 동작하려면 PgvectorDocumentStore와 인덱싱된 문서가 필요해요.

import os
from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
from haystack_integrations.components.retrievers.pgvector import (
    PgvectorEmbeddingRetriever,
)

os.environ["PG_CONN_STR"] = "postgresql://postgres:***@localhost:5432/postgres"

document_store = PgvectorDocumentStore()
retriever = PgvectorEmbeddingRetriever(document_store=document_store)

# using a fake vector to keep the example simple
retriever.run(query_embedding=[0.1] * 768)

파이프라인에서 사용하기

import os
from haystack.document_stores.types import DuplicatePolicy
from haystack import Document, Pipeline
from haystack_integrations.components.embedders.sentence_transformers import (
    SentenceTransformersTextEmbedder,
    SentenceTransformersDocumentEmbedder,
)
from haystack_integrations.document_stores.pgvector import PgvectorDocumentStore
from haystack_integrations.components.retrievers.pgvector import (
    PgvectorEmbeddingRetriever,
)

os.environ["PG_CONN_STR"] = "postgresql://postgres:***@localhost:5432/postgres"
document_store = PgvectorDocumentStore(
    embedding_dimension=768,
    vector_function="cosine_similarity",
    recreate_table=True,
)

documents = [
    Document(content="There are over 7,000 languages spoken around the world today."),
    Document(
        content="Elephants have been observed to behave in a way that indicates a high level of self-awareness, such as recognizing themselves in mirrors.",
    ),
    Document(
        content="In certain parts of the world, like the Maldives, Puerto Rico, and San Diego, you can witness the phenomenon of bioluminescent waves.",
    ),
]

document_embedder = SentenceTransformersDocumentEmbedder()
documents_with_embeddings = document_embedder.run(documents)
document_store.write_documents(
    documents_with_embeddings.get("documents"),
    policy=DuplicatePolicy.OVERWRITE,
)

query_pipeline = Pipeline()
query_pipeline.add_component("text_embedder", SentenceTransformersTextEmbedder())
query_pipeline.add_component(
    "retriever",
    PgvectorEmbeddingRetriever(document_store=document_store),
)
query_pipeline.connect("text_embedder.embedding", "retriever.query_embedding")

query = "How many languages are there?"
result = query_pipeline.run({"text_embedder": {"text": query}})
print(result["retriever"]["documents"][0])

더 알아보기 (Learn more)

• pgvector 저장소
• PgvectorDocumentStore
• API reference: Pgvector