

TITLE
Getting Bert Embeddings in Batch

COMMENTS
1

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
Hi all

I have a list of sentences(a batch during training) and for every word in each sentence, I need an aligned bert embedding which should be the mean of every word-piece that word was split into.

Right now, I am doing it sentence by sentence and obtain the aligned embedding for every word by reiterating over the sentence, tokenise the individual word, note the number of word-pieces it was split into and lookup into the Bert embedding matrix to average out those rows of the matrix. Following is the code to get the aligned embeddings:

`
  
       def get_bert_aligned_embeddings(self, last_hidden_states, tokens):     

            count = 0
            aligned_embeddings = []

            for i in tokens:

                 tokenisation_length = len(self.tokenizer.tokenize(i))

                 emb = torch.mean(last_hidden_states[count:count+tokenisation_length], axis = 0)

                 count += tokenisation_length

                 aligned_embeddings.append(emb)


            aligned_embeddings = torch.stack(aligned_embeddings)

            return aligned_embeddings
`

tokens is the list of words in a sentence, last_hidden_states are the embeddings I obtained from Bert. 

So, the above function runs for every sentence and every sentence is passed one by one to Bert Model. 

I want to know if there is any faster way of doing this? Can this entire process be done in batches? Any suggestions that could help me speed up this process would be great. Thanks