[Vector DB] 3. Milvus 튜토리얼 (1) - 설치, 변수 정의, Collection 생성하기

이전 글들에서는 Vector Database의 전체적인 그림을 그리는 내용이었습니다. 이번 글부터는 여러 Vector Database 들 중 에서 "Milvus"에서 기초적인 내용들(Collection 생성, 데이터 및 임베딩 값 적재, 검색 등)에 대해 튜토리얼 형식으로 글 포스팅을 하고자 합니다.

 

2023.06.10 - [AI/Vector Database] - [Vector DB] 1. Vector Database 배경 & 필요성

2023.07.24 - [AI/Vector Database] - [Vector DB] 2. Vector Database 종류 & 한계점

 

Milvus 설치방법

---

설치법은 영상이나 아래 Documents를 참고하시는 것을 권장합니다. 가장 많이 설치하는 방식으로는 Milvus Standardalone으로 Docker Compose를 구성하고 Pymilvus를 설치하는 형태가 가장 일반적입니다. 설치하는 방법은 따로 Docker 선수 지식만 있으면 충분히 설치할 수 있으니 이번 글에서 다루지 않겠습니다.

• Install Milvus - Standalone
  • https://milvus.io/docs/install_standalone-docker.md
• Install SDKs - Pymilvus
  • https://milvus.io/docs/install-pymilvus.md

 

Load library & Parameters 정의하기

---

pymilvus와 임베딩모델을 활용할 transformer와 sentence_transformers가 설치되어있어야 합니다. 실제로 내가 검색을 하기 위한 메타 데이터 또한 준비되어있어야 합니다. 관련 라이브러리들을 부른다고 가정하겠습니다.

import torch
import pandas as pd
import numpy as np

from sentence_transformers import SentenceTransformer
from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection, utility

 

Python에서 Sqlite3 을 사용할 때처럼 유사하게, Milvus 또한 마찬가지로 여러 값들을 초기에 세팅을 합니다. 일반적으로 자주 사용하는 부분들과 주석을 달면 다음과 같습니다.

MODEL = 'snunlp/KR-SBERT-V40K-klueNLI-augSTS'  # 내가 사용하고자 하는 임베딩 모델
COLLECTION_NAME = 'BalancedNewsCorpus_db'  # 저장할 Vector Database의 Collection name
DIMENSION = 768  # 내가 사용하는 임베딩 모델의 dimension
LIMIT = 10  # topk와 같이 내가 검색했을 때 최대 몇개를 볼 것 인지
MILVUS_HOST = '000.000.00.00' # 설치 Host
MILVUS_PORT = '8888' # 설치 Port
INDEX_TYPE = "IVF_FLAT" # 다양한 Index 알고리즘들을 활용 가능 "HNSW" , "ANNOY" , "IVF_FLAT" , "IVF_SQ8" , "IVF_PQ"

 

여기서 가장 중요한 부분은 임베딩 품질에 영향을 주는 Model입니다. 영어나 Multilingual 인 경우에 대해서는 HuggingFace에 존재하는 임베딩 모델을 활용을 하는 것이 일반적입니다. MTEB 리더보드에서 본인 여건에 맞게 임베딩 모델을 선정해서 사용합니다. MTEB 리더보드를 처음 보면 어떤 모델을 선택해야 할지 막막할 수 있습니다.

 

만약에 품질이 중요하다면, dimension이 높은 모델을 고려할 수 있고 속도가 중요하다면 dimension이 낮은 모델을 선택할 수 있습니다. 또한 사용하는 사람의 도메인(벤치마크 데이터셋 참고)이나 Task(retrieval, Reranking 등)에 따라서 다른 모델들을 선택할 수 있습니다. 하지만 본인이 한국어 임베딩 모델을 사용한다면 단순히 PLM 모델을 Embedding 값을 추출 용도로 사용하는 것은 옳지 않고 성능이 상당히 떨어집니다. 문장 단위의 임베딩 벡터 추출은 KR-SBERT처럼 STS, NLI 계열 특화된 모델을 사용하셔야 합니다. 여기서 더 본인이 다루는 언어가 Domain Specific 하다면, 따로 데이터를 수집해서 임베딩 모델을 굽는 것을 추천합니다. 

 

Collection 생성하기

---

실제로 Milvus 내부에서 collection을 생성하고 검색을 빠르게 하기 위한 색인(Index)를 생성해야 합니다. 아래 코드는 DB에 Insert 되기 직전 테이블이 생성되는 부분입니다.

 

