하이브리드 검색 파이프라인

하이브리드 검색 파이프라인 (Hybrid Retrieval)

키워드 기반 검색과 임베딩(벡터) 기반 검색을 함께 쓰는 하이브리드 검색 파이프라인을 만들어 볼게요. 각 방식의 장점을 모두 활용하는 구조라, 어느 한쪽만 쓸 때보다 더 정확하고 다양한 결과를 얻을 수 있어요.

출처: 공식문서

개요

하이브리드 검색은 키워드 기반과 임베딩 기반 검색 기법을 결합해요. 대체로 밀집 임베딩(dense embedding) 은 질문의 문맥적 뉘앙스를 잘 잡고, 키워드 기반(BM25 같은) 은 키워드 매칭에 강해요. 도메인이 특화된 경우(예: 의료)에는 밀집 모델이 해당 데이터로 학습돼야 하기 때문에, 단순 키워드 검색이 더 잘 나오는 경우도 있어요. 의료 데이터인 PubMed 초록으로 실습해 보겠어요.

설치

pip install haystack-ai
pip install "datasets>=2.6.1"
pip install sentence-transformers-haystack

문서 스토어 초기화

from haystack.document_stores.in_memory import InMemoryDocumentStore

document_store = InMemoryDocumentStore()

문서 가져와서 처리하기

PubMed 데이터셋(anakin87/medrag-pubmed-chunk)에서 각 항목을 Document로 만들어요. contents(초록+제목)는 검색에 쓰고, 나머지(title, abstract, pmid)는 메타데이터로 저장해 검색 결과 표현이나 메타데이터 필터링에 써요.

from datasets import load_dataset
from haystack import Document

dataset = load_dataset("anakin87/medrag-pubmed-chunk", split="train")

docs = []
for doc in dataset:
    docs.append(
        Document(content=doc["contents"], meta={"title": doc["title"], "abstract": doc["content"], "pmid": doc["id"]})
    )

문서 색인 파이프라인

문서를 512 단어 단위로 나누는 DocumentSplitter, 밀집 검색용 임베딩을 만드는 SentenceTransformersDocumentEmbedder, 그리고 문서를 저장하는 DocumentWriter로 색인 파이프라인을 만들어요.

from haystack.components.writers import DocumentWriter
from haystack_integrations.components.embedders.sentence_transformers import SentenceTransformersDocumentEmbedder
from haystack.components.preprocessors.document_splitter import DocumentSplitter
from haystack import Pipeline
from haystack.utils import ComponentDevice

document_splitter = DocumentSplitter(split_by="word", split_length=512, split_overlap=32)
document_embedder = SentenceTransformersDocumentEmbedder(
    model="BAAI/bge-small-en-v1.5", device=ComponentDevice.from_str("cuda:0")
)
document_writer = DocumentWriter(document_store)

indexing_pipeline = Pipeline()
indexing_pipeline.add_component("document_splitter", document_splitter)
indexing_pipeline.add_component("document_embedder", document_embedder)
indexing_pipeline.add_component("document_writer", document_writer)

indexing_pipeline.connect("document_splitter", "document_embedder")
indexing_pipeline.connect("document_embedder", "document_writer")

indexing_pipeline.run({"document_splitter": {"documents": docs}})

하이브리드 검색 파이프라인 만들기

하이브리드 검색은 키워드 검색과 벡터 검색을 모두 실행한 뒤, 결과를 크로스-인코더 모델로 재정렬해요.

1) 리트리버와 임베더 초기화

밀집 검색용 InMemoryEmbeddingRetriever와 키워드 검색용 InMemoryBM25Retriever를 준비해요. 밀집 검색에는 색인 파이프라인과 같은 임베딩 모델(BAAI/bge-small-en-v1.5)을 쓰는 질문 임베더가 필요해요.

from haystack.components.retrievers.in_memory import InMemoryBM25Retriever, InMemoryEmbeddingRetriever
from haystack_integrations.components.embedders.sentence_transformers import SentenceTransformersTextEmbedder

text_embedder = SentenceTransformersTextEmbedder(
    model="BAAI/bge-small-en-v1.5", device=ComponentDevice.from_str("cuda:0")
)
embedding_retriever = InMemoryEmbeddingRetriever(document_store)
bm25_retriever = InMemoryBM25Retriever(document_store)

2) 결과 재정렬

크로스-인코더 모델로 검색된 문서들의 관련성을 점수화하는 SentenceTransformersSimilarityRanker를 써요.

from haystack_integrations.components.rankers.sentence_transformers import SentenceTransformersSimilarityRanker

ranker = SentenceTransformersSimilarityRanker(model="BAAI/bge-reranker-base")

3) 파이프라인 구성

BM25와 임베딩 리트리버 모두 같은 ranker 입력으로 연결해요.

from haystack import Pipeline

hybrid_retrieval = Pipeline()
hybrid_retrieval.add_component("text_embedder", text_embedder)
hybrid_retrieval.add_component("embedding_retriever", embedding_retriever)
hybrid_retrieval.add_component("bm25_retriever", bm25_retriever)
hybrid_retrieval.add_component("ranker", ranker)

hybrid_retrieval.connect("text_embedder", "embedding_retriever")
hybrid_retrieval.connect("bm25_retriever", "ranker")
hybrid_retrieval.connect("embedding_retriever", "ranker")

하이브리드 검색 테스트

질문을 text_embedder, bm25_retriever, ranker 세 곳에 모두 넘겨서 실행해요.

query = "apnea in infants"

result = hybrid_retrieval.run(
    {"text_embedder": {"text": query}, "bm25_retriever": {"query": query}, "ranker": {"query": query}}
)

결과를 '검색 페이지'처럼 보기 좋게 출력하는 함수도 만들 수 있어요.

def pretty_print_results(prediction):
    for doc in prediction["documents"]:
        print(doc.meta["title"], "\t", doc.score)
        print(doc.meta["abstract"])
        print("\n", "\n")

pretty_print_results(result["ranker"])

더 알아보기