Mistral AI와 LlamaIndex로 Sub-Question Query Engine 만들기

Mistral AI와 LlamaIndex로 Sub-Question Query Engine 만들기

복잡한 사용자 질의를 하위 질의(sub-queries)로 쪼개 처리하는 SubQuestionQueryEngine을 LlamaIndex와 MistralAI로 구현하는 노트북이에요. 여러 문서에서 문맥을 가져와야 하는 질의를 더 정확하게 답할 수 있습니다.

출처: 문서

본문

VectorStoreIndex는 단일 문서나 문서 모음 안의 특정 문맥과 관련된 질의를 잘 처리해요. 그런데 실제 사용자 질의는 복잡해서 답을 내기 위해 여러 문서에서 문맥을 가져와야 하는 경우가 많죠. 그럴 때 단순한 VectorStoreIndex로는 부족할 수 있어요. 대신 복잡한 질의를 하위 질의로 쪼개면 더 정확한 응답을 얻을 수 있습니다.

설치와 설정

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

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

import nest_asyncio
nest_asyncio.apply()

필요한 모듈을 임포트하고 LLM·임베딩 모델을 설정합니다.

from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
from llama_index.llms.mistralai import MistralAI
from llama_index.embeddings.mistralai import MistralAIEmbedding
from llama_index.core import Settings
from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core.query_engine import SubQuestionQueryEngine

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

Settings.llm = llm
Settings.embed_model = embed_model
Settings.chunk_size = 512

로깅 (Logging)

참고: 이 설정은 Jupyter 노트북에서만 필요해요. Jupyter는 백그라운드에서 이벤트 루프를 돌리는데, 비동기 질의를 위해 이벤트 루프를 새로 시작하면 중첩된 이벤트 루프가 생깁니다. 보통은 허용되지 않지만 편의를 위해 nest_asyncio로 허용해요.

# NOTE: This is ONLY necessary in jupyter notebook.
# Details: Jupyter runs an event-loop behind the scenes.
#          This results in nested event-loops when we start an event-loop to make async queries.
#          This is normally not allowed, we use nest_asyncio to allow it for convenience.
import nest_asyncio
nest_asyncio.apply()

import logging
import sys

# Set up the root logger
logger = logging.getLogger()
logger.setLevel(logging.INFO)  # Set logger level to INFO

# Clear out any existing handlers
logger.handlers = []

# Set up the StreamHandler to output to sys.stdout (Colab's output)
handler = logging.StreamHandler(sys.stdout)
handler.setLevel(logging.INFO)  # Set handler level to INFO

# Add the handler to the logger
logger.addHandler(handler)

from IPython.display import display, HTML

데이터 다운로드·로드

Uber·Lyft 10K SEC 보고서와 Paul Graham 에세이 문서를 사용합니다.

!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'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O './paul_graham_essay.txt'

# Uber docs
uber_docs = SimpleDirectoryReader(input_files=["./uber_2021.pdf"]).load_data()

# Lyft docs
lyft_docs = SimpleDirectoryReader(input_files=["./lyft_2021.pdf"]).load_data()

# Paul Graham Essay
paul_graham_docs = SimpleDirectoryReader(input_files=["./paul_graham_essay.txt"]).load_data()

Index와 Query Engine 생성

각 문서에 대해 벡터 인덱스와 쿼리 엔진을 만듭니다.

# Index on uber docs
uber_vector_index = VectorStoreIndex.from_documents(uber_docs)

# Index on lyft docs
lyft_vector_index = VectorStoreIndex.from_documents(lyft_docs)

# Index on Paul Graham docs
paul_graham_vector_index = VectorStoreIndex.from_documents(paul_graham_docs)

# Query Engine over Index with uber docs
uber_vector_query_engine = uber_vector_index.as_query_engine(similarity_top_k = 5)

# Query Engine over Index with lyft docs
lyft_vector_query_engine = lyft_vector_index.as_query_engine(similarity_top_k = 5)

# Query Engine over Index with Paul Graham Essay
paul_graham_vector_query_engine = paul_graham_vector_index.as_query_engine(similarity_top_k = 5)

도구 만들기 (Create Tools)

각 쿼리 엔진을 도구로 감싸 이름과 설명을 붙입니다.

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

Sub-Question Query Engine 생성

sub_question_query_engine = SubQuestionQueryEngine.from_defaults(query_engine_tools=query_engine_tools)

질의하기 (Querying)

여기서 복잡한 질의를 답하기 위해 생성된 하위 질의를 확인할 수 있어요. 위 질의는 Uber와 Lyft 관련 하위 질의 두 개를 만듭니다.

response = sub_question_query_engine.query("Compare the revenue of uber and lyft?")
display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

Uber와 Paul Graham 에세이 관련 하위 질의 두 개를 만듭니다.

response = sub_question_query_engine.query("What is the revenue of uber and why did paul graham start YC?")
display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

Uber, Lyft, Paul Graham 에세이 모두와 관련된 하위 질의를 만듭니다.

response = sub_question_query_engine.query("Compare revenue of uber with lyft and why did paul graham start YC?")
display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

더 알아보기 (Learn more)

• LlamaIndex SubQuestionQueryEngine 문서 — 하위 질의 분해 쿼리 엔진
• QueryEngineTool / ToolMetadata — 쿼리 엔진을 도구로 감싸는 구성 요소
• mistral-large — 이 예제에서 사용한 LLM 모델
• nest_asyncio — Jupyter에서 중첩 이벤트 루프 처리용 라이브러리