Milvus를 사용한 HDBSCAN 클러스터링

원본 데이터의 의미 있는 표현을 포착하는 딥러닝 모델을 사용하여 데이터를 임베딩으로 변환할 수 있습니다. 비지도 클러스터링 알고리즘을 적용하면 고유한 패턴에 따라 유사한 데이터 포인트를 함께 그룹화할 수 있습니다. HDBSCAN(계층적 밀도 기반 노이즈가 있는 애플리케이션의 공간 클러스터링)은 데이터 포인트의 밀도와 거리를 분석하여 데이터 포인트를 효율적으로 그룹화하는 널리 사용되는 클러스터링 알고리즘입니다. 특히 다양한 모양과 크기의 클러스터를 발견하는 데 유용합니다. 이 노트북에서는 고성능 벡터 데이터베이스인 Milvus와 함께 HDBSCAN을 사용하여 데이터 요소를 임베딩에 따라 별개의 그룹으로 클러스터링합니다.

HDBSCAN(계층 밀도 기반 노이즈가 있는 애플리케이션의 공간 클러스터링)은 임베딩 공간에서 데이터 포인트 간의 거리 계산에 의존하는 클러스터링 알고리즘입니다. 딥 러닝 모델에 의해 생성된 이러한 임베딩은 데이터를 고차원적인 형태로 표현합니다. 유사한 데이터 포인트를 그룹화하기 위해 HDBSCAN은 근접성과 밀도를 결정하지만, 특히 대규모 데이터 세트의 경우 이러한 거리를 효율적으로 계산하는 것이 어려울 수 있습니다.

고성능 벡터 데이터베이스인 Milvus는 임베딩을 저장하고 색인화하여 이 프로세스를 최적화함으로써 유사한 벡터를 빠르게 검색할 수 있도록 합니다. HDBSCAN과 Milvus를 함께 사용하면 임베딩 공간에서 대규모 데이터 세트를 효율적으로 클러스터링할 수 있습니다.

이 노트북에서는 BGE-M3 임베딩 모델을 사용하여 뉴스 헤드라인 데이터 세트에서 임베딩을 추출하고, Milvus를 활용하여 임베딩 간의 거리를 효율적으로 계산하여 HDBSCAN의 클러스터링을 지원한 다음, UMAP 방법을 사용하여 분석 결과를 시각화합니다. 이 노트북은 딜런 카스티요의 글을 Milvus로 각색한 것입니다.

준비

https://www.kaggle.com/datasets/dylanjcastillo/news-headlines-2024/ 에서 뉴스 데이터 세트 다운로드

$ pip install "pymilvus[model]"
$ pip install hdbscan
$ pip install plotly
$ pip install umap-learn

데이터 다운로드

https://www.kaggle.com/datasets/dylanjcastillo/news-headlines-2024/ 에서 뉴스 데이터셋을 다운로드하고 news_data_dedup.csv 을 추출하여 현재 디렉토리에 넣습니다.

Milvus로 임베딩 추출

Milvus를 사용해 컬렉션을 만들고, BGE-M3 모델을 사용해 밀도가 높은 임베딩을 추출합니다.

import pandas as pd
from dotenv import load_dotenv
from pymilvus.model.hybrid import BGEM3EmbeddingFunction
from pymilvus import FieldSchema, Collection, connections, CollectionSchema, DataType

load_dotenv()

df = pd.read_csv("news_data_dedup.csv")


docs = [
    f"{title}\n{description}" for title, description in zip(df.title, df.description)
]
ef = BGEM3EmbeddingFunction()

embeddings = ef(docs)["dense"]

connections.connect(uri="milvus.db")

fields = [
    FieldSchema(
        name="id", dtype=DataType.INT64, is_primary=True, auto_id=True
    ),  # Primary ID field
    FieldSchema(
        name="embedding", dtype=DataType.FLOAT_VECTOR, dim=1024
    ),  # Float vector field (embedding)
    FieldSchema(
        name="text", dtype=DataType.VARCHAR, max_length=65535
    ),  # Float vector field (embedding)
]

schema = CollectionSchema(fields=fields, description="Embedding collection")

collection = Collection(name="news_data", schema=schema)

for doc, embedding in zip(docs, embeddings):
    collection.insert({"text": doc, "embedding": embedding})
    print(doc)

index_params = {"index_type": "FLAT", "metric_type": "L2", "params": {}}

collection.create_index(field_name="embedding", index_params=index_params)

collection.flush()

HDBSCAN의 거리 매트릭스 구성하기

HDBSCAN은 클러스터링을 위해 지점 간 거리를 계산해야 하므로 연산 집약적일 수 있습니다. 거리가 먼 포인트는 클러스터링 할당에 미치는 영향이 적기 때문에 가장 가까운 상위 k개의 이웃을 계산하여 효율성을 높일 수 있습니다. 이 예에서는 FLAT 인덱스를 사용했지만 대규모 데이터 세트의 경우 Milvus는 검색 프로세스를 가속화하기 위해 더 고급 인덱싱 방법을 지원합니다. 먼저 이전에 만든 Milvus 컬렉션을 반복할 이터레이터를 가져와야 합니다.

import hdbscan
import numpy as np
import pandas as pd
import plotly.express as px
from umap import UMAP
from pymilvus import Collection

