Generative AI on Google Cloud

40 / 80·notebook

Chapter 40

Visualizing embedding similarity from text documents using t-SNE plots

NotebookPython 337 cells

In [ ]python · cell 1

python

# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Visualizing embedding similarity from text documents using t-SNE plots


Author(s)	Gabe Rives-Corbett

This notebook demonstrates how vector similarity is relevant to LLM-generated embeddings. You will embed a collection of labelled documents and then plot the embeddings on a two-dimensional t-SNE plot to observe how similar documents tend to cluster together based on their embeddings.

Getting started

Install libraries

In [ ]python · cell 7

python

%pip install --upgrade google-genai scikit-learn pandas seaborn

Authenticate your notebook environment (Colab only)

If you are running this notebook on Google Colab, you will need to authenticate your environment. To do this, run the new cell below. This step is not required if you are using Agent Platform Workbench.

In [ ]python · cell 9

python

import sys

if "google.colab" in sys.modules:
    from google.colab import auth

    auth.authenticate_user()

Set Google Cloud project information and initialize Agent Platform SDK

To get started using Agent Platform, you must have an existing Google Cloud project and enable the Agent Platform API.

Learn more about setting up a project and a development environment.

In [ ]python · cell 11

python

# Use the environment variable if the user doesn't provide Project ID.
import os

from google import genai

# fmt: off
PROJECT_ID = "[your-project-id]"  # @param {type: "string", placeholder: "[your-project-id]", isTemplate: true}
# fmt: on
if not PROJECT_ID or PROJECT_ID == "[your-project-id]":
    PROJECT_ID = str(os.environ.get("GOOGLE_CLOUD_PROJECT"))

LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")

client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)

Import libraries

In [ ]python · cell 13

python

import re

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from google.api_core import retry
from sklearn.datasets import fetch_20newsgroups
from sklearn.manifold import TSNE
from sklearn.model_selection import train_test_split
from tqdm.auto import tqdm

tqdm.pandas()

Fetch and clean the data

In this example, you will use the open source 20 Newsgroups dataset, a collection of approximately 20,000 newsgroup documents, partitioned (nearly) evenly across 20 different newsgroups

In [ ]python · cell 15

python

newsgroups = fetch_20newsgroups(
    categories=["comp.graphics", "sci.space", "sci.med", "rec.sport.hockey"]
)

In [ ]python · cell 16

python

raw_data = pd.DataFrame(
    {
        "text": newsgroups.data,
        "target": [newsgroups.target_names[x] for x in newsgroups.target],
    }
)

Because of the 8k input token limit, in this example you will exclude all documents that have a length outside this limit.

Even though tokens typically are >=1 characters, for simplicity, you can just filter for documents that have <= 8000 characters.

In [ ]python · cell 18

python

raw_data = raw_data.loc[raw_data["text"].str.len() <= 8000]

Subsample the dataset into 500 data points, stratified on the label

In [ ]python · cell 20

python

x_subsample, _, y_subsample, _ = train_test_split(
    raw_data["text"], raw_data["target"], stratify=raw_data["target"], train_size=500
)

Clean out the text removing by emails, names, etc. This will help improve the data that will then be converted into embeddings.

In [ ]python · cell 22

python

x_subsample = [re.sub(r"[\w\.-]+@[\w\.-]+", "", d) for d in x_subsample]  # Remove email
x_subsample = [re.sub(r"\([^()]*\)", "", d) for d in x_subsample]  # Remove names
x_subsample = [d.replace("From: ", "") for d in x_subsample]  # Remove "From: "
x_subsample = [
    d.replace("\nSubject: ", "") for d in x_subsample
]  # Remove "\nSubject: "

In [ ]python · cell 23

python

df = pd.DataFrame({"text": x_subsample, "target": list(y_subsample)})

You now have 500 data points roughly evenly distributed across the categories:

In [ ]python · cell 25

python

df["target"].value_counts()

Create and visualize the embeddings using a t-SNE plot

