[AI] Vector database chroma 사용해보기

벡터 데이터베이스는 비정형 데이터를 벡터 (크기와 방향으로 정의되는 수치) 형태로 저장하고 쿼리하는데 특화된 DB입니다.

대표적으로 언급되는 DB가 Pinecone, Zillis, Chroma, Weaviate 등이 있습니다.

 

일반적으로 DB라 하면 Oracle, Postgresql 과 같은 관계형 데이터베이스나 mongo DB와 같은 NoSQL 정도를 떠올릴 수 있는데요.

AI 시장에서는 데이터를 효율적으로 저장하고 소비하기 위해 벡터 데이터베이스가 자주 활용됩니다.

LLM 자체가 입력받은 정보를 매번 저장하는게 아니다보니 이 저장의 영역을 보완하는게 벡터 데이터베이스라고 볼 수 있을 것 같습니다.

 

1. 벡터 임베딩이란?

방금 말한 벡터 임베딩은 데이터에 대응하는 벡터를 만들어서 결국 기계가 이해할 수 있도록 표현하는 것입니다. 이 임베딩은 data scientist들이 직접 정의하는 것이 아니라 각 데이터에서 추출된 embedding 간 거리를 바탕으로 패턴과 관계를 발견하는 방식입니다. 

 

2. chroma 설치

https://docs.trychroma.com/getting-started

여기 들어가면 Python 과 JavaScript를 지원하는데요. 저는 Python을 활용해보겠습니다.

 

Chroma 공식 사이트 들어가면 colab demo가 있는데요. 저는 이 데모를 보고 따라해봤습니다.

여기서 사용하는 dataset은 huggingface에서 제공하는 SciQ 데이터셋을 활용했습니다.

이 데이터셋은 13,679 개의 과학 시험 4지선다 문제와 정답에 대한 근거를 모아둔 걸로 보입니다. 

 

해당 예시의 시나리오는 해당 데이터셋에서 과학 문제에 대한 근거를 저장해서 각 문제별 관련이 높은 근거를 조회해보는 내용이네요.

 

예시 그대로 따라해볼게요!

🔑 Getting Started | Chroma

pip install chromadb
pip install numpy datasets

 

3. HuggingFace sciq 데이터셋 로드하기

# Get the SciQ dataset from HuggingFace
from datasets import load_dataset

dataset = load_dataset("sciq", split="train")

# Filter the dataset to only include questions with a support
dataset = dataset.filter(lambda x: x["support"] != "")

print("Number of questions with support: ", len(dataset))

결과를 조회해보면 아래와 같이 조회됩니다.

HuggingFace SciQ dataset load

실제 해당 데이터를 보시면 질문, 정답, 근거 등이 있습니다. 

 

 

4. chroma DB에 데이터 저장

import chromadb
client = chromadb.Client()

collection = client.create_collection(name="my_collection")

 

여기서 Client()는 기본 설정인 Memory에 저장하는 EphemeralClient라고 보시면 됩니다.

chroma db에 100개의 근거를 저장해봅니다.

 

# Embed and store the first 100 supports for this demo
collection.add(
    ids=[str(i) for i in range(0, 100)],  # IDs are just strings
    documents=dataset["support"][:100],
    metadatas=[{"type": "support"} for _ in range(0, 100)
    ],
)

 

 

5. 쿼리 날리기

chroma db에 앞서 저장했던 근거를 바탕으로 10가지의 질문에 대한 관련성이 높은 근거를 조회해봅니다.

results = collection.query(
    query_texts=dataset["question"][:10],
    n_results=1)

# Print the question and the corresponding support
for i, q in enumerate(dataset['question'][:10]):
    print(f"Question: {q}")
    print(f"Retrieved support: {results['documents'][i][0]}")
    print()

chroma db에 저장후 print

 

chroma DB에서 제공해주는 기본적인 사용법을 따라해보았습니다.

추후에는 chatgpt와 연계해서 사용할 수 있는 방법을 확인해봐야겠네요.

 

감사합니다.