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

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)

完整的代码如下：

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 = model.encode(sentences)

data['vector'] = sentence_embeddings