Updating a train/val/test set

In case of imbalanced datasets, we can either use stratify while splitting the data, or use Cross validation, or both. Stratifying/CV upfront helps mitigate this data leakage by developer confirmation bias.

Stratify

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y,
                                                    stratify=y, 
                                                    test_size=0.25)

Cross validation

>>> from sklearn.model_selection import cross_val_score
>>> clf = svm.SVC(kernel='linear', C=1, random_state=42)
>>> scores = cross_val_score(clf, X, y, cv=5)
>>> scores
    array([0.96..., 1. , 0.96..., 0.96..., 1. ])

In this example there are 5 random splits, 5 models score.

RepeatedStratifiedKFold helps.

from sklearn.model_selection import cross_val_score
from sklearn.model_selection import RepeatedStratifiedKFold

cv = RepeatedStratifiedKFold(n_splits=5, n_repeats=2, random_state=42)
scores = cross_val_score(model, X, y, scoring='mean_squared_error', cv=cv, n_jobs=-1)
print('Mean MSE: %.4f' % mean(scores))

In this example there are 5 splits and two iteration, total of 10 models aggregated score.

There are two really different scenarios:

The training and test data are obtained from the same dataset

If the data has been randomly split between training and test set, this is extremely unlikely to happen in the first place. If the data contains some small groups/classes, then the split can be made not only randomly but also with stratified sampling: this prevents that 100% of one group ends up on either side of the split by chance. And if there are some classes or groups which still appear too rarely (only once or twice), then these should usually be discarded or replaced by some generic value at some preprocessing stage.

The training and test data are obtained independently

This is the serious case. There are cases where external constraints make it unavoidable to have slightly different distributions in the training and test set, even though this breaches the main assumption of supervised learning. Note that if the distributions differ too much, it's very likely a lost cause. In this scenario one doesn't have a choice: if the training and test set are provided separately,then it's expected that the performance is measure on the test set "as is". So it's a matter of working with this constraint: some specific preprocessing may be necessary (e.g. introduce a special group/class 'unknown' in the training set), a robust method may be preferable, possibly some plan a default prediction (majority class) for invalid instances, etc.

Mistakes happen

In any case, if one realizes that there's something wrong in the design of the splitting process or any other problem which makes it necessary to re-shuffle the data, well, it's sometimes better to redo the whole process despite the risk of data leakage. Of course it's better if this can be avoided, but it's not as if the ML police is going to arrest you ;)