Recursive Retriever + Node References

Recursive Retriever + Node References

재귀 검색을 사용해 노드 관계를 탐색하고 '참조(reference)'에 기반해 노드를 가져오는 방법을 보여드려요. 청크 참조와 메타데이터 참조라는 두 가지 노드 참조 활용법을 실제 코드로 확인해 볼게요.

출처: 문서

본문

이 가이드는 재귀 검색을 사용해 노드 관계를 탐색하고 '참조'에 기반해 노드를 가져오는 방법을 보여줍니다.

노드 참조는 강력한 개념입니다. 처음 검색을 수행할 때 원시 텍스트 대신 참조를 가져오고 싶을 수 있습니다. 여러 참조가 같은 노드를 가리킬 수 있습니다.

이 가이드에서는 노드 참조의 몇 가지 다른 용법을 살펴봅니다:

• 청크 참조(Chunk references): 더 큰 청크를 가리키는 서로 다른 청크 크기
• 메타데이터 참조(Metadata references): 더 큰 청크를 가리키는 요약 + 생성된 질문

%pip install llama-index-llms-openai
%pip install llama-index-readers-file

%load_ext autoreload
%autoreload 2
%env OPENAI_API_KEY=YOUR_OPENAI_KEY

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

!pip install llama-index pypdf

데이터 로드 + 설정 (Load Data + Setup)

이 섹션에서는 Llama 2 논문을 다운로드하고 초기 노드 집합(chunk size 1024)을 생성합니다.

!mkdir -p 'data/'
!wget --user-agent "Mozilla" "https://arxiv.org/pdf/2307.09288.pdf" -O "data/llama2.pdf"

from pathlib import Path
from llama_index.readers.file import PDFReader
from llama_index.core.response.notebook_utils import display_source_node
from llama_index.core.retrievers import RecursiveRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core import VectorStoreIndex
from llama_index.llms.openai import OpenAI
import json

loader = PDFReader()
docs0 = loader.load_data(file=Path("./data/llama2.pdf"))

from llama_index.core import Document


doc_text = "\n\n".join([d.get_content() for d in docs0])
docs = [Document(text=doc_text)]

from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.schema import IndexNode

node_parser = SentenceSplitter(chunk_size=1024)

base_nodes = node_parser.get_nodes_from_documents(docs)
# set node ids to be a constant
for idx, node in enumerate(base_nodes):
    node.id_ = f"node-{idx}"

from llama_index.core.embeddings import resolve_embed_model


embed_model = resolve_embed_model("local:BAAI/bge-small-en")
llm = OpenAI(model="gpt-3.5-turbo")

기준(Baseline) 리트리버

임베딩 유사도로 상위 k개 원시 텍스트 노드를 단순히 가져오는 기준 리트리버를 정의합니다.

base_index = VectorStoreIndex(base_nodes, embed_model=embed_model)
base_retriever = base_index.as_retriever(similarity_top_k=2)

retrievals = base_retriever.retrieve(
    "Can you tell me about the key concepts for safety finetuning"
)

for n in retrievals:
    display_source_node(n, source_length=1500)

query_engine_base = RetrieverQueryEngine.from_args(base_retriever, llm=llm)

response = query_engine_base.query(
    "Can you tell me about the key concepts for safety finetuning"
)
print(str(response))

청크 참조: 더 작은 자식 청크가 더 큰 부모 청크를 가리키기

이 활용 예시에서는 더 작은 청크가 더 큰 부모 청크를 가리키는 그래프를 만드는 방법을 보여줍니다.

쿼리 시점에는 더 작은 청크를 검색하지만, 참조를 따라 더 큰 청크에 접근합니다. 이를 통해 합성에 더 많은 컨텍스트를 사용할 수 있습니다.

sub_chunk_sizes = [128, 256, 512]
sub_node_parsers = [
    SentenceSplitter(chunk_size=c, chunk_overlap=20) for c in sub_chunk_sizes
]


all_nodes = []
for base_node in base_nodes:
    for n in sub_node_parsers:
        sub_nodes = n.get_nodes_from_documents([base_node])
        sub_inodes = [
            IndexNode.from_text_node(sn, base_node.node_id) for sn in sub_nodes
        ]
        all_nodes.extend(sub_inodes)


    # also add original node to node
    original_node = IndexNode.from_text_node(base_node, base_node.node_id)
    all_nodes.append(original_node)

