how to get the correct embedding from Roberta transformers?

user_8787039 0
Viewed 2011

I got confused by which hidden state should I use as the output of fine-tuned Roberta transformer models.

from transformers import AutoConfig, AutoModelForMaskedLM, AutoTokenizer
config = AutoConfig.from_pretrained("roberta-base")
config.output_hidden_states = True

tok = AutoTokenizer.from_pretrained("roberta-base")
model = AutoModelForMaskedLM.from_pretrained("roberta-base", config=config)

inp = "alright let s do this  "

sentence = tok.encode(inp, padding='max_length', max_length=512, truncation=True, return_tensors='pt')

output = model(sentence)

According to Huggingface documentation for RobertaForMaskedLM:

Returns a tuple of:

  1. masked_lm_loss (optional)
  2. prediction_scores
  3. hidden_states (optional)
  4. attentions (optional)

By passing the config to enable hidden_states output, the output is a tuple of (prediction_scores, hidden_states)

My question is: should I use output[-1][0] or output[-1][-1] as the final output embedding from the fine-tuned Roberta Model? My understanding is that output[-1][0] is the initial embedding feeding into the Roberta Model, and output[-1][-1] is the final embedding output.

bert-language-model huggingface-transformers
1 Answers

output[-1][-1] is correct if you are looking for the output of the last encoding layer. You can figure this out by looking at the source code and validate it by comparing the output:

import torch

print(len(output[-1]))

outputEmbeddings = model.roberta.embeddings(sentence)

#the first tensor is the output of the embedding layer
print(torch.equal(output[1][0],  outputEmbeddings))

#the second tensor is the output of the first encoding layer
print(torch.equal(output[1][1], model.roberta.encoder.layer[0](outputEmbeddings)[0]))

previousLayer = outputEmbeddings
for x in range(12):
    #it is now the current layer
    previousLayer = model.roberta.encoder.layer[x](previousLayer)[0]
    print(torch.equal(output[1][1+x], previousLayer))

Output:

13
True
True
True
True
True
True
True
True
True
True
True
True
True
True
user_66648720
Related