Visualising Word Embeddings Clusters in Jupyter Notebook (python) using t-SNE and k-means methods

Co-authored by Tobi Okusanya.

This article shows how to take a locally stored Excel (or .csv/.txt) file and extract one column of words (entities) and convert each entry into a word embedding using a pre-trained model (GloVe). It then clusters these embeddings using the t-SNE algorithm which visualises the results on a 2D plot.

A silhouette score method is then used to determine the optimal number of clusters, k.

Finally, the t-SNE plot is then updated with the new optimised number of clusters k, with colour-coded clusters. Meanwhile the entries within each cluster are identified and printed out.

Your requirements file should contain the following:

nltk
pandas 
numpy==1.26.3 
scikit-learn
matplotlib
utils
seaborn

Install it from the Terminal by running the following command:

pip install -r requirements.txt

Then we add our dependencies:

from nltk.tokenize import sent_tokenize, word_tokenize
import warnings
import os, json
from glob import glob
import numpy as np
np.set_printoptions(precision=4, linewidth=100)
from matplotlib import pyplot as plt
%matplotlib inline
import utils
from utils import *
import csv
import string
import random
import pandas as pd

import sklearn
from sklearn.manifold import TSNE
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
from sklearn.decomposition import PCA
import seaborn as sns

from collections import defaultdict

warnings.filterwarnings(action='ignore')

We are using the GloVe word embeddings model. These are available to download from here:

GloVe model link

You must download the GloVe model and save it into the same directory as your Notebook for it to work.

You could also use another word embeddings model instead of GloVe if desired.

There are several GloVe models you can use, varying in size, how they were trained and the number of dimensions. We define the number of dimensions in the code below. We chose a 300D model but this can be changed if desired. If you use a model with a different number of dimensions, you must change the code parameter accordingly.

You will probably not want to upload the GloVe model to git as the .txt file is very large so remember to delete it if so. Make a copy and save it elsewhere on your machine and copy it back over when you need to run the code again.

num_dimensions_emb = 300

Now we define a function load_glove_model that reads in each word in the GloVe model and its corresponding word embedding and populates a dictionary with these two items. If you use a different GloVe model then you should change the local file name appropriately.

def load_glove_model(File):
    print("Loading Glove Model")
    glove_model = {}
    with open(File,'r') as f:
        for line in f:
            split_line = line.split(' ')
            word = split_line[0]
            embeddings_array = np.array(split_line[1:], dtype=np.float64)
            glove_model[word] = embeddings_array
    print(f"{len(glove_model)} words loaded!")
    return glove_model

glove_model = load_glove_model('glove.840B.300d.txt')

Now we define a function generate_embeddings that takes in a list of entries and returns the same list of entries and a list of the word embeddings of those entries.

def generate_embeddings(entries_list):

    # We need this random seed to enable the np.random.rand() method further down this function
    np.random.seed(7)

    entry_embeddings_list = []

    for entry in entries_list:

        entry_embedding = np.zeros(num_dimensions_emb)
        entry = entry.replace('_' , ' ')
        entry = entry.split(" ")
        for word in entry:
            try:
                entry_embedding_temp = glove_model[word]
            except:
                # If the word in entry doesn't exist in the GloVe model. assign it a random embedding
                entry_embedding_temp = np.random.rand(num_dimensions_emb)
            entry_embedding+=entry_embedding_temp
        entry_embedding = entry_embedding/len(entry)
        #The previous two lines of code find the mean of the composite entry embedding
        entry_embeddings_list.append(entry_embedding)

    return entries_list, entry_embeddings_list

Now we read in a locally stored Excel file that contains our entries and assign it to a pandas dataframe df_tsne.

You need to ensure this file is stored in the same folder as your Notebook, and that you edit the file name in the cell below accordingly.

filename = <insert name of file here>

Note: the code below is for reading in an Excel file. This could easily be modified to ingest e.g. a .csv or .txt file.

df_tsne = pd.read_excel(filename)

If the file contains more than one column, you will need to run the cell below to grab only that column from the dataframe.

You should change the column name ‘entry’ in the cell below to the one you need in your local file for your use case.

df_entry_tsne = df_tsne['entry']

Then we want to only grab the unique entries from this column, since we don’t want repeat entries.

unique_entries_tsne = df_entry_tsne.unique().tolist()

