[DB]. 벡터 데이터베이스(Vector database) ChromaDB 사용해보기

최근에 RAG, LangChain 등
ChatGPT의 단점을 보완해주는 알고리즘, 프레임워크들이 등장함에 따라
벡터 데이터 베이스(Vector Database)역시 주목을 받고 있다.

What is Vector Database?

벡터 DB는 벡터 데이터를 저장하고 검색하기 위한 데이터베이스입니다. 일반적인 데이터베이스는 텍스트, 숫자, 날짜 등의 형태의 데이터를 저장하고 검색하는 데에 사용되지만, 벡터 DB는 벡터 형태의 데이터를 저장하고 유사성을 기반으로 검색하는 데에 특화되어 있습니다.

벡터 DB는 대규모의 고차원 벡터 데이터를 효율적으로 저장하고 빠르게 검색할 수 있는 기능을 제공합니다.

기존에 알고 사용해보았던 Vector database는 Milvus, pinecorn이 있었다, 하지만, Milvus는 사용방법이 너무 어려웠고 공식문서가 친절하지 않았다. pinecorn은 사용하기 편리했지만 무료 사용량 제한이 있어 사용하지 않았다.

그래서 그냥 인메모리로 벡터를 사용하거나, 벡터를 문자열 형태로 저장해서 변환하는 방법을 사용했었는데..

최근에 무료인 ChromaDB를 알게되어 사용해보았다.

사용법이 매우 간단하고, http 방식, 디스크 저장방식, 인메모리 방식을 선택할 수 있어서 사용자 입장에서 매우 괜찮았다.

현재 토이 프로젝트로 진행하고 있는
"위로봇 오복이"의 검색 방식을 ChromaDB로 구현해보았다.

1. 설치

pip install chromadb

2. 클라이언트 연결

import chromadb
client = chromadb.PersistentClient()

PersistentClient는 데이터를 파일에 저장합니다. 그 외에 메모리에 저장하는 EphmeralClient, 네트워크를 통해서 접속하는 HttpClient가 있습니다. 실제 서비스에서는 HttpClient가 권장됩니다.

3. 컬렉션 생성

answers = client.create_collection(
    name="answers"
)

4. 데이터 불러오기

import pandas as pd
from tqdm import tqdm

df = pd.read_excel("./오복이_데이터.xlsx")
df.sample(5)

5. 모델 불러오기

from sentence_transformers import SentenceTransformer
model = SentenceTransformer('snunlp/KR-SBERT-V40K-klueNLI-augSTS')

6. 데이터 삽입하기

ids = []
metadatas = []
embeddings = []

for row in tqdm(df.iterrows()):
    index = row[0]
    query = row[1].user
    answer = row[1].answer
    
    metadata = {
        "query": query,
        "answer": answer
    }
    
    embedding = model.encode(query, normalize_embeddings=True)
    
    ids.append(str(index))
    metadatas.append(metadata)
    embeddings.append(embedding)
    
chunk_size = 1024  # 한 번에 처리할 chunk 크기 설정
total_chunks = len(embeddings) // chunk_size + 1  # 전체 데이터를 chunk 단위로 나눈 횟수
embeddings = [ e.tolist() for e in tqdm(embeddings)]  

for chunk_idx in tqdm(range(total_chunks)):
    start_idx = chunk_idx * chunk_size
    end_idx = (chunk_idx + 1) * chunk_size
    
    # chunk 단위로 데이터 자르기
    chunk_embeddings = embeddings[start_idx:end_idx]
    chunk_ids = ids[start_idx:end_idx]
    chunk_metadatas = metadatas[start_idx:end_idx]
    
    # chunk를 answers에 추가
    answers.add(embeddings=chunk_embeddings, ids=chunk_ids, metadatas=chunk_metadatas)
    

7. 검색

result = answers.query(
    query_embeddings=model.encode("어제 여자친구랑 헤어졌다", normalize_embeddings=True).tolist(),
    n_results=3
)


print(result)

8. 결과

{
   "ids":[
      [
         "11271",
         "11317",
         "11210"
      ]
   ],
   "distances":[
      [
         0.4837903380393982,
         0.4841890037059784,
         0.49379977583885193
      ]
   ],
   "metadatas":[
      [
         {
            "answer":"이유를 모른다면 그 이유를 찾기 위해 노력해야 해요. 곰곰히 생각해보세요 여자친구는 충분히 표현해왔을 거예요.",
            "query":"여자친구가 왜 나랑 헤어졌는지 아직도 모르겠어"
         },
         {
            "answer":"아직도 많이 힘드시군요. 시간이 약이라는 말이 있듯이 조금만 더 힘내시길 바랍니다.",
            "query":"여자친구와 헤어진지 8일"
         },
         {
            "answer":"연애는 서로 존중하는 마음이 있어야 오래갈 수 있습니다. 이번 기회를 통해 자신의 행동을 되돌아보고, 다음 연애에서는 같은 실수를 반복하지 않도록 노력해보세요.",
            "query":"여자친구 무시했다가 이별 당했습니다."
         }
      ]
   ],
   "embeddings":"None",
   "documents":[
      [
         "None",
         "None",
         "None"
      ]
   ]
}