ColossalAI/examples/language/mixtral/data_utils.py