Hi !
You can follow the instructions here 4 and use this batch iterator instead:
def batch_iterator(batch_size=1000):
    batch = []
    for example in dataset:
        batch.append(example["text"])
        if len(batch) == batch_size:
            yield batch
            batch = []
    if batch:  # yield last batch
        yield batch