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愛好者。歡迎與譯站一起，學習新知，分享成長。

VirMach 便宜 VPS

QNews

QNews