Agrupación HDBSCAN con Milvus

Los datos pueden transformarse en incrustaciones utilizando modelos de aprendizaje profundo, que capturan representaciones significativas de los datos originales. Aplicando un algoritmo de clustering no supervisado, podemos agrupar puntos de datos similares basándonos en sus patrones inherentes. HDBSCAN (Hierarchical Density-Based Spatial Clustering of Applications with Noise) es un algoritmo de clustering ampliamente utilizado que agrupa eficientemente puntos de datos analizando su densidad y distancia. Resulta especialmente útil para descubrir conglomerados de formas y tamaños diversos. En este cuaderno, utilizaremos HDBSCAN con Milvus, una base de datos vectorial de alto rendimiento, para agrupar puntos de datos en distintos grupos en función de sus incrustaciones.

HDBSCAN (Hierarchical Density-Based Spatial Clustering of Applications with Noise) es un algoritmo de agrupación que se basa en el cálculo de distancias entre puntos de datos en el espacio de incrustación. Estas incrustaciones, creadas por modelos de aprendizaje profundo, representan los datos en una forma de alta dimensión. Para agrupar puntos de datos similares, HDBSCAN determina su proximidad y densidad, pero calcular eficientemente estas distancias, especialmente para grandes conjuntos de datos, puede ser un reto.

Milvus, una base de datos vectorial de alto rendimiento, optimiza este proceso almacenando e indexando incrustaciones, lo que permite recuperar rápidamente vectores similares. Cuando se utilizan conjuntamente, HDBSCAN y Milvus permiten una agrupación eficaz de conjuntos de datos a gran escala en el espacio de incrustación.

En este cuaderno, utilizaremos el modelo de incrustación BGE-M3 para extraer incrustaciones de un conjunto de datos de titulares de noticias, utilizaremos Milvus para calcular eficientemente las distancias entre incrustaciones para ayudar a HDBSCAN en la agrupación y, a continuación, visualizaremos los resultados para su análisis utilizando el método UMAP. Este cuaderno es una adaptación a Milvus del artículo de Dylan Castillo.

Preparación

descargar el conjunto de datos de noticias de https://www.kaggle.com/datasets/dylanjcastillo/news-headlines-2024/

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

Descarga de datos

Descargue el conjunto de datos de noticias de https://www.kaggle.com/datasets/dylanjcastillo/news-headlines-2024/, extraiga news_data_dedup.csv y póngalo en el directorio actual.

Extraer las incrustaciones a Milvus

Crearemos una colección utilizando Milvus, y extraeremos las incrustaciones densas utilizando el modelo 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()

Construir la matriz de distancias para HDBSCAN

HDBSCAN requiere el cálculo de distancias entre puntos para la agrupación, lo que puede ser computacionalmente intensivo. Dado que los puntos distantes tienen menos influencia en las asignaciones de clustering, podemos mejorar la eficiencia calculando los top-k vecinos más cercanos. En este ejemplo, utilizamos el índice FLAT, pero para conjuntos de datos a gran escala, Milvus soporta métodos de indexación más avanzados para acelerar el proceso de búsqueda. En primer lugar, necesitamos obtener un iterador para iterar la colección Milvus que hemos creado previamente.

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

Recorreremos todas las incrustaciones de la colección Milvus. Para cada incrustación, buscaremos sus vecinos top-k en la misma colección, obtendremos sus ids y distancias. A continuación, también tenemos que construir un diccionario para asignar ID original a un índice continuo en la matriz de distancia. Cuando hayamos terminado, tendremos que crear una matriz de distancias inicializada con todos los elementos en infinito y rellenar los elementos que hayamos buscado. De esta forma, la distancia entre puntos lejanos será ignorada. Finalmente usamos la librería HDBSCAN para agrupar los puntos usando la matriz de distancia que hemos creado. Debemos establecer la métrica en "precomputed" para indicar que los datos son matrices de distancia en lugar de incrustaciones originales.

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)

Después de esto, el HDBSCAN clustering está terminado. Podemos obtener algunos datos y mostrar su cluster. Observe que algunos datos no serán asignados a ningún cluster, lo que significa que son ruido, porque están localizados en alguna región dispersa.

Visualización de clusters usando UMAP

Ya hemos agrupado los datos usando HDBSCAN y podemos obtener las etiquetas para cada punto de datos. Sin embargo, utilizando algunas técnicas de visualización, podemos obtener la imagen completa de los clusters para un análisis intuitivo. Ahora vamos a utilizar UMAP para visualizar los clusters. UMAP es un método eficaz para reducir la dimensionalidad, preservando la estructura de los datos de alta dimensión y proyectándolos en un espacio de menor dimensión para su visualización o análisis posterior. De nuevo, iteramos los puntos de datos y obtenemos el id y el texto de los datos originales, luego usamos ploty para representar los puntos de datos con estos metainfo en una figura, y usamos diferentes colores para representar diferentes clusters.

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

imagen

Aquí, demostramos que los datos están bien agrupados, y se puede pasar el ratón sobre los puntos para comprobar el texto que representan. Con este cuaderno, esperamos que aprenda a utilizar HDBSCAN para agrupar incrustaciones con Milvus de forma eficiente, lo que también puede aplicarse a otros tipos de datos. Combinado con grandes modelos lingüísticos, este enfoque permite un análisis más profundo de sus datos a gran escala.