How to split the data as per the label name count using sklearn

Question

Before training a model, I would like to split the data in a 80/20 ratio.

For example, in my dataset of 3000 rows, I have different labels. Out of which A is one label name, has 100 records in the dataset. How can I divide the dataset so that 80% of the label A data is in training and 20% is in the test set.

Do we have any pre defined functions for this?

Answer 1

Use a stratified test split.

import numpy as np
from sklearn.model_selection import StratifiedShuffleSplit
X = np.array([[1, 3], [3, 7], [2, 4], [4, 8]])
y = np.array([0, 1, 0, 1])
stratSplit = StratifiedShuffleSplit(y, n_iter=1, test_size=0.5, random_state=42)
for train_idx, test_idx in stratSplit:
    X_train=X[train_idx]
    y_train=y[train_idx]

print(X_train)
# [[3 7]
#  [2 4]]
print(y_train)
# [1 0]

Source: https://stackoverflow.com/a/35473089/9524424