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

실제로 데이터를 docker로 설치한 ChromaDB에 저장하고 Retrieval한뒤 model에 넣어서 정제된 답변을 추출하는 방법이다.

이번에 사용하는 데이터는 txt 파일이고, 추후 pdf도 테스트 해 볼 예정이다.

• 코드 참고 블로그

1. Data Load

from langchain_community.document_loaders import TextLoader

loader = TextLoader('data.txt')
documents = loader.load()

langchain의 TextLoader를 사용해서 txt 파일을 불러온다. documents의 type은 list이다.

2. Chunk Split

Chunk는 임베딩을 해서 벡터DB에 저장하는 단위이다. Chunk의 사이즈, 얼마나 겹치게 분리할 지 등의 파라미터를 정해줘야 하는데, 생각보다 이 부분이 나중에 결과에 큰 영향을 미치는 것 같다. 필자는 langchain에서 제공하는 두가지의 Chunk spliter를 테스트 해봤다.

CharacterTextSpliter

• 공식문서
  구분자를 지정해서 split하는 방법이다. 예를 들어, 엔터를 기준으로 나누고자 한다면, '\n'등과 같은 문자를 지정해주면 된다. 형식이 정해져있거나 구분자가 명확할 때 사용할 수 있을 것 같다.

import uuid
from tqdm import tqdm
# from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.text_splitter import CharacterTextSplitter
from chromadb.utils import embedding_functions
from chromadb.utils.embedding_functions import OllamaEmbeddingFunction

def insert(documents):
    client = chromadb.HttpClient(host="localhost", port=8000)
    ollama_ef = OllamaEmbeddingFunction(url="http://localhost:11434/api/embeddings", model_name="gemma")
    collection = client.get_collection(name="collection1", embedding_function=ollama_ef)
    text_splitter = CharacterTextSplitter(separator="==============", chunk_size=500, chunk_overlap=0, length_function = len)
    docs = text_splitter.split_text(documents[0].page_content)

    for doc in tqdm(docs, desc="Processing", total=len(docs), leave=True):
        uuid_val = uuid.uuid1()
        #print("Insert documents for ", uuid_val)
        collection.add(ids=[str(uuid_val)], documents = doc)

chunk split을 한 뒤 임베딩해서 collection에 저장하는 것 까지 한번에 함수로 만들었다. 이 함수는 아래 3. Embedding, Store에서 하나씩 살펴보고, 지금은 text_splitter 부분만 보면 된다. text_splitter는 CharacterTextSplitter를 지정했고 seperator에 원하는 문자를 넣어준다.

RecursiveCharacterTextSplitter

• 공식문서

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

def insert(documents):
    client = chromadb.HttpClient(host="localhost", port=8000)
    ollama_ef = OllamaEmbeddingFunction(url="http://localhost:11434/api/embeddings", model_name="llama3:70b")
    collection = client.get_collection(name="collection1", embedding_function=ollama_ef)
    text_splitter = RecursiveCharacterTextSplitter(chunk_size = 120, chunk_overlap = 50)
    docs = text_splitter.split_documents(documents)

    for doc in tqdm(docs, desc="Processing", total=len(docs), leave=True):
        uuid_val = uuid.uuid1()
        #print("Insert documents for ", uuid_val)
        collection.add(ids=[str(uuid_val)], documents = doc.page_content)

RecursiveCharacterTextSplitter는 여러 문자에 대해서 재귀적으로 분할을 하고 기본적으로 '\n\n', '\n', 공백, 빈문자열 등을 기준으로 분리한다. 문맥적으로 관련된 텍스트를 유지할 수 있다는 장점이 있지만 속도가 느리다는 단점이 있다. 이때는 구분자를 지정해주지 않고 chunk_size와 chunk_overlap을 설정해서 어떤 길이로 자를 것인지, 얼만큼 문장을 겹치게 만들 것인지는 지정해준다. 이 값은 input text의 특성에 따라 적절한 값이 달라지기 때문에 사실 여러번 테스트를 해보면서 가장 성능이 좋은 경우를 찾는 방법 밖에 없는 것 같다. Chunk Size Stretegy라는 글이 있던데, 나중에 참고해보면 좋을 것 같다.

3. Embedding, Store

이제 위에 코드를 하나씩 살펴보자.
client는 저장할 DB이다. 현재 docker로 localhost:8000 으로 열어줬기 때문에 그대로 설정해준다.
ollama_ef부분은 임베딩을 진행하는 모델을 설정해주는 것이다. 이때는 Ollamadml llama3:7b를 설정해줬는데, HuggingFace 모델을 사용한거나 OpenAI의 API를 사용한다면 아래처럼 코드를 수정해주면 된다.

