재밌는 토이 프로젝트! RAG 기술을 활용하여 pdf 파일에서 내용을 검색, 질의응답(QA)을 구현하는 챗봇 시스템을 만들어 보았습니다.
* 본 포스팅 활용하시는 경우 출처를 밝혀주세요 :-)
* 전체 코드 .py파일을 원하시는 분들께서는 포스팅 하단에 github 링크 참고 바랍니다.
* 사용한 라이브러리
pdfplumber, pytesseract, langchain, ollama, chromadb, gradio (전부 무료)
* 프로젝트 개요
- PDF 파일의 텍스트를 pdfplumber, pytesseract로 추출합니다.
- 추출한 내용을 langchain을 이용해서 split, 임베딩하여 벡터화한 다음 Chroma 벡터 저장소에 저장합니다.
- 벡터 저장소에서 질문에 해당하는 내용을 검색하여 context로 준비합니다.
- Ollama 라이브러리를 이용해서 LLaMA3 모델에게 질문과 context를 프롬프트로 제공하고 답변을 받습니다.
- 모든 과정을 gradio와 연동하여 GUI로 실행할 수 있도록 설정했습니다.
1. 환경설정 및 라이브러리 준비
먼저 새로운 콘다 환경을 만들고 해당 환경 안에서 프로젝트를 진행하도록 하겠습니다.
conda create -n ragchatbot python=3.11 -y # 환경 생성
conda activate ragchatbot # 활성화
필요한 라이브러리를 설치 후 임포트합니다.
# 설치 (터미널에서 실행)
pip install pdfplumber pytesseract ollama gradio langchain
# 임포트
import gradio as gr
from langchain.docstore.document import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import OllamaEmbeddings
import ollama
import pdfplumber
import pytesseract
from PIL import Image
2. 텍스트 추출, 임베딩, chroma에 저장
텍스트 추출, 임베딩, 벡터 데이터베이스 저장에 필요한 2가지 함수를 작성하겠습니다.
# PDF page에서 텍스트 추출하는 함수 작성
def extract_text_with_ocr(page):
text = page.extract_text()
if not text: # 만약 추출할 텍스트가 없다면
# PDF page를 이미지로 변환
image = page.to_image()
# 이미지에서 OCR 재실행하여 텍스트 추출
text = pytesseract.image_to_string(image)
return text
# PDF 파일을 열어서 extract_text_with_ocr 함수 실행 -> 벡터 데이터베이스에 저장하는 함수 작성
def load_and_retrieve_docs(file):
text = ""
try:
with pdfplumber.open(file) as pdf:
for page in pdf.pages:
page_text = extract_text_with_ocr(page)
# 페이지에서 추출한 텍스트가 있을 때마다 text에 누적해서 저장
if page_text:
text += page_text
except Exception as e:
return f"Error reading PDF file: {e}"
# 만약 추출한 텍스트가 하나도 없는 경우 아래와 같은 메세지 출력하고 함수 종료
if not text:
return "No text found in the PDF file."
# 추출한 텍스트가 있는 경우
docs = [Document(page_content=text)]
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
splits = text_splitter.split_documents(docs)
embeddings = OllamaEmbeddings(model="mxbai-embed-large")
vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings)
return vectorstore.as_retriever()
- extract_text_with_ocr
- 본 함수는 pdf 개별 페이지를 Input으로 받아서 페이지 내의 텍스트를 추출, 반환하는 함수입니다.
- 텍스트 추출에는 pdfplumber 라이브러리의 extract_text() 함수를 사용했습니다. 만약 페이지 내에서 추출할 텍스트가 없는 경우 pdfplumber 라이브러리의 to_image() 함수를 사용해서 pdf파일을 pdfplumber가 핸들할 수 있는 PageImage object객체화(이미지화)합니다. 이후 해당 이미지에서 pytesseract의 image_to_string() 함수를 사용해서 이미지로부터 텍스트를 추출하도록 했습니다.
- pytesseract docs : https://pypi.org/project/pytesseract/
- pdfplumber docs : https://github.com/jsvine/pdfplumber
- load_and_retrieve_docs
- 본 함수는 pdf파일을 열어서 모든 페이지마다 extract_text_with_ocr함수를 반복, text에 누적해서 더해주어 하나의 string으로 완성합니다.
- pdf파일을 열기 위해서(load) pdfplumber 라이브러리의 open()함수를 사용했습니다.
- 완성된 text를 langchain의 Document 객체로 만든 다음, RecursiveCharacterTextSplitter 객체를 생성한 뒤 만들어 둔 Document 객체를 집어넣어 적당히 청킹을 해주었습니다. 이후 청킹이 완료된 텍스트들을 벡터화하기 위한 OllamaEmbeddings 객체를 생성해 주었습니다.
- langchain에 연동된 Chroma 라이브러리를 불러온 다음, from_documents() 함수를 이용해 Chroma에 load를 해 줍니다. 이 때, 저장할 텍스트 청킹 객체와 임베딩 객체를 지정해 주어 무엇을 어떻게 임베딩하여 저장할지를 지정해주면 자동으로 임베딩을 실행하면서 vectorstore에 저장을 완료하기 때문에 사용이 매우 간편합니다.
- 마지막으로 load_and_retrieve_docs 함수는 결과값으로 vectorstore에 as_retriever() 함수를 적용한 객체를 반환합니다.
- langchain docs(Chroma) : click
- langchain docs(RecursiveCharacterTextSplitter) : click
3. 검색 실행
먼저 터미널에서 LLaMA3 모델을 다운로드(pull) 해줍니다.
참고로 docs에 따르면 llama3 기본 모델보다 chatqa 모델이 대화형 QA와 RAG 기능을 훨씬 잘 수행한다 - excels at conversational question answering (QA) and retrieval-augmented generation (RAG) - 라고 기재가 되어 있어서 처음에는 LLaMA3-chatqa모델을 선택하였는데, 실험 결과 LLaMA3-chatqa보다 기본 LLaMA3 모델이 더 답변이 정상적으로 생성되는 것을 확인할 수 있었습니다. 그 이유에 대해서는 좀더 살펴봐야 할 것 같아요.
참고로 모델 크기가 커서 pull에 시간이 좀 소요되는 편입니다.
ollama pull llama3
# 리스트 안의 모든 document 객체 내용을 추출해서 string으로 이어붙여 반환
def format_docs(docs):
return "\n\n".join(doc.page_content for doc in docs)
# Function that defines the RAG chain
def rag_chain(file, question):
retriever = load_and_retrieve_docs(file)
if isinstance(retriever, str): # 리턴받은 값이 string인 경우 에러를 의미하므로 함수 중단
return retriever
retrieved_docs = retriever.get_relevant_documents(question) # Use get_relevant_documents method
formatted_context = format_docs(retrieved_docs)
formatted_prompt = f"Question: {question}\n\nContext: {formatted_context}"
response = ollama.chat(model='llama3',
messages=[
{"role": "system",
"content": "You are a helpful assistant. Check the pdf content and answer the question."
},
{"role": "user", "content": formatted_prompt}])
return response['message']['content']
- rag_chain
- 위에서 작성했던 load_and_retrieve_docs 함수는 pdf 내용을 저장해 둔 vectorstore에 as_retriever() 함수를 적용한 객체를 반환합니다. 먼저 이 객체에 retriever라는 이름을 할당한 다음, 만약 retriever의 객체 타입이 string인 경우 에러 메세지(추출한 텍스트 없을 때)가 리턴된 것을 의미하므로 해당 에러를 프린트하면서 함수를 종료합니다.
- docs에 따르면 get_relevant_documents() 함수는 주어진 쿼리(질문)에 가장 관련성이 높은 문서를 벡터 저장소에서 검색하는 기능을 합니다. 예를 들어 제가 lloco_paper.pdf 파일로 테스트를 했을 때, RecursiveCharacterTextSplitters를 이용해서 split한 document 객체는 총 55개였고, 그 중에서 검색을 통해 반환하는 document 객체는 항상 4개인 것으로 확인이 되었습니다. 이 수치를 변경할 수 있는지는 나중에 확인을 해 보도록 하겠습니다.
- 반환받은 document 객체 리스트인 retrieved_docs를 format_docs 함수를 이용해서 하나의 string으로 변환하고, 해당 string은 prompt의 context로 주어지게 됩니다.
- ollama 라이브러리의 chat() 함수를 이용해서 LLaMA 3 모델에 question과 context를 제공하고, 답변(response)을 받습니다. chat
만약 답변을 한국어로 받고 싶은 경우 messages의 context 안에 Generate your answer in Korean 등과 같이 한국어로 답변을 내놓으라는 한 문장을 추가해서 프롬프팅을 해 주시면 됩니다.
4. Gradio GUI 인터페이스 할당
# Gradio interface
iface = gr.Interface(
fn=rag_chain,
inputs=["file", "text"],
outputs="text",
title="[LLAMA 3] RAG 검색 활용 챗봇 시스템",
description="PDF파일을 업로드하고 질문을 입력하면 답변을 생성해 드립니다. (영어로!)"
)
# app 실행
iface.launch()
주피터 파일로 실행하는 경우 바로 인터페이스가 보이실 것이고, 만약 .py파일로 실행하신다면 아래와 같이 할당되는 로컬 URL 주소로 접속을 하시면 인터페이스를 확인하실 수 있습니다.
URL로 접속하면 위와 같은 화면이 보이고, pdf파일 업로드 + question 작성 후 submit을 누르면 output에 답변이 생성됩니다.
저는 제가 현재 읽고 있는 LLoCO 논문(2024년 4월 발행)을 업로드해서 4.1장과 4.5장을 요약해달라는 요청을 해 보았구요
위와 같이 답변을 생성받았습니다 :-) 이모지 넣어달라고 했더니 야무지게 2개 넣어줬네요. 제공한 pdf 파일이 15장의 긴 파일임을 생각하면, 간단히 진행해본 것과 비교했을 때 성능이 꽤 잘 나온 것으로 생각하고 있습니다.
https://github.com/surisurikim/deep_learning/tree/main/RAG
해당 토이 프로젝트는 위 깃허브에서 코드 파일을 확인하실 수 있습니다. 감사합니다 :)
'Data Science > DL 딥러닝' 카테고리의 다른 글
딥러닝 | RAG(2021) 논문 리뷰 (0) | 2024.08.05 |
---|---|
딥러닝 | 효율적인 파인튜닝에 관한 고찰 - LoRA(2021) 논문 리뷰, peft, unsloth (0) | 2024.08.01 |
딥러닝 | BERT(2019) 논문 리뷰 (7) | 2024.07.22 |
딥러닝 | 트랜스포머 positional encoding 코드 구현 (문제 해결) (0) | 2024.07.06 |
딥러닝 | Microsoft 테이블 트랜스포머 PubTables-1m(2021) 논문 리뷰 (0) | 2024.07.04 |