llm-cluster-optimizer 1.0.4

llm-cluster-optimizer

A drop-in replacement for sklearn clustering models, this package optimizes sklearn clusters using LLMs

Installation

By default, this package relies on OpenAI for LLM interactions. This requires having copying your API key from here and putting it in the OPENAI_API_KEY envvar.

Otherwise, you can use any other LLM provider, as described below.

To install llm-cluster-optimizer, run the following command in your terminal:

pip install llm-cluster-optimizer

Usage

Here is a simple example of usage that provides two cluster models as options to be optimized. labels will be an np.ndarray of shape (len(input_texts),) representing which cluster each input text falls into. Use verbose=True in LLMClusterOptimizer.__init__() to get more information on progress.

from llm_cluster_optimizer import LLMClusterOptimizer
from sklearn.cluster import AffinityPropagation, HDBSCAN
import numpy as np
from langchain_openai import OpenAIEmbeddings

def embed_text_openai(text):
    return np.array(OpenAIEmbeddings(model="text-embedding-3-large").embed_query(text))

model = HDBSCAN(min_cluster_size=2, min_samples=1),
optimizer = LLMClusterOptimizer(model, embed_text_openai, 0.3)

input_texts = []  # your list of input texts
labels = optimizer.fit_predict_text(input_texts)

To use a clustering algorithm that takes the number of clusters as an input, you likely want to optimize n using the silhouette score or a similar metric. This requires access to the embeddings. You can pass the model wrapped in a function to achieve this:

from sklearn.cluster import KMeans
from util import guess_optimal_n_clusters

def get_kmeans_model(embeddings: np.ndarray):
    n_clusters = guess_optimal_n_clusters(embeddings, lambda n: KMeans(n_clusters=n))
    return KMeans(n_clusters=n_clusters)


optimizer = LLMClusterOptimizer(get_kmeans_model, embed_text_openai, 0.3)

Docs

init()

fit_predict_text()

Mimics sklearn.base.ClusterMixin.fit_predict(). Given a sequence of texts, returns the cluster labels for each text. Chooses best of cluster_models passed in __init__() based on calculate_clustering_score() Post-processes clusters so that:

• large clusters are broken up
• small clusters are combined based on similarity