MilvusによるHDBSCANクラスタリング

データはディープラーニングモデルを使って埋め込みに変換することができ、元のデータの意味のある表現を捉えることができる。教師なしクラスタリングアルゴリズムを適用することで、固有のパターンに基づいて類似のデータポイントをグループ化することができる。HDBSCAN(Hierarchical Density-Based Spatial Clustering of Applications with Noise)は、データ点の密度と距離を分析することで効率的にデータ点をグループ化する、広く使われているクラスタリングアルゴリズムである。様々な形や大きさのクラスタを発見するのに特に有用です。このノートブックでは、HDBSCANと高性能ベクトルデータベースMilvusを使って、データ点をその埋め込みに基づく明確なグループにクラスタリングします。

HDBSCAN（Hierarchical Density-Based Spatial Clustering of Applications with Noise）は、埋め込み空間におけるデータ点間の距離計算に依存するクラスタリングアルゴリズムです。ディープラーニングモデルによって作成されたこれらの埋め込みは、高次元形式でデータを表現する。類似したデータ点をグループ化するために、HDBSCANはそれらの近接性と密度を決定するが、これらの距離を効率的に計算することは、特に大規模なデータセットでは困難な場合がある。

高性能ベクトルデータベースであるMilvusは、埋め込みデータを格納しインデックスを付けることでこのプロセスを最適化し、類似ベクトルの高速検索を可能にします。HDBSCANとMilvusを併用することで、埋め込み空間における大規模データセットの効率的なクラスタリングが可能になります。

このノートブックでは、BGE-M3埋め込みモデルを使ってニュースのヘッドラインデータセットから埋め込みを抽出し、Milvusを使って埋め込み間の距離を効率的に計算し、HDBSCANのクラスタリングを支援します。このノートブックは、Dylan Castilloの論文を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ライブラリを使って、作成した距離行列を使って点をクラスタリングする。このとき、埋め込みデータではなく距離行列であることを示すために、metricを'precomputed'に設定する必要がある。

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は、次元削減のために使用される効率的な手法であり、高次元データの構造を保持しながら、可視化やさらなる分析のために低次元空間に投影する。UMAPを用いることで、元の高次元データを2次元または3次元空間に可視化し、クラスターを明確に見ることができる。 ここでも、データ点を反復処理し、元データのidとテキストを取得する。次に、plotyを用いて、これらのメタ情報を持つデータ点を図にプロットし、異なるクラスターを表すために異なる色を用いる。

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()

イメージ

ここでは、データがうまくクラスタ化されていることを示し、点の上にカーソルを置くと、それらが表すテキストを確認することができます。このノートブックで、HDBSCANを使ってMilvusで埋め込みデータを効率的にクラスタリングする方法を学んでいただければと思います。大規模な言語モデルと組み合わせることで、大規模なデータ分析が可能になります。