New to Nutbox?

옵시디언 노트 기반 RAG 시스템 구축하기 4 (마지막)

7 comments

anpigon

13 days ago3 min read

이전글

이 글은 로컬 기반 AI 모델을 활용해 비용 없이 옵시디언에 RAG(Retrieval-Augmented Generation) 시스템을 구축하는 방법을 소개합니다.

사용 모델 및 라이브러리

벡터DB: Chroma
임베딩 모델: HuggingFace의 BAAI/bge-m3
Rerank 모델: Cohere의 rerank-multilingual-v3.0
AI 모델: Ollama EEVE-Korean-10.8B

구현 과정

임베딩 모델 불러오기

from langchain_community.embeddings import HuggingFaceEmbeddings

model_name = "BAAI/bge-m3"
model_kwargs = {'device': 'mps'} # 'device': 'cuda' for GPU, 'device': 'cpu' for CPU, 'device': 'mps' for Apple Silicon Macs

embedding = HuggingFaceEmbeddings(
    model_name=model_name, 
    model_kwargs=model_kwargs, 
    encode_kwargs=model_kwargs,
    show_progress=True
)

벡터DB 불러오기

from langchain_community.vectorstores import Chroma

# 벡터DB를 불러옵니다.
vector_store = Chroma(persist_directory="db", embedding_function=embedding)

# 벡터DB를 검색기로 변환합니다.
retriever = vector_store.as_retriever(search_kwargs={"k": 20})

Rerank 모델을 불러오기

from langchain_cohere import CohereRerank

# Cohere의 Rerank 모델을 불러옵니다.
compressor = CohereRerank(model="rerank-multilingual-v3.0")

LLM 모델 불러오기

from langchain_community.llms import Ollama

# Ollma에서 AI 모델을 불러옵니다.
llm = Ollama(model="EEVE-Korean-10.8B")

필요한 함수 작성

from langchain.chains import RetrievalQA
from langchain.retrievers.contextual_compression import ContextualCompressionRetriever


def query_docs(query: str) -> dict:
    compression_retriever = ContextualCompressionRetriever(
        base_compressor=compressor, 
        base_retriever=retriever
    )

    chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=compression_retriever,
        return_source_documents=True,
        verbose=True,
    )

    return chain({"query": query})


def chatbot(query):
    response = query_docs(query)
    if response and "result" in response:
        answer = response["result"].strip()
        source_docs_string = ""
        if response and "result" in response and "source_documents" in response:
            answer = response["result"]
            sources = response["source_documents"]
            source_info = "\n출처:\n" + "\n".join(
                [f"- {src.metadata['path']}" for src in sources]
            )
            return f"{answer.strip()}\n{source_info}"
    else:
        return "No relevant documents found."


while True:
    user_query = input("User: ")
    response = chatbot(user_query)
    print("Chatbot:", response)