FAISS + SBERT實現(xiàn)的十億級語義相似性搜索

本文作者： AI研習社-譯站

2020-11-14 10:30

導語：這是一個基本的實現(xiàn)，在語言模型部分和索引部分仍然需要做很多工作。

譯者：AI研習社（FIONAbiubiu）

雙語原文鏈接：Billion-scale semantic similarity search with FAISS+SBERT

FAISS + SBERT實現(xiàn)的十億級語義相似性搜索

介紹

語義搜索是一種關注句子意義而不是傳統(tǒng)的關鍵詞匹配的信息檢索系統(tǒng)。盡管有許多文本嵌入可用于此目的，但將其擴展到構建低延遲api以從大量數(shù)據(jù)集合中獲取數(shù)據(jù)是很少討論的。在本文中，我將討論如何使用SOTA語句嵌入（語句轉換器）和FAISS來實現(xiàn)最小語義搜索引擎。

句子Transformers

它是一個框架或一組模型，給出句子或段落的密集向量表示。這些模型是transformer網絡（BERT、RoBERTa等），它們專門針對語義文本相似性的任務進行了微調，因為BERT在這些任務中執(zhí)行得不是很好。下面給出了不同模型在STS基準測試中的性能。

FAISS + SBERT實現(xiàn)的十億級語義相似性搜索

圖片來源：句子 transformers

我們可以看到句子transformer模型比其他模型有很大的優(yōu)勢。

但是如果你用代碼和GLUE來看看排行榜，你會看到很多的模型超過90。為什么我們需要句子transformers？

在這些模型中，語義文本相似度被視為一個回歸任務。這意味著，每當我們需要計算兩個句子之間的相似度得分時，我們需要將它們一起傳遞到模型中，然后模型輸出它們之間的數(shù)值分數(shù)。雖然這對于基準測試很有效，但是對于實際的用例來說，它的伸縮性很差，原因如下。

1.當你需要搜索大約10k個文檔時，你需要進行10k個獨立的推理計算，不可能單獨計算嵌入量而只計算余弦相似度。見作者的解釋。
2.最大序列長度（模型一次可以接受的單詞/標記的總數(shù)）在兩個文檔之間共享，這會導致的表示的含義由于分塊而被稀釋

FAISS

Faiss是一個基于C++的庫，由FacebookAI構建，在Python中有完整的包裝器，用于索引矢量化數(shù)據(jù)并對其進行有效的搜索。Faiss基于以下因素提供了不同的索引。

搜索時間
搜索質量
每個索引向量使用的內存
訓練時間
無監(jiān)訓練需要外部數(shù)據(jù)

因此，選擇合適的指數(shù)將是這些因素之間的權衡。

加載模型并對數(shù)據(jù)集執(zhí)行推理

首先，讓我們安裝并加載所需的庫

!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

加載一個包含一百萬個數(shù)據(jù)點的數(shù)據(jù)集

我使用了一個來自Kaggle的數(shù)據(jù)集，其中包含了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)

為數(shù)據(jù)集編制索引

我們可以根據(jù)我們的用例通過參考指南來選擇不同的索引選項。

讓我們定義索引并向其添加數(shù)據(jù)

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')

執(zhí)行語義相似性搜索

讓我們首先為搜索構建一個包裝函數(shù)

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]]

執(zhí)行搜索

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

CPU中的結果

現(xiàn)在讓我們看看搜索結果和響應時間

FAISS + SBERT實現(xiàn)的十億級語義相似性搜索

只需1.5秒，就可以在僅使用CPU后端的百萬文本文檔的數(shù)據(jù)集上執(zhí)行基于意義的智能搜索。

GPU中的結果

首先讓我們關閉CPU版本的Faiss并重啟GPU版本

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

之后執(zhí)行相同步驟，但是最后將索引移到GPU上。

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

現(xiàn)在讓我們轉移這個搜索方法并用GPU執(zhí)行這個搜索

FAISS + SBERT實現(xiàn)的十億級語義相似性搜索

很好，你可以在0.02秒內得到結果，使用GPU（在這個實驗中使用了Tesla T4），它比CPU后端快75倍

但是為什么我不能僅僅序列化編碼數(shù)據(jù)的NumPy數(shù)組而不是索引它們呢？如果我能等幾秒鐘的話，使用余弦相似性呢？

因為NumPy沒有序列化函數(shù)，因此唯一的方法是將其轉換為JSON，然后保存JSON對象，但是大小將增加五倍。例如，在768維向量空間中編碼的一百萬個數(shù)據(jù)點具有正常的索引，大約為3GB，將其轉換為JSON將使其成為15GB，而普通機器無法保存它的RAM。因此，每次執(zhí)行搜索時，我們都要運行一百萬次計算推理，這是不實際的。

最后的想法

這是一個基本的實現(xiàn)，在語言模型部分和索引部分仍然需要做很多工作。有不同的索引選項，應該根據(jù)用例、數(shù)據(jù)大小和可用的計算能力選擇正確的索引選項。另外，這里使用的句子嵌入只是對一些公共數(shù)據(jù)集進行了微調，在特定領域的數(shù)據(jù)集上對它們進行微調可以改進，從而提高搜索結果。