Agrupamento HDBSCAN com Milvus

Os dados podem ser transformados em embeddings utilizando modelos de aprendizagem profunda, que captam representações significativas dos dados originais. Ao aplicar um algoritmo de agrupamento não supervisionado, podemos agrupar pontos de dados semelhantes com base nos seus padrões inerentes. O HDBSCAN (Hierarchical Density-Based Spatial Clustering of Applications with Noise) é um algoritmo de agrupamento amplamente utilizado que agrupa eficazmente os pontos de dados através da análise da sua densidade e distância. É particularmente útil para descobrir clusters de formas e tamanhos variados. Neste caderno, utilizaremos o HDBSCAN com o Milvus, uma base de dados vetorial de elevado desempenho, para agrupar pontos de dados em grupos distintos com base nas suas incorporações.

O HDBSCAN (Hierarchical Density-Based Spatial Clustering of Applications with Noise) é um algoritmo de agrupamento que se baseia no cálculo de distâncias entre pontos de dados no espaço de incorporação. Estas incorporações, criadas por modelos de aprendizagem profunda, representam os dados numa forma altamente dimensional. Para agrupar pontos de dados semelhantes, o HDBSCAN determina a sua proximidade e densidade, mas o cálculo eficiente dessas distâncias, especialmente para grandes conjuntos de dados, pode ser um desafio.

O Milvus, uma base de dados vetorial de elevado desempenho, optimiza este processo através do armazenamento e indexação de embeddings, permitindo uma recuperação rápida de vectores semelhantes. Quando usados em conjunto, o HDBSCAN e o Milvus permitem o agrupamento eficiente de conjuntos de dados de grande escala no espaço de incorporação.

Neste caderno, usaremos o modelo de embedding BGE-M3 para extrair embeddings de um conjunto de dados de manchetes de notícias, utilizaremos o Milvus para calcular eficientemente as distâncias entre embeddings para auxiliar o HDBSCAN no agrupamento e, em seguida, visualizaremos os resultados para análise usando o método UMAP. Este caderno é uma adaptação de Milvus do artigo de Dylan Castillo.

Preparação

descarregar conjunto de dados de notícias de https://www.kaggle.com/datasets/dylanjcastillo/news-headlines-2024/

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

Descarregar dados

Descarregue o conjunto de dados de notícias de https://www.kaggle.com/datasets/dylanjcastillo/news-headlines-2024/, extraia news_data_dedup.csv e coloque-o no diretório atual.

Extrair Embeddings para Milvus

Vamos criar uma coleção utilizando o Milvus e extrair embeddings densos utilizando o 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 a matriz de distância para HDBSCAN

O HDBSCAN requer o cálculo de distâncias entre pontos para o agrupamento, o que pode ser computacionalmente intensivo. Uma vez que os pontos distantes têm menos influência nas atribuições de agrupamento, podemos melhorar a eficiência calculando os top-k vizinhos mais próximos. Neste exemplo, usamos o índice FLAT, mas para conjuntos de dados de grande escala, o Milvus suporta métodos de indexação mais avançados para acelerar o processo de pesquisa. Em primeiro lugar, precisamos de obter um iterador para iterar a coleção Milvus que criámos anteriormente.

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

Vamos iterar todos os embeddings da coleção Milvus. Para cada embedding, vamos procurar os seus top-k vizinhos na mesma coleção, obter os seus ids e distâncias. Depois, também precisamos de construir um dicionário para mapear o ID original para um índice contínuo na matriz de distâncias. Quando terminarmos, precisamos de criar uma matriz de distâncias inicializada com todos os elementos como infinito e preencher os elementos que procurámos. Desta forma, a distância entre pontos distantes será ignorada. Finalmente, utilizamos a biblioteca HDBSCAN para agrupar os pontos utilizando a matriz de distâncias que criámos. Temos de definir a métrica como 'precomputed' para indicar que os dados são a matriz de distâncias e não os embeddings originais.

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)

Depois disto, o agrupamento HDBSCAN está concluído. Podemos obter alguns dados e mostrar o respetivo agrupamento. Note-se que alguns dados não serão atribuídos a nenhum cluster, o que significa que são ruído, porque estão localizados numa região esparsa.

Visualização de clusters usando UMAP

Já agrupámos os dados utilizando o HDBSCAN e podemos obter as etiquetas para cada ponto de dados. No entanto, usando algumas técnicas de visualização, podemos obter a imagem completa dos clusters para uma análise intuitiva. Agora vamos utilizar o UMAP para visualizar os clusters. O UMAP é um método eficiente utilizado para a redução da dimensionalidade, preservando a estrutura de dados de elevada dimensão enquanto os projecta num espaço de dimensão inferior para visualização ou análise posterior. Aqui, mais uma vez, iteramos os pontos de dados e obtemos o id e o texto dos dados originais, depois utilizamos o ploty para representar os pontos de dados com estas metainformações numa figura e utilizamos cores diferentes para representar os 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()

imagem

Aqui, demonstramos que os dados estão bem agrupados e pode passar o rato sobre os pontos para verificar o texto que representam. Com este caderno de notas, esperamos que aprenda a utilizar o HDBSCAN para agrupar embeddings com o Milvus de forma eficiente, o que também pode ser aplicado a outros tipos de dados. Combinada com modelos de linguagem de grande dimensão, esta abordagem permite uma análise mais profunda dos seus dados em grande escala.