Quick tip: Visualise OpenAI Vector Embeddings using Plotly Express

Abstract

This article demonstrates how to visualise OpenAI vector embeddings for a search term using t-SNE and Plotly Express. We build on the work from a previous article, where we showed how to adapt an OpenAI example to work with SingleStoreDB. With some minor code modifications, we can use the same example to visualise vector embeddings for a search term.

The notebook file used in this article is available on GitHub.

Introduction

In a great article, the author demonstrates how to visualise vector embeddings using several technologies. We can use a previous OpenAI example dataset, simplify the code, and use Plotly Express to render a similar visualisation. Let's see how.

As described in a previous article, we'll follow the instructions to create a Notebook.

Fill out the Notebook

First, we'll install some libraries:

!pip install openai --quiet !pip install scikit-learn --quiet !pip install wget --quiet 

and imports:

import ast import getpass import numpy as np import os import pandas as pd import plotly.express as px import wget from openai import OpenAI from scipy.spatial.distance import cdist from sklearn.manifold import TSNE from sklearn.preprocessing import MinMaxScaler 

Next, we'll set our OpenAI API Key:

os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:") 

We'll now download a CSV file from OpenAI that contains text and embeddings related to the Winter Olympics 2022:

embeddings_path = "https://cdn.openai.com/API/examples/data/winter_olympics_2022.csv" file_path = "winter_olympics_2022.csv" if not os.path.exists(file_path): wget.download(embeddings_path, file_path) print("File downloaded successfully") else: print("File already exists in the local file system") 

Now we'll read the file into a Dataframe and convert the data to a NumPy Array:

df = pd.read_csv(file_path) # Convert embeddings from CSV str type to NumPy Array embedding_array = np.array( df["embedding"].apply(ast.literal_eval).to_list() ) 

Our search term is "curling gold medal", and we'll get the vector embeddings for this from OpenAI:

client = OpenAI() query = "curling gold medal" query_embedding_response = client.embeddings.create( input = query, model = "text-embedding-ada-002" ).data[0].embedding 

Now we'll find and store the Euclidean Distance between the search term and the vector embeddings we previously loaded:

df["distance"] = cdist( embedding_array, [query_embedding_response] ) 

and scale the values between 0 and 1, then store them, as follows:

df["normalised"] = MinMaxScaler().fit_transform(df[["distance"]]) 

Finally, we'll create a t-SNE model and plot the data using Plotly Express:

# Create a t-SNE model tsne_model = TSNE( n_components = 2, perplexity = 15, random_state = 42, init = "random", learning_rate = 200 ) tsne_embeddings = tsne_model.fit_transform(embedding_array) # Create a Dataframe for visualisation visualisation_data = pd.DataFrame({ "x": tsne_embeddings[:, 0], "y": tsne_embeddings[:, 1], "Similarity": df["normalised"] }) # Create the scatter plot using Plotly Express plot = px.scatter( visualisation_data, x = "x", y = "y", color = "Similarity", color_continuous_scale = "rainbow", opacity = 0.3, title = f"Similarity to '{query}' visualised using t-SNE" ) plot.update_layout( width = 650, height = 650 ) plot.show() 

The output should be as shown in Figure 1.

Colours specify similarity to the search term. In this case, we can see red areas on the plot that are closer and blue areas that are further away.

Summary

Data visualisation can be used to gain insights into the distribution of data, as seen in a previous article. In this article, we saw how to use vector embeddings and a search term to create a t-SNE model and visualise it using Plotly Express. This simple example showed how to use data visualisation to identify patterns and trends.