Pengelompokan HDBSCAN dengan Milvus

Data dapat diubah menjadi embeddings menggunakan model pembelajaran mendalam, yang menangkap representasi yang bermakna dari data asli. Dengan menerapkan algoritme pengelompokan tanpa pengawasan, kita dapat mengelompokkan titik-titik data yang serupa berdasarkan pola yang melekat. HDBSCAN (Hierarchical Density-Based Spatial Clustering of Applications with Noise) adalah algoritme pengelompokan yang banyak digunakan untuk mengelompokkan titik-titik data secara efisien dengan menganalisis kepadatan dan jaraknya. Algoritma ini sangat berguna untuk menemukan cluster dengan berbagai bentuk dan ukuran. Dalam buku catatan ini, kita akan menggunakan HDBSCAN dengan Milvus, database vektor berkinerja tinggi, untuk mengelompokkan titik data ke dalam kelompok-kelompok yang berbeda berdasarkan penyematannya.

HDBSCAN (Hierarchical Density-Based Spatial Clustering of Applications with Noise) adalah algoritme pengelompokan yang mengandalkan penghitungan jarak antara titik data dalam ruang embedding. Ruang embedding ini, yang dibuat oleh model pembelajaran mendalam, merepresentasikan data dalam bentuk dimensi tinggi. Untuk mengelompokkan titik data yang serupa, HDBSCAN menentukan kedekatan dan kepadatannya, tetapi menghitung jarak ini secara efisien, terutama untuk kumpulan data yang besar, dapat menjadi tantangan.

Milvus, basis data vektor berkinerja tinggi, mengoptimalkan proses ini dengan menyimpan dan mengindeks penyematan, sehingga memungkinkan pengambilan vektor yang serupa dengan cepat. Ketika digunakan bersama-sama, HDBSCAN dan Milvus memungkinkan pengelompokan yang efisien untuk set data berskala besar dalam ruang embedding.

Dalam buku catatan ini, kami akan menggunakan model embedding BGE-M3 untuk mengekstrak embedding dari dataset berita utama, memanfaatkan Milvus untuk menghitung jarak antar embedding secara efisien untuk membantu HDBSCAN dalam pengelompokan, dan kemudian memvisualisasikan hasilnya untuk analisis menggunakan metode UMAP. Buku catatan ini merupakan adaptasi Milvus dari artikel Dylan Castillo.

Persiapan

unduh dataset berita dari https://www.kaggle.com/datasets/dylanjcastillo/news-headlines-2024/

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

Unduh Data

Unduh dataset berita dari https://www.kaggle.com/datasets/dylanjcastillo/news-headlines-2024/, ekstrak news_data_dedup.csv dan letakkan di direktori saat ini.

Mengekstrak sematan ke Milvus

Kita akan membuat koleksi menggunakan Milvus, dan mengekstrak embedding padat menggunakan model 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()

Membangun Matriks Jarak untuk HDBSCAN

HDBSCAN memerlukan penghitungan jarak antar titik untuk pengelompokan, yang bisa jadi sangat intensif secara komputasi. Karena titik-titik yang jauh memiliki pengaruh yang lebih kecil pada penugasan pengelompokan, kita dapat meningkatkan efisiensi dengan menghitung k tetangga terdekat. Dalam contoh ini, kita menggunakan indeks FLAT, tetapi untuk kumpulan data berskala besar, Milvus mendukung metode pengindeksan yang lebih canggih untuk mempercepat proses pencarian. Pertama, kita perlu mendapatkan iterator untuk mengulang koleksi Milvus yang telah kita buat sebelumnya.

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 = []

Kita akan mengulang semua sematan di dalam koleksi Milvus. Untuk setiap embedding, kita akan mencari top-k tetangganya dalam koleksi yang sama, mendapatkan id dan jaraknya. Kemudian kita juga perlu membuat sebuah kamus untuk memetakan ID asli ke sebuah indeks kontinu di dalam matriks jarak. Setelah selesai, kita perlu membuat matriks jarak yang diinisialisasi dengan semua elemen sebagai tak terhingga dan mengisi elemen-elemen yang kita cari. Dengan cara ini, jarak antara titik-titik yang jauh akan diabaikan. Terakhir, kita menggunakan library HDBSCAN untuk mengelompokkan titik-titik menggunakan matriks jarak yang telah kita buat. Kita perlu mengatur metrik ke 'precomputed' untuk menunjukkan bahwa data tersebut merupakan matriks jarak dan bukannya embedding asli.

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)

Setelah itu, pengelompokan HDBSCAN selesai. Kita bisa mendapatkan beberapa data dan menunjukkan klasternya. Perhatikan bahwa beberapa data tidak akan dimasukkan ke dalam klaster mana pun, yang berarti data tersebut adalah noise, karena terletak di wilayah yang jarang.

Visualisasi Cluster menggunakan UMAP

Kita telah mengelompokkan data menggunakan HDBSCAN dan mendapatkan label untuk setiap titik data. Namun dengan menggunakan beberapa teknik visualisasi, kita bisa mendapatkan gambaran keseluruhan dari cluster untuk analisis intuitif. Sekarang kita akan menggunakan UMAP untuk memvisualisasikan cluster. UMAP adalah metode yang efisien yang digunakan untuk pengurangan dimensi, mempertahankan struktur data dimensi tinggi sambil memproyeksikannya ke dalam ruang dimensi yang lebih rendah untuk visualisasi atau analisis lebih lanjut. Dengannya, kita dapat memvisualisasikan data dimensi tinggi asli dalam ruang 2D atau 3D, dan melihat cluster dengan jelas. Di sini sekali lagi, kita mengulang titik data dan mendapatkan id dan teks untuk data asli, kemudian kita menggunakan ploty untuk memplot titik-titik data dengan metainfo ini dalam sebuah gambar, dan menggunakan warna yang berbeda untuk merepresentasikan cluster yang berbeda.

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

gambar

Di sini, kami mendemonstrasikan bahwa data dikelompokkan dengan baik, dan Anda dapat mengarahkan kursor ke titik-titik tersebut untuk memeriksa teks yang diwakilinya. Dengan buku catatan ini, kami harap Anda dapat mempelajari cara menggunakan HDBSCAN untuk mengelompokkan embedding dengan Milvus secara efisien, yang juga dapat diterapkan pada jenis data lainnya. Dikombinasikan dengan model bahasa yang besar, pendekatan ini memungkinkan analisis yang lebih dalam pada data Anda dalam skala besar.