0Pricing
AI Engineering Academy · Lesson

Clustering and Visualizing Embeddings

Apply k-means clustering to a set of embeddings and visualize them in 2D using UMAP to discover natural topic groupings in your data.

Clustering and Visualizing Embeddings is a free AI Engineering Academy lesson on CoddyKit — lesson 4 of 4. You can read the complete lesson below for free — then practise it hands-on in the browser with a built-in code editor and a 24/7 AI tutor. It is part of the AI Engineering Academy learning path, one of 4 lessons in the course, and your progress syncs across the web and the CoddyKit app.

Why Cluster Embeddings?

When you have hundreds or thousands of documents, you often want to discover what topics exist without manually reading everything. Clustering embeddings groups semantically similar documents together automatically, revealing the natural structure of your data.

Common applications include: auto-tagging support tickets, discovering content categories, finding redundant documents, and understanding what users ask about most.

K-Means Clustering Overview

K-means partitions n data points into k clusters by iteratively assigning each point to the nearest centroid, then recomputing centroids as the mean of assigned points. It converges when assignments stop changing.

For embeddings, k-means finds documents that are close together in the high-dimensional embedding space, effectively grouping them by semantic similarity.

from sklearn.cluster import KMeans
import numpy as np

# corpus_embeddings: (n_docs, 1536) — pre-computed
corpus_embeddings = np.random.randn(200, 1536)  # placeholder

kmeans = KMeans(n_clusters=5, random_state=42, n_init='auto')
kmeans.fit(corpus_embeddings)

labels = kmeans.labels_
print(f'Cluster assignments: {labels[:10]}')
print(f'Unique clusters: {set(labels)}')

Choosing the Right Number of Clusters

Choosing k (the number of clusters) is the hardest part. The elbow method plots inertia (sum of squared distances to centroids) for different values of k and looks for the point where adding more clusters stops reducing inertia significantly.

The silhouette score measures how much better a point fits its own cluster versus the nearest other cluster — scores closer to 1 are better. Try k values from 3 to 20 and pick the best score.

from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
import numpy as np

corpus_embeddings = np.random.randn(200, 50)  # placeholder (50D for speed)

best_k, best_score = 2, -1
for k in range(2, 11):
    km = KMeans(n_clusters=k, random_state=42, n_init='auto')
    labels = km.fit_predict(corpus_embeddings)
    score = silhouette_score(corpus_embeddings, labels, sample_size=100)
    print(f'k={k}: silhouette={score:.3f}')
    if score > best_score:
        best_score, best_k = score, k

print(f'Best k: {best_k}')

Labeling Clusters with an LLM

After clustering, you can ask an LLM to generate a human-readable label for each cluster by passing a few representative documents from that cluster to the model. This turns raw cluster numbers into meaningful category names automatically.

from openai import OpenAI

client = OpenAI()

def label_cluster(cluster_docs, n_examples=3):
    examples = '\n'.join(f'- {d}' for d in cluster_docs[:n_examples])
    prompt = f'Given these documents, provide a 3-word category label:\n{examples}\nLabel:'
    response = client.chat.completions.create(
        model='gpt-4o-mini',
        messages=[{'role': 'user', 'content': prompt}],
        max_tokens=20
    )
    return response.choices[0].message.content.strip()

# Example usage:
# cluster_0_docs = [documents[i] for i in range(len(documents)) if labels[i] == 0]
# print(label_cluster(cluster_0_docs))

Dimensionality Reduction with UMAP

1536-dimensional vectors cannot be plotted directly. UMAP (Uniform Manifold Approximation and Projection) reduces high-dimensional data to 2D or 3D while preserving local neighborhood structure. Unlike PCA, UMAP is non-linear and much better at preserving clusters.

UMAP is the standard choice for visualizing text embeddings because similar documents remain close together in the 2D plot.

import umap
import numpy as np

corpus_embeddings = np.random.randn(200, 1536)  # placeholder

reducer = umap.UMAP(
    n_components=2,
    metric='cosine',   # important for text embeddings
    random_state=42
)

embeddings_2d = reducer.fit_transform(corpus_embeddings)
print(f'Reduced shape: {embeddings_2d.shape}')  # (200, 2)

Visualizing Clusters with Matplotlib

Once you have 2D coordinates from UMAP and cluster labels from k-means, a simple scatter plot makes the cluster structure visible. Color each point by its cluster label and add representative document titles as text annotations to make the chart interpretable.

import matplotlib.pyplot as plt
import numpy as np

# Assume: embeddings_2d (n, 2), labels (n,), documents list
# Placeholder data:
np.random.seed(42)
embeddings_2d = np.random.randn(50, 2)
labels = np.random.randint(0, 5, 50)

fig, ax = plt.subplots(figsize=(10, 7))
scatter = ax.scatter(
    embeddings_2d[:, 0],
    embeddings_2d[:, 1],
    c=labels,
    cmap='tab10',
    s=60,
    alpha=0.8
)
plt.colorbar(scatter, label='Cluster')
ax.set_title('Document Embedding Clusters (UMAP)')
plt.tight_layout()
plt.savefig('/tmp/clusters.png', dpi=150)
print('Saved cluster plot')

