나만의 RAG 구현하기: OpenAI API와 FAISS로 PDF 문서 기반 질문-응답 시스템 만들기

안녕하세요 오늘은 BESPIN GLOBAL Innovate AI실 이규민님이 작성해주신 ‘나만의 RAG 구현하기: OpenAI API와 FAISS로 PDF 문서 기반 질문-응답 시스템 만들기’ 대해 소개해드리도록 하겠습니다.

목차

  1. 들어가며
  2. 나만의 RAG 구현하기
  3. 소스의 내부 동작을 이해해보기 위해 추가적으로 해본 작업
  4. 마치며

1. 들어가며

최근 자연어 처리 기술의 발전으로, 다양한 방식으로 정보 검색과 질문-응답 시스템을 구현할 수 있습니다. 그 중 하나가 RAG(Retrieval-Augmented Generation) 모델로, 이는 정보 검색(Information Retrieval)과 텍스트 생성(Generation)을 결합한 방식입니다. 

RAG는 일반적으로 사용자가 던진 질문에 대해, 관련 문서를 먼저 검색하고 그 문서를 기반으로 답변을 생성하는 방식으로 동작합니다. 이를 통해 모델은 보다 정확하고 구체적인 답변을 제공할 수 있습니다.

이번 글에서는 OpenAI API와 FAISS 라이브러리를 사용하여 RAG 모델을 구현하는 방법을 소개하겠습니다. OpenAI는 강력한 언어 모델인 GPT를 제공하여 텍스트 생성의 고도화를 가능하게 하며, FAISS는 벡터 데이터베이스로 효율적인 검색을 지원하는 라이브러리입니다. 

이 두 가지를 결합하여 PDF 문서에서 질문을 받아 답변을 생성해보는 샘플 코드를 구현해보았습니다.

2. 나만의 RAG 구현하기

개발 환경은 Google Colab 입니다. 구현에 대해 Step by Step으로 알아보겠습니다.

  • Cell 1: 필요한 라이브러리 설치
# Cell 1: 필요한 라이브러리 설치
# OpenAI와 LangChain을 사용하기 위해 필요한 라이브러리를 설치합니다.
# PDF 처리를 위해 pypdf를, 벡터 데이터베이스를 위해 faiss-cpu를 설치합니다.
!pip install openai langchain pypdf faiss-cpu langchain-community tiktoken

필요한 라이브러리를 설치합니다. 

여기서는 OpenAI와 LangChain을 사용하여 GPT 기반의 질문-응답 시스템을 구축하며, PDF 파일 처리를 위해 pypdf 라이브러리, 벡터 데이터베이스 구축을 위해 faiss-cpu를 설치합니다. 

langchain-community와 tiktoken은 LangChain의 확장 기능을 제공하며, 텍스트를 보다 효율적으로 처리해주는 라이브러리 입니다.

  • Cell 2: OpenAI API 키 설정 및 초기화
# Cell 2: OpenAI API 키 설정 및 초기화
import os
from google.colab import userdata
from langchain.embeddings import OpenAIEmbeddings  # OpenAI 임베딩 생성
from langchain.vectorstores import FAISS  # FAISS를 사용한 벡터 데이터베이스 생성
from langchain.document_loaders import PyPDFLoader  # PDF 문서 로더
from langchain.chains import RetrievalQA  # 질문-응답 체인 생성
from langchain.chat_models import ChatOpenAI  # OpenAI GPT 모델

# OpenAI API 키를 환경 변수에 설정합니다.
os.environ["OPENAI_API_KEY"] = userdata.get('OPENAI_API_KEY')

OpenAI API 키를 환경 변수로 설정합니다.

OpenAIEmbeddings를 사용하여 문서에서 임베딩을 생성하고, FAISS 라이브러리를 사용하여 벡터 데이터베이스를 생성합니다.  PyPDFLoader를 통해 PDF 파일을 로드하고, RetrievalQA를 사용해 질문과 답변을 처리할 수 있는 체인을 생성합니다.

  • Cell 3: PDF 파일 업로드 및 벡터 DB 생성
# Cell 3: PDF 파일 업로드 및 벡터 DB 생성
from google.colab import files
import shutil
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import FAISS
from langchain.document_loaders import PyPDFLoader

def upload_and_create_vector_db():
    """
    PDF 파일을 업로드하고 벡터 DB를 생성하는 함수.
    파일 업로드가 없을 경우, 사용자에게 메시지를 출력하고 종료.
    """
    # 이전 업로드 파일 및 벡터 DB 초기화
    shutil.rmtree('/content/uploads', ignore_errors=True)
    os.makedirs('/content/uploads', exist_ok=True)

    # 파일 업로드 요청
    print("문서를 업로드하고 해당 문서 내용에 대해 질문하세요.")
    uploaded = files.upload()

    # 파일 업로드 여부 확인
    if not uploaded:
        print("파일 업로드가 필요합니다.")
        return None, None

    # 업로드된 파일 처리
    for filename in uploaded.keys():
        file_path = f'/content/uploads/{filename}'
        with open(file_path, 'wb') as f:
            f.write(uploaded[filename])

        # PDF 로드 및 벡터 DB 생성
        loader = PyPDFLoader(file_path)
        documents = loader.load()
        embeddings = OpenAIEmbeddings()
        vector_db = FAISS.from_documents(documents, embeddings)

        print(f"파일 '{filename}'이 벡터 DB에 저장되었습니다.")
        return filename, vector_db  # 파일명과 생성된 벡터 DB 반환

