Clustering Music Tracks

Introduction: Why Cluster Tracks?

Welcome back! So far, you have learned how to represent music tracks and user preferences as vectors, and how to use cosine similarity to recommend tracks to users. In this lesson, we will take a new step: grouping similar tracks together using a technique called clustering . Clustering helps us organize our music library by finding groups of tracks that are similar to each other. This is useful for many reasons. For example, you can use clusters to create playlists, suggest new genres to users, or simply explore your music collection in a more structured way. In this lesson, you will learn how to use the KMeans algorithm to cluster tracks based on their embeddings.

Recap: Data and Embedding Setup

Before we start clustering, let’s quickly remind ourselves how we get the data and embeddings for our tracks. You have already seen how to load track data and generate embeddings in previous lessons. Here is a short code block that shows the basic setup: Python from src.user_model import get_track_embeddings from src.database import get_all_tracks # Get track IDs and their embeddings track_ids, track_embeddings_matrix = get_track_embeddings() # Load all track details tracks_df = get_all_tracks() This code gives us two important things: track_ids and track_embeddings_matrix : These are the unique IDs for each track and their corresponding embedding vectors. tracks_df : This is a DataFrame containing all the details about each track. We will use these as the starting point for clustering. To confirm everything is wired up correctly, we run some isolated tests on our clustering system (you’ll see these in the practice inside the src/tests/test_clustering.py). Here's the initial diagnostic output: text ✅ Loaded 5 tracks for clustering. ✅ Retrieved embeddings for 5 tracks. Embedding dimensions: 11

What Are Clusters

Clustering is a type of unsupervised learning, where we don’t start with any labels or categories. Instead, we ask the algorithm to find natural groupings in the data — that is, clusters. Think of a cluster as a “cloud” of similar items in a multidimensional space. Each track is represented as a point in this space (based on its embedding), and clustering algorithms try to group nearby points together. The assumption is: if points are close, they’re likely to be similar in meaningful ways (e.g., mood, tempo, instrumentation). Unlike classification, where we already know categories (like genre) and assign items to them, clustering figures out the categories for us . It answers: "What kinds of groups naturally exist in my data?" For example, without knowing any genres up front, clustering might still discover groups like “slow instrumental tracks,” “fast electronic tracks,” or “melancholic acoustic songs” — purely based on numerical similarities. This makes clustering especially useful for: Exploring datasets you don’t fully understand yet Discovering unexpected patterns Creating structure from messy or unlabeled data KMeans is one of the simplest and most widely used clustering methods, which is why we start with it here.

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

KMeans Clustering Explained

KMeans is a popular clustering algorithm. It works by dividing your data into a set number of groups, called clusters. Each cluster contains tracks that are similar to each other based on their embeddings. Two tracks can be placed in the same cluster even if their genres differ, as long as their embeddings are close in vector space.

Here’s how KMeans works in simple terms:

You choose how many clusters you want (for example, 3).
The algorithm tries to group the tracks so that tracks in the same cluster are as similar as possible.
Each track is assigned a cluster label (like 0, 1, or 2).

In the context of music tracks, this means that tracks with similar features (like genre, tempo, or mood) will end up in the same cluster. This makes it easier to find and recommend similar music.

Under the hood, KMeans follows a simple but powerful iterative process to find good cluster groupings:

Initialization: It randomly picks k points from the dataset to act as the first "centroids" (these are like the centers of each cluster).
Assignment Step: Every data point (in our case, every track embedding) is assigned to the nearest centroid — using Euclidean distance (straight-line distance in vector space).
Update Step: Once all points are assigned, each centroid is moved to the center of the cluster — that is, it’s recalculated as the average of all points in that group.
Repeat: Steps 2 and 3 are repeated until the centroids stop moving (or only move very little). This means the algorithm has converged to a stable solution.

This helps you understand that KMeans doesn’t magically “know” what a cluster is. Instead, it’s optimizing a mathematical goal: minimize the sum of squared distances within each cluster. It’s entirely based on numerical similarity, which is why having good embeddings is so important — bad embeddings = bad clusters, no matter how good your KMeans implementation is.

🤔 But How Many Clusters Should You Use? Choosing the number of clusters (n_clusters) can be tricky. There's no perfect answer—it depends on the data. For a small number of tracks, 2–5 clusters might work. In real-world applications, you might:

Experiment with different values and evaluate results visually (e.g., with PCA or t-SNE).
Use metrics like the elbow method or silhouette score to guide your choice.
Start small and scale up once you have more data and a better understanding of your embedding space.

Clustering Tracks in Code

Let’s walk through the main function that clusters tracks using KMeans: assign_track_clusters. Here is the code, broken down into key steps:

