data_source.py

# data_source.py
from __future__ import annotations

import os

import pandas as pd
from huggingface_hub import hf_hub_download


def get_token() -> "str | None":
    return os.environ.get("HF_TOKEN")


def load_slices(cfg: dict, token: "str | None" = None):
    """Download medians + metros CSVs from the PRIVATE HF Dataset and return
    them as DataFrames. Token comes from the HF_TOKEN Space secret."""
    d = cfg["data"]
    token = token or get_token()
    medians_path = hf_hub_download(
        repo_id=d["dataset_repo"], filename=d["medians_file"],
        repo_type="dataset", token=token,
    )
    metros_path = hf_hub_download(
        repo_id=d["dataset_repo"], filename=d["metros_file"],
        repo_type="dataset", token=token,
    )
    medians = pd.read_csv(medians_path, dtype={"procedure_code": str})
    metros = pd.read_csv(metros_path, dtype={"h3_cell": str})
    return medians, metros