python - 使用 RoBERTa 加速嵌入 200 万个句子

Question

我有大约 200 万个句子，我想使用 Facebook AI 的 RoBERTa-large 将它们转换为向量，在 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)