我想使用 Sentence Transoformer 中的 BERT 嵌入(尤其是 bert-base-nli-mean 标记)对我的文档进行聚类,并且我想使用 kmeans 聚类对该嵌入进行聚类,但我有一个问题,我可以使用余弦距离使用 kmeans 聚类吗?
这个问题的解决方案和代码?
是的,您可以将 K-Means 聚类与从 Sentence Transformers 获得的 BERT 嵌入结合使用,例如 bert-base-nli-mean-tokens。然而,scikit-learn 等库中 K-Means 的标准实现使用欧几里得距离,而不是余弦距离。要使用余弦距离对嵌入进行聚类,您有几种选择。
!pip install sentence-transformers scikit-learn
from sentence_transformers import SentenceTransformer
from sklearn.preprocessing import normalize
from sklearn.cluster import KMeans
# Load BERT model
model = SentenceTransformer('bert-base-nli-mean-tokens')
# Your document texts
documents = ["Document 1 text...", "Document 2 text...", "..."]
# Generate BERT embeddings
embeddings = model.encode(documents)
# Normalize the embeddings
normalized_embeddings = normalize(embeddings)
# Define the K-Means model
num_clusters = 5 # Adjust the number of clusters
kmeans = KMeans(n_clusters=num_clusters, random_state=42)
# Fit the model
kmeans.fit(normalized_embeddings)
# Get cluster labels
labels = kmeans.labels_
# Output the cluster labels for your documents
print(labels)