基于向量表示的聚类(K-Means)与相似度查询

root
abc abc
  • 23 May

下载模型 文本向量表示模型:

git clone https://gitcode.com/hf_mirr...

pip install faiss-cpu

或者

pip install faiss-gpu

```python
import os
from sentence_transformers import SentenceTransformer
import faiss
import numpy

model = SentenceTransformer(r'D:\text2vec\text2vec-base-chinese')

corpus = ['书', '笔', '电脑', '桌子', '椅子', '窗户', '门', '花', '树', '草',
'天空', '云', '太阳', '月亮', '星星', '海', '山', '河流', '桥', '路',
'汽车', '自行车', '飞机', '火车', '船', '手机', '钥匙', '钱包', '衣服',
'鞋子', '帽子', '手套', '眼镜', '手表', '相机', '电视', '冰箱',
'橡皮', '铅笔', '笔记本', '日历', '地图', '地球仪', '乐器', '吉他',
'钢琴', '鼓', '小提琴', '笛子', '喇叭', '口琴', '画', '雕塑', '照片',
'相框', '邮票', '硬币', '纸币', '信用卡', '身份证', '护照', '钥匙扣',
'雨伞', '扇子', '风筝', '球', '篮球', '足球', '乒乓球', '羽毛球',
'网球', '棒球', '高尔夫球', '滑板', '溜冰鞋', '滑雪板', '冲浪板',
'帐篷', '睡袋', '登山杖', '望远镜', '显微镜', '放大镜', '温度计',
'血压计', '体重秤', '钟表', '计算器', '游戏机', '耳机', '音响',
'麦克风', '投影仪', '打印机', '扫描仪', '路由器', '电源插座', '电池',
'充电器', '数据线', '存储卡', '硬盘', '鼠标', '键盘', '显示器',
'音箱', '耳机', '麦克风', '摄像头']
corpus=list(set(corpus))

print(f'len=={len(corpus)}')

corpus_embeddings = model.encode(corpus)

def vector_search(query = "阅读", k=5):

# 通过embedding模型将查询短语向量化
query_embedding = model.encode([query])

# 初始化 Faiss 索引
dimension = corpus_embeddings.shape[1]  # 向量维度
print("dimension:", dimension)
index = faiss.IndexFlatL2(dimension)  # 使用 L2 距离
index.add(corpus_embeddings)  # 添加语料库向量到索引

# 检索与查询向量最接近的前 k 个结果
# 返回前k = 5(默认为5) 个最相似的结果
distances, indices = index.search(query_embedding, k)

# 打印检索结果
print("Query:", query)
print(f"Top {k} Results:")
for i, idx in enumerate(indices[0]):
    print(
        f"Rank {i + 1}: {corpus[idx]} (Distance: {distances[0][i]:.4f})")
    prompt=prompt+corpus[idx]

def cluster(embb=None, k=15):
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs

X=numpy.array(embb,dtype=float) #得到文本向量数据

# 创建 KMeans 实例并指定要形成的簇的数量
kmeans = KMeans(n_clusters=k, random_state=170)

# 训练模型
kmeans.fit(X)

# 预测每个样本所属的簇
y_pred = kmeans.predict(X)
labels=set(y_pred)
labels=list(labels)
labels=sorted(labels)
clustered_data=[]
for i in range(len(labels)):
    clustered_data.append([])

for i, label in enumerate(y_pred):
    for x in labels:
        if x==label:
            clustered_data[x].append(corpus[i])

for x in labels: print("聚类-", x ,  ": \n", clustered_data[x])

cluster(corpus_embeddings)
vector_search("阅读", k=5)
```python