Phi-3.5와 PGVector 벡터 DB를 이용한 검색증강생성(RAG) 시스템 구축하기

 
Phi 모델 구동 방법 참고:
2024.09.22 - [🟣 AI & ML] - Microsoft의 Phi-3.5 모델 Mac Silicon 환경에서 구동하기
PGVector 설치 방법 참고:
2024.05.13 - [🟣 AI & ML] - 검색증강생성(RAG) - LangChain과 PGVector를 이용한 간단한 RAG 시스템 구축해보기
 
지난 포스트에서 Microsoft의 Phi-3.5 모델을 구동하는 법을 살펴보았습니다. 이번 포스트에서는, 이 Phi-3.5가 특정 문서의 정보를 참고하여 답변할 수 있게하는 RAG 시스템을 PGVector를 이용하여 구축해보겠습니다. 시스템의 대략적인 구조는 다음과 같습니다. RAG 시스템 구축 프로세스는 크게 두 단계로 구성됩니다.
 

• Step 1: 사전 작업 단계로, Phi 모델을 다운로드하고 문서를 벡터로 변환하여 벡터 DB(PGVector)에 저장합니다.
• Step 2: RAG 시스템 구축 단계로, Phi 모델과 벡터 저장소 기반 검색기(retriever) 결합한 후 프롬프트를 적절히 설정하여 RAG 시스템을 완성합니다.

 

 
 

프로젝트 구성 및 환경 변수 설정(.env)

 
프로젝트 예제는 다음과 같이 구성합니다. 모델은 models 폴더 내에 저장되며, flashattention.pdf는 예제로 사용할 PDF 파일입니다. env 파일에 각종 환경 변수를 저장합니다. 각자 본인의 환경에 맞게 토큰이나 키를 설정해주면 됩니다.
 
 

myllm
└───models
│   └───phi-3.5-mini
│       │   file01
│       │   file02
│       │    ...
└───phi-3.5-rag.ipynb  
└───flashattention.pdf
└───.env

 

# .env file
HF_TOKEN="YOUR_TOKEN"
OPENAI_API_KEY ="sk-YOUR_KEY"
DB_CONNECTION = "postgresql+psycopg://{user}:{password}@{host}:{port}/{db}"

 
 

Step 1: 모델 다운로드 및 벡터 생성

 

모듈 불러오기

 
필요한 모듈을 불러오고, Huggingface 인증도 수행해줍니다.
 
 

import os
import torch

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from huggingface_hub import snapshot_download
from huggingface_hub import login
login(token=os.getenv("HF_TOKEN"), add_to_git_credential=True)

from langchain_openai import OpenAIEmbeddings
from langchain_postgres.vectorstores import PGVector
from langchain_community.document_loaders import UnstructuredPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

from IPython.display import Markdown, display

 
 

Phi 모델 다운로드

 
이전 포스트에서 설명한 대로, snapshot_download를 이용해 Huggingface 모델을 로컬의 지정한 경로에 다운받을 수 있습니다.
 
 

snapshot_download(repo_id="microsoft/Phi-3.5-mini-instruct",
                  local_dir="models/Phi-3.5-mini-instruct")

 
 
 

문서 벡터 변환 및 PGVector 저장

 
다음은 예제로 사용할 PDF 파일입니다. 어떤 PDF를 사용해도 상관 없습니다만, Phi-3.5가 지원하는 언어를 참고하여 문서를 설정해야 합니다.
 

 
 
먼저 Unstructured 모듈을 이용해 PDF로부터 텍스트를 추출합니다. PDF 데이터 추출 모듈은 다양하므로 상황에 따라 바꾸어 사용해도 됩니다.
 
 

file_path = "flashattention.pdf" # https://arxiv.org/abs/2205.14135
loader = UnstructuredPDFLoader(file_path)
loaded = loader.load()

 
 
벡터 검색이 용이하도록 불러온 문서를 분할해줍니다(chunking). 분할 길이나 문서 간 겹침 정도(chunk_overlap)는 문서에 따라 적당히 조정해 줍니다.
문서 간 겹침 정도를 설정하는 이유는 어떤 특정 정보의 양이 많을 때, 하나의 정보를 나타내는 연속된 분할 문서를 함께 검색해 올 확률을 높이기 위해서입니다.
 
 

