时间戳

首页 > 解决方案 > 如何将句子-Bert 输出向量保存到文件中？

embedding - 如何将句子-Bert 输出向量保存到文件中？

问题描述

我正在使用 Bert 来获取多词之间的相似性。这是我用于嵌入的代码：

from sentence_transformers import SentenceTransformer
model = SentenceTransformer('bert-large-uncased-whole-word-masking')
words = [
"Artificial intelligence",
"Data mining",
"Political history",
"Literature book"]

我还有一个包含 540000 个单词的数据集。

Vocabs = [
"Winter flooding",
"Cholesterol diet", ....]

问题是当我想将 Vocabs 嵌入到向量中时，它永远需要时间。

words_embeddings = model.encode(words)
Vocabs_embeddings = model.encode(Vocabs)

有没有办法让它更快？或者我想在 for 循环中嵌入词汇并将输出向量保存在一个文件中，这样我就不必每次需要时都嵌入 540000 个词汇。有没有办法将嵌入保存到文件中并再次使用它？我将非常感谢您花时间帮助我。

标签： embeddingbert-language-modelword-embeddingsentence-transformers

解决方案

您可以像这样腌制您的语料库和嵌入，也可以腌制字典，或者以您喜欢的任何其他格式将它们写入文件。

import pickle
with open("my-embeddings.pkl", "wb") as fOut:
    pickle.dump({'sentences': words, 'embeddings': word_embeddings},fOut)

或更一般地，如下所示，因此您在嵌入不存在时进行编码，但之后您需要从文件中加载它们，而不是重新编码您的语料库：

if not os.path.exists(embedding_cache_path):
    # read your corpus etc
    corpus_sentences = ...
    print("Encoding the corpus. This might take a while")
    corpus_embeddings = model.encode(corpus_sentences, show_progress_bar=True, convert_to_numpy=True)
    corpus_embeddings = corpus_embeddings / np.linalg.norm(corpus_embeddings, axis=1, keepdims=True)

    print("Storing file on disc")
    with open(embedding_cache_path, "wb") as fOut:
        pickle.dump({'sentences': corpus_sentences, 'embeddings': corpus_embeddings}, fOut)

else:
    print("Loading pre-computed embeddings from disc")
    with open(embedding_cache_path, "rb") as fIn:
        cache_data = pickle.load(fIn)
        corpus_sentences = cache_data['sentences']
        corpus_embeddings = cache_data['embeddings']

推荐阅读