Adds unit tests for the dataloader. Adds segment level training pipeline test. PiperOrigin-RevId: 445277327
PiperOrigin-RevId: 402314322
