42 lines
1.2 KiB
Python
42 lines
1.2 KiB
Python
from typing import Any
|
|
|
|
from numpy.typing import ArrayLike
|
|
from sklearn.model_selection import train_test_split
|
|
|
|
|
|
def split_train_test(
|
|
*data: Any,
|
|
test_size: float | None = None,
|
|
train_size: float | None = None,
|
|
random_state: int | None = None,
|
|
shuffle: bool = True,
|
|
stratify: ArrayLike | None = None,
|
|
) -> tuple[Any, Any, Any, Any]:
|
|
"""
|
|
Split arrays or matrices into random train and test subsets.
|
|
|
|
Wrapper for sklearn.model_selection.train_test_split.
|
|
|
|
Args:
|
|
*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,
|
|
test_size=test_size,
|
|
train_size=train_size,
|
|
random_state=random_state,
|
|
shuffle=shuffle,
|
|
stratify=stratify,
|
|
)
|
|
return X_train, X_test, y_train, y_test
|