如何基于LangChain FAISS向量存储实现KMeans聚类？

方法一：用Scikit-learn KMeans（适合中小数据集） #

• 先从FAISS向量存储里提取所有向量，再用sklearn的KMeans做聚类，最后把标签和原文档关联起来

from sklearn.cluster import KMeans
import numpy as np
from collections import Counter

# 提取所有向量
total_vecs = vector_store.index.ntotal
all_embeddings = np.array([vector_store.index.reconstruct(i) for i in range(total_vecs)])

# 训练KMeans，这里假设聚成5类，可按需调整
kmeans = KMeans(n_clusters=5, random_state=42)
cluster_labels = kmeans.fit_predict(all_embeddings)

# 把聚类标签和原文档绑定
all_docs = vector_store.get_all_documents()
clustered_docs = []
for doc, label in zip(all_docs, cluster_labels):
    clustered_docs.append({
        "content": doc.page_content,
        "metadata": doc.metadata,
        "cluster_label": label
    })

# 查看各聚类的文档数量
print(Counter(cluster_labels))

方法二：用Faiss原生KMeans（适合大数据集，速度更快） #

Faiss自带的KMeans做了性能优化，适合大规模向量聚类，还能支持GPU加速：

import faiss
import numpy as np

# 提取所有向量，转成Faiss要求的float32格式
total_vecs = vector_store.index.ntotal
all_embeddings = np.array([vector_store.index.reconstruct(i) for i in range(total_vecs)]).astype('float32')

# 初始化Faiss KMeans
n_clusters = 5
kmeans = faiss.Kmeans(d=all_embeddings.shape[1], k=n_clusters, niter=20, verbose=True)

# 训练模型
kmeans.train(all_embeddings)

# 获取每个向量的聚类标签
cluster_labels, _ = kmeans.index.search(all_embeddings, 1)
cluster_labels = cluster_labels.squeeze()  # 压缩维度

# 关联文档与聚类标签
all_docs = vector_store.get_all_documents()
clustered_docs = [
    {"content": doc.page_content, "metadata": doc.metadata, "cluster_label": label}
    for doc, label in zip(all_docs, cluster_labels)
]

# 查看聚类分布
print(np.bincount(cluster_labels))

注意点 #

• n_clusters要根据你的数据集调整，可通过肘部法则确定最优聚类数
• 若向量量极大，优先用Faiss原生实现，安装faiss-gpu还能开启GPU加速
• 向量和文档的顺序是一一对应的，不用担心匹配错误

内容的提问来源于stack exchange，提问作者Yining Wang