main.py 文件源码

def balancedSplit(X, y, seed, test_sz=1000):
    stratSplit = StratifiedShuffleSplit(
        y, 1, test_size=test_sz, random_state=seed
    )
    for train_idx, test_idx in stratSplit:
        X_train = X[train_idx]
        y_train = y[train_idx]
        X_test = X[test_idx]
        y_test = y[test_idx]
        break
    return X_train, y_train, X_test, y_test