HDBSCAN Clustering avec Milvus

Les données peuvent être transformées en embeddings à l'aide de modèles d'apprentissage profond, qui capturent des représentations significatives des données d'origine. En appliquant un algorithme de clustering non supervisé, nous pouvons regrouper des points de données similaires sur la base de leurs modèles inhérents. HDBSCAN (Hierarchical Density-Based Spatial Clustering of Applications with Noise) est un algorithme de clustering largement utilisé qui regroupe efficacement les points de données en analysant leur densité et leur distance. Il est particulièrement utile pour découvrir des grappes de formes et de tailles variées. Dans ce carnet, nous utiliserons HDBSCAN avec Milvus, une base de données vectorielles haute performance, pour regrouper des points de données en groupes distincts sur la base de leurs encastrements.

HDBSCAN (Hierarchical Density-Based Spatial Clustering of Applications with Noise) est un algorithme de clustering qui repose sur le calcul des distances entre les points de données dans l'espace d'intégration. Ces embeddings, créés par des modèles d'apprentissage profond, représentent les données sous une forme hautement dimensionnelle. Pour regrouper des points de données similaires, HDBSCAN détermine leur proximité et leur densité, mais le calcul efficace de ces distances, en particulier pour les grands ensembles de données, peut s'avérer difficile.

Milvus, une base de données vectorielles haute performance, optimise ce processus en stockant et en indexant les encastrements, ce qui permet de retrouver rapidement des vecteurs similaires. Utilisés conjointement, HDBSCAN et Milvus permettent un regroupement efficace d'ensembles de données à grande échelle dans l'espace d'intégration.

Dans ce bloc-notes, nous utiliserons le modèle d'intégration BGE-M3 pour extraire les intégrations d'un ensemble de données de titres de journaux, nous utiliserons Milvus pour calculer efficacement les distances entre les intégrations afin d'aider HDBSCAN dans le regroupement, puis nous visualiserons les résultats pour l'analyse à l'aide de la méthode UMAP. Ce bloc-notes est une adaptation Milvus de l'article de Dylan Castillo.

Préparation

Télécharger le jeu de données des actualités sur https://www.kaggle.com/datasets/dylanjcastillo/news-headlines-2024/

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

Télécharger les données

Télécharger le jeu de données d'actualités à partir de https://www.kaggle.com/datasets/dylanjcastillo/news-headlines-2024/, extraire news_data_dedup.csv et le placer dans le répertoire courant.

Extraire les Embeddings vers Milvus

Nous allons créer une collection à l'aide de Milvus, et extraire des embeddings denses à l'aide du modèle 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()

Construire la matrice de distance pour HDBSCAN

HDBSCAN nécessite le calcul des distances entre les points pour le regroupement, ce qui peut nécessiter un calcul intensif. Comme les points éloignés ont moins d'influence sur les affectations de regroupement, nous pouvons améliorer l'efficacité en calculant les k premiers voisins les plus proches. Dans cet exemple, nous utilisons l'index FLAT, mais pour les ensembles de données à grande échelle, Milvus prend en charge des méthodes d'indexation plus avancées afin d'accélérer le processus de recherche. Tout d'abord, nous devons obtenir un itérateur pour parcourir la collection Milvus que nous avons précédemment créée.

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

Nous allons itérer tous les encastrements dans la collection Milvus. Pour chaque encastrement, nous rechercherons ses principaux voisins dans la même collection, nous obtiendrons leurs identifiants et leurs distances. Ensuite, nous devons également construire un dictionnaire pour faire correspondre l'ID original à un index continu dans la matrice de distance. Une fois la recherche terminée, nous devons créer une matrice de distance dont tous les éléments sont initialisés à l'infini et remplir les éléments que nous avons recherchés. De cette manière, la distance entre les points éloignés sera ignorée. Enfin, nous utilisons la bibliothèque HDBSCAN pour regrouper les points à l'aide de la matrice de distance que nous avons créée. Nous devons définir la métrique sur "précalculée" pour indiquer que les données sont des matrices de distance plutôt que des encastrements originaux.

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)

Après cela, le regroupement HDBSCAN est terminé. Nous pouvons obtenir des données et afficher leur cluster. Notez que certaines données ne seront assignées à aucun cluster, ce qui signifie qu'il s'agit de bruit, parce qu'elles sont situées dans une région peu dense.

Visualisation des grappes à l'aide de l'UMAP

Nous avons déjà regroupé les données à l'aide de HDBSCAN et nous pouvons obtenir les étiquettes pour chaque point de données. Cependant, en utilisant certaines techniques de visualisation, nous pouvons obtenir une image complète des grappes pour une analyse intuitive. Nous allons maintenant utiliser UMAP pour visualiser les grappes. UMAP est une méthode efficace utilisée pour la réduction de la dimensionnalité, préservant la structure des données à haute dimension tout en les projetant dans un espace à plus faible dimension pour la visualisation ou une analyse plus approfondie. Ici encore, nous itérons les points de données et obtenons l'identifiant et le texte des données originales, puis nous utilisons ploty pour représenter les points de données avec ces méta-informations dans une figure, et nous utilisons différentes couleurs pour représenter les différents groupes.

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

Image

Ici, nous démontrons que les données sont bien regroupées, et vous pouvez survoler les points pour vérifier le texte qu'ils représentent. Avec ce carnet, nous espérons que vous apprendrez à utiliser HDBSCAN pour regrouper efficacement des embeddings avec Milvus, ce qui peut également être appliqué à d'autres types de données. Combinée à de grands modèles de langage, cette approche permet une analyse plus approfondie de vos données à grande échelle.