Milvus 빠른 시작 — Milvus Lite로 벡터 검색하기

Milvus 빠른 시작 — Milvus Lite로 벡터 검색하기

벡터는 신경망 모델이 출력하는 데이터 형식이에요. Milvus는 이 벡터를 저장하고 의미 기반 검색을 가능하게 해 줍니다. 가장 가벼운 시작은 로컬 파이썬 라이브러리인 Milvus Lite 로, 몇 분 안에 전체 흐름을 체험할 수 있어요.

출처: https://milvus.io/docs/quickstart.md

설치

pymilvus 하나로 클라이언트 라이브러리와 Milvus Lite를 함께 설치합니다. 공식 문서는 Python 3.8+ 환경을 권장해요.

$ pip install -U pymilvus

데이터베이스와 컬렉션 준비

파일명을 지정해 MilvusClient를 만들면 로컬 벡터 데이터베이스가 준비돼요. 이어서 컬렉션을 만들 때는 벡터 필드의 차원만 지정하면 됩니다.

from pymilvus import MilvusClient

client = MilvusClient("milvus_demo.db")

if client.has_collection(collection_name="demo_collection"):
    client.drop_collection(collection_name="demo_collection")
client.create_collection(
    collection_name="demo_collection",
    dimension=768,  # The vectors we will use in this demo has 768 dimensions
)

이 기본 설정에서 기본 키와 벡터 필드는 각각 "id", "vector" 라는 기본 이름을 쓰고, 거리 측정값(metric)은 기본값인 COSINE(코사인 유사도)을 사용합니다.

데이터 준비와 삽입

텍스트를 벡터로 바꾸려면 pymilvus[model]의 기본 임베딩 함수를 쓰면 돼요. Milvus는 데이터를 딕셔너리 목록(각 딕셔너리가 엔티티 하나) 형태로 받습니다.

from pymilvus import model

# This will download a small embedding model "paraphrase-albert-small-v2" (~50MB).
embedding_fn = model.DefaultEmbeddingFunction()

docs = [
    "Artificial intelligence was founded as an academic discipline in 1956.",
    "Alan Turing was the first person to conduct substantial research in AI.",
    "Born in Maida Vale, London, Turing was raised in southern England.",
]

vectors = embedding_fn.encode_documents(docs)
print("Dim:", embedding_fn.dim, vectors[0].shape)  # Dim: 768 (768,)

data = [
    {"id": i, "vector": vectors[i], "text": docs[i], "subject": "history"}
    for i in range(len(vectors))
]

네트워크 문제로 모델을 못 받았다면 무작위 벡터로 흉내낼 수 있어요. 이때는 벡터가 진짜가 아니라 의미 유사도가 반영되지 않는 점만 주의하면 됩니다. 삽입은 insert 한 번입니다.

res = client.insert(collection_name="demo_collection", data=data)
print(res)
# {'insert_count': 3, 'ids': [0, 1, 2], 'cost': 0}

의미 검색

질의 텍스트를 임베딩으로 바꾼 뒤 search를 호출하면 가장 가까운 벡터부터 돌려줍니다.

query_vectors = embedding_fn.encode_queries(["Who is Alan Turing?"])

res = client.search(
    collection_name="demo_collection",  # target collection
    data=query_vectors,  # query vectors
    limit=2,  # number of returned entities
    output_fields=["text", "subject"],  # specifies fields to be returned
)

print(res)

검색 결과는 질의마다 묶여서 옵니다. 각 결과에는 엔티티 기본 키와 질의 벡터까지의 distance, 그리고 output_fields로 지정한 세부 정보가 담겨요.

data: ["[{'id': 2, 'distance': 0.5859944820404053, 'entity': {'text': 'Born in Maida Vale, London, Turing was raised in southern England.', 'subject': 'history'}}, {'id': 1, 'distance': 0.5118255615234375, 'entity': {'text': 'Alan Turing was the first person to conduct substantial research in AI.', 'subject': 'history'}}]"] , extra_info: {'cost': 0}

메타데이터 필터링과 기타 검색

벡터 검색을 하면서 스칼라 필드(메타데이터) 조건을 함께 적용할 수 있어요. filter 표현식으로 특정 주제만 골라내는 식입니다.

res = client.search(
    collection_name="demo_collection",
    data=embedding_fn.encode_queries(["tell me AI related information"]),
    filter="subject == 'biology'",
    limit=2,
    output_fields=["text", "subject"],
)

query()는 필터 표현식이나 ID로 조건에 맞는 엔티티를 통째로 가져오는 연산이고, delete()는 기본 키나 필터로 엔티티를 지웁니다.

res = client.query(
    collection_name="demo_collection",
    filter="subject == 'history'",
    output_fields=["text", "subject"],
)

res = client.delete(collection_name="demo_collection", ids=[0, 2])

검색 성능을 위해 큰 데이터셋에서 스칼라 필터를 자주 쓴다면 고정 스키마를 사용하고 스칼라 필드 인덱스를 켜는 걸 권장해요. Milvus Lite는 모든 데이터를 로컬 파일에 저장하므로, 같은 파일로 MilvusClient를 만들면 프로그램이 끝난 뒤에도 데이터가 그대로 복원됩니다.

더 알아보기