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.
Top comments (0)