Mistral AI로 시작하는 RAG 기초

Mistral AI로 시작하는 RAG 기초 (RAG Basics with Mistral AI)

LLM과 정보 검색 시스템의 능력을 결합한 RAG(검색 증강 생성)의 가장 기본적인 예시를 배우는 문서예요. Mistral만으로 처음부터 만드는 방법과, LangChain·LlamaIndex·Haystack 같은 인기 프레임워크를 활용하는 네 가지 구현을 다룹니다.

출처: 문서

본문

RAG(Retrieval-augmented generation, 검색 증강 생성)는 LLM과 정보 검색 시스템의 능력을 시너지하는 AI 프레임워크예요. 외부 지식을 활용해 질문에 답하거나 콘텐츠를 생성할 때 유용하죠. RAG에는 두 가지 주요 단계가 있어요:

  1. retrieval(검색): 벡터 스토어에 저장된 텍스트 임베딩으로 지식 베이스에서 관련 정보를 검색해요.
  2. generation(생성): 관련 정보를 프롬프트에 삽입해 LLM이 정보를 생성하게 해요.

이 가이드에서는 네 가지 구현으로 RAG의 아주 기본적인 예시를 살펴볼게요:

  • Mistral로 처음부터 RAG 만들기 (RAG from scratch)
  • Mistral과 LangChain으로 RAG
  • Mistral과 LlamaIndex로 RAG
  • Mistral과 Haystack으로 RAG

처음부터 RAG 만들기 (RAG from scratch)

이 섹션은 기본 RAG를 처음부터 구축하는 과정을 안내해요. 두 가지 목표가 있어요: 첫째, RAG의 내부 작동 원리에 대한 포괄적 이해를 제공해 그 메커니즘을 밝히는 것. 둘째, 최소한의 의존성만으로 RAG를 구축하는 데 필요한 기초 역량을 갖추게 하는 것.

필요 패키지 가져오기

mistralai와 faiss-cpu 패키지를 설치하고 필요한 패키지를 가져와요.

! pip install faiss-cpu==1.7.4 mistralai
from mistralai.client import Mistral
import requests
import numpy as np
import faiss
import os
from getpass import getpass

api_key= getpass("Type your API Key")
client = Mistral(api_key=api_key)

데이터 가져오기

이 간단한 예시에서는 Paul Graham이 쓴 에세이에서 데이터를 가져와요.

response = requests.get('https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt')
text = response.text

에세이를 로컬 파일로 저장할 수도 있어요.

f = open('essay.txt', 'w')
f.write(text)
f.close()
len(text)

문서를 청크로 나누기 (Split document into chunks)

RAG 시스템에서는 문서를 더 작은 청크로 나누는 것이 중요해요. 그래야 나중에 검색 과정에서 가장 관련 있는 정보를 더 효과적으로 식별·검색할 수 있거든요. 이 예시에서는 텍스트를 문자 단위로 나누고, 각 청크에 2048자를 결합해 37개의 청크를 얻어요.

chunk_size = 2048
chunks = [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]
len(chunks)

고려사항:

  • Chunk size (청크 크기): 특정 사용 사례에 따라 최적 성능을 위해 다른 청크 크기와 오버랩을 조정·실험해야 할 수 있어요. 예를 들어 작은 청크는 검색 과정에서 더 유리한데, 큰 텍스트 청크는 의미 표현을 흐릿하게 만들 수 있는 채움 텍스트(filler text)를 자주 포함하기 때문이에요. 다만 작은 청크는 처리 시간과 컴퓨팅 자원 증가 같은 트레이드오프가 있어요.
  • How to split (분할 방법): 가장 단순한 방법은 문자 단위 분할이지만, 사용 사례와 문서 구조에 따라 다른 옵션이 있어요. 예를 들어 API 호출의 토큰 한도를 피하려면 토큰 단위로 나누고, 청크의 응집력을 유지하려면 문장·문단·HTML 헤더 단위로 나누는 게 좋아요. 코드를 다룬다면 AST(Abstract Syntax Tree) 파서로 의미 있는 코드 청크로 나누는 걸 권장해요.

각 텍스트 청크의 임베딩 생성 (Create embeddings)

각 텍스트 청크에 대해 텍스트 임베딩을 만들어요. 이는 벡터 공간에서 텍스트의 숫자 표현이에요. 의미가 비슷한 단어들은 벡터 공간에서 더 가까이 있거나 거리가 짧을 것으로 기대돼요.

임베딩을 만들려면 Mistral의 임베딩 API 엔드포인트와 임베딩 모델 mistral-embed를 사용해요. 단일 텍스트 청크의 임베딩을 얻는 get_text_embedding을 만들고, 리스트 컴프리헨션으로 모든 텍스트 청크의 임베딩을 얻어요.

def get_text_embedding(input):
    embeddings_batch_response = client.embeddings.create(
        model="mistral-embed",
        inputs=input
    )
    return embeddings_batch_response.data[0].embedding
text_embeddings = np.array([get_text_embedding(chunk) for chunk in chunks])
text_embeddings.shape
text_embeddings

벡터 데이터베이스에 로드 (Load into a vector database)

