问用于极端多标签分类的测试分割

EN

from sklearn import datasets
from sklearn.model_selection import train_test_split

iris = datasets.load_iris()
X = pd.DataFrame(iris.data)
y = iris.target
# these labels will not cause any problems
X['cat'] = np.random.choice(['label1','label2','label3','label4'],len(X))

# but these ones will, because they are unique
X.loc[37, 'cat'] = 'label5'
X.loc[137, 'cat'] = 'label6'

# this row will raise an exception if uncommented
#X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, 
#                                                    random_state=43, stratify = X.cat)

# let's drop rows with unique labels
from collections import Counter
unique_labels = [lab for lab, count in Counter(X.cat).items() if count == 1]

print(f"unique labels to be dropped: {unique_labels}")

# drop rows with unique labels
X = X[~X.cat.isin(unique_labels)]
y = y[X.index]

# now datasets X and y can be used in train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, 
                                                    random_state=43, stratify = X.cat)