Load the text embedding model from Agent Platform (documentation).

Since we are using these embeddings for visualization, we will set the task type to clustering.

In [ ]python · cell 27

python

MODEL_ID = "gemini-embedding-001"

In [ ]python · cell 28

python

# Retrieve embeddings from the specified model with retry logic

from google.genai import types

def get_embeddings():
    @retry.Retry(timeout=300.0)
    def embed_fn(contents: str) -> list[float]:
        response = client.models.embed_content(
            model=MODEL_ID,
            contents=contents,
            config=types.EmbedContentConfig(output_dimensionality=768),
        )
        return response.embeddings[0].values

    return embed_fn

Create the embeddings. This may take a minute or two.

In [ ]python · cell 30

python

df["embeddings"] = df["text"].progress_apply(get_embeddings())

In [ ]python · cell 31

python

df.head()

The vectors generate by our model are 768 dimensions, and so visualizing across 768 dimensions is impossible. Instead, you can use t-SNE to reduce to 2 dimensions.

In [ ]python · cell 33

python

tsne = TSNE(random_state=0, max_iter=1000)
tsne_results = tsne.fit_transform(
    np.array(df["embeddings"].to_list(), dtype=np.float32)
)

In [ ]python · cell 34

python

df_tsne = pd.DataFrame(tsne_results, columns=["TSNE1", "TSNE2"])
df_tsne["target"] = df["target"]  # Add labels column from df_train to df_tsne

In [ ]python · cell 35

python

df_tsne.head()

Plot the data points. It should now be visually clear how the documents from the same newsgroup show up close to each other in the vector space with text embeddings.

In [ ]python · cell 37

python

fig, ax = plt.subplots(figsize=(8, 6))  # Set figsize
sns.set_style("darkgrid", {"grid.color": ".6", "grid.linestyle": ":"})
sns.scatterplot(data=df_tsne, x="TSNE1", y="TSNE2", hue="target", palette="hls")
sns.move_legend(ax, "upper left", bbox_to_anchor=(1, 1))
plt.title("Scatter plot of news using t-SNE")
plt.xlabel("TSNE1")
plt.ylabel("TSNE2")
plt.axis("equal")

← 返回课程

Generative AI on Google Cloud

40 / 80·notebook

Chapter 40

Visualizing embedding similarity from text documents using t-SNE plots

NotebookPython 337 cells

In [ ]python · cell 1

python

# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Visualizing embedding similarity from text documents using t-SNE plots


Author(s)	Gabe Rives-Corbett

Getting started

Install libraries

In [ ]python · cell 7

python

%pip install --upgrade google-genai scikit-learn pandas seaborn

Authenticate your notebook environment (Colab only)

In [ ]python · cell 9

python

import sys

if "google.colab" in sys.modules:
    from google.colab import auth

    auth.authenticate_user()

Set Google Cloud project information and initialize Agent Platform SDK

To get started using Agent Platform, you must have an existing Google Cloud project and enable the Agent Platform API.

Learn more about setting up a project and a development environment.

In [ ]python · cell 11

python

# Use the environment variable if the user doesn't provide Project ID.
import os

from google import genai

# fmt: off
PROJECT_ID = "[your-project-id]"  # @param {type: "string", placeholder: "[your-project-id]", isTemplate: true}
# fmt: on
if not PROJECT_ID or PROJECT_ID == "[your-project-id]":
    PROJECT_ID = str(os.environ.get("GOOGLE_CLOUD_PROJECT"))

LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")

client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)

Import libraries

In [ ]python · cell 13

python

import re

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from google.api_core import retry
from sklearn.datasets import fetch_20newsgroups
from sklearn.manifold import TSNE
from sklearn.model_selection import train_test_split
from tqdm.auto import tqdm

tqdm.pandas()

Fetch and clean the data

In this example, you will use the open source 20 Newsgroups dataset, a collection of approximately 20,000 newsgroup documents, partitioned (nearly) evenly across 20 different newsgroups

