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

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’

[mrscp]

dataset = load_dataset(‘csv’, data_files=[‘files/datasets/dataset.csv’]) dataset = dataset[‘train’] dataset = dataset.train_test_split(test_size=0.1)

[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. 🙂