It seems that using BertForNextSentencePrediction with TextDatasetForNextSentencePrediction and DataCollatorForLanguageModeling would be equivalent to the BERT objective (except static masking part). And for dataset, we can use datasets.concatenate_datasets() method for BookCorpus and Wikipedia. This might be close right ? Any additional details ?