FAISS + SBERT实现的十亿级语义相似性搜索

本文作者： AI研习社-译站

2020-11-14 10:30

导语：这是一个基本的实现，在语言模型部分和索引部分仍然需要做很多工作。

译者：AI研习社（FIONAbiubiu）

双语原文链接：Billion-scale semantic similarity search with FAISS+SBERT

FAISS + SBERT实现的十亿级语义相似性搜索

介绍

语义搜索是一种关注句子意义而不是传统的关键词匹配的信息检索系统。尽管有许多文本嵌入可用于此目的，但将其扩展到构建低延迟api以从大量数据集合中获取数据是很少讨论的。在本文中，我将讨论如何使用SOTA语句嵌入（语句转换器）和FAISS来实现最小语义搜索引擎。

句子Transformers

它是一个框架或一组模型，给出句子或段落的密集向量表示。这些模型是transformer网络（BERT、RoBERTa等），它们专门针对语义文本相似性的任务进行了微调，因为BERT在这些任务中执行得不是很好。下面给出了不同模型在STS基准测试中的性能。

FAISS + SBERT实现的十亿级语义相似性搜索

图片来源：句子 transformers

我们可以看到句子transformer模型比其他模型有很大的优势。

但是如果你用代码和GLUE来看看排行榜，你会看到很多的模型超过90。为什么我们需要句子transformers？

在这些模型中，语义文本相似度被视为一个回归任务。这意味着，每当我们需要计算两个句子之间的相似度得分时，我们需要将它们一起传递到模型中，然后模型输出它们之间的数值分数。虽然这对于基准测试很有效，但是对于实际的用例来说，它的伸缩性很差，原因如下。

1.当你需要搜索大约10k个文档时，你需要进行10k个独立的推理计算，不可能单独计算嵌入量而只计算余弦相似度。见作者的解释。
2.最大序列长度（模型一次可以接受的单词/标记的总数）在两个文档之间共享，这会导致的表示的含义由于分块而被稀释

FAISS

Faiss是一个基于C++的库，由FacebookAI构建，在Python中有完整的包装器，用于索引矢量化数据并对其进行有效的搜索。Faiss基于以下因素提供了不同的索引。

搜索时间
搜索质量
每个索引向量使用的内存
训练时间
无监训练需要外部数据

因此，选择合适的指数将是这些因素之间的权衡。

加载模型并对数据集执行推理

首先，让我们安装并加载所需的库

!pip install faiss-cpu
!pip install -U sentence-transformersimport numpy as np
import torch
import os
import pandas as pd
import faiss
import time
from sentence_transformers import SentenceTransformer

加载一个包含一百万个数据点的数据集

我使用了一个来自Kaggle的数据集，其中包含了17年来出版的新闻标题。

df=pd.read_csv("abcnews-date-text.csv")
data=df.headline_text.to_list()

加载预训练模型并且进行推断

model = SentenceTransformer('distilbert-base-nli-mean-tokens')encoded_data = model.encode(data)

为数据集编制索引

我们可以根据我们的用例通过参考指南来选择不同的索引选项。

让我们定义索引并向其添加数据

index = faiss.IndexIDMap(faiss.IndexFlatIP(768))index.add_with_ids(encoded_data, np.array(range(0, len(data))))

序列化索引

faiss.write_index(index, 'abc_news')

将序列化的索引导出到托管搜索引擎的任何计算机中

反序列化索引

index = faiss.read_index('abc_news')

执行语义相似性搜索

让我们首先为搜索构建一个包装函数

def search(query):
t=time.time()
query_vector = model.encode([query])
k = 5
top_k = index.search(query_vector, k)
print('totaltime: {}'.format(time.time()-t))
return [data[_id] for _id in top_k[1].tolist()[0]]

执行搜索

query=str(input())
results=search(query)
print('results :')
for result in results:
print('\t'

CPU中的结果

现在让我们看看搜索结果和响应时间

FAISS + SBERT实现的十亿级语义相似性搜索

只需1.5秒，就可以在仅使用CPU后端的百万文本文档的数据集上执行基于意义的智能搜索。

GPU中的结果

首先让我们关闭CPU版本的Faiss并重启GPU版本

!pip uninstall faiss-cpu
!pip install faiss-gpu

之后执行相同步骤，但是最后将索引移到GPU上。

res = faiss.StandardGpuResources()
gpu_index = faiss.index_cpu_to_gpu(res, 0, index)

现在让我们转移这个搜索方法并用GPU执行这个搜索

FAISS + SBERT实现的十亿级语义相似性搜索

很好，你可以在0.02秒内得到结果，使用GPU（在这个实验中使用了Tesla T4），它比CPU后端快75倍

但是为什么我不能仅仅序列化编码数据的NumPy数组而不是索引它们呢？如果我能等几秒钟的话，使用余弦相似性呢？

因为NumPy没有序列化函数，因此唯一的方法是将其转换为JSON，然后保存JSON对象，但是大小将增加五倍。例如，在768维向量空间中编码的一百万个数据点具有正常的索引，大约为3GB，将其转换为JSON将使其成为15GB，而普通机器无法保存它的RAM。因此，每次执行搜索时，我们都要运行一百万次计算推理，这是不实际的。