ChromaDB로 RAG(Retrieval-Augmented Generation)방식의 챗봇 구현하기

 

이전 WIZnet AI Chatbot의 동작 원리글에서 설명했듯이, 기업에서 AI 챗봇을 개발할 때, 주로 두 가지 방법을 고민합니다. 

Fine-tuning과 Retrieval Augmented Generation(RAG). 이 두 방식은 AI 챗봇의 성능을 극대화하고, 사용자 경험을 향상하는 데 중요한 역할을 합니다. 

 

RAG (Retrieval Augmented Generation) : 검색 증강 생성이란?

RAG는 대규모 언어 모델에게 추가적인 외부 데이터를 제공하는 방식입니다. 이 방식은 먼저 필요한 데이터를 임베딩 형태로 변환하여 벡터 데이터베이스에 저장합니다. 사용자의 질문이 들어오면, 이 질문과 관련된 데이터를 데이터베이스에서 검색하여 언어 모델에 제공합니다. 이렇게 검색된 데이터는 언어 모델이 사용자의 질문에 대한 더 정확하고 사실적인 답변을 생성하는 데 사용됩니다.

RAG 구조 (Retriever와 Generator)

 

왜 RAG를 사용하는가?

RAG의 핵심 장점은 최신 정보에 대한 접근성과 답변의 정확성 향상입니다. 특히, 지속적으로 변화하는 데이터나 최신 트렌드에 대해 챗봇이 신속하게 반응할 수 있도록 도와줍니다. 이를 통해 기업은 고객의 다양한 질문에 보다 신속하고 정확하게 대응할 수 있는 AI 챗봇을 구축할 수 있습니다.

 

---

 

이번 프로젝트에서는 RAG 방식을 실제 코드로 구현했습니다. 구현 과정에서는 외부 데이터의 임베딩, 이를 데이터베이스에 저장하는 과정, 그리고 사용자 질문에 대응하여 관련 데이터를 검색하고 이를 언어 모델에 입력하는 과정을 포함했습니다. 이러한 구현을 통해, AI 챗봇은 사용자 질문에 대한 보다 정확하고 신뢰할 수 있는 답변을 생성할 수 있게 되었습니다.

 

Diagram

 

우선, ChromaDB 클라이언트를 초기화하고 OpenAI의 임베딩 기능을 활용해 데이터를 처리합니다. 이 과정은 아래의 코드로 구현되었습니다.

import chromadb
from chromadb.utils import embedding_functions
import openai

# ChromaDB 클라이언트 및 컬렉션 초기화
client = chromadb.Client()
client.heartbeat()

# OpenAI 임베딩 함수 초기화
openai_ef = embedding_functions.OpenAIEmbeddingFunction(api_key="OPENAI_API_KEY", model_name="text-embedding-ada-002")

# 새 컬렉션 생성 및 데이터 저장
collection = client.create_collection(name="my_collection", embedding_function=openai_ef)

 

다음으로, WIZnet 관련 example data를 Sample로 넣어주었습니다.

example_texts = [
    "WIZnet Document Wiki: 이 사이트는 WIZnet 제품 및 오픈 소스 하드웨어에 대한 정보를 호스팅하는 오픈 정보 수집 플랫폼입니다. 계정을 등록하여 내용을 개선해 주시면 감사하겠습니다.",
    "Products Home: WIZnet 제품에 대한 정보 모음으로, 데이터시트, 소스 파일, 응용 노트 및 관련 응용 정보를 포함합니다.",
    "WIZwiki-W7500, WIZwiki-W7500P, WIZwiki-W7500ECO, WIZwiki-W7500 mbed Starter Kit, W7500, W7500P, W6100: 하드와이어드 TCP/IP 칩(IPv4/6 듀얼 스택), W5100S: 하드와이어드 TCP/IP 칩, W5500: 하드와이어드 TCP/IP 칩, W5500-EVB: W5500 평가 보드, WIZ550io: 플러그인 네트워크 모듈, WIZ550S2E: 시리얼 투 이더넷, WIZ550web: 임베디드 웹 서버, WizFi250: Wi-Fi 모듈, WizFi360: Wi-Fi 모듈(2019년 1월 8일 출시), WIZ510SSL: S2E 모듈(2021년 5월 25일 출시)",
    "Open Source Hardware: WIZnet 오픈 소스 하드웨어 제품에 대한 자료 및 정보를 제공합니다.",
    "VAR Products: 가치 추가 리셀러 제품으로, Ethernet Shield 및 WiFi Shield를 포함합니다.",
    "Tech Support: 기술 지원에 대한 질문은 https://forum.wiznet.io 에서 해주세요."
]


