LlamaIndex로 RAG 파이프라인 구축하기

LlamaIndex로 RAG 파이프라인 구축하기

LlamaIndex에서 MistralAI LLM과 임베딩 모델을 함께 사용해 RAG(검색 증강 생성) 파이프라인을 만드는 노트북이에요. 기본 RAG 파이프라인을 한 단계씩 쌓아 가는 것과 함께, 만든 Index를 Retriever로도 쓸 수 있는 방법까지 다룹니다.

출처: 문서

본문

이 노트북에서는 MistralAI LLM과 임베딩 모델로 LlamaIndex를 이용해 RAG를 구축하는 법을 살펴봅니다. 두 가지를 다룰게요.

  • 기본 RAG 파이프라인 (Basic RAG pipeline)
  • Index를 Retriever로 사용하기 (Index as Retriever)

필요한 패키지를 설치하고 API 키를 설정합니다.

!pip install llama-index
!pip install llama-index-embeddings-mistralai
!pip install llama-index-llms-mistralai

import os
os.environ['MISTRAL_API_KEY'] = '<YOUR MISTRALAI API KEY>'

기본 RAG 파이프라인

기본 RAG 파이프라인을 만드는 단계는 다음과 같아요.

  1. LLM과 임베딩 모델 설정
  2. 데이터 다운로드
  3. 데이터 로드
  4. 노드 생성
  5. Index 생성
  6. Query Engine 생성
  7. 쿼리 실행

Query Engine은 Retrieval과 Response Synthesis 모듈을 결합해서 주어진 쿼리에 대한 응답을 생성해 줍니다.

먼저 LLM과 임베딩 모델을 설정합니다.

from llama_index.llms.mistralai import MistralAI
from llama_index.embeddings.mistralai import MistralAIEmbedding

llm = MistralAI(model='mistral-large')
embed_model = MistralAIEmbedding()

from llama_index.core import Settings

Settings.llm = llm
Settings.embed_model = embed_model

이번 데모에서는 Uber 2021 10K SEC 보고서를 사용할게요.

!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/uber_2021.pdf' -O './uber_2021.pdf'

데이터를 로드하고 노드로 분할합니다.

from llama_index.core import SimpleDirectoryReader

documents = SimpleDirectoryReader(input_files=["./uber_2021.pdf"]).load_data()

from llama_index.core.node_parser import TokenTextSplitter

splitter = TokenTextSplitter(
    chunk_size=512,
    chunk_overlap=0,
)

nodes = splitter.get_nodes_from_documents(documents)

노드로 Index를 만들고 Query Engine을 만들어 쿼리를 실행합니다.

from llama_index.core import VectorStoreIndex

index = VectorStoreIndex(nodes)

query_engine = index.as_query_engine(similarity_top_k=2)

response = query_engine.query("What is the revenue of Uber in 2021?")

print(response)

Index를 Retriever로 사용하기

만든 index를 그대로 Retriever로 활용할 수도 있어요. Retriever는 사용자 쿼리에 대해 관련성 높은 청크/노드를 검색해 반환합니다.

retriever = index.as_retriever(similarity_top_k = 2)

retrieved_nodes = retriever.retrieve("What is the revenue of Uber in 2021?")

from llama_index.core.response.notebook_utils import display_source_node

for node in retrieved_nodes:
    display_source_node(node, source_length=1000)

더 알아보기 (Learn more)

  • LlamaIndex 공식 문서 — RAG 파이프라인 구축 프레임워크
  • llama-index-llms-mistralai — Mistral LLM용 LlamaIndex 통합
  • llama-index-embeddings-mistralai — Mistral 임베딩 모델용 통합
  • VectorStoreIndex.as_retriever() — Index를 검색기로 변환하는 API