Drop columns after loading samples in prepare_tf_dataset
Another super-small fix to `prepare_tf_dataset()` - this time we apply the same fix we applied to `to_tf_dataset()`, and keep columns until after samples have been loaded from the dataset. This ensures that columns that are needed to compute the transform aren't dropped.
_The documentation is not available anymore as the PR was closed or merged._<|||||>I think this PR is ready to go, but it's waiting on https://github.com/huggingface/datasets/pull/4553 to be merged and a release in Datasets. Tests will fail until that makes it through, so I won't merge until then!