Generative AI Google Brain

embeddings-embedding-similarity-visualization

3 min read

Visualizing embedding similarity from text documents using t-SNE plots

Source notebook

Repo path: embeddings/embedding-similarity-visualization.ipynb · Open on GitHub · intro

Embeds 20 Newsgroups text with gemini-embedding-001 and visualizes similarity clusters using t-SNE.

Summary

This notebook teaches how vector similarity appears in LLM-generated text embeddings. It fetches labeled 20 Newsgroups documents, filters and cleans the text, creates 768-dimensional embeddings with gemini-embedding-001, reduces them to two dimensions with t-SNE, and plots category clusters with seaborn.

Key code patterns

Initialize GenAI client for Vertex AI

from google import genai
 
PROJECT_ID = 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)

Configures google-genai to call Vertex AI-hosted models in a Google Cloud project and region.

Create text embeddings with retry

from google.api_core import retry
from google.genai import types
 
@retry.Retry(timeout=300.0)
def embed_fn(contents: str) -> list[float]:
    response = client.models.embed_content(
        model="gemini-embedding-001",
        contents=contents,
        config=types.EmbedContentConfig(output_dimensionality=768),
    )
    return response.embeddings[0].values

Adds retry handling around embedding calls and requests 768-dimensional vectors.

Reduce embeddings with t-SNE

tsne = TSNE(random_state=0, max_iter=1000)
tsne_results = tsne.fit_transform(
    np.array(df["embeddings"].to_list(), dtype=np.float32)
)
df_tsne = pd.DataFrame(tsne_results, columns=["TSNE1", "TSNE2"])
df_tsne["target"] = df["target"]

Turns high-dimensional embedding vectors into two coordinates for visual inspection.

Plot labeled clusters

fig, ax = plt.subplots(figsize=(8, 6))
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.axis("equal")

Visualizes whether documents with the same newsgroup label appear close in embedding space.

Models & APIs used

  • Models: gemini-embedding-001
  • APIs / services: Vertex AI
  • SDKs / libraries: google-genai, scikit-learn, pandas, seaborn, matplotlib, numpy, google-api-core, tqdm

When to use this

Use this pattern to sanity-check whether text embeddings group labeled documents by semantic similarity before building clustering, retrieval, or classification workflows.

Gotchas & caveats

  • Requires an existing Google Cloud project and the Agent Platform API enabled.
  • Colab requires explicit user authentication; Agent Platform Workbench does not require that Colab auth cell.
  • The notebook filters documents to 8000 characters because of the stated 8k input token limit.
  • Embedding generation over 500 documents may take a minute or two.
  • LOCATION defaults to us-central1 when GOOGLE_CLOUD_REGION is not set.

Best practices

  • Use environment variables for GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_REGION when notebook parameters are not provided.
  • Subsample with stratification so labels are roughly evenly distributed.
  • Clean emails, names, From headers, and Subject markers before embedding text.
  • Use retry logic around embedding requests with a 300 second timeout.
  • Set output_dimensionality to 768 before applying t-SNE for visualization.

Graph View

Backlinks