使用嵌入向量进行聚类分析：Google Colab 教程

本文以 Gemini API 生成的文本嵌入为核心，演示从 20 Newsgroups 数据集采样、清洗、生成嵌入，到用 t-SNE 降维可视化、KMeans 聚类并借助混淆矩阵评估的完整流程。教程强调嵌入模型选择、task_type 设置、维度与降维取舍、聚类与分类评估差异，并给出环境准备、代码解析、常见错误排查与生产优化建议，帮助开发者把嵌入聚类落地到真实文本主题发现场景。

技术背景

在真实业务中，大量文本没有标签，人工归类成本高且难以扩展。传统 TF-IDF 或词袋模型只能捕捉词面重合，无法理解同义、改写和上下文语义，导致主题发现效果差。嵌入模型把文本映射到高维语义空间，使语义相近的文本距离更近，从而让聚类、检索、分类等下游任务可以直接在向量空间中进行。Gemini 嵌入模型提供 CLUSTERING、SEMANTIC_SIMILARITY、CLASSIFICATION、RETRIEVAL_QUERY、RETRIEVAL_DOCUMENT 等任务类型，不同任务类型会改变向量分布，因此必须与下游目标匹配。本教程解决的是“无监督文本主题发现”问题，适用于新闻聚合、工单归类、评论洞察、知识库去重等场景。主要限制包括：嵌入模型有速率与计费限制，不同模型向量不兼容，高维向量需要降维才能可视化，聚类结果需要人工解释，且 t-SNE 可能收敛慢或陷入局部最优。

核心原理

核心数据流是：原始文本 -> 清洗 -> 采样 -> 调用 embed_content 生成 3072 维向量 -> 转为 float32 矩阵 -> t-SNE 降到二维 -> KMeans 聚类 -> 用多数投票把簇映射回类别 -> 混淆矩阵评估。嵌入模型把每个文本编码为固定长度向量，语义相似文本在向量空间中距离更近。task_type='clustering' 会针对聚类目标调整向量分布，使簇结构更明显。t-SNE 通过在高维空间构建邻居概率分布，再在低维空间优化相似分布，从而保留局部邻接关系，适合可视化但不适合作为最终聚类特征。KMeans 迭代寻找 k 个中心点并最小化簇内平方和，直接作用于原始嵌入矩阵，因此其聚类质量反映嵌入本身的语义可分性。多数投票映射把无监督簇标签转换为可解释类别，再用混淆矩阵观察错分去向，比单一准确率更能暴露类别混淆。

架构分析

整体架构分为四层：数据层负责加载 20 Newsgroups、正则清洗邮箱/姓名/From/Subject 并过滤超长文本；嵌入层封装 make_embed_text_fn，通过 retry.Retry 增加超时重试，并用 tqdm.progress_apply 批量生成向量；分析层把嵌入列表转为 np.float32 矩阵，先做 t-SNE 降维用于可视化，再直接对原始高维向量做 KMeans；评估层用 get_majority_cluster_per_group 建立类别到簇的映射，生成 Predicted 列并绘制混淆矩阵。关键设计决策是：可视化与聚类使用不同输入，t-SNE 结果只用于画图，KMeans 使用原始嵌入，避免降维信息损失；采样时按 Label 分组各取 150 条并只保留 sci 类别，保证类别均衡且规模可控；模型 ID 固定为 gemini-embedding-001，避免混用不同模型导致向量空间不一致。

完整实战教程

##### 版权所有 2026 Google LLC。

# @title Licensed under the Apache License, Version 2.0 (the "License");
#
# 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.

使用嵌入进行聚类

<a target="_blank" href="https://colab.research.google.com/github/google-gemini/cookbook/blob/main/examples/clustering_with_embeddings.ipynb">

</a>

注意： 本笔记本需要付费层级的速率限制才能正常运行。（更多详情请参阅定价）。

概述

本教程演示如何使用 Gemini API 生成的嵌入向量进行可视化和聚类。你将使用 t-SNE{:.external} 对 20 Newsgroup 数据集的一个子集进行可视化，并使用 KMeans 算法对该子集进行聚类。

