ColossalAI/examples/language/deepseek/data_utils.py