Docker에서 ChromaDB 사용하기(3) - pdf 데이터 기반 RAG 모델 만들기

이전글은 txt 데이터를 ChromaDB에 저장하고 retrieval하는 내용을 다뤘다. 이번에는 똑같은 프로세스로 진행하지만, PDF파일을 불러와서 진행하는 방법이다. PDF파일 로더를 사용해서 가져오는 방법을 제외하면 크게 다르지 않다.

PDF file Loader

pip install pymupdf

from langchain_community.document_loaders import PyMuPDFLoader
loader = PyMuPDFLoader("data.pdf")
pages = loader.load()
len(pages)

langchain 공식문서에 보면 pdf를 불러오는 여러가지 방법이 있다. 처음에는 PyPDFLoader를 사용했는데, PyMulPDFLoader가 속도도 더 빠르고 인코딩 문제가 없는 것 같아서 해당 방법을 선택했다.
불러온 pages의 type은 list이다.

Create Collection

import chromadb
chromadb_client = chromadb.HttpClient(host="localhost", port=8000)
print(chromadb_client.list_collections())

collection = chromadb_client.create_collection(name="pdf_test_240531")
print(chromadb_client.list_collections())

collection을 만드는 방법은 이전 포스트에서 설명했으니 생략하겠다. 만들어준 collection 이름은 'pdf_test_240531'이다.

Chunk Split, Embedding, Vector Store

import chromadb
import uuid
from langchain.text_splitter import RecursiveCharacterTextSplitter
from chromadb.utils import embedding_functions

def insert(content):
    client = chromadb.HttpClient(host="localhost", port=8000)
    openai_ef = embedding_functions.OpenAIEmbeddingFunction(
                api_key='',
                model_name="text-embedding-ada-002"
            )
    collection = client.get_collection(name="pdf_test_240531", embedding_function=openai_ef)

    text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=50)
    docs = text_splitter.split_documents(documents)

    for doc in docs:
        uuid_val = uuid.uuid1()
        print("Inserted documents for ", uuid_val)
        collection.add(ids=[str(uuid_val)], documents=doc.page_content)

이 과정 역시 txt 파일 때와 동일하다. 다만 이번에는 임베딩 모델을 Ollama가 아닌 OpenAI API의 'text-embedding-ada-002'를 사용하였다. chunk_size와 chunk_overlap도 임의로 설정해줬다. 실행이 끝난뒤 collection.peek()를 실행해보면 정상적으로 잘 들어간 것을 확인할 수 있다.

Query, Retrieval

import chromadb
from langchain.vectorstores import Chroma
from langchain.embeddings.openai import OpenAIEmbeddings

def queryDB(query):
    client = chromadb.HttpClient(host="localhost", port=8000)
    embedding_function = OpenAIEmbeddings(api_key='')
    db4 = Chroma(client=client, collection_name="pdf_test_240531", embedding_function=embedding_function)
    docs = db4.similarity_search_with_score(query) ## 유사도 점수까지 확인
   
    return docs 

query=""
result = queryDB(query)
result

max_similarity_doc = min(result, key=lambda x: x[1])[0].page_content

가장 유사한 문장을 가져온 뒤, ChatGPT-4o를 활용하여 답변을 생성한다.

from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI

# 프롬프트 생성
prompt = f"""Answer the question based only on the following context, which can include texts:
{max_similarity_doc}
If you can't find the answer in the context, answer "I can't answer the question.".
Please answer in Korean.
"""

# LLM
model = ChatOpenAI(temperature=0, model="gpt-4o")

# RAG pipeline
chain = (
    RunnablePassthrough(lambda x: prompt)  # 함수를 직접 전달
    | model
    | StrOutputParser()
)

# 파이프라인 실행
output = chain.invoke(query)
print(output)

몇 개 질문을 테스트해봤는데, 성능이 꽤 괜찮다.