Python

import pandas as pd
from sklearn.cluster import KMeans
from src.user_model import get_track_embeddings
from src.database import get_all_tracks

_clustered_tracks_df = None

def assign_track_clusters(n_clusters_requested: int = 3):
    """
    Clusters tracks based on their embeddings using KMeans.
    Returns a DataFrame of tracks with an added 'cluster' column.
    Caches the result.
    """
    global _clustered_tracks_df
    if _clustered_tracks_df is not None:
        # Check if n_clusters matches cached version, simplistic check
        if 'cluster' in _clustered_tracks_df and \
           _clustered_tracks_df['cluster'].nunique() == n_clusters_requested:
            return _clustered_tracks_df

    track_ids, track_embeddings_matrix = get_track_embeddings()
    
    if not track_ids or track_embeddings_matrix.shape[0] == 0:
        _clustered_tracks_df = pd.DataFrame() # Cache empty df
        return _clustered_tracks_df

    # Adjust n_clusters if it's more than the number of samples
    num_samples = track_embeddings_matrix.shape[0]
    actual_n_clusters = min(n_clusters_requested, num_samples)
    
    if actual_n_clusters <= 0: # Should not happen if num_samples > 0
        _clustered_tracks_df = get_all_tracks().copy()
        if not _clustered_tracks_df.empty:
             _clustered_tracks_df['cluster'] = -1 # Indicate no clustering done
        return _clustered_tracks_df

    kmeans = KMeans(n_clusters=actual_n_clusters, random_state=42, n_init='auto')
    try:
        cluster_labels = kmeans.fit_predict(track_embeddings_matrix)
    except ValueError as e: # Catch potential errors if num_samples is too small for n_init
        print(f"KMeans fitting error: {e}. Assigning default cluster.")
        cluster_labels = [-1] * num_samples # Default cluster if KMeans fails

    # Combine cluster labels with original track data
    tracks_df = get_all_tracks().copy() # Get a fresh copy of all track details
    
    # Create a mapping from track_id to cluster_label
    # This assumes track_ids from get_track_embeddings are in the same order as rows in tracks_df
    # A more robust way is to merge based on ID if orders can mismatch.
    if len(track_ids) == len(tracks_df) and list(tracks_df['id']) == track_ids:
        tracks_df['cluster'] = cluster_labels
    else:
        # Fallback: create a temporary df for merging if orders/lengths differ
        cluster_df = pd.DataFrame({'id': track_ids, 'cluster': cluster_labels})
        tracks_df = pd.merge(tracks_df, cluster_df, on='id', how='left')
        tracks_df['cluster'] = tracks_df['cluster'].fillna(-1).astype(int) # Fill missing clusters

    _clustered_tracks_df = tracks_df
    return _clustered_tracks_df

import pandas as pd
from sklearn.cluster import KMeans
from src.user_model import get_track_embeddings
from src.database import get_all_tracks

_clustered_tracks_df = None

def assign_track_clusters(n_clusters_requested: int = 3):
    """
    Clusters tracks based on their embeddings using KMeans.
    Returns a DataFrame of tracks with an added 'cluster' column.
    Caches the result.
    """
    global _clustered_tracks_df
    if _clustered_tracks_df is not None:
        # Check if n_clusters matches cached version, simplistic check
        if 'cluster' in _clustered_tracks_df and \
           _clustered_tracks_df['cluster'].nunique() == n_clusters_requested:
            return _clustered_tracks_df

    track_ids, track_embeddings_matrix = get_track_embeddings()
    
    if not track_ids or track_embeddings_matrix.shape[0] == 0:
        _clustered_tracks_df = pd.DataFrame() # Cache empty df
        return _clustered_tracks_df

    # Adjust n_clusters if it's more than the number of samples
    num_samples = track_embeddings_matrix.shape[0]
    actual_n_clusters = min(n_clusters_requested, num_samples)
    
    if actual_n_clusters <= 0: # Should not happen if num_samples > 0
        _clustered_tracks_df = get_all_tracks().copy()
        if not _clustered_tracks_df.empty:
             _clustered_tracks_df['cluster'] = -1 # Indicate no clustering done
        return _clustered_tracks_df

    kmeans = KMeans(n_clusters=actual_n_clusters, random_state=42, n_init='auto')
    try:
        cluster_labels = kmeans.fit_predict(track_embeddings_matrix)
    except ValueError as e: # Catch potential errors if num_samples is too small for n_init
        print(f"KMeans fitting error: {e}. Assigning default cluster.")
        cluster_labels = [-1] * num_samples # Default cluster if KMeans fails

    # Combine cluster labels with original track data
    tracks_df = get_all_tracks().copy() # Get a fresh copy of all track details
    
    # Create a mapping from track_id to cluster_label
    # This assumes track_ids from get_track_embeddings are in the same order as rows in tracks_df
    # A more robust way is to merge based on ID if orders can mismatch.
    if len(track_ids) == len(tracks_df) and list(tracks_df['id']) == track_ids:
        tracks_df['cluster'] = cluster_labels
    else:
        # Fallback: create a temporary df for merging if orders/lengths differ
        cluster_df = pd.DataFrame({'id': track_ids, 'cluster': cluster_labels})
        tracks_df = pd.merge(tracks_df, cluster_df, on='id', how='left')
        tracks_df['cluster'] = tracks_df['cluster'].fillna(-1).astype(int) # Fill missing clusters

    _clustered_tracks_df = tracks_df
    return _clustered_tracks_df