有关如何开始使用 Gemini API 生成的嵌入向量的更多信息，请查看快速入门。

前提条件

你可以在 Google Colab 中运行此快速入门。

要在你自己的开发环境中完成此快速入门，请确保你的环境满足以下要求：

设置

首先，下载并安装 Gemini API Python 库。

%pip install -U -q google-genai seaborn scikit-learn pandas matplotlib

import re
import tqdm
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

from google import genai
from google.genai import types

# Used to securely store your API key
from google.colab import userdata

from sklearn.datasets import fetch_20newsgroups
from sklearn.manifold import TSNE
from sklearn.cluster import KMeans
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

获取 API 密钥

在使用 Gemini API 之前，你必须先获取一个 API 密钥。如果你还没有，可以在 Google AI Studio 中一键创建密钥。

<a class="button button-primary" href="https://aistudio.google.com/app/apikey" target="_blank" rel="noopener noreferrer">获取 API 密钥</a>

在 Colab 中，将密钥添加到左侧面板“🔑”下的密钥管理器中。将其命名为 GEMINI_API_KEY。

获取 API 密钥后，将其传递给 SDK。你可以通过两种方式完成：

* 将密钥放入 GEMINI_API_KEY 环境变量中（SDK 会自动从那里读取）。

* 将密钥传递给 genai.Client(api_key=...)

# Or use `os.getenv('GEMINI_API_KEY')` to fetch an environment variable.
GEMINI_API_KEY=userdata.get('GEMINI_API_KEY')

client = genai.Client(api_key=GEMINI_API_KEY)

关键点：接下来，你需要选择一个模型。任何嵌入模型都适用于本教程，但在实际应用中，重要的是选择一个特定的模型并坚持使用。不同模型的输出彼此不兼容。

注意：目前，Gemini API 仅在特定地区可用。

for m in client.models.list():
  if 'embedContent' in m.supported_actions:
    print(m.name)

选择要使用的模型

MODEL_ID = "gemini-embedding-001" # @param ["gemini-embedding-2-preview", "gemini-embedding-001"] {"allow-input":true, isTemplate: true}

数据集

20 Newsgroups 文本数据集{:.external} 包含 18,000 篇新闻组帖子，涵盖 20 个主题，分为训练集和测试集。训练集与测试集的划分基于帖子发布时间是否早于或晚于某个特定日期。在本教程中，你将使用训练子集。

newsgroups_train = fetch_20newsgroups(subset='train')

# View list of class names for dataset
newsgroups_train.target_names

这是训练集中的第一个示例。

idx = newsgroups_train.data[0].index('Lines')
print(newsgroups_train.data[0][idx:])

# Apply functions to remove names, emails, and extraneous words from data points in newsgroups.data
newsgroups_train.data = [re.sub(r'[\w\.-]+@[\w\.-]+', '', d) for d in newsgroups_train.data] # Remove email
newsgroups_train.data = [re.sub(r"\([^()]*\)", "", d) for d in newsgroups_train.data] # Remove names
newsgroups_train.data = [d.replace("From: ", "") for d in newsgroups_train.data] # Remove "From: "
newsgroups_train.data = [d.replace("\nSubject: ", "") for d in newsgroups_train.data] # Remove "\nSubject: "

# Put training points into a dataframe
df_train = pd.DataFrame(newsgroups_train.data, columns=['Text'])
df_train['Label'] = newsgroups_train.target
# Match label to target name index
df_train['Class Name'] = df_train['Label'].map(newsgroups_train.target_names.__getitem__)
# Retain text samples that can be used in the gecko model.
df_train = df_train[df_train['Text'].str.len() < 10000]

df_train

接下来，你将通过从训练数据集中抽取 100 个数据点来对部分数据进行采样，并删除一些类别以运行本教程。选择要比较的科学类别。

# Take a sample of each label category from df_train
SAMPLE_SIZE = 150
df_train = (df_train.groupby('Label', as_index = False)
                    .apply(lambda x: x.sample(SAMPLE_SIZE))
                    .reset_index(drop=True))

# Choose categories about science
df_train = df_train[df_train['Class Name'].str.contains('sci')]

