Ускорьте внедрение 2M предложений с помощью RoBERTa - PullRequest
Новая
       95

Ускорьте внедрение 2M предложений с помощью RoBERTa

0 голосов
/ 04 мая 2020

У меня есть примерно 2 миллиона предложений, которые я хочу превратить в векторы, используя большой RoBERTa AI AI, настроенный на NLI и STSB, для сходства предложений (используя потрясающий пакет-преобразователь предложений пакет).

У меня уже есть фрейм данных с двумя столбцами: «высказывание», содержащее каждое предложение из корпуса, и «отчет», содержащий для каждого предложения заголовок документа, из которого он взят.

Оттуда мой код выглядит следующим образом: 

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 часов, что больше, чем я могу сэкономить.

Можно ли как-нибудь ускорить это, изменив мой код? Является ли создание огромного списка в памяти, а затем добавление его в фрейм данных лучшим способом для продолжения? (Я подозреваю, что нет).

Большое спасибо заранее!

1 Ответ

1 голос
/ 05 мая 2020

Я обнаружил нелепое ускорение, используя этот пакет, подавая высказывания в виде списка, а не перебирая список. Я предполагаю, что происходит некоторая хорошая внутренняя векторизация. 

%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)
...