Dataset loading issue, integration with Colossal-AI

[Edenzzzz]

Hi, Thanks for the good work. I'm trying to intergrate this into Colossal-AI(https://github.com/hpcaitech/ColossalAI), compatible with Tensor Parallel and ZeRO. However, I had trouble loading the dataset; seems they updated the dataset to remove the json schema. Could you share your dataset version and how you're able to load it? Thanks!