Dividing data set into (almost) equal batches based on characteristics
If I got this right you want to split the dataset into train and test sets in a way that preserves the same proportions of examples in each class as observed in the original dataset?
This is called a stratified train-test split. See the stratify argument here sklearn split