FAISS + SBERT实现的十亿级语义相似性搜索
- 2020 年 11 月 13 日
- AI
译者：AI研习社（FIONAbiubiu）
双语原文链接：Billion-scale semantic similarity search with FAISS+SBERT
介绍
语义搜索是一种关注句子意义而不是传统的关键词匹配的信息检索系统。尽管有许多文本嵌入可用于此目的，但将其扩展到构建低延迟api以从大量数据集合中获取数据是很少讨论的。在本文中，我将讨论如何使用SOTA语句嵌入（语句转换器）和FAISS来实现最小语义搜索引擎。
句子Transformers
它是一个框架或一组模型，给出句子或段落的密集向量表示。这些模型是transformer网络（BERT、RoBERTa等），它们专门针对语义文本相似性的任务进行了微调，因为BERT在这些任务中执行得不是很好。下面给出了不同模型在STS基准测试中的性能。
图片来源：句子 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中的结果
现在让我们看看搜索结果和响应时间
只需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执行这个搜索
很好，你可以在0.02秒内得到结果，使用GPU（在这个实验中使用了Tesla T4），它比CPU后端快75倍
但是为什么我不能仅仅序列化编码数据的NumPy数组而不是索引它们呢？如果我能等几秒钟的话，使用余弦相似性呢？
因为NumPy没有序列化函数，因此唯一的方法是将其转换为JSON，然后保存JSON对象，但是大小将增加五倍。例如，在768维向量空间中编码的一百万个数据点具有正常的索引，大约为3GB，将其转换为JSON将使其成为15GB，而普通机器无法保存它的RAM。因此，每次执行搜索时，我们都要运行一百万次计算推理，这是不实际的。
最后的想法
这是一个基本的实现，在语言模型部分和索引部分仍然需要做很多工作。有不同的索引选项，应该根据用例、数据大小和可用的计算能力选择正确的索引选项。另外，这里使用的句子嵌入只是对一些公共数据集进行了微调，在特定领域的数据集上对它们进行微调可以改进，从而提高搜索结果。
参考文献
[1] Nils Reimers and Iryna Gurevych. “Making Monolingual Sentence Embeddings Multilingual using Knowledge Distillation.” arXiv (2020): 2004.09813.
[2]Johnson, Jeff and Douze, Matthijs and J{\’e}gou, Herv{\’e}. “Billion-scale similarity search with GPUs” arXiv preprint arXiv:1702.08734.
AI研习社是AI学术青年和AI开发者技术交流的在线社区。我们与高校、学术机构和产业界合作，通过提供学习、实战和求职服务，为AI学术青年和开发者的交流互助和职业发展打造一站式平台，致力成为中国最大的科技创新人才聚集地。
如果，你也是位热爱分享的AI爱好者。欢迎与译站一起，学习新知，分享成长。
