HDBSCAN-Clustering mit Milvus

Daten können mithilfe von Deep-Learning-Modellen in Einbettungen umgewandelt werden, die aussagekräftige Repräsentationen der ursprünglichen Daten erfassen. Durch die Anwendung eines unüberwachten Clustering-Algorithmus können wir ähnliche Datenpunkte auf der Grundlage ihrer inhärenten Muster zusammenfassen. HDBSCAN (Hierarchical Density-Based Spatial Clustering of Applications with Noise) ist ein weit verbreiteter Clustering-Algorithmus, der Datenpunkte durch Analyse ihrer Dichte und ihres Abstands effizient gruppiert. Er ist besonders nützlich für die Entdeckung von Clustern unterschiedlicher Form und Größe. In diesem Notizbuch werden wir HDBSCAN mit Milvus, einer leistungsstarken Vektordatenbank, verwenden, um Datenpunkte auf der Grundlage ihrer Einbettungen in verschiedene Gruppen zu gruppieren.

HDBSCAN (Hierarchical Density-Based Spatial Clustering of Applications with Noise) ist ein Clustering-Algorithmus, der auf der Berechnung von Abständen zwischen Datenpunkten im Einbettungsraum beruht. Diese Einbettungen, die von Deep-Learning-Modellen erstellt werden, stellen die Daten in einer hochdimensionalen Form dar. Um ähnliche Datenpunkte zu gruppieren, bestimmt HDBSCAN ihre Nähe und Dichte, aber die effiziente Berechnung dieser Abstände, insbesondere bei großen Datensätzen, kann eine Herausforderung darstellen.

Milvus, eine leistungsstarke Vektordatenbank, optimiert diesen Prozess, indem sie Einbettungen speichert und indiziert und so ein schnelles Auffinden ähnlicher Vektoren ermöglicht. Im Zusammenspiel ermöglichen HDBSCAN und Milvus eine effiziente Clusterung großer Datensätze im Einbettungsraum.

In diesem Notebook verwenden wir das BGE-M3 Einbettungsmodell, um Einbettungen aus einem Datensatz von Nachrichtenschlagzeilen zu extrahieren, Milvus zur effizienten Berechnung von Distanzen zwischen Einbettungen zu verwenden, um HDBSCAN beim Clustering zu unterstützen, und die Ergebnisse anschließend für die Analyse mit der UMAP-Methode zu visualisieren. Dieses Notizbuch ist eine Milvus-Anpassung des Artikels von Dylan Castillo.

Vorbereitung

Download des Nachrichtendatensatzes von https://www.kaggle.com/datasets/dylanjcastillo/news-headlines-2024/

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

Daten herunterladen

Laden Sie den Nachrichtendatensatz von https://www.kaggle.com/datasets/dylanjcastillo/news-headlines-2024/ herunter, extrahieren Sie news_data_dedup.csv und legen Sie ihn im aktuellen Verzeichnis ab.

Einbettungen in Milvus extrahieren

Wir erstellen eine Sammlung mit Milvus und extrahieren dichte Einbettungen mit dem BGE-M3-Modell.

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

Konstruieren Sie die Abstandsmatrix für HDBSCAN

HDBSCAN erfordert die Berechnung von Distanzen zwischen Punkten für das Clustering, was sehr rechenintensiv sein kann. Da entfernte Punkte einen geringeren Einfluss auf die Clusterzuordnung haben, können wir die Effizienz durch die Berechnung der nächsten Top-k-Nachbarn verbessern. In diesem Beispiel verwenden wir den FLAT-Index, aber für große Datensätze unterstützt Milvus fortschrittlichere Indizierungsmethoden, um den Suchprozess zu beschleunigen. Zunächst benötigen wir einen Iterator, um die Milvus-Sammlung zu iterieren, die wir zuvor erstellt haben.

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

Wir werden alle Einbettungen in der Milvus-Sammlung iterieren. Für jede Einbettung suchen wir ihre Top-k-Nachbarn in der gleichen Sammlung, erhalten ihre IDs und Entfernungen. Dann müssen wir auch ein Wörterbuch erstellen, um die ursprüngliche ID auf einen kontinuierlichen Index in der Abstandsmatrix abzubilden. Anschließend müssen wir eine Distanzmatrix erstellen, die mit allen Elementen als unendlich initialisiert wird, und die gesuchten Elemente füllen. Auf diese Weise wird der Abstand zwischen weit entfernten Punkten ignoriert. Schließlich verwenden wir die HDBSCAN-Bibliothek, um die Punkte anhand der erstellten Abstandsmatrix zu clustern. Wir müssen die Metrik auf "vorberechnet" setzen, um anzugeben, dass es sich bei den Daten um eine Abstandsmatrix und nicht um ursprüngliche Einbettungen handelt.

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)

Danach ist das HDBSCAN-Clustering abgeschlossen. Wir können einige Daten abrufen und deren Cluster anzeigen. Beachten Sie, dass einige Daten keinem Cluster zugeordnet werden können, was bedeutet, dass es sich um Rauschen handelt, da sie sich in einer dünn besiedelten Region befinden.

Clustervisualisierung mit UMAP

Wir haben die Daten bereits mit HDBSCAN geclustert und können die Bezeichnungen für jeden Datenpunkt erhalten. Mit Hilfe einiger Visualisierungstechniken können wir uns jedoch ein Gesamtbild der Cluster für eine intuitive Analyse verschaffen. Jetzt werden wir UMAP verwenden, um die Cluster zu visualisieren. UMAP ist eine effiziente Methode zur Dimensionalitätsreduzierung, bei der die Struktur hochdimensionaler Daten erhalten bleibt, während sie zur Visualisierung oder weiteren Analyse in einen niedrigdimensionalen Raum projiziert werden. Damit können wir hochdimensionale Originaldaten im 2D- oder 3D-Raum visualisieren und die Cluster klar erkennen. Auch hier iterieren wir die Datenpunkte und erhalten die ID und den Text für die Originaldaten, dann verwenden wir Ploty, um die Datenpunkte mit diesen Metainformationen in einer Abbildung darzustellen, und verwenden verschiedene Farben, um verschiedene Cluster darzustellen.

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

Bild

Hier zeigen wir, dass die Daten gut geclustert sind, und Sie können den Mauszeiger über die Punkte bewegen, um den Text zu sehen, den sie darstellen. Wir hoffen, dass Sie mit diesem Notizbuch lernen, wie Sie HDBSCAN zum effizienten Clustern von Einbettungen mit Milvus verwenden können. In Kombination mit großen Sprachmodellen ermöglicht dieser Ansatz eine tiefere Analyse Ihrer Daten in großem Maßstab.