| from pandas import DataFrame, Series, read_csv |
| from sklearn.model_selection import train_test_split |
|
|
| from src.config import DATASET_FILE_PATH |
|
|
|
|
| def get_dataset() -> DataFrame: |
| """ |
| Get the dataset |
| |
| Returns: |
| DataFrame: The dataset as a DataFrame |
| """ |
| try: |
| return DataFrame(data=read_csv(DATASET_FILE_PATH)) |
| except FileNotFoundError: |
| return DataFrame(data={}) |
|
|
|
|
| def get_features_target(df: DataFrame) -> tuple[DataFrame, Series]: |
| """ |
| Get the feature and target from the dataset |
| |
| Args: |
| df (DataFrame): The dataset as a DataFrame |
| |
| Returns: |
| tuple[DataFrame, Series]: The features and target as a tuple |
| """ |
| return df.drop(columns=["TARGET"], axis=1), df["TARGET"] |
|
|
|
|
| def get_train_test_sets( |
| X: DataFrame, y: Series |
| ) -> tuple[DataFrame, Series, DataFrame, Series]: |
| """ |
| Get the train and test sets from the features and target |
| |
| Args: |
| features (DataFrame): The features as a DataFrame |
| target (Series): The target as a Series |
| |
| Returns: |
| tuple[DataFrame, Series, DataFrame, Series]: The train and test sets as a tuple |
| """ |
| X_train, X_test, y_train, y_test = train_test_split( |
| X, y, test_size=0.2, random_state=42 |
| ) |
| return X_train, y_train, X_test, y_test |
|
|
