Stop dropping columns in to_tf_dataset() before we load batches
`to_tf_dataset()` dropped unnecessary columns before loading batches from the dataset, but this is causing problems when using a transform, because the dropped columns might be needed to compute the transform. Since there's no real way to check which columns the transform might need, we skip dropping columns and instead drop keys from the batch after we load it.

cc @amyeroberts and https://github.com/huggingface/notebooks/pull/202