How to use AutoModel

Beginners
1

Hello,

I am trying to implement Bert with multiple heads. However, im currently facing problem with loading BERT model and BertOnlyMLMHead separately to build a functional model.

When I use the pre-implemented BertForMaskedLM model I get reasonable MASK predictions, when I use my own implementation with AutoModel.from_pretrained or BertModel.from_pretrained I dont get reasonable predictions.

Does anyone have any idea what I am doing wrong?

from transformers import AutoModel, AutoTokenizer, AutoConfig, BertModel, 
BertForMaskedLM, BertTokenizer
import torch
from transformers.modeling_bert import BertOnlyMLMHead

print('Auto Model')
tok_auto = AutoTokenizer.from_pretrained("bert-base-uncased")
bert_auto = AutoModel.from_pretrained("bert-base-uncased")
cls_auto = BertOnlyMLMHead(bert_auto.config)
bert_auto.eval()
cls_auto.eval()

input_idx_auto = tok_auto.encode(f"Paris is the capital of {tok_auto.mask_token}.")
bert_auto_output = bert_auto(torch.tensor([input_idx_auto]))
logits = cls_auto(bert_auto_output[0])
prediction_auto = logits[0].argmax(dim=1)
res = []
for i in range(len(prediction_auto)):
    res.append(tok_auto.ids_to_tokens[prediction_auto[i].numpy().tolist()])
print(res)


print('Bert + mlmHead')
tok = BertTokenizer.from_pretrained('bert-base-uncased')
bert = BertModel.from_pretrained("bert-base-uncased")
cls = BertOnlyMLMHead(bert.config)
bert.eval()
cls.eval()

input_idx = tok.encode(f"Paris is the capital of {tok.mask_token}.")
bert_output = bert(torch.tensor([input_idx]))

print('MLM_BERT')
tok_mlm = BertTokenizer.from_pretrained('bert-base-uncased')
bert_mlm = BertForMaskedLM.from_pretrained("bert-base-uncased")
bert_mlm.eval()

input_idx_mlm = tok_mlm.encode(f"Paris is the capital of {tok_mlm.mask_token}.")
logits = bert_mlm(torch.tensor([input_idx_mlm]))[0]
prediction = logits[0].argmax(dim=1)
res = []
for i in range(len(prediction)):
     res.append(tok_mlm.ids_to_tokens[prediction[i].numpy().tolist()])
print(res)



Outputs:
Auto Model
['webster', '##lier', '[unused20]', 'forcing', 'bulls', '##zia', 'terry', '##den', 'reference']
Bert + mlmHead
['æ°·', 'illnesses', 'puppet', 'intended', 'intended', 'intended', 'niccolo', '##42', 'dinosaur']
MLM_BERT
['.', 'paris', 'is', 'the', 'capital', 'of', 'france', '.', '.']