src/data/ragbench_loader.py

"""RAGBench dataset loading with parquet caching.

Same loader as Notebook 5/6's `load_ragbench_config`: pulls
`rungalileo/ragbench/<config_name>` from the Hugging Face Hub once and caches
the concatenated splits as a parquet file under `data/raw/`.
"""

import pandas as pd
from datasets import load_dataset

from src import settings


def load_ragbench_config(config_name: str) -> pd.DataFrame:
    # v2 cache includes source_split and source_row_idx for cross-referencing HF dataset rows.
    cache_path = settings.DATA_DIR / "raw" / f"{config_name}_v2.parquet"
    if cache_path.exists():
        return pd.read_parquet(cache_path)

    ds = load_dataset("rungalileo/ragbench", config_name)
    df = pd.concat(
        [
            ds[split].to_pandas().assign(source_split=split)
            for split in ds.keys()
        ],
        ignore_index=True,
    )
    cache_path.parent.mkdir(parents=True, exist_ok=True)
    df.to_parquet(cache_path)
    return df