How can I get RoBERTa word embeddings?

Viewed 2133

Given a sentence of the type 'Roberta is a heavily optimized version of BERT.', I need to get the embeddings for each of the words in this sentence with RoBERTa. I have tried to look at the sample codes online, failing to find a definite answer.

My take is the following:

tokens = roberta.encode(headline)
all_layers = roberta.extract_features(tokens, return_all_hiddens=True)
embedding = all_layers[0]
n = embedding.size()[1] - 1
embedding = embedding[:,1:n,:]

where embedding[:,1:n,:] is used to extract only the embeddings for the words in the sentence, without the start and end tokens.

Is it correct?

1 Answers
TOKENIZER_PATH = "../input/roberta-transformers-pytorch/roberta-base"
ROBERTA_PATH = "../input/roberta-transformers-pytorch/roberta-base"

text= "How are you? I am good."
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH)

##how the words are broken into tokens
print(tokenizer.tokenize(text))

##the format of a encoding
print(tokenizer.batch_encode_plus([text]))

##op wants the input id
print(tokenizer.batch_encode_plus([text])['input_ids'])

##op wants the input id without first and last token
print(tokenizer.batch_encode_plus([text])['input_ids'][0][1:-1])

Output:

['How', 'Ġare', 'Ġyou', '?', 'ĠI', 'Ġam', 'Ġgood', '.']

{'input_ids': [[0, 6179, 32, 47, 116, 38, 524, 205, 4, 2]], 'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]}

[[0, 6179, 32,47, 116, 38, 524, 205, 4, 2]]

[6179, 32, 47, 116, 38, 524, 205, 4]

And dont worry about the "Ġ" character. It just indicates that there is a space in front of the word.

Related