# host와 port에 연결
connections.connect(host=MILVUS_HOST, port=MILVUS_PORT)

# 내가 collection 이름이 겹치는 경우에는 삭제하겠다는 뜻
if utility.has_collection(COLLECTION_NAME):
    utility.drop_collection(COLLECTION_NAME)

# 직접 밀어넣을 필드 & 스키마 정의하기
fields = [
    FieldSchema(name='file_id', dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name='filename', dtype=DataType.VARCHAR, max_length=20),
    FieldSchema(name='date', dtype=DataType.INT64, max_length=10),
    FieldSchema(name='NewsPaper', dtype=DataType.VARCHAR, max_length=1000),
    FieldSchema(name='Topic', dtype=DataType.VARCHAR, max_length=1000),
    FieldSchema(name='News', dtype=DataType.VARCHAR, max_length=1000),
    FieldSchema(name='News_embedding', dtype=DataType.FLOAT_VECTOR, dim=DIMENSION)
]

schema = CollectionSchema(fields=fields)
collection = Collection(name=COLLECTION_NAME, schema=schema)

# 매우 중요한 인덱스 부분
index_params = {
    'metric_type':'L2',
    'index_type': INDEX_TYPE,
    'params':{"nlist":1536}
}

# colletion에서 색인 부분을 생성하는 부분을 명시하고 부르겠다는 뜻
collection.create_index(field_name="News_embedding", index_params=index_params)
collection.load()

 

• 기본적으로 id는 auto_id = True를 설정하면 자동으로 생성되며 is_primary = True로 설정하면 Primary key로 설정됩니다.
• 나머지는 가지고 있는 데이터 필드들에 대해 type들을 설정합니다.
  • 일반적인 int, float, bool 등과 같은 여러 datatype에 대해 설정이 가능합니다.
  • varchar 같은 경우에는 최대 길이는 65,535 입니다. 가지고 있는 데이터의 텍스트 길이를 insert 하기 이전에 살펴볼 필요가 있습니다. 
  • 큰 차이점은 역시 임베딩 스키마 입니다. 추후 색인을 하기 위해 참고를 하기 때문에 제일 중요한 부분입니다.
  • schema에 더 depth 있게 파고 싶다면 여기 공식 문서를 참고해주세요. 
• 인덱스 부분은 내가 검색을 하기 위한 알고리즘을 선택하는 부분입니다. 거리기반(L2 등), 유사도 기반(IP, Inner Product 등)을 선정하고 여러 알고리즘을 적용하고 파라미터들을 설정합니다. 인덱스 부분은 다룰 내용이 많아서 추후에 글 하나 포스팅할 예정입니다. 

Collection을 실제로 조회하면 아래와 같이 구성이 되어있음을 확인 가능합니다.

<Collection>:
-------------
<name>: BalancedNewsCorpus_db
<description>: 
<schema>: 
    {'auto_id': True, 'description': '', 
    
    'fields': [
        {'name': 'file_id', 'description': '', 
        'type': <DataType.INT64: 5>, 
        'is_primary': True, 
        'auto_id': True}, 

        {'name': 'filename', 
        'description': '', 
        'type': <DataType.VARCHAR: 21>, 
        'params': {'max_length': 20}}, 

        {'name': 'date', 
        'description': '', 
        'type': <DataType.INT64: 5>}, 

        {'name': 'NewsPaper', 
        'description': '', 
        'type': <DataType.VARCHAR: 21>, 
        'params': {'max_length': 100}}, 

        {'name': 'Topic', 
        'description': '', 
        'type': <DataType.VARCHAR: 21>, 
        'params': {'max_length': 100}}, 

        {'name': 'News', 
        'description': '', 
        'type': <DataType.VARCHAR: 21>, 
        'params': {'max_length': 30000}}, 

        {'name': 'News_embedding', 
        'description': '', 
        'type': <DataType.FLOAT_VECTOR: 101>, 
        'params': {'dim': 768}}
        ]
    }

 

다음 글에서는 실제 샘플 데이터를 바탕으로 데이터를 삽입하고 인덱스를 생성하고, 검색하는 부분에 대해서 다루어보겠습니다. 

 

---

아래는 블로그 주인장의 토스 익명 후원 링크입니다. 글이 도움되거나 흡족스러웠다면 후원해주시면 감사하겠습니다.

https://toss.me/hotorch

hotorch님에게 보내주세요