Sentence-transformer + FAISS 语义搜索
1. 关于FAISS
在上一篇博客中介绍了SBERT系列的各种应用，其中包括了语义搜索。但是在进行语义搜索的时候采用的是全局搜索的方法，时间代价较高，因而在数据库很大的情况下，需要采取一些近似计算的策略，其中FAISS就是一个非常有效的工具。
FAISS是由 Facebook AI Research开发一款基于C++的对稠密向量的语义检索和聚类工具，并提供了完整的python的封装。
由于其对稠密向量直接进行计算，这就使之可以具有非常多样的下游任务，理论上任何可以向量化的特征都可以利用这一工具实现快速检索，例如图片，文本等。
sentence-transformer模块提供的样例semantic_search_quora_faiss.py中，就包含了对结合FAISS使用的介绍，这篇博客基本内容就是对这个样例进行搬运，并做简单的注释。
2. 安装FAISS
打开pypi搜索FAISS，出现的第一条FAISS 1.5.1不要安装，那个项目停在两年前不维护了，也没有在py3.8以上的版本上编译，所以直接下载cpu版本或者gpu版本的1.7.1.
3. 下载预训练模型和数据集
在样例中，数据集采用的是quora的重复问句数据集，点此下载。这个数据集原本是用来做句对匹配训练的，每一条中有两个question，以及这两个question是否是重复提问的label，但是在这个应用场景中，语义检索希望把用户输入问句的最相关的问句返回，并不关心重复的问题，于是将question1和question2一视同仁的看待。这个数据集里一共有500k的语句。
关于预训练模型，样例中给出的是在quora上训练的’quora-distilbert-multilingual’，而我使用的是mpnet-base，实验效果也是不错的，预训练模型的下载和使用可以参考上一篇博客。
from sentence_transformers import SentenceTransformer, util
import os
import csv
import pickle
import time
import faiss
import numpy as np
4. 获取embedding
由于数据量比较大，不可能每次检索的时候都对库里的所有语句现场编码，所以需要把所有候选语句全都编码完，保存到本地。我在实验时使用3090对全部500k的候选用mpnet-base编码完，总共耗时约5分钟。
首先设置一下embedding相关的参数：
# 下载的数据集的路径
dataset_path = "your_path_to_/quora_duplicate_questions.tsv"
# 在500k的数据集里选择多少条构建候选语料库
max_corpus_size = 100000
# 候选语料保存在本地文件的名称
embedding_cache_path = 'quora-embeddings-{}-size-{}.pkl'.format(model_name.replace('/', '_'), max_corpus_size)
# 预训练模型的编码输出特征维度
embedding_size = 768
然后开始编码，并保存到本地：
扫描二维码关注公众号，回复：
15817732 查看本文章
#Check if embedding cache path exists
if not os.path.exists(embedding_cache_path):
# Check if the dataset exists. If not, download and extract
# Download dataset if needed
# if not os.path.exists(dataset_path):
# print("Download dataset")
# util.http_get(url, dataset_path)
# Get all unique sentences from the file
corpus_sentences = set()
with open(dataset_path, encoding='utf8') as fIn:
reader = csv.DictReader(fIn, delimiter='\t', quoting=csv.QUOTE_MINIMAL)
for row in reader:
corpus_sentences.add(row['question1'])
if len(corpus_sentences) >= max_corpus_size:
break
corpus_sentences.add(row['question2'])
if len(corpus_sentences) >= max_corpus_size:
break
corpus_sentences = list(corpus_sentences)
print("Encode the corpus. This might take a while")
corpus_embeddings = model.encode(corpus_sentences, show_progress_bar=True, convert_to_numpy=True)
print("Store file on disc")
with open(embedding_cache_path, "wb") as fOut:
pickle.dump({
'sentences': corpus_sentences, 'embeddings': corpus_embeddings}, fOut)
else:
print("Load pre-computed embeddings from disc")
with open(embedding_cache_path, "rb") as fIn:
cache_data = pickle.load(fIn)
corpus_sentences = cache_data['sentences']
corpus_embeddings = cache_data['embeddings']
然后这个pkl会保存到本地，100k的编码文件大概有300M大小。
5. 创建索引并训练
首先设置一下相关的参数。
# 查询前K个结果
top_k_hits = 10
# 聚类的数量，这个数量一般介于4*sqrt(N) 到 16*sqrt(N)，N是语料库的大小
n_clusters = 1024
# 在最相关的多少个簇中搜索答案，这个参数越大查的越全，消耗的时间也就越多
_nprobe = 3
然后创建索引对象，并传入训练数据。
# 创建
quantizer = faiss.IndexFlatIP(embedding_size)
index = faiss.IndexIVFFlat(quantizer, embedding_size, n_clusters, faiss.METRIC_INNER_PRODUCT)
index.nprobe = _nprobe
# 训练
# 因为对向量做点积计算以进行召回，先对所有语料的编码进行normalize
corpus_embeddings = corpus_embeddings / np.linalg.norm(corpus_embeddings, axis=1)[:, None]
index.train(corpus_embeddings)
index.add(corpus_embeddings)
6. 应用
训练完成之后就可以应用了。
while True:
inp_question = input("Please enter a question: ")
start_time = time.time()
question_embedding = model.encode(inp_question)
#FAISS works with inner product (dot product). When we normalize vectors to unit length, inner product is equal to cosine similarity
question_embedding = question_embedding / np.linalg.norm(question_embedding)
question_embedding = np.expand_dims(question_embedding, axis=0)
## 使用FAISS进行检索
# Search in FAISS. It returns a matrix with distances and corpus ids.
distances, corpus_ids = index.search(question_embedding, top_k_hits)
# We extract corpus ids and scores for the first query
hits = [{
'corpus_id': id, 'score': score} for id, score in zip(corpus_ids[0], distances[0])]
hits = sorted(hits, key=lambda x: x['score'], reverse=True)
end_time = time.time()
print("Input question:", inp_question)
print("Results (after {:.3f} seconds):".format(end_time-start_time))
for hit in hits[0:top_k_hits]:
print("\t{:.3f}\t{}".format(hit['score'], corpus_sentences[hit['corpus_id']]))
## 为了评估FAISS的效果，下面这部分是为了对照在全局检索的结果中，FAISS成功命中了多少
## 在实际应用的时候要删除下面的内容
# Approximate Nearest Neighbor (ANN) is not exact, it might miss entries with high cosine similarity
# Here, we compute the recall of ANN compared to the exact results
correct_hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k_hits)[0]
correct_hits_ids = set([hit['corpus_id'] for hit in correct_hits])
ann_corpus_ids = set([hit['corpus_id'] for hit in hits])
if len(ann_corpus_ids) != len(correct_hits_ids):
print("Approximate Nearest Neighbor returned a different number of results than expected")
recall = len(ann_corpus_ids.intersection(correct_hits_ids)) / len(correct_hits_ids)
print("\nApproximate Nearest Neighbor Recall@{}: {:.2f}".format(top_k_hits, recall * 100))
if recall < 1:
print("Missing results:")
for hit in correct_hits[0:top_k_hits]:
if hit['corpus_id'] not in ann_corpus_ids:
print("\t{:.3f}\t{}".format(hit['score'], corpus_sentences[hit['corpus_id']]))
print("\n\n========\n")
以上就是本期全部内容了，如果这篇博客对你有帮助的话，还请一键三连支持一下up主，我们下期再见。