pytorch - 提高 Huggingface 标记器输出的速度

Question

我需要使用 HuggingFace 从 BERT 模型中获取最后一层嵌入。以下代码有效，但是速度极慢，如何提高速度？

这是一个玩具示例，我的真实数据由数千个带有长文本的示例组成。

import transformers
import pandas as pd
from transformers import BertModel, BertTokenizer

tokenizer = BertTokenizer.from_pretrained("bert-base-uncased") 
model = BertModel.from_pretrained("bert-base-uncased")

def getPrediction(text): 
  encoded_input = tokenizer(text, return_tensors='pt')
  outputs = model(**encoded_input)
  embedding = outputs[0][:, -1]  
  embedding_1ist = embedding.cpu().detach().tolist()[0]
  return embedding_1ist

df = pd.DataFrame({'text':['First text', 'Second text']})
results =  pd.DataFrame(df.apply(lambda x: getPrediction(x.text), axis=1))