RAG 실습을 위해 가장 먼저 해야할 일은 Vector Database를 선택하는 일 입니다. Vector DB의 종류로는 Milvus, Weaviate, ChromaDB 등이 있으며 서로 다른 장단점을 가지고 있습니다. 그리고 이번 실습에서는 Milvus를 사용할 예정입니다.
Milvus의 장점으로는 가장 기본에 충실하다고 할 수 있습니다. Vector DB의 존재 이유는 데이터를 안전하게 저장하고 빠르게 vector similarity search를 하기 위함입니다. 특히 빠른 검색 부분에서 Milvus는 다른 Open Source에 비해 우월한 성능을 보입니다.

그림 1. Various Vector DBs
© Zilliz

Open Source Vector DB 중에서는 Milvus와 Weaviate가 자주 비교되는데, 운영적인 측면을 고려했을 땐 Weaviate가 좀 더 유리한 부분이 있습니다. Milvus의 경우는 합의(Consensus) 알고리즘을 etcd같은 외부 서비스에 의존하는 등의 이유로 운영 복잡성이 좀 더 높을 수 있습니다.

Milvus

이번 실습에서는 standalone(단일 시스템)환경에서 docker를 이용하여 Milvus를 실행시킬 예정입니다. 하지만 Milvus는 기본적으로 분산 시스템에서 동작하는 것을 가정하고 제작되었기에, standalone일지라도 docker-compose를 이용해서 etcd와 minio를 함께 실행하는 것을 권장하고 있습니다.

Install Milvus

wget https://github.com/milvus-io/milvus/releases/download/v2.5.12/milvus-standalone-docker-compose.yml -O docker-compose.yml
sudo docker compose up -d

도커가 설치되어있는 환경에서 위의 명령어만 입력하면 간단하게 milvus를 구축할 수 있습니다.

Stop and delete Milvus

sudo docker compose down
sudo rm -rf volumes

마찬가지로 실행시킨 도커를 정지시키고 볼륨을 지우는 것 만으로 milvus를 완전히 제거할 수 있습니다.

Create Collection

A collection is a two-dimensional table with fixed columns and variant rows. Each column represents a field, and each row represents an entity.

Milvus에서 데이터를 삽입하기 위해서는 기존 관계형 DB처럼 Schema를 정의하고 Collection을 생성해야 합니다. Documents 상에서는 fixed columns라는 표현을 사용하고 있지만, Milvus의 schema는 기존 관계형 DB와 다르게 dynamic schema를 지원하는 등 비교적 덜 엄격하게 관리되고 있습니다. 

from pymilvus import MilvusClient, DataType

client = MilvusClient(
    uri="http://localhost:19530",
    token="root:Milvus" # Default user and password
)

client.create_collection(
    collection_name="collection_1", # 컬렉션 이름 (string)
    schema=schema,                  # MilvusClient.create_schema를 이용해 정의된 스키마
    index_params=index_params       # 빠른 검색을 위한 index 정의
)

가장 먼저 pymilvus의 MilvusClient를 이용해 DB와의 connection을 생성해줍니다. 그리고 schema와 index와 함께 collection을 정의해 주면 컬렉션을 생성할 수 있습니다.

Create Schema
schema = MilvusClient.create_schema(
    auto_id=False, # timestamp를 기반으로 자동으로 id 값을 생성할지 여부
    enable_dynamic_field=True,
)

schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True)
schema.add_field(field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=1024)
schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=1024)

스키마의 경우 보통의 경우 primary key로 사용할 id 필드와 vector 값을 저장할 vector 필드는 필수적으로 필요합니다. 임베딩 전의 원본 텍스트를 보존하기 위해 text 필드를 하나 더 만들어주겠습니다.

Create Index
index_params = client.prepare_index_params()

index_params.add_index(
    field_name="id",
    index_type="AUTOINDEX" # 대부분의 scalar 타입은 자동 지정시 Inverted Index 사용
)

index_params.add_index(
    field_name="vector",
    index_type="HNSW",
    index_name="vector_index",
    metric_type="COSINE",
    params={
        "M": 64,
        "efConstruction": 100,
    }
)

id 필드의 경우 AUTOINDEX를 이용해 인덱스 타입을 자동으로 지정해주었고, vector 필드의 경우 HNSW을 이용하겠습니다.

Schema, Index를 정의하고 Collection을 생성하는 방법까지 진행해 보았습니다. 이제 남은 건 실제 데이터를 DB에 삽입(insert)하고 검색(search)하는 일 입니다. 이 때 데이터를 Vector DB에 입력하기 위해서는 임베딩 모델을 이용해서 텍스트 데이터를 벡터 데이터로 임베딩하는 과정을 거쳐야합니다.

Insert & Search

Embedding
from openai import OpenAI
client = OpenAI()

response = client.embeddings.create(
    input="Your text string goes here",
    model="text-embedding-3-small"
)

ebedded = response.data[0].embedding)

Open AI의 모델을 이용하거나, Huggingface 모델을 이용할 경우 text-embeddings-inference를 이용하는 것을 추천드립니다. 그 외에도 다양한 방법을 이용해서 텍스트를 벡터로 변환하기만 하시면 됩니다.

Insert
data=[
    {"id": 0, "vector": embedded, "text": "Your text string goes here"},
]

res = client.insert(
    collection_name="collection_1",
    data=data
)

print(res)

해당 콜렉션에 알맞은 json list 형태로 data를 넣어주면 됩니다.

search_params = {
    "params": {
        "ef": 10, # Number of neighbors to consider during the search
    }
}

res = client.search(
    collection_name="collection_1",
    anns_field="vector",    # Vector field name
    data=[embedded],        # Query vector
    limit=3,                # TopK results to return
    search_params=search_params
)

for hits in res:
    for hit in hits:
        print(hit)

마지막으로 index_type에 따른 search_params을 정의해 주고, Query vector 또한 기존에 이용했던 동일한 방법으로 임베딩해준 후 검색을 진행하면 벡터 유사도 검색 결과를 받아 볼 수 있습니다.

Reference

[1] Milvus | Milvus Docs