Datasets: AttributeError: 'DatasetDict' object has no attribute 'train_test_split'

[david-waterworth]

The following code fails with "'DatasetDict' object has no attribute 'train_test_split'" - am I doing something wrong?

from datasets import load_dataset
dataset = load_dataset('csv', data_files='data.txt')
dataset = dataset.train_test_split(test_size=0.1)

AttributeError: 'DatasetDict' object has no attribute 'train_test_split'

[SBrandeis]

Hi @david-waterworth!

As indicated in the error message, load_dataset("csv") returns a DatasetDict object, which is mapping of str to Dataset objects. I believe in this case the behavior is to return a train split with all the data.
train_test_split is a method of the Dataset object, so you will need to do something like this:

dataset_dict = load_dataset(`'csv', data_files='data.txt')
dataset = dataset_dict['split name, eg train']
dataset.train_test_split(test_size=0.1)

Please let me know if this helps. 🙂

[david-waterworth]

Thanks, that's working - the same issue also tripped me up with training.

I also agree https://github.com/huggingface/datasets/issues/767 would be a useful addition.

[SBrandeis]

Closing this now