all_nodes_dict = {n.node_id: n for n in all_nodes}

vector_index_chunk = VectorStoreIndex(all_nodes, embed_model=embed_model)

vector_retriever_chunk = vector_index_chunk.as_retriever(similarity_top_k=2)

retriever_chunk = RecursiveRetriever(
    "vector",
    retriever_dict={"vector": vector_retriever_chunk},
    node_dict=all_nodes_dict,
    verbose=True,
)

nodes = retriever_chunk.retrieve(
    "Can you tell me about the key concepts for safety finetuning"
)
for node in nodes:
    display_source_node(node, source_length=2000)

query_engine_chunk = RetrieverQueryEngine.from_args(retriever_chunk, llm=llm)

response = query_engine_chunk.query(
    "Can you tell me about the key concepts for safety finetuning"
)
print(str(response))

메타데이터 참조: 더 큰 청크를 가리키는 요약 + 생성된 질문

이 활용 예시에서는 소스 노드를 가리키는 추가 컨텍스트를 정의하는 방법을 보여줍니다.

이 추가 컨텍스트에는 요약과 생성된 질문이 포함됩니다.

쿼리 시점에는 더 작은 청크를 검색하지만, 참조를 따라 더 큰 청크에 접근합니다. 이를 통해 합성에 더 많은 컨텍스트를 사용할 수 있습니다.

import nest_asyncio


nest_asyncio.apply()

from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.schema import IndexNode
from llama_index.core.extractors import (
    SummaryExtractor,
    QuestionsAnsweredExtractor,
)

extractors = [
    SummaryExtractor(summaries=["self"], show_progress=True),
    QuestionsAnsweredExtractor(questions=5, show_progress=True),
]

# run metadata extractor across base nodes, get back dictionaries
node_to_metadata = {}
for extractor in extractors:
    metadata_dicts = extractor.extract(base_nodes)
    for node, metadata in zip(base_nodes, metadata_dicts):
        if node.node_id not in node_to_metadata:
            node_to_metadata[node.node_id] = metadata
        else:
            node_to_metadata[node.node_id].update(metadata)

# cache metadata dicts
def save_metadata_dicts(path, data):
    with open(path, "w") as fp:
        json.dump(data, fp)




def load_metadata_dicts(path):
    with open(path, "r") as fp:
        data = json.load(fp)
    return data

save_metadata_dicts("data/llama2_metadata_dicts.json", node_to_metadata)

metadata_dicts = load_metadata_dicts("data/llama2_metadata_dicts.json")

# all nodes consists of source nodes, along with metadata
import copy


all_nodes = copy.deepcopy(base_nodes)
for node_id, metadata in node_to_metadata.items():
    for val in metadata.values():
        all_nodes.append(IndexNode(text=val, index_id=node_id))

all_nodes_dict = {n.node_id: n for n in all_nodes}

## Load index into vector index
from llama_index.core import VectorStoreIndex
from llama_index.llms.openai import OpenAI


llm = OpenAI(model="gpt-3.5-turbo")


vector_index_metadata = VectorStoreIndex(all_nodes)

vector_retriever_metadata = vector_index_metadata.as_retriever(
    similarity_top_k=2
)

retriever_metadata = RecursiveRetriever(
    "vector",
    retriever_dict={"vector": vector_retriever_metadata},
    node_dict=all_nodes_dict,
    verbose=False,
)

nodes = retriever_metadata.retrieve(
    "Can you tell me about the key concepts for safety finetuning"
)
for node in nodes:
    display_source_node(node, source_length=2000)

query_engine_metadata = RetrieverQueryEngine.from_args(
    retriever_metadata, llm=llm
)

response = query_engine_metadata.query(
    "Can you tell me about the key concepts for safety finetuning"
)
print(str(response))

평가 (Evaluation)

재귀 검색 + 노드 참조 방식이 얼마나 잘 작동하는지 평가합니다. 청크 참조와 메타데이터 참조를 모두 평가하며, 참조 노드를 검색하기 위해 임베딩 유사도 조회를 사용합니다.

두 방식을 원시 노드를 직접 가져오는 기준 리트리버와 비교합니다.