splitter = RecursiveCharacterTextSplitter(chunk_size = 400,
                                          chunk_overlap = 50)
splitted = splitter.split_documents(loaded)

 
 
분할된 문서를 벡터로 변환할 임베딩 모델을 설정합니다.
 
 

embedding = OpenAIEmbeddings(model="text-embedding-3-small",
                             openai_api_key=os.getenv('OPENAI_API_KEY'))

 
 
분할된 문서와 임베딩 모델, PGVector가 설치된 PostgresSQL 연결 정보를 PGVector 인스턴스에 입력하여 실행하면 자동으로 문서를 벡터로 변환 후 PGVector에 지정한 collection_name으로 저장합니다.
 
 

vectorstore = PGVector.from_documents(
     documents = splitted,
     embedding = embedding,
     collection_name = "flashattention",
     connection=os.getenv("DB_CONNECTION")
     )

 
 
이제 RAG 시스템을 구축하기 위한 사전 작업이 완료되었습니다. 이제 Step 2에서 RAG 시스템을 만들어봅시다.
 
 

Step 2: RAG 시스템 구축하기

 

Phi 모델 불러오기

 
Step 1에서 다운받은 Phi 모델을 불러오고, 텍스트 생성 파이프라인을 구성합니다.
 
 

device = "mps" if torch.backends.mps.is_available() else "cpu" # if mac silicon
# device = "cuda" if torch.cuda.is_available() else "cpu" # if cuda
model = AutoModelForCausalLM.from_pretrained(
    "models/Phi-3.5-mini-instruct",
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    device_map=device,
    )
tokenizer = AutoTokenizer.from_pretrained("models/Phi-3.5-mini-instruct")

 

pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
)

generation_args = {
    "max_new_tokens": 800,
    "return_full_text": False
    }

 
 

PGVector에서 특정 문서에 대한 벡터 인덱스 가져오기

 
Step 1에서 저장해놓은 벡터 인덱스를 from_existing_index()를 이용해 불러옵니다. 필요하다면 임베딩 모델을 다시 설정하여 입력해줍니다.
불러온 벡터 인덱스가 지정된 PGVector 객체는 as_retriever()를 이용해 언어 모델의 검색기(retriever)로 지정하여 사용할 수 있습니다.
 
 

embedding = OpenAIEmbeddings(model="text-embedding-3-small",
                             openai_api_key=os.getenv('OPENAI_API_KEY'))
connection = os.getenv("DB_CONNECTION")
vectorstore = PGVector.from_existing_index(collection_name="flashattention",
                                           embedding=embedding,
                                           connection=connection)
retriever = vectorstore.as_retriever(search_kwargs={'k': 3})

 
 
 

프롬프트 설정

 
프롬프트를 설정하는 파트입니다. 언어 모델이 답변에 검색한 문서를 참고하도록 하는 프롬프트 설정이 조금 복잡해보일 수 있습니다. 다음과 같이 각 요소의 의미를 확실히 구분하면 됩니다. 특히, 검색한 문서를 포함하여 모델의 입력을 관리하는 messages 객체 대화 기록만을 저장하는 chat history 객체를 혼동하지 않도록 합시다.

• query는 사용자가 매 턴(turn) 입력하는 질문입니다.
• context는 PGVector를 이용해 검색해온, query와 유사한 정보를 담은 문서 집합들입니다.
• messages 객체는 모델로부터 결과를 얻기위해 사용하는 지시문(rag_system_prompt), 포맷팅된 쿼리(query_format) 두 개의 값만을 가집니다. 포맷팅된 쿼리는 대화 기록(chat_history)과 현재 사용자 질문(query), 질문에 대해 검색한 문서(context)의 정보를 모두 담고 있는 객체입니다.
• chat history 객체는 매 턴 생성되는 사용자의 질문과 모델의 답변만을 저장하는 객체입니다. 검색한 문서 정보나, 시스템 프롬프트는 저장되지 않습니다. 매 턴마다 값이 추가되며, 대화 기록은 다시 사용자 입력에 반영됩니다. 

 

