diff --git a/model_manager/sientia/utils.py b/model_manager/sientia/utils.py index 38d5582..fd6b1c4 100644 --- a/model_manager/sientia/utils.py +++ b/model_manager/sientia/utils.py @@ -15,15 +15,20 @@ def split_train_test( """ Split arrays or matrices into random train and test subsets. + Wrapper for sklearn.model_selection.train_test_split. + Args: - *data: data to be splitted. + *data: data to be split. test_size: size of test subset. train_size: size of train subset. random_state: Seed applied to the data before applying the split. shuffle: Whether or not to shuffle the data before splitting. stratify: If not None, data is split in a stratified fashion, using this as the class labels. + Returns: X_train, X_test, y_train, y_test + + Thread-safe: This function is stateless and thread-safe. """ X_train, X_test, y_train, y_test = train_test_split( *data,