?

在這一期中，我們延續(xù)上一期 ?Bert 中文短句相似度計(jì)算 Docker CPU鏡像，繼續(xù)使用 huggingface transformer 和 sentence-transformer 類庫，并將英語句子生成 bert embedding，然后引入 faiss 類庫來建立索引，最后查詢最接近的句子。

?

Docker 鏡像獲取方式

本期 docker 鏡像獲取方式為，關(guān)注 MyEncyclopedia 公眾號(hào)后回復(fù) docker-faiss-transformer 即可獲取如下完整命令。

docker?run?-p?8888:8888?myencyclopedia/faiss-demo?bash?-c?'jupyter?notebook?--allow-root?--port?8888?--NotebookApp.token=?--ip?0.0.0.0'

然后打開瀏覽器，輸入 http://localhost:8888/notebooks/faiss_demo.ipynb

faiss 簡介

Faiss 的全稱是Facebook AI Similarity Search，是由 Facebook 開發(fā)的適用于稠密向量匹配的開源庫，作為向量化檢索開山鼻祖，F(xiàn)aiss 提供了一套查詢海量高維數(shù)據(jù)集的解決方案，它從兩個(gè)方面改善了暴力搜索算法存在的問題：降低空間占用和加快檢索速度。此外，F(xiàn)aiss 提供了若干種方法實(shí)現(xiàn)數(shù)據(jù)壓縮，包括 PCA、Product-Quantization等。

Faiss 主要特性：

• 支持相似度檢索和聚類；
• 支持多種索引方式；
• 支持CPU和GPU計(jì)算；
• 支持Python和C++調(diào)用；

Faiss 使用流程

使用 faiss 分成兩部，第一步需要對(duì)原始向量建立索引文件，第二步再對(duì)索引文件進(jìn)行向量 search 操作。

在第一次建立索引文件的時(shí)候，需要經(jīng)過 train 和 add 兩個(gè)過程；后續(xù)如果有新的向量需要被添加到索引文件，只需要一個(gè) add 操作來實(shí)現(xiàn)增量索引更新，但是如果增量的量級(jí)與原始索引差不多的話，整個(gè)向量空間就可能發(fā)生了一些變化，這個(gè)時(shí)候就需要重新建立整個(gè)索引文件，也就是再用全部的向量來走一遍 train 和 add，至于具體是如何 train 和 add的，就和特定的索引類型有關(guān)了。

1. IndexFlatL2 indexFlatIP

對(duì)于精確搜索，例如歐式距離 faiss.indexFlatL2 或 內(nèi)積距離 faiss.indexFlatIP，沒有 train 過程，add 完直接可以 search。

import?faiss?

#?建立索引,?定義為dimension?d?=?128
index?=?faiss.IndexFlatL2(d)

?#?add?vectors,?xb?為?(100000,128)大小的numpy
index.add(xb)?????????????????
print(index.ntotal)?
#?索引中向量的數(shù)量,?輸出100000

#?求4-近鄰
k?=?4
#?xq為query?embedding,?大小為(10000,128)
D,?I?=?index.search(xq,?k)?????
##?D?shape?(10000,4)，表示每個(gè)返回點(diǎn)的embedding?與?query?embedding的距離,
##?I?shape?(10000,4)，表示和query?embedding最接近的k個(gè)物品id，
print(I[:5])

2. IndexIVFFlat

IndexFlatL2 的結(jié)果雖然精確，但當(dāng)數(shù)據(jù)集比較大的時(shí)候，暴力搜索的時(shí)間復(fù)雜度很高，因此我們一般會(huì)使用其他方式的索引來加速。比如 IndexIVFFlat，將數(shù)據(jù)集在 train 階段分割為幾部分，技術(shù)術(shù)語為 Voronoi Cells，每個(gè)數(shù)據(jù)向量只能落在一個(gè)cell中。Search 時(shí)只需要查詢query向量落在cell中的數(shù)據(jù)了，降低了距離計(jì)算次數(shù)。這個(gè)過程本質(zhì)就是高維 KNN 聚類算法。search 階段使用倒排索引來。

IndexIVFFlat 需要一個(gè)訓(xùn)練的階段，其與另外一個(gè)索引 quantizer 有關(guān)，通過 quantizer 來判斷屬于哪個(gè)cell。IndexIVFFlat 在搜索階段，引入了nlist(cell的數(shù)量)與nprob(執(zhí)行搜索的cell數(shù))參數(shù)。增大nprobe可以得到與brute-force更為接近的結(jié)果，nprobe就是速度與精度的調(diào)節(jié)器。

import?faiss
nlist?=?100
k?=?4

#?建立索引,?定義為dimension?d?=?128
quantizer?=?faiss.IndexFlatL2(d)

#?使用歐式距離 L2 建立索引。
index?=?faiss.IndexIVFFlat(quantizer,?d,?nlist,?faiss.METRIC_L2)

##?xb:?(100000,128)
index.train(xb)?
index.add(xb)????????????????
index.nprobe?=?10??#?默認(rèn)?nprobe?是?1?,可以設(shè)置的大一些試試
D,?I?=?index.search(xq,?k)
print(I[-5:])???#?最后五次查詢的結(jié)果

3. IndexIVFPQ