collection = Collection(name="news_data")
collection.load()

iterator = collection.query_iterator(
    batch_size=10, expr="id > 0", output_fields=["id", "embedding"]
)

search_params = {
    "metric_type": "L2",
    "params": {"nprobe": 10},
}  # L2 is Euclidean distance

ids = []
dist = {}

embeddings = []

Milvus 컬렉션의 모든 임베딩을 반복할 것입니다. 각 임베딩에 대해 동일한 컬렉션에서 상위 k 개의 이웃을 검색하고 그 ID와 거리를 가져옵니다. 그런 다음 원본 ID를 거리 행렬의 연속 인덱스에 매핑하는 사전도 구축해야 합니다. 완료되면 모든 요소를 무한대로 초기화한 거리 행렬을 생성하고 검색한 요소를 채워야 합니다. 이렇게 하면 멀리 떨어진 지점 사이의 거리가 무시됩니다. 마지막으로 HDBSCAN 라이브러리를 사용하여 생성한 거리 행렬을 사용하여 점을 클러스터링합니다. 데이터가 원래의 임베딩이 아닌 거리 행렬임을 나타내기 위해 메트릭을 '미리 계산됨'으로 설정해야 합니다.

while True:
    batch = iterator.next()
    batch_ids = [data["id"] for data in batch]
    ids.extend(batch_ids)

    query_vectors = [data["embedding"] for data in batch]
    embeddings.extend(query_vectors)

    results = collection.search(
        data=query_vectors,
        limit=50,
        anns_field="embedding",
        param=search_params,
        output_fields=["id"],
    )
    for i, batch_id in enumerate(batch_ids):
        dist[batch_id] = []
        for result in results[i]:
            dist[batch_id].append((result.id, result.distance))

    if len(batch) == 0:
        break

ids2index = {}

for id in dist:
    ids2index[id] = len(ids2index)

dist_metric = np.full((len(ids), len(ids)), np.inf, dtype=np.float64)

for id in dist:
    for result in dist[id]:
        dist_metric[ids2index[id]][ids2index[result[0]]] = result[1]

h = hdbscan.HDBSCAN(min_samples=3, min_cluster_size=3, metric="precomputed")
hdb = h.fit(dist_metric)

이렇게 하면 HDBSCAN 클러스터링이 완료됩니다. 이제 일부 데이터를 가져와서 클러스터를 표시할 수 있습니다. 일부 데이터는 희박한 영역에 위치하기 때문에 어떤 클러스터에도 할당되지 않으며, 이는 노이즈가 있음을 의미합니다.

UMAP을 사용한 클러스터 시각화

이미 HDBSCAN을 사용하여 데이터를 클러스터링했으며 각 데이터 포인트에 대한 레이블을 얻을 수 있습니다. 그러나 몇 가지 시각화 기법을 사용하면 직관적인 분석을 위해 클러스터의 전체 그림을 얻을 수 있습니다. 이제 클러스터를 시각화하기 위해 UMAP을 사용하겠습니다. UMAP은 고차원 데이터의 구조를 보존하면서 시각화 또는 추가 분석을 위해 저차원 공간에 투영하는 차원 축소에 사용되는 효율적인 방법입니다. 이를 통해 원본 고차원 데이터를 2D 또는 3D 공간에서 시각화하고 클러스터를 명확하게 볼 수 있습니다. 여기서도 데이터 포인트를 반복하여 원본 데이터의 ID와 텍스트를 얻은 다음, 플로티를 사용하여 이러한 메타정보를 가진 데이터 포인트를 그림으로 플로팅하고 다른 색상을 사용하여 서로 다른 클러스터를 나타냅니다.

import plotly.io as pio

pio.renderers.default = "notebook"

umap = UMAP(n_components=2, random_state=42, n_neighbors=80, min_dist=0.1)

df_umap = (
    pd.DataFrame(umap.fit_transform(np.array(embeddings)), columns=["x", "y"])
    .assign(cluster=lambda df: hdb.labels_.astype(str))
    .query('cluster != "-1"')
    .sort_values(by="cluster")
)
iterator = collection.query_iterator(
    batch_size=10, expr="id > 0", output_fields=["id", "text"]
)

ids = []
texts = []

while True:
    batch = iterator.next()
    if len(batch) == 0:
        break
    batch_ids = [data["id"] for data in batch]
    batch_texts = [data["text"] for data in batch]
    ids.extend(batch_ids)
    texts.extend(batch_texts)

show_texts = [texts[i] for i in df_umap.index]

df_umap["hover_text"] = show_texts
fig = px.scatter(
    df_umap, x="x", y="y", color="cluster", hover_data={"hover_text": True}
)
fig.show()

이미지

여기에서는 데이터가 잘 클러스터링되어 있음을 보여드리며, 포인트 위로 마우스를 가져가서 포인트가 나타내는 텍스트를 확인할 수 있습니다. 이 노트북을 통해 Milvus로 임베딩을 효율적으로 클러스터링하는 방법을 배우고 다른 유형의 데이터에도 적용할 수 있기를 바랍니다. 대규모 언어 모델과 결합하여 이 접근 방식을 사용하면 대규모로 데이터를 심층적으로 분석할 수 있습니다.