I am trying to add noise to the embeddings from the 5th layer of bert model before it is forwarded to the next layer. The final layer output after the noise is added is then used for classification task. How can I do this in huggingface bert model?
class MainModel(BertPreTrainedModel):
def __init__(self, config, loss_fn):
super(MainModel,self).__init__(config)
self.num_labels = 2
self.loss_fn = loss_fn
config.output_hidden_states = True
self.bert = AutoModel.from_pretrained("bert-base-uncased",config = config)
self.classifier = nn.Linear(768, self.num_labels)
def forward(self, input_ids, attention_mask, token_type_ids, labels, noise, device):
output = self.bert(input_ids, attention_mask = attention_mask, token_type_ids = token_type_ids)
hidden_emb = output.hidden_states[5]
hidden_emb = hidden_emb + noise
#Forward these hidden_emb to next layer of bert model
#Write the code here
output = output.last_hidden_state #Output from final layer of bert model after adding noise
output = output[:,0,:]
classifier_out = self.classifier(output)
main_prob = F.softmax(classifier_out, dim = 1)
main_gold_prob = torch.gather(main_prob, 1, labels)
loss_main = self.loss_fn.forward(main_gold_prob)
return loss_main,main_prob