使用RoBERTa加速2M句子的嵌入

Question

[我想使用Facebook AI的RoBERTa-large将大约200万个句子转换为矢量，并在NLI和STSB上进行微调以实现句子相似性（使用令人敬畏的sentence-transformers程序包。

[我已经有一个包含两列的数据框：“话语”包含来自语料库的每个句子，“报告”包含每个句子的文档的标题。

从那里，我的代码如下：

from sentence_transformers import SentenceTransformer
from tqdm import tqdm

model = SentenceTransformer('roberta-large-nli-stsb-mean-tokens')

print("Embedding sentences")

data = pd.read_csv("data/sentences.csv")

sentences = data['utterance'].tolist()

sentence_embeddings = []

for sent in tqdm(sentences):
    embedding = model.encode([sent])
    sentence_embeddings.append(embedding[0])

data['vector'] = sentence_embeddings

现在，tqdm估计整个过程将在我的计算机上花费大约160个小时，这比我可以节省的时间还多。

有什么方法可以通过更改代码来加快速度？是在内存中创建一个巨大的列表，然后将其追加到数据框，这是进行此操作的最佳方法吗？ （我怀疑不是）。

非常感谢！

Answer 1

%timeit utterances_enc = model.encode(utterances[:10])                                                                                                                                                                                                                 
3.07 s ± 53.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit utterances_enc = [model.encode(utt) for utt in utterances[:10]]
4min 1s ± 8.08 s per loop (mean ± std. dev. of 7 runs, 1 loop each)