Notebooks
G
Google Gemini Clustering With Embeddings

Clustering With Embeddings

gemini-cookbookgemini-apiexamplesgemini
alph-notebooks/gemini-cookbook / clustering_with_embeddings.ipynb
Copyright 2025 Google LLC.
[ ]

Clustering with embeddings

⚠️

This notebook requires paid tier rate limits to run properly.
(cf. pricing for more details).

Overview

This tutorial demonstrates how to visualize and perform clustering with the embeddings from the Gemini API. You will visualize a subset of the 20 Newsgroup dataset using t-SNE{:.external} and cluster that subset using the KMeans algorithm.

For more information on getting started with embeddings generated from the Gemini API, check out the Get Started.

Prerequisites

You can run this quickstart in Google Colab.

To complete this quickstart on your own development environment, ensure that your envirmonement meets the following requirements:

  • Python 3.11+
  • An installation of jupyter to run the notebook.

Setup

First, download and install the Gemini API Python library.

[1]
[ ]

Grab an API Key

Before you can use the Gemini API, you must first obtain an API key. If you don't already have one, create a key with one click in Google AI Studio.

Get an API key

In Colab, add the key to the secrets manager under the "🔑" in the left panel. Give it the name GEMINI_API_KEY.

Once you have the API key, pass it to the SDK. You can do this in two ways:

  • Put the key in the GEMINI_API_KEY environment variable (the SDK will automatically pick it up from there).
  • Pass the key to genai.Client(api_key=...)
[ ]

Key Point: Next, you will choose a model. Any embedding model will work for this tutorial, but for real applications it's important to choose a specific model and stick with it. The outputs of different models are not compatible with each other.

Note: At this time, the Gemini API is only available in certain regions.

[3]
models/embedding-001
models/text-embedding-004
models/gemini-embedding-exp-03-07
models/gemini-embedding-exp
models/gemini-embedding-001

Select the model to be used

[4]
MODEL_ID

Dataset

The 20 Newsgroups Text Dataset{:.external} contains 18,000 newsgroups posts on 20 topics divided into training and test sets. The split between the training and test datasets are based on messages posted before and after a specific date. For this tutorial, you will be using the training subset.

[5]
['alt.atheism',
, 'comp.graphics',
, 'comp.os.ms-windows.misc',
, 'comp.sys.ibm.pc.hardware',
, 'comp.sys.mac.hardware',
, 'comp.windows.x',
, 'misc.forsale',
, 'rec.autos',
, 'rec.motorcycles',
, 'rec.sport.baseball',
, 'rec.sport.hockey',
, 'sci.crypt',
, 'sci.electronics',
, 'sci.med',
, 'sci.space',
, 'soc.religion.christian',
, 'talk.politics.guns',
, 'talk.politics.mideast',
, 'talk.politics.misc',
, 'talk.religion.misc']

Here is the first example in the training set.

[6]
Lines: 15

 I was wondering if anyone out there could enlighten me on this car I saw
the other day. It was a 2-door sports car, looked to be from the late 60s/
early 70s. It was called a Bricklin. The doors were really small. In addition,
the front bumper was separate from the rest of the body. This is 
all I know. If anyone can tellme a model name, engine specs, years
of production, where this car is made, history, or whatever info you
have on this funky looking car, please e-mail.

Thanks,
- IL
   ---- brought to you by your neighborhood Lerxst ----
[7]
[8]

Next, you will sample some of the data by taking 100 data points in the training dataset, and dropping a few of the categories to run through this tutorial. Choose the science categories to compare.

[9]
/tmp/ipykernel_200977/406673449.py:4: FutureWarning: DataFrameGroupBy.apply operated on the grouping columns. This behavior is deprecated, and in a future version of pandas the grouping columns will be excluded from the operation. Either pass `include_groups=False` to exclude the groupings or explicitly select the grouping columns after groupby to silence this warning.
  .apply(lambda x: x.sample(SAMPLE_SIZE))
[10]
Class Name
,sci.crypt          150
,sci.electronics    150
,sci.med            150
,sci.space          150
,Name: count, dtype: int64

Generate the embeddings

In this section, you will see how to generate embeddings for the different texts in the dataframe using the embeddings from the Gemini API.

The Gemini embedding model supports several task types, each tailored for a specific goal. Here’s a general overview of the available types and their applications:

Task TypeDescription
RETRIEVAL_QUERYSpecifies the given text is a query in a search/retrieval setting.
RETRIEVAL_DOCUMENTSpecifies the given text is a document in a search/retrieval setting.
SEMANTIC_SIMILARITYSpecifies the given text will be used for Semantic Textual Similarity (STS).
CLASSIFICATIONSpecifies that the embeddings will be used for classification.
CLUSTERINGSpecifies that the embeddings will be used for clustering.
[11]
/home/lucianomartins/Documents/projects/gemini-launches/gemini-embedding/ge-env/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm
100%|██████████| 600/600 [05:15<00:00,  1.90it/s]

Dimensionality reduction

The dimension of the document embedding vector is 3072. In order to visualize how the embedded documents are grouped together, you will need to apply dimensionality reduction as you can only visualize the embeddings in 2D or 3D space. Contextually similar documents should be closer together in space as opposed to documents that are not as similar.

[12]
3072
[13]
(600, 3072)

You will apply the t-Distributed Stochastic Neighbor Embedding (t-SNE) approach to perform dimensionality reduction. This technique reduces the number of dimensions, while preserving clusters (points that are close together stay close together). For the original data, the model tries to construct a distribution over which other data points are "neighbors" (e.g., they share a similar meaning). It then optimizes an objective function to keep a similar distribution in the visualization.

[15]
[16]
[17]
(np.float64(-45.10566120147705),
, np.float64(47.539869117736814),
, np.float64(-41.033332443237306),
, np.float64(37.77369651794434))
Output

Compare results to KMeans

KMeans clustering{:.external} is a popular clustering algorithm and used often for unsupervised learning. It iteratively determines the best k center points, and assigns each example to the closest centroid. Input the embeddings directly into the KMeans algorithm to compare the visualization of the embeddings to the performance of a machine learning algorithm.

[18]
[19]
[20]
(np.float64(-45.10566120147705),
, np.float64(47.539869117736814),
, np.float64(-41.033332443237306),
, np.float64(37.77369651794434))
Output
[21]
[22]
{'sci.crypt': np.int32(1),
, 'sci.electronics': np.int32(0),
, 'sci.med': np.int32(3),
, 'sci.space': np.int32(2)}

Get the majority of clusters per group, and see how many of the actual members of that group are in that cluster.

[23]
Class Name
,sci.space          0.993333
,sci.med            0.986667
,sci.crypt          0.973333
,sci.electronics    0.960000
,Name: count, dtype: float64
[24]

To better visualize the performance of the KMeans applied to your data, you can use a confusion matrix. The confusion matrix allows you to assess the performance of the classification model beyond accuracy. You can see what misclassified points get classified as. You will need the actual values and the predicted values, which you have gathered in the dataframe above.

[25]
Output

Next steps

You've now created your own visualization of embeddings with clustering! Try using your own textual data to visualize them as embeddings. You can perform dimensionality reduction in order to complete the visualization step. Note that TSNE is good at clustering inputs, but can take a longer time to converge or might get stuck at local minima.

There are other clustering algorithms outside of KMeans as well, such as density-based spatial clustering (DBSCAN).

To learn how to use other services in the Gemini API, see the Get started guide.