In [ ]python · cell 15

python

newsgroups = fetch_20newsgroups(
    categories=["comp.graphics", "sci.space", "sci.med", "rec.sport.hockey"]
)

In [ ]python · cell 16

python

raw_data = pd.DataFrame(
    {
        "text": newsgroups.data,
        "target": [newsgroups.target_names[x] for x in newsgroups.target],
    }
)

Because of the 8k input token limit, in this example you will exclude all documents that have a length outside this limit.

Even though tokens typically are >=1 characters, for simplicity, you can just filter for documents that have <= 8000 characters.

In [ ]python · cell 18

python

raw_data = raw_data.loc[raw_data["text"].str.len() <= 8000]

Subsample the dataset into 500 data points, stratified on the label

In [ ]python · cell 20

python

x_subsample, _, y_subsample, _ = train_test_split(
    raw_data["text"], raw_data["target"], stratify=raw_data["target"], train_size=500
)

Clean out the text removing by emails, names, etc. This will help improve the data that will then be converted into embeddings.

In [ ]python · cell 22

python

x_subsample = [re.sub(r"[\w\.-]+@[\w\.-]+", "", d) for d in x_subsample]  # Remove email
x_subsample = [re.sub(r"\([^()]*\)", "", d) for d in x_subsample]  # Remove names
x_subsample = [d.replace("From: ", "") for d in x_subsample]  # Remove "From: "
x_subsample = [
    d.replace("\nSubject: ", "") for d in x_subsample
]  # Remove "\nSubject: "

In [ ]python · cell 23

python

df = pd.DataFrame({"text": x_subsample, "target": list(y_subsample)})

You now have 500 data points roughly evenly distributed across the categories:

In [ ]python · cell 25

python

df["target"].value_counts()

Create and visualize the embeddings using a t-SNE plot

Load the text embedding model from Agent Platform (documentation).

Since we are using these embeddings for visualization, we will set the task type to clustering.

In [ ]python · cell 27

python

MODEL_ID = "gemini-embedding-001"

In [ ]python · cell 28

python

# Retrieve embeddings from the specified model with retry logic

from google.genai import types

def get_embeddings():
    @retry.Retry(timeout=300.0)
    def embed_fn(contents: str) -> list[float]:
        response = client.models.embed_content(
            model=MODEL_ID,
            contents=contents,
            config=types.EmbedContentConfig(output_dimensionality=768),
        )
        return response.embeddings[0].values

    return embed_fn

Create the embeddings. This may take a minute or two.

In [ ]python · cell 30

python

df["embeddings"] = df["text"].progress_apply(get_embeddings())

In [ ]python · cell 31

python

df.head()

The vectors generate by our model are 768 dimensions, and so visualizing across 768 dimensions is impossible. Instead, you can use t-SNE to reduce to 2 dimensions.

In [ ]python · cell 33

python

tsne = TSNE(random_state=0, max_iter=1000)
tsne_results = tsne.fit_transform(
    np.array(df["embeddings"].to_list(), dtype=np.float32)
)

In [ ]python · cell 34

python

df_tsne = pd.DataFrame(tsne_results, columns=["TSNE1", "TSNE2"])
df_tsne["target"] = df["target"]  # Add labels column from df_train to df_tsne

In [ ]python · cell 35

python

df_tsne.head()

Plot the data points. It should now be visually clear how the documents from the same newsgroup show up close to each other in the vector space with text embeddings.

In [ ]python · cell 37

python

fig, ax = plt.subplots(figsize=(8, 6))  # Set figsize
sns.set_style("darkgrid", {"grid.color": ".6", "grid.linestyle": ":"})
sns.scatterplot(data=df_tsne, x="TSNE1", y="TSNE2", hue="target", palette="hls")
sns.move_legend(ax, "upper left", bbox_to_anchor=(1, 1))
plt.title("Scatter plot of news using t-SNE")
plt.xlabel("TSNE1")
plt.ylabel("TSNE2")
plt.axis("equal")