保存 BERT 句子嵌入

Question

我目前正在從事信息檢索任務。 我正在使用 SBERT 執行語義搜索。 我已經按照這里的文檔

我使用的 model

model = SentenceTransformer('sentence-transformers/paraphrase-multilingual-mpnet-base-v2')

大綱是

1. 你有一個這樣的語料庫列表：

    data = ['A man is eating food.',
          'A man is eating a piece of bread.',
          'The girl is carrying a baby.',
          'A man is riding a horse.',
          'A woman is playing violin.',
          'Two men pushed carts through the woods.',
          'A man is riding a white horse on an enclosed ground.',
          'A monkey is playing drums.',
          'A cheetah is running behind its prey.'
          ]

2. 您有這樣的查詢：

queries = ['A man is eating pasta.']

3. 使用查詢和語料庫執行編碼

query_embedding = model.encode(query)
doc_embedding = model.encode(data)

編碼 function 輸出 numpy.ndarray 就像model.encode(data) 的輸出

4. 並像這樣使用余弦相似度計算相似度

similarity = util.cos_sim(query_embedding, doc_embedding)

5. 如果你打印相似度，你會得到 torch.Tensor 包含這樣的相似度分數

tensor([[0.4389, 0.4288, 0.6079, 0.5571, 0.4063, 0.4432, 0.5467, 0.3392, 0.4293]])

它工作得很好而且很快。 但當然它只是使用少量的語料庫。 當使用大量語料庫時，編碼工作需要時間。

注意：query的編碼不需要時間，因為它只有一個句子，但是語料庫的編碼需要一些時間

那么，問題是我們可以將 doc_embedding 保存在本地，然后再次使用它嗎？ 尤其是在使用大型語料庫時

是否有任何內置的類/功能可以從變壓器中完成？

Answer 1

在處理大型語料庫時，您需要使用矢量數據庫，我在這里和這里寫了一些關於 Faiss 的指南，您可能會發現它們很有用。 Faiss 確實需要大量學習才能獲得合理的性能，並且只存儲向量（而不是任何其他信息，如 ID、文本等），因此您需要設置另一個數據庫，如 SQL 來處理它。

根據經驗，這可能非常煩人。 所以我建議研究像Pinecone這樣的托管矢量數據庫，你可以在大約 10 分鍾內從現在的位置啟動並運行它——它最多可以釋放大約 100 萬個矢量，並且性能令人難以置信。

Answer 2

將它們保存為泡菜文件並稍后加載它們 = ]

import pickle

with open('doc_embedding.pickle', 'wb') as pkl:
    pickle.dump(doc_embedding, pkl)

with open('doc_embedding.pickle', 'rb') as pkl:
    doc_embedding = pickle.load(pkl)