📢 转载信息
原文链接：https://machinelearningmastery.com/document-clustering-with-llm-embeddings-in-scikit-learn/
原文作者：Jason Brownlee
本文将介绍如何使用 LLM 嵌入（例如来自 OpenAI 的嵌入模型）与 Scikit-Learn 库相结合，对文档进行聚类。
传统的文档聚类方法（如 TF-IDF）依赖于词频或词袋模型，这些方法通常无法很好地捕捉语义信息。而 LLM 嵌入能够将文本转换为密集的向量表示，这些向量捕捉了文本的深层语义含义，因此非常适合用于聚类。
设置环境和先决条件
在开始之前，我们需要安装一些必要的 Python 库。
我们将使用 Scikit-Learn 进行聚类，使用 OpenAI 库来生成嵌入。
pip install scikit-learn openai
您还需要设置您的 OpenAI API 密钥。请确保它已设置为环境变量，以便代码可以访问它。
准备示例数据
我们将使用一组简短的示例文本进行演示。在实际应用中，您将用自己的文档数据集替换它们。
# 示例文档列表
documents = [
"The cat sat on the mat.",
"The dog chased the ball in the park.",
"A feline rested upon the rug.",
"The canine pursued the toy outdoors.",
"I love programming in Python.",
"Machine learning is a fascinating field."
]
生成 LLM 嵌入
我们将使用 OpenAI 的 text-embedding-ada-002 模型（或更新的模型，如 text-embedding-3-small）将文本转换为嵌入向量。
以下是一个用于生成嵌入的函数：
import os
from openai import OpenAI
import numpy as np
# 确保设置了 OPENAI_API_KEY 环境变量
# client = OpenAI()
def get_embeddings(texts, model="text-embedding-3-small"):
"""生成给定文本列表的嵌入向量"""
try:
client = OpenAI()
response = client.embeddings.create(input=texts, model=model)
# 提取嵌入向量
return [item.embedding for item in response.data]
except Exception as e:
print(f"Error generating embeddings: {e}")
return None
# 获取嵌入
embeddings = get_embeddings(documents)
if embeddings:
print(f"Generated {len(embeddings)} embeddings.")
print(f"Embedding dimension: {len(embeddings[0])}")
# 将列表转换为 NumPy 数组，以便与 Scikit-Learn 兼容
X = np.array(embeddings)
print(f"Shape of embedding matrix: {X.shape}")
else:
print("Failed to generate embeddings.")
对于上述示例数据，输出应显示我们生成了 6 个嵌入，每个嵌入的维度取决于所选模型（text-embedding-3-small 为 1536 维）。
使用 Scikit-Learn 进行聚类
有了嵌入向量（X），我们现在可以将其用作 Scikit-Learn 聚类算法的输入。我们将使用 K-Means 算法，因为它简单高效，适用于大多数聚类任务。
确定聚类数量 (K)
在实际操作中，K 的选择通常需要通过肘部法则（Elbow Method）或轮廓系数（Silhouette Score）等方法来确定。对于我们的示例，我们知道我们有三组不同的主题（猫/狗、编程/ML），所以我们姑且设定 K=3。
from sklearn.cluster import KMeans
# 设定聚类数量
k = 3
# 初始化 K-Means 模型
kmeans = KMeans(n_clusters=k, random_state=42, n_init=10)
# 训练模型
kmeans.fit(X)
# 获取每个文档的聚类标签
labels = kmeans.labels_
print("Cluster Labels:")
print(labels)
注意：在 Scikit-Learn 版本 1.2 及更高版本中，KMeans 构造函数中 n_init 的默认值已更改。将其明确设置为 10（或您选择的值）是推荐的做法，以避免未来版本的警告。
分析聚类结果
现在我们将原始文档与其分配的聚类标签关联起来，看看结果如何：
results = zip(documents, labels)
print("\nDocument Clustering Results:")
for doc, label in results:
print(f"Cluster {label}: {doc}")
针对示例数据运行此代码，我们应该会看到如下分组（标签可能因随机初始化而异，但分组应保持一致）：
Cluster 0: The cat sat on the mat. (猫/猫科动物)
Cluster 0: A feline rested upon the rug. (猫/猫科动物)
Cluster 1: The dog chased the ball in the park. (狗/犬科动物)
Cluster 1: The canine pursued the toy outdoors. (狗/犬科动物)
Cluster 2: I love programming in Python. (编程/ML)
Cluster 2: Machine learning is a fascinating field. (编程/ML)
如您所见，即使原始文本的词汇差异很大（例如“cat”和“feline”），由于 LLM 嵌入的语义捕捉能力，它们仍被正确地分到了同一组。
评估聚类（可选）
为了客观地评估聚类的质量，我们可以计算轮廓系数（Silhouette Score）。该分数范围从 -1 到 1，越接近 1 表示聚类越好。
from sklearn.metrics import silhouette_score
# 确保至少有两个聚类才能计算轮廓系数
if len(np.unique(labels)) > 1:
score = silhouette_score(X, labels)
print(f"\nSilhouette Score: {score:.4f}")
else:
print("\nCannot calculate silhouette score with only one cluster.")
对于这种非常清晰的小数据集，您应该会得到一个很高的轮廓系数，表明模型成功地找到了语义紧密的群组。
总结
将 LLM 嵌入与 Scikit-Learn 结合是现代文档聚类的一个强大范式。它利用了预训练模型的丰富语义知识，显著提高了聚类结果的质量，特别是对于具有复杂或隐含语义关系的数据集。
此方法的核心优势在于将复杂的文本表示问题简化为经典的向量聚类问题，这使得我们可以利用成熟且高效的 ML 算法。
🚀 想要体验更好更全面的AI调用？
欢迎使用青云聚合API，约为官网价格的十分之一，支持300+全球最新模型，以及全球各种生图生视频模型，无需翻墙高速稳定，文档丰富，小白也可以简单操作。
评论区