使用 RoBERTa 加速嵌入 2M 句子

Question

我有大约 200 万个句子，我想在 NLI 和 STSB 上使用 Facebook AI 的 RoBERTa-large、fine-tuned 将其转换为向量以获取句子相似性（使用很棒的 sentence-transformers 包） .

我已经有一个包含两列的数据框："utterance" 包含语料库中的每个句子，"report" 包含每个句子的文档标题。

从那里开始，我的代码如下：

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

Answer 2

在同一型号中。如果你想有更高效的编码方式。

您可以将 sentence-transoformer 模型转换为 onnxruntime 的 onnx 模型或 plantensorrt 模型。但是sentence-transformer的作者没有提供转换的方法。

我找到了一个演示转换步骤的教程。 quick_sentence_transformers

使用 RoBERTa 加速嵌入 2M 句子

Speed up embedding of 2M sentences with RoBERTa

python

nlp

transformer

word-embedding