먼저, 검색한 문서 정보인 context를 참고하여 답변할 수 있도록 시스템 프롬프트를 구성합니다.
 
 

rag_system_prompt = """You are an assistant for question-answering tasks.
Use the following pieces of retrieved context and chat history to answer the question.
If you don't know the answer, just say that you don't know.
keep the answer brief and concise."""

 
 
시스템 프롬프트를 messages 객체에 담고, chat_history 객체를 생성해줍니다.
 
 

messages = [{"role" : "system", "content" : rag_system_prompt}]
chat_history = []

 
 
질문을 입력하고, 질문에 포함된 정보와 유사한 정보를 담은 문서를 벡터 DB인 PGVector에서 검색합니다. 그리고, 검색해 온 정보를 간단히 전처리 해줍니다.
 
 

query = input()
relavant_docs = retriever.invoke(query)
context = ' \n\n'.join([i.page_content for i in relavant_docs])
display(Markdown(context))

 
 
검색해온 정보의 예시는 다음과 같습니다. 여기서는 3개의 문서를 검색해오도록 설정했습니다.
 
 

 
 
사용자의 질문을 chat_history에 저장합니다.
 
 

chat_history.append({"role" : "user", "content" : query})

 
 
이제, 대화 기록(chat_history)을 모델 답변에 참고할 수 있도록 간단히 전처리해줍니다. 대화 기록이 너무 길어지면 답변 생성 성능이 떨어지거나 시간이 너무 오래 걸릴 수 있으므로, max_last_history를 설정하여 마지막 몇 개의 대화 기록만 참고하도록 합시다.

길어지는 대화 기록을 관리하는 또 다른 방법으로, 대화 기록을 매 턴마다 별도 설정한 LLM을 이용하여 요약하여 이용하기도 합니다.
 

def chat_history_to_string(chat_history, max_last_history):
    result = ''
    max_last_history = -(max_last_history) -1
    for i in chat_history[max_last_history:-1]:
        result += i['role'].upper() + ': ' + i['content'] + ' \n'
    return result

 
 
이제 설정한 chat_history, query, context를 합쳐 모델 입력(query_format)을 구성합니다. 
 
 

query_format  = f"""chat_history: {chat_history_to_string(chat_history, max_last_history=4)} \n
question: {query} \n
context:\n {context} \n
answer: """
display(Markdown(query_format))

 
 
query_format 예시는 다음과 같습니다.
 
 

 
 
이제 query_format도 messages에 더해줍니다. 이제 입력이 완성되었습니다. 모델은 지시문(system_prompt), 사용자 질문(query), 마지막 몇 개의 대화 기록(chat_history), 질문과 유사한 문서 정보(context)를 참고하여 답변을 수행할 것입니다.
 
 

messages.append({"role" : "user", "content" : query_format})

 
 
 
대화 기록을 출력할 간단한 함수를 구성합니다.
 
 

def display_messages(chat_history):
    for i in chat_history:
        if i['role'] == 'user':
            display_text = "**[User]** \n" + i['content']
            display(Markdown(display_text))
        elif i['role'] == 'assistant':
            display_text = "**[Phi]** \n" + i['content']
            display(Markdown(display_text))

 
 
 
이제 답변 생성을 수행합니다. 출력된 모델 답변은 chat_history에 저장하여 업데이트하고, messages에서는 또 다음 모델 입력을 구성하여 입력해야 하므로 pop()을 이용해 마지막 query_format을 삭제해줍니다.
 
 

output = pipe(messages, **generation_args)
response = output[0]['generated_text']
chat_history.append({"role" : "assistant", "content" : response})
messages.pop()
display_messages(chat_history)

 
 
다음은 RAG 시스템을 이용해 몇 차례 질문과 답변을 수행하는 예입니다. 문서를 참고하여 답변이 잘 생성되고, 대화 기록도 잘 반영되고 있는 것 같습니다.