지표로는 hit-rate와 MRR을 모두 사용합니다.

데이터셋 생성 (Dataset Generation)

먼저 텍스트 청크 집합에서 질문 데이터셋을 생성합니다.

from llama_index.core.evaluation import (
    generate_question_context_pairs,
    EmbeddingQAFinetuneDataset,
)
from llama_index.llms.openai import OpenAI


import nest_asyncio


nest_asyncio.apply()

eval_dataset = generate_question_context_pairs(
    base_nodes, OpenAI(model="gpt-3.5-turbo")
)

eval_dataset.save_json("data/llama2_eval_dataset.json")

# optional
eval_dataset = EmbeddingQAFinetuneDataset.from_json(
    "data/llama2_eval_dataset.json"
)

결과 비교 (Compare Results)

각 리트리버에 대해 hit rate와 MRR을 측정하는 평가를 실행합니다.

노드 참조(청크든 메타데이터든)가 있는 리트리버가 원시 청크를 검색하는 것보다 더 좋은 성과를 내는 경향이 있음을 확인할 수 있습니다.

import pandas as pd
from llama_index.core.evaluation import (
    RetrieverEvaluator,
    get_retrieval_results_df,
)


# set vector retriever similarity top k to higher
top_k = 10




def display_results(names, results_arr):
    """Display results from evaluate."""


    hit_rates = []
    mrrs = []
    for name, eval_results in zip(names, results_arr):
        metric_dicts = []
        for eval_result in eval_results:
            metric_dict = eval_result.metric_vals_dict
            metric_dicts.append(metric_dict)
        results_df = pd.DataFrame(metric_dicts)


        hit_rate = results_df["hit_rate"].mean()
        mrr = results_df["mrr"].mean()
        hit_rates.append(hit_rate)
        mrrs.append(mrr)


    final_df = pd.DataFrame(
        {"retrievers": names, "hit_rate": hit_rates, "mrr": mrrs}
    )
    display(final_df)

vector_retriever_chunk = vector_index_chunk.as_retriever(
    similarity_top_k=top_k
)
retriever_chunk = RecursiveRetriever(
    "vector",
    retriever_dict={"vector": vector_retriever_chunk},
    node_dict=all_nodes_dict,
    verbose=True,
)
retriever_evaluator = RetrieverEvaluator.from_metric_names(
    ["mrr", "hit_rate"], retriever=retriever_chunk
)
# try it out on an entire dataset
results_chunk = await retriever_evaluator.aevaluate_dataset(
    eval_dataset, show_progress=True
)

vector_retriever_metadata = vector_index_metadata.as_retriever(
    similarity_top_k=top_k
)
retriever_metadata = RecursiveRetriever(
    "vector",
    retriever_dict={"vector": vector_retriever_metadata},
    node_dict=all_nodes_dict,
    verbose=True,
)
retriever_evaluator = RetrieverEvaluator.from_metric_names(
    ["mrr", "hit_rate"], retriever=retriever_metadata
)
# try it out on an entire dataset
results_metadata = await retriever_evaluator.aevaluate_dataset(
    eval_dataset, show_progress=True
)

base_retriever = base_index.as_retriever(similarity_top_k=top_k)
retriever_evaluator = RetrieverEvaluator.from_metric_names(
    ["mrr", "hit_rate"], retriever=base_retriever
)
# try it out on an entire dataset
results_base = await retriever_evaluator.aevaluate_dataset(
    eval_dataset, show_progress=True
)

full_results_df = get_retrieval_results_df(
    [
        "Base Retriever",
        "Retriever (Chunk References)",
        "Retriever (Metadata References)",
    ],
    [results_base, results_chunk, results_metadata],
)
display(full_results_df)

	retrievers	hit_rate	mrr
0	Base Retriever	0.778351	0.563103
1	Retriever (Chunk References)	0.896907	0.691114
2	Retriever (Metadata References)	0.891753	0.718440

위 표에서 볼 수 있듯이, 노드 참조를 사용한 두 리트리버(청크·메타데이터 참조) 모두 hit-rate와 MRR에서 기준(Baseline) 리트리버보다 더 높은 성능을 보입니다.

더 알아보기 (Learn more)

• LlamaIndex RecursiveRetriever
• LlamaIndex Evaluators