保存 BERT 句子嵌入

Question

我目前正在从事一项信息检索任务。我正在使用 SBERT 执行语义搜索。我已经按照文档here

我使用的模型

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

1. 您有这样的查询：

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

1. 使用查询和语料库执行编码

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

encode 函数像这样输出一个 numpy.ndarray outputs of 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 here 和 here 上写了一些您可能会觉得有用的指南。 Faiss 确实需要大量学习才能获得合理的性能，并且只存储向量（而不是任何其他信息，如 ID、文本等），因此您需要设置另一个数据库（如 SQL）来处理。

根据经验，这可能非常烦人。因此，我建议您研究像 Pinecone 这样的托管矢量数据库，您可以在大约 10 分钟内从现在的位置启动并运行它 - 它最多可以免费提供大约 100 万个矢量，并且性能令人难以置信。

Answer 2

将它们保存为 pickle 文件并稍后加载它们 = ]

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)