Maybe explore flatten_indices API in datasets when doing offline dataset processing

[dushyantbehl]

Is your feature request related to a problem? Please describe.

Explore the API and use it if it makes more sense to get better performance
https://huggingface.co/docs/datasets/en/package_reference/main_classes#datasets.Dataset.flatten_indices

Describe the solution you'd like

A clear and concise description of what you want to happen.

Describe alternatives you've considered

A clear and concise description of any alternative solutions or features you've considered.

Additional context

Add any other context about the feature request here.

[dushyantbehl]

Fixed here - #494