AstraEmbeddingRetriever

AstraEmbeddingRetriever

Astra Document Store와 호환되는 임베딩 기반 Retriever예요.

파이프라인에서 가장 흔한 위치:

  1. RAG 파이프라인에서 Text Embedder 다음, PromptBuilder 앞
  2. 의미 검색(semantic search) 파이프라인의 마지막 컴포넌트
  3. 추출형 QA 파이프라인에서 Text Embedder 다음, TransformersExtractiveReader 앞

필수 init 변수: document_store — AstraDocumentStore 인스턴스 필수 run 변수: query_embedding — float 목록 출력 변수: documents — 문서 목록 API reference: Astra GitHub link: https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/astra Package name: astra-haystack

출처: 문서

본문

Overview

AstraEmbeddingRetriever는 쿼리와 문서 임베딩을 비교하고, 그 결과를 바탕으로 AstraDocumentStore에서 쿼리와 가장 관련 있는 문서를 가져와요.

NLP 시스템에서 AstraEmbeddingRetriever를 쓸 때는 쿼리 임베딩과 문서 임베딩이 준비돼 있어야 해요. 인덱싱 파이프라인에 Document Embedder를, 쿼리 파이프라인에 Text Embedder를 추가하면 됩니다.

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

Setup and installation

AstraDB 계정을 만들고 데이터베이스를 생성했다면 astra-haystack 통합을 설치하세요:

pip install astra-haystack

AstraDB 웹 UI의 설정에서 데이터베이스 ID와 생성된 토큰이 필요해요.

추가로 컬렉션 이름과 네임스페이스가 필요해요. 컬렉션 이름을 만들 때는 임베딩 차원과 유사도 지표도 설정해야 합니다. 네임스페이스는 데이터베이스 안의 데이터를 조직화하며 Apache Cassandra에서는 키스페이스(keyspace)라고 불러요.

그리고 아래 예시를 실행하려면 선택적으로 sentence-transformers-haystack 패키지도 설치하세요:

pip install sentence-transformers-haystack

Usage

인증 데이터는 환경 변수로 전달하는 걸 강력히 권장합니다. 아래 예시를 실행하기 전에 ASTRA_DB_API_ENDPOINT와 ASTRA_DB_APPLICATION_TOKEN 환경 변수를 채워두세요.

In a pipeline

쿼리 파이프라인에서 이 Retriever를 이렇게 사용해요:

from haystack import Document, Pipeline
from haystack.document_stores.types import DuplicatePolicy
from haystack_integrations.components.embedders.sentence_transformers import (
    SentenceTransformersTextEmbedder,
    SentenceTransformersDocumentEmbedder,
)
from haystack_integrations.components.retrievers.astra import AstraEmbeddingRetriever
from haystack_integrations.document_stores.astra import AstraDocumentStore

document_store = AstraDocumentStore()
model = "sentence-transformers/all-mpnet-base-v2"
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(model=model)
documents_with_embeddings = document_embedder.run(documents)
document_store.write_documents(
    documents_with_embeddings.get("documents"),
    policy=DuplicatePolicy.SKIP,
)

query_pipeline = Pipeline()
query_pipeline.add_component(
    "text_embedder",
    SentenceTransformersTextEmbedder(model=model),
)
query_pipeline.add_component(
    "retriever",
    AstraEmbeddingRetriever(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])

예시 출력:

Document(id=cfe93bc1c274908801e6670440bf2bbba54fad792770d57421f85ffa2a4fcc94, content: 'There are over 7,000 languages spoken around the world today.', score: 0.8929937, embedding: vector of size 768)

더 알아보기 (Learn more)

🧑‍🍳 쿡북: Using AstraDB as a data store in your Haystack pipelines