# Reset the index
df_train = df_train.reset_index()
df_train

df_train['Class Name'].value_counts()

生成嵌入

在本节中，你将了解如何使用 Gemini API 的嵌入模型为数据框中的不同文本生成嵌入。

Gemini 嵌入模型支持多种任务类型，每种类型都针对特定目标而设计。以下是可用类型及其应用场景的总体概览：

任务类型 | 描述

from tqdm.auto import tqdm
tqdm.pandas()

from google.api_core import retry

def make_embed_text_fn(model):

  @retry.Retry(timeout=300.0)
  def embed_fn(text: str) -> list[float]:
    # Set the task_type to CLUSTERING.
    result = client.models.embed_content(model=model,
                                            contents=text,
                                            config=types.EmbedContentConfig(
                                                task_type="clustering"))
    return np.array(result.embeddings[0].values)

  return embed_fn

def create_embeddings(df):
  df['Embeddings'] = df['Text'].progress_apply(make_embed_text_fn(MODEL_ID))
  return df

df_train = create_embeddings(df_train)

降维

文档嵌入向量的维度是 3072。为了可视化嵌入文档是如何分组在一起的，你需要应用降维，因为你只能在 2D 或 3D 空间中可视化嵌入。上下文相似的文档在空间中应该比不相似的文档更接近。

len(df_train['Embeddings'][0])

# Convert df_train['Embeddings'] Pandas series to a np.array of float32
X = np.array(df_train['Embeddings'].to_list(), dtype=np.float32)
X.shape

你将应用 t-分布随机邻域嵌入（t-SNE）方法来进行降维。该技术减少了维度数量，同时保留了聚类（彼此靠近的点仍然保持靠近）。对于原始数据，模型尝试构建一个关于哪些其他数据点是“邻居”的分布（例如，它们具有相似的含义）。然后，它优化一个目标函数，以在可视化中保持相似的分布。

tsne = TSNE(random_state=0)
tsne_results = tsne.fit_transform(X)

df_tsne = pd.DataFrame(tsne_results, columns=['TSNE1', 'TSNE2'])
df_tsne['Class Name'] = df_train['Class Name'] # Add labels column from df_train to df_tsne
df_tsne

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='Class Name', 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')

Notebook output

将结果与 KMeans 进行比较

KMeans 聚类{:.external}是一种流行的聚类算法，常用于无监督学习。它迭代地确定最佳的 k 个中心点，并将每个样本分配给最近的中心点。将嵌入直接输入 KMeans 算法，以将嵌入的可视化与机器学习算法的性能进行比较。

# Apply KMeans
kmeans_model = KMeans(n_clusters=4, random_state=1, n_init='auto').fit(X)
labels = kmeans_model.fit_predict(X)

df_tsne['Cluster'] = labels
df_tsne

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='Cluster', palette='magma')
sns.move_legend(ax, "upper left", bbox_to_anchor=(1, 1))
plt.title('Scatter plot of news using KMeans Clustering')
plt.xlabel('TSNE1')
plt.ylabel('TSNE2')
plt.axis('equal')

Notebook output

def get_majority_cluster_per_group(df_tsne_cluster, class_names):
  class_clusters = dict()
  for c in class_names:
    # Get rows of dataframe that are equal to c
    rows = df_tsne_cluster.loc[df_tsne_cluster['Class Name'] == c]
    # Get majority value in Cluster column of the rows selected
    cluster = rows.Cluster.mode().values[0]
    # Populate mapping dictionary
    class_clusters[c] = cluster
  return class_clusters

classes = df_tsne['Class Name'].unique()
class_clusters = get_majority_cluster_per_group(df_tsne, classes)
class_clusters

获取每组中的多数聚类，并查看该组实际成员中有多少位于该聚类中。

# Convert the Cluster column to use the class name
class_by_id = {v: k for k, v in class_clusters.items()}
df_tsne['Predicted'] = df_tsne['Cluster'].map(class_by_id.__getitem__)

# Filter to the correctly matched rows
correct = df_tsne[df_tsne['Class Name'] == df_tsne['Predicted']]

