Fix train_test_split output format
There was an issue in the `transmit_format` wrapper that returned bad formats when using train_test_split.
This was due to `column_names` being handled as a List[str] instead of Dict[str, List[str]] when the dataset transform (train_test_split) returns a DatasetDict (one set of column names per split).

This should fix @timothyjlaurent 's issue in #620 and fix #676 

I added tests for `transmit_format` so that it doesn't happen again