training/data_split.py

"""Train/test split."""

from sklearn.model_selection import train_test_split

from config.constants import LABEL_COLUMN, RANDOM_STATE, TEST_SIZE, TEXT_COLUMN


def split_data(df, text_column=TEXT_COLUMN, label_column=LABEL_COLUMN,
               test_size=TEST_SIZE, random_state=RANDOM_STATE):
    """Split the cleaned DataFrame into train/test text and labels."""
    X = df[text_column]
    y = df[label_column]
    return train_test_split(X, y, test_size=test_size, shuffle=True,
                            random_state=random_state)