PDF 파일을 업로드하고, 업로드된 파일을 처리하여 벡터 데이터베이스를 생성하는 작업을 합니다. 

PyPDFLoader를 사용하여 PDF 파일을 로드한 후, OpenAIEmbeddings로 임베딩을 생성하고 이를 FAISS로 저장합니다. FAISS 라이브러리를 통해 사용자가 업로드한 파일에 대해 해당 문서를 벡터화하고, 이를 기반으로 검색을 수행할 수 있는 벡터 데이터베이스를 간편하게 생성할 수 있습니다.

  • Cell 4: 사용자 질문 및 응답 생성
# Cell 4: 사용자 질문 및 응답 생성
from langchain.chains import RetrievalQA
from langchain.chat_models import ChatOpenAI

def ask_question_and_get_answer(filename, vector_db):
    """
    질문을 받고 업로드된 PDF 내용을 기반으로 답변을 생성하는 함수.
    """
    # 벡터 DB가 없는 경우 함수 실행 방지
    if not vector_db:
        print("질문을 받기 전에 PDF 파일을 업로드하세요.")
        return

    # GPT 모델 선택
    model_name = "gpt-3.5-turbo"

    # 사용자 질문 입력
    print("질문을 입력하세요: ")
    user_question = input()

    # 질문-응답 체인 설정 및 실행
    qa_chain = RetrievalQA.from_chain_type(
        llm=ChatOpenAI(model=model_name, temperature=0.1),  # 선택한 모델 사용, temperature를 0에 가깝게 설정하여 파일 내용에 더 적합한 답변을 생성하도록 함
        retriever=vector_db.as_retriever(),  # 벡터 DB 검색
        return_source_documents=False
    )
    # 답변 생성
    answer = qa_chain.run(user_question)
    print(f"업로드한 파일 '{filename}'을 기반으로한 답변입니다:\n {answer}")

사용자가 입력한 질문에 대해 벡터 DB를 사용하여 답변을 생성하는 부분입니다. RetrievalQA 체인을 사용하여 PDF 파일에서 가장 관련성 높은 문서(text)를 검색하고, 이를 기반으로 GPT 모델인 gpt-3.5-turbo를 사용하여 답변을 생성합니다. 

temperature 값을 0에 가깝게 설정하여 모델이 보다 구체적이고 정확한 답변을 생성할 수 있도록 합니다. (temperaturer값이 1에 가까울 수록 창의적인 답변을 생성)

  • Cell 5: 테스트

Cell 5: 테스트

최종적으로 위 함수들을 테스트 해보겠습니다.

제가 최근에 요약했던 AWS Re:Invent 2024 세션 요약 파일로 테스트 해본 결과 위와 같이 파일 내용을 기반으로 답변을 생성하는 것을 확인 할 수 있었습니다.

3. 소스의 내부 동작을 이해해보기 위해 추가적으로 해본 작업

OpenAI의 API를 활용하여 간편하게 구현해볼 수 있었지만, 모델의 동작이나 임베딩 등을 더 잘 이해해보기 위해서 오픈소스 모델과 라이브러리로도 구현해보았습니다.

  • Cell 1: 필요한 라이브러리 설치
# Cell 1: 라이브러리 설치
!pip install faiss-cpu sentence-transformers langchain langchain_community pypdf transformers
  • Cell 2: Model과 Tokenizer 생성

# Hugging Face 모델 로드
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

if torch.cuda.is_available():
    print("GPU 사용 가능:", torch.cuda.get_device_name(0))
else:
    print("GPU를 사용할 수 없습니다.")

_model_name = "EleutherAI/gpt-neo-1.3B"

def load_generation_model(model_name="EleutherAI/gpt-neo-1.3B"):
    """
    Hugging Face의 GPT-Neo 모델 로드.
    """
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    tokenizer.add_special_tokens({'pad_token': '[PAD]'})

    return tokenizer, model

# GPT-Neo 모델 로드
generation_tokenizer, generation_model = load_generation_model(_model_name)

OpenAI 라이브러리에서 생성해주는 model 대신 오픈소스 모델을 통해 직접 생성해보았습니다.

토큰 길이를 자동으로 채우기 위해 tokenizer에 padding을 추가하였습니다.

  • Cell 3: Embedding 생성
# Hugging Face 임베딩 모델 및 FAISS 설정
from sentence_transformers import SentenceTransformer
from langchain_community.vectorstores import FAISS
from langchain.embeddings.base import Embeddings
import numpy as np

