src/data.py

"""Loading and light exploratory helpers for the Banking77 dataset."""

from pathlib import Path

import pandas as pd

REQUIRED_COLUMNS = {"text", "category"}


def load_data(train_path: str | Path, test_path: str | Path) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Load train/test CSVs and validate they have the expected shape."""
    train_df = pd.read_csv(train_path)
    test_df = pd.read_csv(test_path)

    for name, df in [("train", train_df), ("test", test_df)]:
        missing = REQUIRED_COLUMNS - set(df.columns)
        if missing:
            raise ValueError(f"{name} set is missing required columns: {missing}")

    return train_df, test_df


def data_quality_report(df: pd.DataFrame) -> pd.Series:
    """Missing values, duplicate rows, and blank-text rows for a quick sanity check."""
    return pd.Series(
        {
            "rows": len(df),
            "missing_text": df["text"].isna().sum(),
            "missing_category": df["category"].isna().sum(),
            "duplicate_rows": df.duplicated(subset=["text", "category"]).sum(),
            "blank_text": (df["text"].str.strip() == "").sum(),
            "num_classes": df["category"].nunique(),
        }
    )


def with_text_length(df: pd.DataFrame) -> pd.DataFrame:
    """Return a copy of df with a word-count `text_length` column."""
    out = df.copy()
    out["text_length"] = out["text"].str.split().str.len()
    return out