This 216 looks relevant. You can see there an example of LM task, you can reuse it/build on it and create your own LM task inside which you will initialize the weights of bert with a pretrained version and then train it with your own data.