Reciprocal Rerank Fusion Retriever

여러 쿼리와 여러 인덱스에서 가져온 검색 결과를 어떻게 결합하는지 살펴보는 예시예요. Reciprocal Rerank Fusion 알고리즘으로 노드를 재정렬하는 방법을 함께 배워볼게요.

출처: 문서

본문

이 예시에서는 여러 쿼리와 여러 인덱스에서 가져온 검색 결과를 결합하는 방법을 살펴봅니다.

검색된 노드는 이 논문에서 소개된 Reciprocal Rerank Fusion 알고리즘에 따라 재정렬됩니다. 이 알고리즘은 과도한 계산이나 외부 모델에 의존하지 않고도 검색 결과를 효율적으로 재정렬하는 방법을 제공합니다.

이 예시 구현은 github의 @Raduaschl이 만든 예시 구현에 크게 기여받았습니다.

%pip install llama-index-llms-openai
%pip install llama-index-retrievers-bm25

import os
import openai


os.environ["OPENAI_API_KEY"] = "sk-..."
openai.api_key = os.environ["OPENAI_API_KEY"]

Setup

이 노트북을 colab에서 여는 경우 LlamaIndex 🦙를 설치해야 할 수 있습니다.

데이터 다운로드

!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

from llama_index.core import SimpleDirectoryReader


documents = SimpleDirectoryReader("./data/paul_graham/").load_data()

다음으로 문서 위에 벡터 인덱스를 구성합니다.

from llama_index.core import VectorStoreIndex
from llama_index.core.node_parser import SentenceSplitter


splitter = SentenceSplitter(chunk_size=256)


index = VectorStoreIndex.from_documents(documents, transformations=[splitter])

Hybrid Fusion Retriever 생성

이 단계에서는 인덱스를 BM25 기반 리트리버와 융합합니다. 이렇게 하면 입력 쿼리에서 의미적 관계와 키워드를 모두 포착할 수 있습니다.

두 리트리버 모두 점수를 계산하므로, 추가 모델이나 과도한 계산 없이 reciprocal rerank 알고리즘을 사용해 노드를 재정렬할 수 있습니다.

이 설정은 또한 쿼리를 4번 수행하는데, 원래 쿼리로 한 번, 그리고 추가로 3개의 쿼리를 생성합니다.

기본적으로 추가 쿼리를 생성하는 데 다음 프롬프트를 사용합니다:

QUERY_GEN_PROMPT = (
    "You are a helpful assistant that generates multiple search queries based on a "
    "single input query. Generate {num_queries} search queries, one on each line, "
    "related to the following input query:\n"
    "Query: {query}\n"
    "Queries:\n"
)

먼저 리트리버를 생성합니다. 각 리트리버는 가장 유사한 상위 2개 노드를 가져옵니다:

from llama_index.retrievers.bm25 import BM25Retriever


vector_retriever = index.as_retriever(similarity_top_k=2)


bm25_retriever = BM25Retriever.from_defaults(
    docstore=index.docstore, similarity_top_k=2
)

다음으로 퓨전 리트리버를 만들 수 있는데, 이 리트리버는 리트리버들이 반환한 4개 노드 중 상위 2개를 반환합니다:

from llama_index.core.retrievers import QueryFusionRetriever


retriever = QueryFusionRetriever(
    [vector_retriever, bm25_retriever],
    similarity_top_k=2,
    num_queries=4,  # set this to 1 to disable query generation
    mode="reciprocal_rerank",
    use_async=True,
    verbose=True,
    # query_gen_prompt="...",  # we could override the query generation prompt here
)

# apply nested async to run in a notebook
import nest_asyncio


nest_asyncio.apply()

nodes_with_scores = retriever.retrieve(
    "What happened at Interleafe and Viaweb?"
)

for node in nodes_with_scores:
    print(f"Score: {node.score:.2f} - {node.text}...\n-----\n")

보시다시피 두 반환 노드 모두 Viaweb과 Interleaf를 정확히 언급합니다!

Query Engine에서 사용하기!

이제 리트리버를 쿼리 엔진에 연결해 자연어 응답을 합성할 수 있습니다.

from llama_index.core.query_engine import RetrieverQueryEngine


query_engine = RetrieverQueryEngine.from_args(retriever)

response = query_engine.query("What happened at Interleafe and Viaweb?")

from llama_index.core.response.notebook_utils import display_response


display_response(response)

Reciprocal Rerank Fusion Retriever