Allow to change the input data format for pretraining data generator
Allow to change the input data format for pretraining data generator

PiperOrigin-RevId: 293917774

google-research/albert,google-research/albert