data/loader.py

"""Reading the raw and the previously cleaned Excel datasets."""

from pathlib import Path

import pandas as pd

from config.constants import LABEL_COLUMN, TEXT_COLUMN
from config.paths import DEFAULT_CLEANED_DATA_PATH, DEFAULT_RAW_DATA_PATH


def load_raw_dataset(path=DEFAULT_RAW_DATA_PATH):
    """Load the raw Excel dataset."""
    return pd.read_excel(path)


def load_cleaned_dataset(path=DEFAULT_CLEANED_DATA_PATH):
    """Load a dataset already written by the preprocessing stage."""
    path = Path(path)
    if not path.is_file():
        raise FileNotFoundError(
            'No cleaned dataset at %s - run the pipeline once without '
            '--from-cleaned to create it.' % path
        )

    df = pd.read_excel(path)

    missing = [c for c in (TEXT_COLUMN, LABEL_COLUMN) if c not in df.columns]
    if missing:
        raise KeyError(
            '%s is missing the column(s) %s; found %s'
            % (path, missing, list(df.columns))
        )
    return df