IndexFlatL2 和 IndexIVFFlat都要存儲(chǔ)所有的向量數(shù)據(jù)。對(duì)于超大規(guī)模數(shù)據(jù)集來說，可能會(huì)不大現(xiàn)實(shí)。因此IndexIVFPQ 索引可以用來壓縮向量，具體的壓縮算法就是 Product-Quantization，注意，由于高維向量被壓縮，因此 search 時(shí)候返回也是近似的結(jié)果。

import?faiss

nlist?=?100
#?每個(gè)向量分8段
m?=?8?
#?求4-近鄰
k?=?4?
quantizer?=?faiss.IndexFlatL2(d)????#?內(nèi)部的索引方式依然不變
index?=?faiss.IndexIVFPQ(quantizer,?d,?nlist,?m,?8)?#?每個(gè)向量都被編碼為8個(gè)字節(jié)大小
index.train(xb)
index.add(xb)
index.nprobe?=?10????????????????
D,?I?=?index.search(xq,?k)??#?檢索
print(I[-5:])

在本期中，我們僅使用基本的 IndexIVFFlat 和 IndexFlatIP 完成 bert embedding 的索引和搜索，后續(xù)會(huì)有篇幅來解讀 Product-Quantization 的論文原理和代碼實(shí)踐。

ag_news 新聞數(shù)據(jù)集

ag_news 新聞數(shù)據(jù)集 3.0 包含了英語新聞標(biāo)題，training 部分包含 120000條數(shù)據(jù)， test 部分包含 7600條數(shù)據(jù)。

ag_news 可以通過 huggingface datasets API 自動(dòng)下載

def?load_dataset(part='test')?->?List[str]:
????ds?=?datasets.load_dataset("ag_news")
????list_str?=?[r['text']?for?r?in?ds[part]]
????return?list_str
????
list_str?=?load_dataset(part='train')
print(f'{len(list_str)}')
for?s?in?list_str[:3]:
????print(s)
????print('
')

顯示前三條新聞標(biāo)題為

120000
Wall St. Bears Claw Back Into the Black (Reuters) Reuters - Short-sellers, Wall Street's dwindlingand of ultra-cynics, are seeing green again.


Carlyle Looks Toward Commercial Aerospace (Reuters) Reuters - Private investment firm Carlyle Group,which has a reputation for making well-timed and occasionallycontroversial plays in the defense industry, has quietly placedits bets on another part of the market.


Oil and Economy Cloud Stocks' Outlook (Reuters) Reuters - Soaring crude prices plus worriesabout the economy and the outlook for earnings are expected tohang over the stock market next week during the depth of thesummer doldrums.

sentence-transformer

和上一期一樣，我們利用sentence-transformer 生成句子級(jí)別的embedding。其原理基于 Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks （https://arxiv.org/abs/1908.10084）這篇論文?；舅枷牒苤苯?，將句子中的每個(gè)詞的 bert embedding ，輸進(jìn)入一個(gè)池化層(pooling)，例如選擇最簡單的平均池化層，將所有token embedding 的均值作為輸出，便得到跟輸入句子長度無關(guān)的一個(gè)定長的 sentence embedding。

結(jié)果展示

數(shù)據(jù)集 train 部分由于包含的樣本比較多，需要一段時(shí)間生成 bert embedding，大家可以使用 load_dataset(part='test') 來快速體驗(yàn)。下面我們演示一個(gè)查詢 how to make money 的最接近結(jié)果。

index?=?load_index('news_train.index')
list_id?=?query(model,?index,?'how?to?make?money')
for?id?in?list_id:
????print(list_str[id])

Profit From That Traffic Ticket Got a traffic ticket? Can't beat 'em? Join 'em by investing in the company that processes those tickets.

Answers in the Margins By just looking at operating margins, investors can find profitable industry leaders.

Types of Investors: Which Are You? Learn a little about yourself, and it may improve your performance.

Target Can Aim High Target can maintain its discount image while offering pricier services and merchandise.

Finance moves Ford into the black US carmaker Ford Motor returns to profit, as the money it makes from lending to customers outweighs losses from selling vehicles.

核心代碼

所有可運(yùn)行代碼和數(shù)據(jù)都已經(jīng)包含在 docker 鏡像中了，下面列出核心代碼

建立索引

def?train_flat(index_name,?id_list,?embedding_list,?num_clusters):
????import?numpy?as?np
????import?faiss

????dim?=?768
????m?=?16
????
????embeddings?=?np.asarray(embedding_list)
????
????quantiser?=?faiss.IndexFlatIP(dim)
????index?=?faiss.IndexIVFFlat(quantiser,?dim,?num_clusters,?faiss.METRIC_INNER_PRODUCT)
????index.train(embeddings)??##?clustering
????
????ids?=?np.arange(0,?len(id_list))
????ids?=?np.asarray(ids.astype('int64'))
????
????index.add_with_ids(embeddings,?ids)
????print(index.is_trained)?
????print("Total?Number?of?Embeddings?in?the?index",?index.ntotal)
????faiss.write_index(index,?index_name)

查詢結(jié)果

def?query(model,?index,?query_str:?str)?->?List[int]:
????topk?=?5
????q_embed?=?model.encode([query_str])
????D,?I?=?index.search(q_embed,?topk)
????print(D)
????print(I)
????return?I[0].tolist()

?

?

審核編輯 ：李倩