# Summarise, as a percentage
acc = correct['Class Name'].value_counts() / SAMPLE_SIZE
acc

# Get predicted values by name
df_tsne['Predicted'] = ''
for idx, rows in df_tsne.iterrows():
  cluster = rows['Cluster']
  # Get key from mapping based on cluster value
  key = list(class_clusters.keys())[list(class_clusters.values()).index(cluster)]
  df_tsne.at[idx, 'Predicted'] = key

df_tsne

为了更好地可视化 KMeans 应用于数据的性能，你可以使用混淆矩阵。混淆矩阵允许你在准确率之外评估分类模型的性能。你可以看到被错误分类的点被归类为什么。你将需要实际值和预测值，这些你已经在上面的数据框中收集到了。

cm = confusion_matrix(df_tsne['Class Name'].to_list(), df_tsne['Predicted'].to_list())
disp = ConfusionMatrixDisplay(confusion_matrix=cm,
                              display_labels=classes)
disp.plot(xticks_rotation='vertical')
plt.title('Confusion Matrix for Actual and Clustered Newsgroups');
plt.grid(False)

Notebook output

后续步骤

现在你已经创建了自己的嵌入可视化，并进行了聚类！尝试使用你自己的文本数据，将其可视化为嵌入。你可以执行降维以完成可视化步骤。请注意，TSNE 擅长对输入进行聚类，但可能需要较长时间才能收敛，或者可能会陷入局部最小值。

除了 KMeans 之外，还有其他聚类算法，例如基于密度的空间聚类（DBSCAN）。

要了解如何使用 Gemini API 中的其他服务，请参阅入门指南。

常见问题

常见错误一：未设置 GEMINI_API_KEY 或未启用付费层级，导致 429/403，应检查环境变量、AI Studio 密钥和可用地区。错误二：混用不同嵌入模型或不同 task_type，导致向量不可比，应固定 MODEL_ID 和 task_type。错误三：把 t-SNE 输出直接喂给 KMeans，损失高维语义信息，应使用原始嵌入做聚类。错误四：文本过长超过模型限制，代码中已用 len<10000 过滤，生产环境应做分块或截断。错误五：KMeans 的 n_clusters 与真实类别数不一致，导致多数投票映射失真，应结合轮廓系数或业务先验选 k。错误六：t-SNE 收敛慢或结果随机，应设置 random_state、perplexity 并多次运行。错误七：混淆矩阵标签顺序错乱，应显式传入 classes 并旋转 x 轴标签。生产优化包括：批量请求与并发控制、缓存嵌入结果、使用向量数据库、对长文本分块后池化、用 UMAP 替代 t-SNE 提升速度、用 HDBSCAN 处理非球形簇和噪声点。

企业应用扩展

可进一步扩展的方向：第一，把嵌入聚类接入 RAG 或语义检索，用 RETRIEVAL_DOCUMENT 与 RETRIEVAL_QUERY 分别编码文档和查询。第二，用分类任务类型训练轻量分类器，把聚类得到的伪标签作为弱监督信号。第三，替换为 DBSCAN 或 HDBSCAN，自动发现簇数量并识别离群点。第四，引入 UMAP 做更快更稳定的降维可视化。第五，对超长文档先分块嵌入再平均或加权池化，提升长文本主题发现质量。第六，把流程封装为可配置管道，支持自定义数据集、模型、采样比例和聚类算法，并加入嵌入缓存与增量更新，便于在生产环境中持续运行。

总结

本文以 Gemini API 生成的文本嵌入为核心，演示从 20 Newsgroups 数据集采样、清洗、生成嵌入，到用 t-SNE 降维可视化、KMeans 聚类并借助混淆矩阵评估的完整流程。教程强调嵌入模型选择、task_type 设置、维度与降维取舍、聚类与分类评估差异，并给出环境准备、代码解析、常见错误排查与生产优化建议，帮助开发者把嵌入聚类落地到真实文本主题发现场景。

内容来源与验证

继续系统学习

站内还有更多完整实战教程

加入会员，解锁全部教程、完整代码与持续更新内容。

加入会员
查看全部API 与工程实践教程