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 维度
# 构建 IndexFlatL2 索引，它是最简单的索引类型，只执行强力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) # 查看数据规模：6417

# 查找 k 近邻
k = 3
q = '我咋没收到交费的通知呢？'
q_emb = bc.encode([q]) # 获得特征
# dis_res:距离矩阵，由小到大排序
# idx_res:数据索引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] 先映射到训练集文本对应的真实意图 label_id，再转换为 label 标签
print(list(id2name[labels[i]] for i in idx_res[0]))
"""['未收到催费短信', '未收到催费短信', '未收到催费短信']"""

# m: 特征拆分子空间的数量
m, d = 8, 768
# 每一个子空间分配的bit大小
n_bits = 8
# 创建乘积量化索引
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]))
"""['租户不交电费是否影响户主信用', '租户不交电费是否影响户主信用', '租户不交电费是否影响户主信用']"""