Github13K！相似搜索百寶箱，文本匹配入門必備！

# CPU-only version
$ conda install -c pytorch faiss-cpu

# GPU(+CPU) version
$ conda install -c pytorch faiss-gpu

# CPU-only version
$ pip install faiss-cpu -i https://pypi.douban.com/simple

# GPU(+CPU) version
$ pip install faiss-gpu -i https://pypi.douban.com/simple

import faiss
d = 768 # embedding 維度
# 構(gòu)建 IndexFlatL2 索引，它是最簡單的索引類型，只執(zhí)行強力L2距離搜索
faiss_l2_index = faiss.IndexFlatL2(d)
print(faiss_l2_index.is_trained) # True

# 在索引中添加原始特征
# faq_texts_emb 是特征向量，維度 [6417, 768]
faiss_l2_index.add(faq_texts_emb)
print(faiss_l2_index.ntotal) # 查看數(shù)據(jù)規(guī)模：6417

# 查找 k 近鄰
k = 3
q = '我咋沒收到交費的通知呢？'
q_emb = bc.encode([q]) # 獲得特征
# dis_res:距離矩陣，由小到大排序
# idx_res:數(shù)據(jù)索引id
dis_res, idx_res = faiss_l2_index.search(q_emb, k)
print(idx_res.shape, dis_res.shape) # [1,3], [1,3]

print(idx_res) # [[107, 111, 75]]
print(dis_res) # [[18.602684 20.919521 21.184694]]
# labels[i] 先映射到訓練集文本對應(yīng)的真實意圖 label_id，再轉(zhuǎn)換為 label 標簽
print(list(id2name[labels[i]] for i in idx_res[0]))
"""['未收到催費短信', '未收到催費短信', '未收到催費短信']"""

# m: 特征拆分子空間的數(shù)量
m, d = 8, 768
# 每一個子空間分配的bit大小
n_bits = 8
# 創(chuàng)建乘積量化索引
faiss_pq_index = faiss.IndexPQ(d, m, n_bits)        
# train 階段
faiss_pq_index.train(faq_texts_emb)
# add 階段
faiss_pq_index.add(faq_texts_emb)

k = 3
q = '我妹夫家的租戶不交電費，這樣子會影響我妹夫的信用不？'
q_emb = bc.encode([q])
dis_res, idx_res = faiss_pq_index.search(q_emb, k)
print(idx_res.shape, dis_res.shape) # [1,3], [1,3]
print(idx_res) # [[2563 2575 2547]]
print(dis_res) # [[11.660806 11.660806 11.660806]]
print(list(id2name[labels[i]] for i in idx_res[0]))
"""['租戶不交電費是否影響戶主信用', '租戶不交電費是否影響戶主信用', '租戶不交電費是否影響戶主信用']"""