Weird example of batching in Dataset.map document

Hi @lhoestq , with that example above, it is just splitting sentence1, but not saving the labels, ie the label and idx columns in the chunked_dataset. How does one split a long sentence and save the labels in a format for the Trainer to understand?