# 사용자 정의 Hugging Face 임베딩 클래스
class HuggingFaceEmbeddings(Embeddings):
    def __init__(self, model_name="sentence-transformers/all-MiniLM-L6-v2"):
        self.model = SentenceTransformer(model_name)
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'

    def embed_documents(self, texts):
        """텍스트를 임베딩 벡터로 변환"""
        return self.model.encode(texts, convert_to_tensor=True, device=self.device).cpu().detach().numpy()

    def embed_query(self, text):
        """질문 텍스트를 임베딩 벡터로 변환"""
        return self.model.encode(text, convert_to_tensor=True, device=self.device).cpu().detach().numpy()

# Hugging Face 임베딩 객체 생성
embedding_model = HuggingFaceEmbeddings()

라이브러리에서 사용할 수 있는 OpenAIEmbeddings 대신 오픈소스를 라이브러리를 사용하여 임베딩을 직접 생성하였습니다. 백터 디비 생성 시 Faiss 를 사용하기 위해서는 embed_documents 메소드가 필요하여 직접 구현하였습니다.

  • Cell 4: Embedding 생성

라해당 코드는 기존 코드와 동일하나 embedding만 직접 구현한 embedding_model로 대체 합니다.

        # FAISS 벡터 DB 생성
        vector_db = FAISS.from_documents(documents, embedding_model)
  • Cell 5: 사용자 질문 및 응답 생성
# 질문 응답 함수
def ask_question_and_generate_answer(filename, vector_db, tokenizer, model):
    """
    GPT-Neo 모델을 사용하여 자연스러운 답변 생성.
    """
    if not vector_db:
        print("PDF 파일을 먼저 업로드하세요.")
        return

    print("질문을 입력하세요: ")
    user_question = input()

    # 벡터 DB에서 문맥 검색
    docs = vector_db.similarity_search(user_question, k=1)
    context = "\n".join([doc.page_content for doc in docs])

    # 프롬프트 구성
    prompt = f"문맥:\n{context}\n\n질문: {user_question}\n\n답변:"
    print(f"\n---\n생성된 프롬프트:\n{prompt}\n---")

    # 텍스트 생성
    # inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True, padding=True)
    inputs = tokenizer(user_question, return_tensors="pt", max_length=512, truncation=True, padding=True)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    inputs = {key: value.to(device) for key, value in inputs.items()}

    # 텍스트 생성
    outputs = model.generate(
        inputs["input_ids"],
        attention_mask=inputs["attention_mask"],  # attention_mask 추가
        # max_length=1024,
        max_new_tokens=256,  # 생성되는 토큰 길이 제한
        do_sample=True,  # 샘플링 활성화
        temperature=0.7,  # 샘플링 시 온도 조정
        top_k=50,  # 상위 k개 단어만 고려
        top_p=0.95,  # 누적 확률이 95%를 넘는 단어 제거
        pad_token_id=tokenizer.eos_token_id  # 패딩 토큰 설정
    )

    # 생성된 답변
    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print(f"\n업로드된 파일 '{filename}'을 기반으로 생성된 답변입니다:")
    # print(f"---\n{answer[len(prompt):].strip()}\n---")  # 프롬프트 제외한 텍스트만 출력
    print(f"---\n{answer}\n---")  # 프롬프트 제외한 텍스트만 출력
  • Cell 5: 테스트

테스트 해본 결과, 오픈소스 무료 모델인 GPT-3 기반 모델이라 토큰 제한도 매우 짧았고, 영어에 대해서는 그나마 답변이 나왔으나 한국어는 잘 이해하지 못하였습니다. 하지만, 토큰 길이 제한이 너무 작아(최대 약 1024 토근 이내) 질문 + PDF 파일 내용을 기반으로 검색한 유사도 높은 context 를 입력으로 하여 사용하기에는 부적절하였습니다.

오픈소스를 통해 직접 구현해보는 과정을 통해 벡터디비를 통해 조회한 문맥 context, input과 output을 조절하기 위한 각 파라미터 등에 대해 더 상세히 알 수 있었습니다.

4. 마치며

OpenAI API와 FAISS를 사용하여 RAG 모델을 구현해 보았지만 다음과 같은 한계점도 알게 되었습니다.

  1. 문서의 길이가 매우 긴 경우나 복잡한 구조를 가진 PDF 파일에서는 검색 성능이나 모델 응답 시간이 저하될 수 있습니다.
  2. 모델이 생성하는 답변의 품질은 입력된 데이터의 품질과 관련이 깊기 때문에, 문서의 내용이 명확하지 않거나 부정확한 경우 올바른 답변을 얻기 어려울 수 있습니다.

이러한 한계점은 특정 도메인에 특화된 모델을 사용하여 문서 전처리 과정을 강화하거나, 질문을 보다 세분화하여 여러 단계로 답변을 생성하거나 다양한 모델을 결합하여 더 정확한 결과를 도출하는 방법을 통해 개선해볼 수 있을 것 같습니다.

여기까지 ‘나만의 RAG 구현하기: OpenAI API와 FAISS로 PDF 문서 기반 질문-응답 시스템 만들기’에 대해 소개해드렸습니다. 유익한 정보가 되셨길 바랍니다. 감사합니다. 

Written by 이 규민/ Innovate AI실

BESPIN GLOBAL