for i, text in enumerate(example_texts): 
    content_metadata = {  
        "id": str(i+1),  # 각 텍스트에 고유 번호를 부여 (리스트의 인덱스에 1을 더함)
        "content": text  # 현재 텍스트의 내용
    }
    collection.add(  
        documents=[text],  # 현재 텍스트를 문서로 추가
        metadatas=[content_metadata],  # 현재 텍스트의 메타데이터 추가
        ids=[str(i+1)]  # 고유 번호를 식별자로 사용하여 추가
    )

 

example_texts 리스트로부터 content_metadata 딕셔너리에 id, content를 넣어주고 collection.add 메소드를 호출하여 ChromaDB 컬렉션에 데이터를 추가하는 과정입니다. 이렇게 저장된 데이터는 추후 RAG 방식의 AI 챗봇에서 검색 및 응답 생성 과정에 사용됩니다.

 

 

다음으로 RAG의 핵심인 사용자의 질문에 대해 관련 데이터를 검색, 이를 바탕으로 답변을 생성하는 코드입니다. 이를 위해 'Retriver'와 'Generator'클래스를 구현하였습니다.

# Retriver : 사용자의 질문과 관련된 데이터를 검색하는 역할 
class Retriver:
    def __init__(self):
        self.collection = client.get_collection(name="my_collection", embedding_function=openai_ef)

    def get_retrieval_results(self, input, k=1):
        retrieval_results = self.collection.query(
            query_texts=[input],
            n_results=k,
        )
        return retrieval_results["documents"]

Retriver 클래스

# 'Generator' : 검색된 데이터를 바탕으로 GPT-4 모델을 사용하여 답변 생성
class Generator:
    def __init__(self, openai_model="gpt-4"):
        self.openai_model = openai_model
        self.prompt_template = """
            You're an expert on Wiznet products. Provide a concise and relevant answer based on the following information:
            {text}
        """

    def generate_response(self, retrieval_results):
        prompts = []
        for result in retrieval_results:
            prompt = self.prompt_template.format(text=result)
            prompts.append(prompt)
        prompts.reverse()

        response = openai.ChatCompletion.create(
            model=self.openai_model,
            messages=[{"role": "assistant", "content": prompt} for prompt in prompts],
            temperature=0,
        )

        return response["choices"][0]["message"]["content"]

 

 

마지막으로 AI Chatbot의 작동 방식을 설정하고 실행합니다.

# Chatbot 클래스
class Chatbot:
    def __init__(self):
        self.retriver = Retriver()
        self.generator = Generator()
    
    def answer(self, input):
        retrieval_results = self.retriver.get_retrieval_results(input)
        return self.generator.generate_response(retrieval_results)

# 챗봇 인스턴스 생성 및 실행
chatbot = Chatbot()

while True:
    user_input = input("You: ")
    response = chatbot.answer(user_input)
    print(f"Chatbot: {response}")

 

결과

처음에 Retriver클래스의 'get_rertrieval_results'에서 k=5로 두니, 상위 5개의 결과가 반환되었습니다. 그러나 제공한 데이터 example_text에 6개의 문장만 구성하였기에, 적절하지 않은 답변이 추론되는 문제가 있었습니다.

 

이에따라 'k=1'로 수정하여 가장 관련성 높은 단 하나의 결과만을 반환하도록 하였고, 답변의 정확성이 크게 향상되었습니다.

 

한편, 데이터셋의 크기가 클 경우에는 k 값을 5에서 15 사이로 설정하여 더 많은 검색 결과를 고려한다고 합니다. 데이터셋이 충분히 크고 다양한 경우, 더 많은 검색 결과를 포함시키는 것이 챗봇이 더 정확하고 유용한 답변을 생성하는 데 도움이 될 수 있습니다.

---

이번 프로젝트를 통해 Retrieval Augmented Generation(RAG) 방식을 직접 구현해 보면서, 대규모 언어 모델(LLM)을 활용한 애플리케이션 개발에 대한 기초를 다질 수 있었습니다. LLM의 복잡한 메커니즘을 이해하고, 실제 애플리케이션에 적용하는 데 약간의 자신감을 얻을 수 있었습니다.

현재, LangChain과 같은 다양한 LLM 애플리케이션을 위한 프레임워크가 개발되고 있으며, 이는 LLM의 활용 범위를 더욱 넓혀가고 있습니다. 또한, LLM 모델의 fine-tuning과 prompt engineering과 같은 방법들을 통해서도 LLM의 성능을 개선하기 위한 연구가 활발히 진행되고 있습니다. 

매일 같이 굵직한 논문들이 발표되는 이 시기에, LLM의 최신 트렌드를 파악하고 이를 적용하는 것이 매우 중요하다고 생각됩니다. 이러한 흐름을 따라가면서, generative AI engineer로서 계속해서 새로운 도전을 추구할 것입니다. 

LLM의 미래는 매우 밝다고 생각되며 계속되는 기술의 발전을 통해 머지않아 보다 쉽게 고성능의 솔루션을 개발할 수 있을 것으로 기대하며 계속 노력하겠습니다.