Call the main embeddings function from above and convert the output to arrays for processing.

entries, embeddings = generate_embeddings(unique_entries_tsne)

entries_array = np.array(entries)
embeddings_array = np.array(embeddings)

NOTE: The plot generated will be in general different each time qualitatively due to the random state value.

Therefore you should run the cell below a few times, to ensure that we get good clustering. This should be done with common sense in mind: are the correct entities being clustered together roughly? If not, run the cell below again with different random state values until they are.

tsne = TSNE(n_components=2, random_state=1, perplexity=1) #edit the random state value to different integer here
Y_tsne = tsne.fit_transform(embeddings_array)

start=0; end=len(embeddings_array)
dat = Y_tsne[start:end]
plt.figure(figsize=(15,15))
plt.scatter(dat[:, 0], dat[:, 1])
for label, x, y in zip(entries_array[start:end], dat[:, 0], dat[:, 1]):
    plt.text(x,y,label, color=np.random.rand(3)*0.7,
                 fontsize=14)
plt.show()

Now let’s find out the optimal number of clusters for the t-SNE method.

In order to do this, we plot the number of clusters, k against the average silhouette score.

The silhoutte score is a metric used to evaluate the quality of the clusters, with it ranging from -1 to +1 whereby +1 conveys that a given cluster is well-defined. We then average each silhouette score over all clusters, for different values of k (typically between 2 and 20 or so). Consequently, the closer the average score is to +1 indicates the most optimal k value.

Hence the peak value of the average silhouette score in this plot corresponds to the optimal number of clusters k to be used in the TSNE method.

k_values = range(2, 20) 
silhouette_scores_emb_tsne = []

for k in k_values:
    kmeans_tsne = KMeans(n_clusters=k, random_state=42)
    cluster_labels_tsne = kmeans_tsne.fit_predict(Y_tsne)
    silhouette_avg_emb_tsne = silhouette_score(Y_tsne, cluster_labels_tsne)

    silhouette_scores_emb_tsne.append(silhouette_avg_emb_tsne)

plt.figure(figsize=(8, 5))
plt.plot(k_values, silhouette_scores_emb_tsne, marker='o')
plt.xlabel('Number of clusters (k)')
plt.ylabel('Silhouette Score Average')
plt.title('Silhouette Method for Optimal k')
plt.savefig("silhouette_plot.png")
plt.show()

# fig2 = plt_silhouette.get_figure() #if you want to save the plot, do it in these two last lines
# fig2.savefig("silhouette_plot.png")

Example of the silhouette plot generated by the cell above

Once you’ve determined the optimal number of clusters conveyed by the silhouette score plot, it is recommended to refer back to the t-SNE plot above in order to verify if it makes sense.

Now we need to input the number of clusters (you get this by simply reading from the silhouette plot generated above the value of k that gives the maximum silhouette score).

num_clusters = <insert_optimised_number_of_clusters_here>

Then we replot the t-SNE data but with the optimal number of clusters information also captured.

kmeans_tsne = KMeans(n_clusters=num_clusters, random_state=42)
cluster_labels_tsne = kmeans_tsne.fit_predict(Y_tsne)

pca = PCA(2)
pca_tsne = pca.fit_transform(Y_tsne)
df_pca_tsne = pd.DataFrame(pca_tsne)

plt.figure(figsize=(5,5))
sns_plot = sns.scatterplot(
    x=dat[:, 0], y=dat[:, 1],
    hue=cluster_labels_tsne,
    palette=sns.color_palette("hls", 10),
    data=df_pca_tsne,       
    legend="full",
    alpha=1
)

fig = sns_plot.get_figure()

Example of the clustering plot generated by the cell above

The next cell will save the plot into the current directory as something appropriate (you must input the file name below).

plotname = <insert_tsne_plot_file_name_here>
fig.savefig(plotname)

This final cell outputs all the entries in each cluster.

cluster_to_entries_tsne = defaultdict(list)
for entry, label in zip(entries_array, cluster_labels_tsne):
    cluster_to_entries_tsne[label].append(entry)
    # print(cluster_labels_tsne)

for cluster_id, entries in cluster_to_entries_tsne.items():
    print(f"Cluster {cluster_id}:")
    for entry in entries:
        print(f"  - {entry}")

And that’s it! Have fun and let me know how you get on.