OpenLab - OpenLLM, RAG 1차 구현

NoOne 24.05.29

1,028 1 0

facebook twitter kakao link

1 0

facebook twitter kakao link

DEVOTEE 요약: 이 코드는 2024년 KBO 야구 규정을 RAG(검색 강화 생성) 방식으로 처리하여 질문에 답변하는 개념을 다룹니다. PyMuPDFLoader로 PDF 문서를 로드하고, RecursiveCharacterTextSplitter로 문서를 분리한 후, HuggingFaceEmbeddings로 임베딩하여 벡터 저장소에 저장합니다. 마지막으로, 검색 쿼리에 따라 관련 문서를 추출하고, 올라마 모델로 질문에 답변합니다.

안녕하세요. 2024년 야구 규정을 RAG 하는 1차 개념 코드를 작성해 봤습니다.

0) ollama 에 모델을 올려두고 사용하는 코드입니다.

$ ollama list
NAME             	ID          	SIZE  	MODIFIED
ggml-q5-km:latest	c3ebb32b93a5	7.7 GB	2 weeks ago
$

1) RAG 는 아래의 순서로 동작합니다.

2) 위의 순서를 코드로 그대로 옮기면 다음과 같습니다.

import os
from langchain.document_loaders import TextLoader

import warnings
warnings.filterwarnings("ignore")

############################################################################
## TextLoader
############################################################################
#loader = TextLoader('little.prince.txt')
#pages = loader.load()

from langchain_community.document_loaders import PyMuPDFLoader
loader = PyMuPDFLoader("2024.kbo.rules.pdf")
pages = loader.load()

text = pages[0].page_content

import tiktoken
from langchain.text_splitter import RecursiveCharacterTextSplitter


# 문장 임베딩 및 벡터 저장소 생성
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma

############################################################################
## TextSplitter
############################################################################
# 문서를 문장으로 분리
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=50,
)
docs = text_splitter.split_documents(pages)
#print(docs[0])
#print(docs[1])


############################################################################
## Text Vector Embedding
############################################################################
# 문장을 임베딩으로 변환하고 벡터 저장소에 저장
embeddings = HuggingFaceEmbeddings(
    #model_name='jhgan/ko-sroberta-nli',
    model_name='BAAI/bge-m3',
    #model_kwargs={'device':'cpu'},
    model_kwargs={'device':'mps'},
    encode_kwargs={'normalize_embeddings':True},
)


############################################################################
## VectorStore(Chroma)
############################################################################
vectorstore = Chroma.from_documents(docs, embeddings)


# 벡터 저장소 경로 설정
vectorstore_path = 'vectorstore'
os.makedirs(vectorstore_path, exist_ok=True)

# 벡터 저장소 생성 및 저장
vectorstore = Chroma.from_documents(docs, embeddings, persist_directory=vectorstore_path)
vectorstore.persist()
print("Vectorstore created and persisted")




############################################################################
## retriever
############################################################################
from langchain_core.runnables import RunnablePassthrough
from langchain_community.chat_models import ChatOllama

model = ChatOllama(model="ggml-q5-km", temperature=0)

# 검색 쿼리
query = "승률계산법을 알려 줘"

retriever = vectorstore.as_retriever(search_kwargs={'k': 3})

docs = retriever.get_relevant_documents(query)
#print(len(docs))
#print(docs[0].page_content)
#print(docs[0].metadata)


from langchain_community.chat_models import ChatOllama
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

# Prompt
template = '''Answer the question in korean based only on the following context:
{context}

Question: {question}
'''

prompt = ChatPromptTemplate.from_template(template)

def format_docs(docs):
    return '\n\n'.join([d.page_content for d in docs])

# RAG Chain 연결
rag_chain = (
    {'context': retriever | format_docs, 'question': RunnablePassthrough()}
    | prompt
    | model
    | StrOutputParser()
)

# Chain 실행
query = "승률계산법을 알려 줘"
answer = rag_chain.invoke(query)

print("Answer:", answer)

3) 실행해 보면, 승률계산법을 정확하게 답변하는 것을 볼 수 있습니다.

$ python rag_concept.py
Answer: 승률을 계산하기 위해서는 다음과 같은 공식을 사용할 수 있습니다:

승률 = (승리수 / 경기수) x 100

여기서 승리수는 팀이 이긴 게임의 수를, 경기수는 총 플레이한 게임의 수를 의미합니다. 예를 들어, 팀이 60경기 중 40경기를 이겼다면, 그 팀의 승률을 계산하려면 다음과 같이 하면 됩니다:

승률 = (40 / 60) x 100 = 67%

이는 해당 팀이 시즌 동안 67%의 경기에서 승리했음을 의미합니다.