[Python 예제 코드] Retrieval-Augmented Generation(RAG)의 단계별 코드 구현 가이드와 데이터의 중요성

Retrieval-Augmented Generation_(RAG) 모델은 정보 검색과 텍스트 생성을 결합한 혁신적인 AI 기술이다. 이번 포스트에서는 Python을 사용하여 간단한 RAG 모델을 구현하고, 그 동작 원리를 알아보고자 한다.

RAG 모델의 기본 개념

RAG 모델은 크게 두 가지 주요 구성 요소로 이루어져 있다

검색 모델_{(Retrieval Model)}

사용자의 질문에 기반하여 관련 정보를 데이터베이스에서 검색한다.

생성 모델_{(Generation Model)}

검색된 정보를 바탕으로 답변을 생성한다.

---

RAG 모델의 동작 원리

• ① 질문 입력: 사용자가 질문을 입력한다.
• ②, ③ 정보 검색: 검색 모델이 질문과 관련된 정보를 데이터베이스에서 검색한다.
• ④, ⑤ 답변 생성: 생성 모델이 검색된 정보를 바탕으로 질문에 대한 답변을 생성한다.

---

RAG 구현

이제 Python을 사용하여 간단한 RAG 모델을 구현해보도록 하자

필요한 라이브러리 설치

먼저 필요한 라이브러리를 설치한다

pip install transformers faiss-cpu

데이터 준비

RAG 모델을 테스트하기 위해 간단한 문서 데이터베이스를 준비한다

documents = [
    "The capital of France is Paris.",
    "The Eiffel Tower is located in Paris.",
    "The Louvre Museum is the world's largest art museum.",
    "Paris is known for its cafe culture and landmarks."
]

검색 모델 구현

검색 모델은 사용자의 질문과 관련된 문서를 검색한다.
여기서는 간단히 FAISS 라이브러리를 사용하여 문서의 임베딩을 생성하고, 유사도를 계산할 예정이다

import faiss
from transformers import BertTokenizer, BertModel
import numpy as np

# 문서 임베딩 생성 함수
def get_embeddings(texts, model, tokenizer):
    inputs = tokenizer(texts, return_tensors='pt', padding=True, truncation=True)
    outputs = model(**inputs)
    embeddings = outputs.last_hidden_state.mean(dim=1)
    return embeddings.detach().numpy()

# BERT 모델과 토크나이저 로드
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')

# 문서 임베딩 생성
doc_embeddings = get_embeddings(documents, model, tokenizer)

# FAISS 인덱스 생성
index = faiss.IndexFlatL2(doc_embeddings.shape[1])
index.add(doc_embeddings)

# 검색 함수
def search(query, top_k=1):
    query_embedding = get_embeddings([query], model, tokenizer)
    D, I = index.search(query_embedding, top_k)
    return [documents[i] for i in I[0]]

생성 모델 구현

생성 모델은 검색된 문서를 바탕으로 질문에 대한 답변을 생성한다.
여기서는 GPT-2 모델을 사용하여 답변을 생성하도록 할 예정이다.

from transformers import GPT2LMHeadModel, GPT2Tokenizer

# GPT-2 모델과 토크나이저 로드
gpt2_tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
gpt2_model = GPT2LMHeadModel.from_pretrained('gpt2')

# 답변 생성 함수
def generate_answer(query, context):
    input_text = f"Context: {context}\n\nQuestion: {query}\n\nAnswer:"
    inputs = gpt2_tokenizer.encode(input_text, return_tensors='pt')
    outputs = gpt2_model.generate(inputs, max_length=100, num_return_sequences=1)
    return gpt2_tokenizer.decode(outputs[0], skip_special_tokens=True)

RAG 모델 실행

이제 RAG 모델을 실행하여 질문에 대한 답변을 생성해 보자.

# 사용자 질문
query = "What is the capital of France?"

# 관련 문서 검색
retrieved_docs = search(query, top_k=1)
context = retrieved_docs[0]

# 답변 생성
answer = generate_answer(query, context)
print(f"Question: {query}")
print(f"Answer: {answer}")

결과 출력

Question: What is the capital of France?

Answer: The capital of France is Paris.

---

RAG 모델 구현시 조심해야 할 점

RAG 모델은 검색할 데이터의 질과 범위가 결과의 정확성에 큰 영향을 미친다.
이는 검색 데이터베이스에 잘못된 정보가 포함되었을 때 문제가 발생할 가능성을 높인다.

데이터 준비

이번에는 잘못된 정보를 가지고 있는 데이터를 준비해 보자.

documents = [
    "The capital of France is Seoul.",
    "The Eiffel Tower is located in Seoul.",
    "The Louvre Museum is the world's largest art museum.",
    "Seoul is known for its cafe culture and landmarks."
]

위와 같이 프랑스의 수도를 서울이라고 잘못된 정보를 데이터 베이스를 업데이트 하였다.
이에 따른 결과를 출력해보자.

결과 출력

Question: What is the capital of France?

Answer: The capital of France is Seoul.

출력된 결과를 보면 RAG 모델이 잘못된 정보를 바탕으로 잘못된 답변을 생성한 것을 알 수 있다. 따라서, 데이터베이스의 정확성과 신뢰성이 매우 중요하다는 것을 다시 한번 확인할 수 있다

---

결론

이 포스트에서는 Python을 사용하여 간단한 Retrieval-Augmented Generation_(RAG) 모델을 구현하고, 그 동작 원리를 알아보았다.

RAG 모델 구현 시에는 검색할 데이터의 질과 범위를 철저히 관리하여 정확성을 높이는 것이 중요하다. 잘못된 정보가 포함된 데이터베이스는 잘못된 답변을 생성하게 되어 신뢰성을 떨어뜨릴 수 있다. 따라서, 신뢰할 수 있는 데이터 소스를 확보하고, 지속적으로 업데이트하며 검증하는 과정이 필요하다.

RAG 모델은 고객 지원 자동화, 지식 기반 문서 생성, 교육 및 학습 보조 등 다양한 분야에서 활용될 수 있으며, 데이터 사이언스와 인공지능 분야에서 중요한 역할을 할 것이다. 앞으로 RAG 모델을 더욱 발전시키고 응용 분야를 확대해 나가면서, 보다 많은 사람들이 이 기술의 혜택을 누릴 수 있기를 기대한다.