Let’s break down what happens here:

The function first checks if it already has a cached result for the requested number of clusters. If so, it returns that to save time.
It gets the track embeddings and checks if there are data to cluster.
It makes sure the number of clusters does not exceed the number of tracks.
It runs KMeans to assign each track to a cluster.
It attaches the cluster labels to the track data, making sure each track has a cluster column.
If something goes wrong (like not enough tracks), it handles the error gracefully.

Additional Clarifications

The get_track_embeddings() function is used twice: once for the embedding matrix and again to ensure that track_ids are aligned with the embeddings. This alignment is crucial because the KMeans output (cluster_labels) will map directly to the rows of the embedding matrix.
The fallback merging logic (pd.merge(...)) is important when there's a mismatch in order or number of track IDs. Without it, cluster labels could be assigned to the wrong tracks, leading to misleading results.
The random_state=42 in KMeans ensures reproducibility. It’s helpful during testing or debugging, especially when clusters appear inconsistent across runs.
Caching with _clustered_tracks_df avoids recomputation, which is important when working with large datasets or limited resources. However, developers should be cautious—if the underlying embeddings change but the number of clusters remains the same, the cache won’t refresh. In production, a more robust cache invalidation strategy should be considered.

Observing the Outputs

Here’s an example of how you might use the assign_track_clusters function and what the output could look like:

Python

clustered_data = assign_track_clusters(n_clusters_requested=2)
print(clustered_data[['id', 'title', 'cluster']].head())

clustered_data = assign_track_clusters(n_clusters_requested=2)
print(clustered_data[['id', 'title', 'cluster']].head())

Here’s a real example showing 5 tracks clustered into 3 groups::

text

         id            title  cluster
0  track001   Sunrise Melody        2
1  track002  Midnight Cruise        0
2  track003      Forest Path        0
3  track004     Quantum Leap        1
4  track005       Neon Skies        2

         id            title  cluster
0  track001   Sunrise Melody        2
1  track002  Midnight Cruise        0
2  track003      Forest Path        0
3  track004     Quantum Leap        1
4  track005       Neon Skies        2

We also verify that each cluster contains a reasonable number of tracks:

text

Cluster distribution (requested=3):
cluster
2    2
0    2
1    1

Cluster distribution (requested=3):
cluster
2    2
0    2
1    1

Even if the number of clusters you request is more than the number of tracks, the system handles it gracefully. This is done because KMeans requires the number of clusters to be less than or equal to the number of samples. If you ask for 10 clusters with only 5 tracks, the system quietly reduces it to 5. Otherwise, scikit-learn would raise a ValueError. This fallback protects the system from crashing and provides a helpful behavior for testing small datasets. For example:

text

Testing edge case: more clusters than tracks...
✅ Assigned 5 clusters (capped by number of tracks).

Testing edge case: more clusters than tracks...
✅ Assigned 5 clusters (capped by number of tracks).

And the final cluster assignment table will look like this:

text

=== Final Cluster Assignment Preview ===
         id            title  cluster
1  track002  Midnight Cruise        0
3  track004     Quantum Leap        1
0  track001   Sunrise Melody        2
2  track003      Forest Path        3
4  track005       Neon Skies        4

=== Final Cluster Assignment Preview ===
         id            title  cluster
1  track002  Midnight Cruise        0
3  track004     Quantum Leap        1
0  track001   Sunrise Melody        2
2  track003      Forest Path        3
4  track005       Neon Skies        4

Introduction: Why Cluster Tracks?

Recap: Data and Embedding Setup

What Are Clusters

Join the 1M+ learners on CodeSignal

KMeans Clustering Explained

Clustering Tracks in Code

Observing the Outputs

Why Clustering Helps in Real Music Applications

Summary And What’s Next