hey @awaiskaleem if i understand correctly, i that what you’re looking for is to fine-tune the language model on your corpus. this will generally produce mask-filling that more accurately captures the relations in your corpus and for BERT you can check out the Masked language modeling section of this tutorial: Google Colaboratory 79
once the language model is fine-tune, you can save the weights and then load them using AutoModel.from_pretrained to generate your embeddings.