UMAP Parameters That Matter

Key UMAP hyperparameters that affect visualization quality:

  • n_neighbors (default 15): larger values capture more global structure, smaller values reveal local clusters
  • min_dist (default 0.1): how tightly points pack together; smaller = tighter clusters but more overlap between them
  • metric: always use 'cosine' for text embeddings — Euclidean distance is inappropriate

Experiment with n_neighbors between 5 and 50 to find the visualization that best reveals your data's structure.

Hierarchical Clustering Alternative

Agglomerative hierarchical clustering builds a tree of nested clusters without requiring you to specify k in advance. You cut the tree at a chosen distance threshold to get your final clusters. This is useful when you do not know the number of natural topics in your data.

from sklearn.cluster import AgglomerativeClustering
import numpy as np

corpus_embeddings = np.random.randn(100, 50)  # placeholder

cluster = AgglomerativeClustering(
    n_clusters=None,
    distance_threshold=1.5,
    metric='cosine',
    linkage='average'
)
labels = cluster.fit_predict(corpus_embeddings)

n_clusters = len(set(labels))
print(f'Found {n_clusters} natural clusters')

Finding Outliers and Near-Duplicates

Embeddings are also excellent for finding near-duplicate documents. Compute pairwise cosine similarity between all documents and flag pairs with similarity above 0.95 as potential duplicates. This is more robust than text diffing because it catches rephrased duplicates.

import numpy as np

def find_near_duplicates(corpus_vecs, threshold=0.95):
    # corpus_vecs assumed L2-normalized
    sim_matrix = corpus_vecs @ corpus_vecs.T  # (n, n)
    duplicates = []
    n = len(corpus_vecs)
    for i in range(n):
        for j in range(i + 1, n):
            if sim_matrix[i, j] >= threshold:
                duplicates.append((i, j, float(sim_matrix[i, j])))
    return duplicates

# pairs = find_near_duplicates(corpus_embeddings)
# print(f'Found {len(pairs)} near-duplicate pairs')

Practical Clustering Workflow

A real clustering workflow follows these steps in order:

  1. Embed all documents with text-embedding-3-small
  2. Optionally reduce to 50D with UMAP before clustering (speeds up k-means)
  3. Cluster with k-means, sweep k to find the best silhouette score
  4. Label each cluster with an LLM using 5 representative documents
  5. Visualize the 2D UMAP scatter plot colored by cluster
  6. Review outliers (points far from any centroid)

Limitations of Clustering

Clustering embeddings has important limitations to keep in mind:

  • K-means assumes spherical clusters — elongated or irregularly shaped topic groups may split incorrectly
  • Embeddings compress meaning — two documents can be semantically close but belong to different actionable categories
  • Cluster labels need verification — always sample documents from each cluster to sanity-check the LLM-generated label

Use clustering as an exploration tool, not as a ground-truth categorization system.

Quick Check

Test your understanding of AI Engineering concepts from this lesson.

Lesson Recap

In this lesson you learned: k-means groups documents by semantic proximity in the embedding space, UMAP reduces high-dimensional embeddings to 2D for visualization using cosine distance, and silhouette score helps you choose the right number of clusters without labeled data. Next up we move beyond in-memory search and explore production vector databases.

Frequently asked questions

Is the “Clustering and Visualizing Embeddings” lesson free?

Yes — the full text of “Clustering and Visualizing Embeddings” is free to read here on the web, and the AI Engineering Academy course includes 4 lessons in total. To practise it interactively (a built-in code editor and a 24/7 AI tutor) and unlock the rest of the AI Engineering Academy course, upgrade to CoddyKit PRO.

What will I learn in “Clustering and Visualizing Embeddings”?

Apply k-means clustering to a set of embeddings and visualize them in 2D using UMAP to discover natural topic groupings in your data. You practise AI Engineering Academy with hands-on code you run directly in the browser, and a 24/7 AI tutor answers your questions as you work through the lesson.

Do I need any experience to start AI Engineering Academy?

No prior experience is required. AI Engineering Academy on CoddyKit is structured for beginners through advanced learners; this is — lesson 4 of 4, so you can start here or from the beginning and move at your own pace.

How long does the “Clustering and Visualizing Embeddings” lesson take?

Most CoddyKit lessons take about 5–10 minutes. Each one is bite-sized and interactive, so you make steady progress and pick up exactly where you left off across the web and the app.

Can I write and run code in this AI Engineering Academy lesson?

Yes. Every AI Engineering Academy lesson includes a built-in code editor, so you write and run real code right in your browser and get instant AI feedback — no local setup required.

All lessons in this course

  1. What Are Vector Embeddings?
  2. Generating Embeddings with OpenAI
  3. Semantic Search with NumPy
  4. Clustering and Visualizing Embeddings
← Back to AI Engineering Academy