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

Ollama에서 실행하는 Mistral 모델과 LlamaIndex로 RAG 파이프라인을 만드는 문서예요. RouterQueryEngine으로 질문을 적절한 인덱스로 라우팅하고, SubQuestionQueryEngine으로 복잡한 질문을 하위 질문으로 쪼개 처리해요.

출처: 문서

본문

이 노트북에서는 Ollama, Mistral 모델, LlamaIndex를 사용해 RAG 파이프라인을 구축하는 방법을 보여줘요. 다루는 주제는 다음과 같아요:

Mistral을 Ollama, LlamaIndex와 통합하기
Mistral 모델로 Ollama와 LlamaIndex를 사용한 RAG 구현하기
RouterQueryEngine으로 쿼리 라우팅하기
SubQuestionQueryEngine으로 복잡한 쿼리 처리하기

이 노트북을 실행하기 전에 Ollama를 설정해야 해요. 여기의 지침을 따라 주세요.

import nest_asyncio
nest_asyncio.apply()
from IPython.display import display, HTML

LLM 설정 (Setup LLM)

Ollama의 mistral:instruct 모델을 LLM으로 설정해요.

from llama_index.llms.ollama import Ollama

llm = Ollama(model="mistral:instruct", request_timeout=60.0)

쿼리 실행 (Querying)

from llama_index.core.llms import ChatMessage

messages = [
    ChatMessage(role="system", content="You are a helpful assistant."),
    ChatMessage(role="user", content="What is the capital city of France?"),
]
response = llm.chat(messages)

display(HTML(f'<p style="font-size:20px">{response}</p>'))

임베딩 모델 설정 (Setup Embedding Model)

HuggingFace의 소형 임베딩 모델을 사용해요.

from llama_index.embeddings.huggingface import HuggingFaceEmbedding

embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")

from llama_index.core import Settings

Settings.llm = llm
Settings.embed_model = embed_model

데이터 다운로드 및 로드 (Download Data / Load Data)

시연에는 Uber와 Lyft의 10K SEC 서류를 사용할 거예요.

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

from llama_index.core import SimpleDirectoryReader

uber_docs = SimpleDirectoryReader(input_files=["./uber_2021.pdf"]).load_data()
lyft_docs = SimpleDirectoryReader(input_files=["./lyft_2021.pdf"]).load_data()

인덱스와 쿼리 엔진 생성 (Create Index and Query Engines)

각 문서에 벡터 인덱스와 쿼리 엔진을 만들어요.

from llama_index.core import VectorStoreIndex
from llama_index.core import SummaryIndex

uber_vector_index = VectorStoreIndex.from_documents(uber_docs)
uber_vector_query_engine = uber_vector_index.as_query_engine(similarity_top_k=2)

lyft_vector_index = VectorStoreIndex.from_documents(lyft_docs)
lyft_vector_query_engine = lyft_vector_index.as_query_engine(similarity_top_k=2)

response = uber_vector_query_engine.query("What is the revenue of uber in 2021 in millions?")
display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

response = lyft_vector_query_engine.query("What is the revenue of lyft in 2021 in millions?")
display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

RouterQueryEngine (쿼리 라우팅)

질문이 Uber 관련인지 Lyft 관련인지에 따라 적절한 인덱스로 사용자 질문을 보내기 위해 RouterQueryEngine을 사용할 거예요.

from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core.query_engine.router_query_engine import RouterQueryEngine
from llama_index.core.selectors.llm_selectors import LLMSingleSelector

query_engine_tools = [
    QueryEngineTool(
        query_engine=lyft_vector_query_engine,
        metadata=ToolMetadata(
            name="vector_lyft_10k",
            description="Provides information about Lyft financials for year 2021",
        ),
    ),
    QueryEngineTool(
        query_engine=uber_vector_query_engine,
        metadata=ToolMetadata(
            name="vector_uber_10k",
            description="Provides information about Uber financials for year 2021",
        ),
    ),
]

query_engine = RouterQueryEngine(
    selector=LLMSingleSelector.from_defaults(),
    query_engine_tools=query_engine_tools,
    verbose = True
)

response = query_engine.query("What are the investments made by Uber?")
display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

response = query_engine.query("What are the investments made by the Lyft in 2021?")

display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

SubQuestionQueryEngine (복잡한 질문 처리)

SubQuestionQueryEngine을 활용하면 복잡한 질문을 하위 질문으로 쪼개 각각 처리함으로써 복잡한 쿼리를 해결할 수 있어요.

from llama_index.core.query_engine import SubQuestionQueryEngine

sub_question_query_engine = SubQuestionQueryEngine.from_defaults(query_engine_tools=query_engine_tools,
                                                                verbose=True)

response = sub_question_query_engine.query("Compare the revenues of Uber and Lyft in 2021?")
display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

response = sub_question_query_engine.query("What are the investments made by Uber and Lyft in 2021?")
display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

이처럼 Ollama에서 로컬로 돌리는 Mistral 모델로도 LlamaIndex를 활용한 RAG 파이프라인을 완성할 수 있고, 라우팅과 하위 질문 분해로 단일·복합 질문을 모두 처리할 수 있어요.

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