-  Migrate the train dataset preprocessing function into the iterative process,
   rather than part of dataset building the client data function. This is necessary
   to enable serialization of datasets for sending to remote workers.

PiperOrigin-RevId: 322691112