텍스트 임베딩을 얻으면 일반적으로 효율적인 처리와 검색을 위해 벡터 데이터베이스에 저장해요. 선택할 수 있는 벡터 데이터베이스가 여러 가지 있는데, 이 간단한 예시에서는 효율적인 유사도 검색을 가능하게 하는 오픈소스 벡터 데이터베이스 Faiss를 사용해요.

Faiss로는 벡터 데이터베이스의 인덱싱 구조를 정의하는 Index 클래스 인스턴스를 만들고, 이 인덱싱 구조에 텍스트 임베딩을 추가해요.

d = text_embeddings.shape[1]
index = faiss.IndexFlatL2(d)
index.add(text_embeddings)

고려사항:

  • Vector database (벡터 데이터베이스): 벡터 데이터베이스를 선택할 때는 속도, 확장성, 클라우드 관리, 고급 필터링, 오픈소스 vs. 클로즈드소스 등 여러 요인을 고려해야 해요.

질문의 임베딩 생성 (Create embeddings for a question)

사용자가 질문할 때마다 이전과 같은 임베딩 모델로 질문의 임베딩도 만들어야 해요.

question = "What were the two main things the author worked on before college?"
question_embeddings = np.array([get_text_embedding(question)])
question_embeddings.shape
question_embeddings

고려사항:

  • HyDE (Hypothetical Document Embeddings): 어떤 경우에는 사용자의 질문이 관련 컨텍스트를 식별하는 가장 적절한 쿼리가 아닐 수 있어요. 대신 사용자 쿼리를 기반으로 가상의 답변이나 가상 문서를 생성하고, 생성된 텍스트의 임베딩으로 유사한 텍스트 청크를 검색하는 게 더 효과적일 수 있어요.

벡터 데이터베이스에서 유사한 청크 검색 (Retrieve similar chunks)

index.search로 벡터 데이터베이스에서 검색할 수 있어요. 이 함수는 두 인자를 받는데, 첫 번째는 질문 임베딩의 벡터, 두 번째는 검색할 유사 벡터의 개수예요. 이 함수는 질문 벡터와 가장 유사한 벡터들의 거리와 인덱스를 반환해요. 그런 다음 반환된 인덱스에 기반해 실제 관련 텍스트 청크를 가져올 수 있어요.

D, I = index.search(question_embeddings, k=2)
print(I)
retrieved_chunk = [chunks[i] for i in I.tolist()[0]]
print(retrieved_chunk)

고려사항:

  • Retrieval methods (검색 방법): 다양한 검색 전략이 있어요. 예시에서는 임베딩을 이용한 간단한 유사도 검색을 보여줘요. 데이터에 메타데이터가 있으면 유사도 검색 전에 먼저 메타데이터로 필터링하는 게 더 나을 때가 있어요. 또한 문서의 용어 빈도·분포를 사용해 관련 청크를 식별하는 TF-IDF, BM25 같은 통계적 검색 방법도 있어요.
  • Retrieved document (검색된 문서): 항상 개별 텍스트 청크를 그대로 검색할까요? 그렇지 않아요. 때로는 실제 검색된 텍스트 청크 주변에 더 많은 컨텍스트를 포함하기를 원해요. 실제 검색된 텍스트 청크를 "child chunk"라고 부르고, 그 "child chunk"가 속한 더 큰 "parent chunk"를 검색하는 게 목표일 수 있어요. 또 검색된 문서에 가중치를 주고 싶을 수도 있어요. 예를 들어 시간 가중치 방식은 가장 최근 문서를 검색하는 데 도움이 돼요.
  • 검색 과정의 흔한 문제 하나는 "lost in the middle" 문제예요. 긴 컨텍스트의 중간에 있는 정보가 사라지는 문제죠. 우리 모델은 이 문제를 완화하려고 노력했어요. 예를 들어 passkey 작업에서 우리 모델은 긴 프롬프트 안에 무작위로 삽입된 passkey를 최대 32k 컨텍스트 길이까지 찾아내는 "바늘 찾기(needle in a haystack)" 능력을 보여줬어요. 다만 가장 관련 있는 청크를 시작과 끝에 배치하면 결과가 개선되는지 실험해 보는 것도 고려할 만해요.

프롬프트에서 컨텍스트와 질문을 결합하고 응답 생성

마지막으로 검색된 텍스트 청크를 프롬프트 안의 컨텍스트 정보로 제공할 수 있어요. 검색된 텍스트와 사용자 질문을 모두 포함하는 프롬프트 템플릿이에요.

prompt = f"""
Context information is below.
---------------------
{retrieved_chunk}
---------------------
Given the context information and not prior knowledge, answer the query.
Query: {question}
Answer:
"""
def run_mistral(user_message, model="mistral-large-latest"):
    messages = [
        {
            "role": "user", "content": user_message
        }
    ]
    chat_response = client.chat.complete(
        model=model,
        messages=messages
    )
    return (chat_response.choices[0].message.content)
run_mistral(prompt)

