# How to use sklearn train_test_split to stratify data for multi-label classification?

I am attempting to mirror a machine learning program by Ahmed Besbes, but scaled up for multi-label classification. It seems that any attempt to stratify the data returns the following error: The least populated class in y has only 1 member, which is too few. The minimum number of labels for any class cannot be less than 2.

In my data set, I have 1 column which contains clean, tokenized text. The other 8 columns are for the classifications based on the content of that text. Just to note, column 1 - 4 have significantly more samples than 5 - 8 (more obscure classifications derived from the text).

Here is a generic sample from my code:

x = data['cleaned_text']
y = data[['car','truck','ford','chevy','black','white','parked', 'driving']]

x_train, x_test, y_train, y_test = train_test_split(x,
y,
test_size=0.1,
random_state=42)

print(x_train.shape, x_test.shape, y_train.shape, y_test.shape)


Output: (6293,) (700,) (6293, 8) (700, 8)

Adding stratify=y to train_test_split returns the error previously mentioned. Even when I limit y to just one column, I still get the error.

How can I stratify the data so that I give the program a fair look in the training set?

• You try to predict more than one class at the same time. It's not a multi-class classification, but a multi-label classification problem. Please add a sample of your dataset since it is not clear what you try to do. – Tasos Feb 6 at 16:57
• Thanks, I edited the title and body of the initial question to reflect multi-label vice multi-class. As for the data, I can give a generic example – Michael Joy Feb 6 at 17:08