1. Ollama

from chromadb.utils.embedding_functions import OllamaEmbeddingFunction
ollama_ef = OllamaEmbeddingFunction(url="http://localhost:11434/api/embeddings", model_name="llama3:70b")

2. OpenAI

openai_ef = embedding_functions.OpenAIEmbeddingFunction(api_key='', model_name="text-embedding-ada-002")

3. HuggingFace
(HuggingFace 코드는 재테스트 필요)

import chromadb.utils.embedding_functions as embedding_functions
    huggingface_ef = embedding_functions.HuggingFaceEmbeddingFunction(
        api_key="",
        model_name="google/gemma-7b"
    )

4. Google GenerativeAI

google_ef  = embedding_functions.GoogleGenerativeAiEmbeddingFunction(api_key="")

collection에는 만들어준 DB의 collection을 설정해준다. 이때 embedding_function에 앞에서 선언한 임베딩 모델 변수를 넣어준다. Collcection을 만드는 방법은 이전 포스트를 참고하길 바란다.
text_splitter는 앞에서 설명한 대로 원하는 split 방법에 따라 설정해주면 되고 분리한 값들을 docs에 담아준다.
이후 docs에 저장된 Chunk를 하나씩 넣어주는데, 이때 Chunk를 구분할 수 있는 고유값을 uuid(Universally Unique Identifier)를 통해서 준다.
이후에 collection.peek()을 했을 때, 다음과 같은 dictionary가 나오면 정상적으로 값이 저장된 것이다.

{'ids': ['000b963c-092d-11ef-8ef7-ea9d1e35f996',
  '0011c3f2-092f-11ef-8ef7-ea9d1e35f996',
  '0017faf0-0930-11ef-8ef7-ea9d1e35f996',
  '001d47dc-092e-11ef-8ef7-ea9d1e35f996',
  '004360c4-092f-11ef-8ef7-ea9d1e35f996',
  '004730ac-092d-11ef-8ef7-ea9d1e35f996',
  '0048a25e-0930-11ef-8ef7-ea9d1e35f996',
  '00635f4c-092e-11ef-8ef7-ea9d1e35f996',
  '0075c474-092f-11ef-8ef7-ea9d1e35f996',
  '0078ef0e-0930-11ef-8ef7-ea9d1e35f996'],
  'embeddings': [[-1.062864899635315,
   1.4597023725509644,
   -0.06015903502702713,
   -0.07649730145931244,
   4.129417896270752,
   -2.124058723449707,
   -2.592519521713257,
   0.8902860283851624,
   -1.3188565969467163,
   0.3617442548274994,
   2.072702646255493,
   -2.419872760772705,
   -3.9154436588287354,......
   }

4. Query, Retrieval

vectorDB에 저장한 값을 기반으로 입력한 Query의 답을 생성할 수 있는 부분을 찾아(retrieval)와서 다시 model이 그 값을 기반으로 답변을 생성하는 부분이다. 이때 사용하는 model 역시 Ollama를 쓸수도, OpenAI API를 쓸 수도 있으며, 원하는 것에 맞게 모델을 설정해주면 된다.

from langchain.vectorstores import Chroma
from langchain_community.embeddings import OllamaEmbeddings
def queryDB(query):
    client = chromadb.HttpClient(host="localhost", port=8000)
    embedding_function = OllamaEmbeddings(model="gemma")
    db = Chroma(client=client, collection_name = "collection1", embedding_function = embedding_function)
    docs = db.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
max_similarity_doc

result를 출력해보면 유사도가 높은 문장(Chunk)와 유사도 값을 보여준다. 거리 기반으로 유사도를 계산하기 때문에 값이 가장 작은 문장이 가장 유사한 문장이라고 볼 수 있다.

import ollama

output = ollama.generate(
    model="gemma",
    prompt=f"Using the data: {max_similarity_doc}. Respond to this prompt: {query}. If data doesn't have an appropriate value to answer the question in the prompt, answer 'I can't answer the question.' 한국어로 대답해줘 "
)

print(output['response'])

출력된 값을 모델에 넣어서 최종 답변을 출력하도록 만드는 부분이다. Ollama를 사용했으며, 모델에게 guide를 미리 준다. prompt부분에는 '출력 데이터를 기반으로 response를 만들어주는데, 만약 적절한 출력값이 없다면 답변을 할 수 없다고 응답해줘'라는 가이드를 줬다. 영어로 가이드를 주는게 더 효과적일것 같아서 영어로 줬지만, 한국어로 줘도 될 것 같다.