Recursive Retriever + Node References + Braintrust

Recursive Retriever + Node References + Braintrust

재귀 검색을 사용해 노드 관계를 탐색하고 '참조(reference)'에 기반해 노드를 가져오는 방법을 보여드려요. 노드 참조의 여러 활용법과 Braintrust를 이용한 평가까지 함께 살펴볼게요.

출처: 문서

본문

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

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

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

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

재귀 검색 + 노드 참조 방식이 얼마나 잘 작동하는지 Braintrust를 사용해 평가합니다. Braintrust는 AI 제품 구축을 위한 엔터프라이즈급 스택입니다. 평가부터 프롬프트 플레이그라운드, 데이터 관리까지 AI를 비즈니스에 통합할 때의 불확실성과 지루함을 제거해 줍니다.

평가 대시보드 예시는 여기에서 볼 수 있습니다:

• base retriever
• recursive metadata retreiver
• recursive chunk retriever

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

%load_ext autoreload
%autoreload 2
# NOTE: Replace YOUR_OPENAI_API_KEY with your OpenAI API Key and YOUR_BRAINTRUST_API_KEY with your BrainTrust API key. Do not put it in quotes.
# Signup for Braintrust at https://braintrustdata.com/ and get your API key at https://www.braintrustdata.com/app/braintrustdata.com/settings/api-keys
%env OPENAI_API_KEY=
%env BRAINTRUST_API_KEY=
%env TOKENIZERS_PARALLELISM=true # This is needed to avoid a warning message from Chroma

%pip install -U llama_hub llama_index braintrust autoevals pypdf pillow transformers torch torchvision

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

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

!mkdir 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) 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))

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

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

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

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

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
metadata_dicts = []
for extractor in extractors:
    metadata_dicts.extend(extractor.extract(base_nodes))

# cache metadata dicts
def save_metadata_dicts(path):
    with open(path, "w") as fp:
        for m in metadata_dicts:
            fp.write(json.dumps(m) + "\n")




def load_metadata_dicts(path):
    with open(path, "r") as fp:
        metadata_dicts = [json.loads(l) for l in fp.readlines()]
        return metadata_dicts

save_metadata_dicts("data/llama2_metadata_dicts.jsonl")

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

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


all_nodes = copy.deepcopy(base_nodes)
for idx, d in enumerate(metadata_dicts):
    inode_q = IndexNode(
        text=d["questions_this_excerpt_can_answer"],
        index_id=base_nodes[idx].node_id,
    )
    inode_s = IndexNode(
        text=d["section_summary"], index_id=base_nodes[idx].node_id
    )
    all_nodes.extend([inode_q, inode_s])

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=True,
)

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)

재귀 검색 + 노드 참조 방식이 얼마나 잘 작동하는지 Braintrust를 사용해 평가합니다. Braintrust는 AI 제품 구축을 위한 엔터프라이즈급 스택입니다. 평가부터 프롬프트 플레이그라운드, 데이터 관리까지 AI를 비즈니스에 통합할 때의 불확실성과 지루함을 제거해 줍니다.

청크 참조와 메타데이터 참조를 모두 평가합니다. 참조 노드를 검색하기 위해 임베딩 유사도 조회를 사용합니다. 두 방식을 원시 노드를 직접 가져오는 기준 리트리버와 비교합니다. 지표로는 hit-rate와 MRR을 모두 사용합니다.

평가 대시보드 예시는 여기에서 볼 수 있습니다:

• base retriever
• recursive metadata retreiver
• recursive chunk retriever

데이터셋 생성 (Dataset Generation)

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

from llama_index.core.evaluation import (
    generate_question_context_pairs,
    EmbeddingQAFinetuneDataset,
)
import nest_asyncio


nest_asyncio.apply()

eval_dataset = generate_question_context_pairs(base_nodes)

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


# 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)

몇 가지 스코어링 함수를 정의하고 데이터셋 데이터 변수를 정의해 봅시다.

queries = eval_dataset.queries
relevant_docs = eval_dataset.relevant_docs
data = [
    ({"input": queries[query], "expected": relevant_docs[query]})
    for query in queries.keys()
]




def hitRateScorer(input, expected, output=None):
    is_hit = any([id in expected for id in output])
    return 1 if is_hit else 0




def mrrScorer(input, expected, output=None):
    for i, id in enumerate(output):
        if id in expected:
            return 1 / (i + 1)
    return 0

import braintrust


# Evaluate the chunk retriever
vector_retriever_chunk = vector_index_chunk.as_retriever(similarity_top_k=10)
retriever_chunk = RecursiveRetriever(
    "vector",
    retriever_dict={"vector": vector_retriever_chunk},
    node_dict=all_nodes_dict,
    verbose=False,
)




def runChunkRetriever(input, hooks):
    retrieved_nodes = retriever_chunk.retrieve(input)
    retrieved_ids = [node.node.node_id for node in retrieved_nodes]
    return retrieved_ids




chunkEval = await braintrust.Eval(
    name="llamaindex-recurisve-retrievers",
    data=data,
    task=runChunkRetriever,
    scores=[hitRateScorer, mrrScorer],
)

# Evaluate the metadata retriever


vector_retriever_metadata = vector_index_metadata.as_retriever(
    similarity_top_k=10
)
retriever_metadata = RecursiveRetriever(
    "vector",
    retriever_dict={"vector": vector_retriever_metadata},
    node_dict=all_nodes_dict,
    verbose=False,
)




def runMetaDataRetriever(input, hooks):
    retrieved_nodes = retriever_metadata.retrieve(input)
    retrieved_ids = [node.node.node_id for node in retrieved_nodes]
    return retrieved_ids




metadataEval = await braintrust.Eval(
    name="llamaindex-recurisve-retrievers",
    data=data,
    task=runMetaDataRetriever,
    scores=[hitRateScorer, mrrScorer],
)

# Evaluate the base retriever
base_retriever = base_index.as_retriever(similarity_top_k=10)




def runBaseRetriever(input, hooks):
    retrieved_nodes = base_retriever.retrieve(input)
    retrieved_ids = [node.node.node_id for node in retrieved_nodes]
    return retrieved_ids




baseEval = await braintrust.Eval(
    name="llamaindex-recurisve-retrievers",
    data=data,
    task=runBaseRetriever,
    scores=[hitRateScorer, mrrScorer],
)

더 알아보기 (Learn more)

• Braintrust
• Braintrust API 키 발급
• LlamaIndex RecursiveRetriever