utils/load_data.py

import pandas as pd
import streamlit as st
from pathlib import Path
from huggingface_hub import snapshot_download
import os

@st.cache_data(show_spinner="Loading data...⏳")
def load_dataset(path: str | Path, category: str | None=None):
    
    DATA_OPTIONS = {
        'Beauty': 'All_Beauty.parquet',
        'Appliances': 'Appliances.parquet',
        'Baby Products': 'Baby_Products.parquet',
        'Electronics': 'Electronics.parquet',
        'Health and Household': 'Health_and_Household.parquet',
        'Movies and TV': 'Movies_and_TV.parquet'
        }
    
    if category:
        data_path = os.path.join(path, DATA_OPTIONS[category])
        df = pd.read_parquet(data_path)
    else:
        df = pd.read_parquet(path)

    return df


@st.cache_resource
def get_data_directory():
    # This downloads the whole review_data folder from your Dataset repo
    data_path = snapshot_download(
        repo_id="tkbarb10/ads505-review-data", 
        repo_type="dataset"
    )
    return Path(data_path)