고려사항:

  • Prompting techniques (프롬프팅 기법): 대부분의 프롬프팅 기법을 RAG 시스템 개발에도 사용할 수 있어요. 예를 들어 몇 가지 예시를 제공해 few-shot learning으로 모델의 답을 안내할 수 있고, 모델에게 특정 방식으로 답을 포맷하도록 명시적으로 지시할 수도 있어요.

다음 섹션들에서는 인기 있는 RAG 프레임워크로 비슷한 기본 RAG를 수행하는 방법을 보여드릴게요.

LangChain

!pip install langchain langchain-mistralai langchain_community mistralai==0.4.2
from langchain_community.document_loaders import TextLoader
from langchain_mistralai.chat_models import ChatMistralAI
from langchain_mistralai.embeddings import MistralAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate
from langchain.chains import create_retrieval_chain

# Load data
loader = TextLoader("essay.txt")
docs = loader.load()
# Split text into chunks
text_splitter = RecursiveCharacterTextSplitter()
documents = text_splitter.split_documents(docs)
# Define the embedding model
embeddings = MistralAIEmbeddings(model="mistral-embed", mistral_api_key=api_key)
# Create the vector store
vector = FAISS.from_documents(documents, embeddings)
# Define a retriever interface
retriever = vector.as_retriever()
# Define LLM
model = ChatMistralAI(mistral_api_key=api_key)
# Define prompt template
prompt = ChatPromptTemplate.from_template("""Answer the following question based only on the provided context:

<context>
{context}
</context>

Question: {input}""")

# Create a retrieval chain to answer questions
document_chain = create_stuff_documents_chain(model, prompt)
retrieval_chain = create_retrieval_chain(retriever, document_chain)
response = retrieval_chain.invoke({"input": "What were the two main things the author worked on before college?"})
print(response["answer"])

LlamaIndex

!pip install llama-index==0.10.55 llama-index-llms-mistralai==0.1.18 llama-index-embeddings-mistralai mistralai==0.4.2
import os
from llama_index.core import Settings, SimpleDirectoryReader, VectorStoreIndex
from llama_index.llms.mistralai import MistralAI
from llama_index.embeddings.mistralai import MistralAIEmbedding

# Load data
reader = SimpleDirectoryReader(input_files=["essay.txt"])
documents = reader.load_data()
# Define LLM and embedding model
Settings.llm = MistralAI(model="mistral-medium", api_key=api_key)
Settings.embed_model = MistralAIEmbedding(model_name='mistral-embed', api_key=api_key)
# Create vector store index
index = VectorStoreIndex.from_documents(documents)
# Create query engine
query_engine = index.as_query_engine(similarity_top_k=2)
response = query_engine.query(
    "What were the two main things the author worked on before college?"
)
print(str(response))

Haystack

!pip install mistral-haystack==0.0.1 mistralai==0.4.2
from haystack import Pipeline
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.dataclasses import ChatMessage
from haystack.utils.auth import Secret

from haystack.components.builders import DynamicChatPromptBuilder
from haystack.components.converters import TextFileToDocument
from haystack.components.preprocessors import DocumentSplitter
from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
from haystack.components.writers import DocumentWriter
from haystack_integrations.components.embedders.mistral import MistralDocumentEmbedder, MistralTextEmbedder
from haystack_integrations.components.generators.mistral import MistralChatGenerator

document_store = InMemoryDocumentStore()

docs = TextFileToDocument().run(sources=["essay.txt"])
split_docs = DocumentSplitter(split_by="passage", split_length=2).run(documents=docs["documents"])
embeddings = MistralDocumentEmbedder(api_key=Secret.from_token(api_key)).run(documents=split_docs["documents"])
DocumentWriter(document_store=document_store).run(documents=embeddings["documents"])

text_embedder = MistralTextEmbedder(api_key=Secret.from_token(api_key))
retriever = InMemoryEmbeddingRetriever(document_store=document_store)
prompt_builder = DynamicChatPromptBuilder(runtime_variables=["documents"])
llm = MistralChatGenerator(api_key=Secret.from_token(api_key),
    model='mistral-small')

chat_template = """Answer the following question based on the contents of the documents.\n
 Question: {{query}}\n

 Documents:
 {% for document in documents %}
 {{document.content}}
 {% endfor%}
 """
messages = [ChatMessage.from_user(chat_template)]

rag_pipeline = Pipeline()
rag_pipeline.add_component("text_embedder", text_embedder)
rag_pipeline.add_component("retriever", retriever)
rag_pipeline.add_component("prompt_builder", prompt_builder)
rag_pipeline.add_component("llm", llm)

rag_pipeline.connect("text_embedder.embedding", "retriever.query_embedding")
rag_pipeline.connect("retriever.documents", "prompt_builder.documents")
rag_pipeline.connect("prompt_builder.prompt", "llm.messages")

question = "What were the two main things the author worked on before college?"

result = rag_pipeline.run(
    {
        "text_embedder": {"text": question},
        "prompt_builder": {"template_variables": {"query": question}, "prompt_source": messages},
        "llm": {"generation_kwargs": {"max_tokens": 225}},
    }
)

print(result["llm"]["replies"][0].content)

더 알아보기 (Learn more)