kirad.dataset_utils

kirad.dataset_utils

Data stage helpers: load metadata CSVs, build file paths, merge rows, collate batches.

broadcast_collate is the DataLoader collate_fn (set in launcher_utils). The Kintsugi dataset-version / GCS bucket helpers need private gsutil URIs from constants.py.

Contains utility functions for processing and manipulating datasets, primarily in pandas DataFrame format.

DistributionData = namedtuple('DistributionData', ['distribution', 'bin_edges', 'categories']) module-attribute

concat_collate = get_collate_fn(torch.concat) module-attribute

concat_collate_to_cpu = get_collate_fn(lambda batch: torch.concat(batch).to('cpu')) module-attribute

load_dataset(dataset, data_sources=None, dataset_splits=None, dataset_fold=None)

Return metadata as a DataFrame: pass-through, read a CSV path, or a Kintsugi version.

Note: DatasetVersion names download private splits with gsutil; use a CSV path.

Source code in kirad/dataset_utils.py

def load_dataset(
    dataset: pd.DataFrame | DatasetVersion | str | Path | Mapping,
    data_sources: Optional[DataSource | list[DataSource]] = None,
    dataset_splits: Optional[DatasetSplit | list[DatasetSplit]] = None,
    dataset_fold: Optional[int] = None,
) -> pd.DataFrame:
    """Return metadata as a DataFrame: pass-through, read a CSV path, or a Kintsugi version.

    Note: DatasetVersion names download private splits with gsutil; use a CSV path.
    """
    if isinstance(dataset, pd.DataFrame):
        return dataset
    elif dataset in typing.get_args(DatasetVersion):
        # User provided a dataset version
        # Versioned datasets need sources, splits and a fold (except test) to pick CSVs.
        if (
            data_sources is None
            or dataset_splits is None
            or (dataset_splits != "test" and dataset_fold is None)
        ):
            raise ValueError(
                "To use a dataset from `kintsugi-ml-datasets`, pass in "
                "`data_sources`, `dataset_splits`, and `dataset_fold`"
            )
        # Normalize sources and splits to lists.
        if isinstance(data_sources, str) or not isinstance(data_sources, Iterable):
            data_sources = [data_sources]
        if isinstance(dataset_splits, str) or not isinstance(dataset_splits, Iterable):
            dataset_splits = [dataset_splits]
        # Fetch split CSVs; concatenate each source/split's first DataFrame.
        full_dataset = load_dataset_version(
            dataset, data_sources, dataset_splits, dataset_fold
        )
        return pd.concat(
            [
                full_dataset[ds][split][0]
                for split in dataset_splits
                for ds in data_sources
            ],
            ignore_index=True,
        )
    elif isinstance(dataset, str) or isinstance(dataset, Path):
        # User provided a string. Assume this is a path to CSV file.
        return pd.read_csv(dataset)
    else:
        raise ValueError(f"Expected a DataFrame or a string; received {type(dataset)}")

load_relative_dataset(dataset, data_sources=None, dataset_splits=None, dataset_fold=None, audio_root=None, text_root=None)

Load metadata and add audio_filename / text_filename paths from the filename column.

A Mapping input {metadata, audio_root, text_root} overrides the roots (default: cwd). Paths: /.wav and /.json.

Source code in kirad/dataset_utils.py

def load_relative_dataset(
    dataset: pd.DataFrame | DatasetVersion | str | Path | Mapping,
    data_sources: Optional[DataSource | list[DataSource]] = None,
    dataset_splits: Optional[DatasetSplit | list[DatasetSplit]] = None,
    dataset_fold: Optional[int] = None,
    audio_root: Optional[str | Path] = None,
    text_root: Optional[str | Path] = None,
) -> pd.DataFrame:
    """Load metadata and add audio_filename / text_filename paths from the filename column.

    A Mapping input {metadata, audio_root, text_root} overrides the roots (default: cwd).
    Paths: <audio_root>/<filename>.wav and <text_root>/<filename>.json.
    """
    if isinstance(dataset, Mapping):
        audio_root = dataset.get("audio_root", audio_root)
        text_root = dataset.get("text_root", text_root)
        dataset = dataset["metadata"]
    # Load metadata; roots default to the current directory.
    df = load_dataset(dataset, data_sources, dataset_splits, dataset_fold)
    audio_root = audio_root or os.getcwd()
    text_root = text_root or os.getcwd()
    # HARDCODED: audio files must be <filename>.wav, transcripts <filename>.json
    df["audio_filename"] = df["filename"].map(
        lambda s: os.path.expanduser(os.path.join(audio_root, str(s) + ".wav"))
    )
    df["text_filename"] = df["filename"].map(
        lambda s: os.path.expanduser(os.path.join(text_root, str(s) + ".json"))
    )
    return df

load_dataset_version(dataset_version, data_sources=None, dataset_splits=None, dataset_folds=None)

Fetch a Kintsugi dataset version (gsutil rsync) and read its split/fold CSVs.

Returns {source: {split: [DataFrame per fold]}}; test has no folds. Hardcoded: METADATA_DIR//splits/..fold_.csv (test: no fold suffix).

Source code in kirad/dataset_utils.py

def load_dataset_version(
    dataset_version: DatasetVersion,
    data_sources: Optional[DataSource | list[DataSource]] = None,
    dataset_splits: Optional[DatasetSplit | list[DatasetSplit]] = None,
    dataset_folds: Optional[int | list[int]] = None,
) -> dict[DataSource, dict[DatasetSplit, list[pd.DataFrame]]]:
    """Fetch a Kintsugi dataset version (gsutil rsync) and read its split/fold CSVs.

    Returns {source: {split: [DataFrame per fold]}}; test has no folds.
    Hardcoded: METADATA_DIR/<version>/splits/<source>.<split>.fold_<k>.csv (test: no fold suffix).
    """
    # Validate the version against constants.py.
    if dataset_version not in typing.get_args(DatasetVersion):
        raise ValueError(
            f"Dataset version {dataset_version} is not registered kirad/constants.py."
        )
    if dataset_version not in DATASET_VERSION_METADATA:
        raise ValueError(
            f"Metadata for dataset version {dataset_version} is not defined in"
            "kirad/constants.py."
        )

    # Normalize to lists; sources/splits: None or "all" = every one; folds: None = all.
    if data_sources is None or data_sources == "all" or "all" in data_sources:
        data_sources = [i for i in typing.get_args(DataSource) if i != "all"]
    elif isinstance(data_sources, str) or not isinstance(data_sources, Iterable):
        data_sources = [data_sources]
    if dataset_splits is None or dataset_splits == "all" or "all" in dataset_splits:
        dataset_splits = [i for i in typing.get_args(DatasetSplit) if i != "all"]
    elif isinstance(dataset_splits, str) or not isinstance(dataset_splits, Iterable):
        dataset_splits = [dataset_splits]
    if dataset_folds is None:
        dataset_folds = list(
            range(0, DATASET_VERSION_METADATA[dataset_version]["n_folds"])
        )
    elif not isinstance(dataset_folds, Iterable):
        dataset_folds = [dataset_folds]

    # Download metadata
    # When the destination directory already exists or when a disk is mounted in a
    # read-only mode don't try to copy the data. It is assumed the data is already
    # there and have the correct directory structure.
    out_dir = METADATA_DIR / dataset_version
    dir_exists = out_dir.exists()
    if not dir_exists:
        out_dir.mkdir(parents=True)
    # Detect a read-only disk (no download possible).
    stat = os.statvfs(out_dir)
    read_only = bool(stat.f_flag & os.ST_RDONLY)
    # Download only into a new writable dir; otherwise reuse the local copy.
    if not dir_exists and not read_only:
        if DATASET_VERSION_METADATA[dataset_version]["uri"] is None:
            raise ValueError(
                f"Need to define URI for dataset version {dataset_version}."
            )
        # HARDCODED: download via gsutil from a private GCS URI in constants.py
        cmd = [
            "gsutil",
            "-m",
            "rsync",
            "-r",
            DATASET_VERSION_METADATA[dataset_version]["uri"],
            out_dir,
        ]
        # Run the rsync, logging gsutil output on failure.
        try:
            subprocess.check_output(
                cmd, stderr=subprocess.STDOUT, universal_newlines=True
            )
        except subprocess.CalledProcessError as e:
            logger.info(e.output)
            raise e
    else:
        if dir_exists:
            logger.info("Dataset directory already exists. Skipping metadata download.")
        else:
            logger.info("Disk is in read-only mode. Skipping metadata download.")

    # Gather the data across the data source and splits
    full_dataset = dict()
    for data_source in data_sources:
        # One list of DataFrames per split (train/val: one per fold).
        full_dataset[data_source] = {"train": [], "val": [], "test": []}
        for dataset_split in dataset_splits:
            # Test has a single file with no fold suffix.
            if dataset_split == "test":
                dataset_fold_strs = [""]
            else:
                # Train and val splits
                dataset_fold_strs = [f".fold_{fold}" for fold in dataset_folds]
            # Read each CSV, drop duplicate recordings, tag split; skip missing files.
            for fold_str in dataset_fold_strs:
                data_csv = (
                    out_dir / "splits" / f"{data_source}.{dataset_split}{fold_str}.csv"
                )
                try:
                    tmp_df = pd.read_csv(data_csv).drop_duplicates(
                        subset="filename", ignore_index=True
                    )
                    tmp_df["split"] = dataset_split
                    full_dataset[data_source][dataset_split].append(tmp_df)
                except FileNotFoundError as e:
                    logger.warning(e)

    return full_dataset

resolve_audio_path(audio_dir_or_feature_store_bucket, is_raw_audio=False, force_download=False)

Return a local audio directory: a given path, or a synced feature-store bucket.

Source code in kirad/dataset_utils.py

def resolve_audio_path(
    audio_dir_or_feature_store_bucket: FeatureStoreBucket | str | Path,
    is_raw_audio: bool = False,
    force_download: bool = False,
) -> Path:
    """Return a local audio directory: a given path, or a synced feature-store bucket."""
    # Bucket name -> synced local folder; any other value is treated as a path.
    if audio_dir_or_feature_store_bucket in typing.get_args(FeatureStoreBucket):
        return resolve_audio_path_from_bucket(
            audio_dir_or_feature_store_bucket,
            is_raw_audio=is_raw_audio,
            force_download=force_download,
        )
    elif isinstance(audio_dir_or_feature_store_bucket, str | Path):
        return Path(audio_dir_or_feature_store_bucket)
    else:
        raise ValueError(
            f"Expected Path or a string; received {type(audio_dir_or_feature_store_bucket)}"
        )

resolve_audio_path_from_bucket(feature_store_bucket, is_raw_audio=False, force_download=False)

Download a feature-store audio bucket with gsutil (if missing) and return its path.

Source code in kirad/dataset_utils.py

def resolve_audio_path_from_bucket(
    feature_store_bucket: FeatureStoreBucket,
    is_raw_audio: bool = False,
    force_download: bool = False,
) -> Path:
    """Download a feature-store audio bucket with gsutil (if missing) and return its path."""
    if feature_store_bucket not in typing.get_args(FeatureStoreBucket):
        raise ValueError(
            f"Received feature store bucket {feature_store_bucket}; "
            f"available dataset versions are {typing.get_args(FeatureStoreBucket)}"
        )

    audio_dir = get_cache_dir(feature_store_bucket)
    # Download only when not cached locally (or when forced).
    if not audio_dir.exists() or force_download:
        if FEATURE_STORE_URI is None:
            raise ValueError("Need to define URI for `FEATURE_STORE_URI`.")
        if RAW_AUDIO_URI is None:
            raise ValueError("Need to define URI for `RAW_AUDIO_URI`.")
        # Raw audio has its own URI; processed features live under one per bucket.
        download_uri = (
            RAW_AUDIO_URI
            if is_raw_audio
            else f"{FEATURE_STORE_URI}/{feature_store_bucket}"
        )
        # Create the local dir and sync the bucket into it.
        audio_dir.mkdir(parents=True, exist_ok=True)
        cmd = ["gsutil", "-m", "rsync", "-r", download_uri, audio_dir]
        subprocess.check_output(cmd, stderr=subprocess.STDOUT, universal_newlines=True)

    return audio_dir

get_cache_dir(location)

Local directory for a bucket name (AUDIO_DIR/) or a given path.

Source code in kirad/dataset_utils.py

def get_cache_dir(location: FeatureStoreBucket | str | Path) -> Path:
    """Local directory for a bucket name (AUDIO_DIR/<bucket>) or a given path."""
    # Bucket name -> AUDIO_DIR/<bucket>; any other value is treated as a path.
    if location in typing.get_args(FeatureStoreBucket):
        return AUDIO_DIR / location
    elif isinstance(location, str | Path):
        return Path(location)
    else:
        raise ValueError(
            f"Received feature store bucket {location}; "
            f"available dataset versions are {typing.get_args(FeatureStoreBucket)}"
        )

calculate_distributions(df, label_column, other_columns=None, existing_label_distribution=None, existing_other_distributions=None)

Label and column distributions (histograms / category masses) used by split_dataset.

Calculate distributions for the labels and specified columns in the data

The distributions will be density-normalized histograms for real-valued data, and probability masses for categorical data. If existing distributions are specified, then the bin edges/categories from those distributions will be used for calculating the distributions.

Parameters:

Name	Type	Description	Default
df	DataFrame	Pandas DataFrame containing the dataset.	required
label_column	str	Column in the DataFrame corresponding to the target labels (eg. "bin_phq"). For now, assumes that the labels are categorical and encoded as non-negative integers.	required
other_columns	list[str]	List of column names in the DataFrame for which to calculate distributions.	None
existing_label_distribution	DistributionData	Contains distribution information about the labels. If provided, the calculated distributions will use the same bin edges/categories.	None
existing_other_distributions	dict[str, DistributionData]	Contains distribution information about various columns in the DataFrame, where the dictionary key specifies the column name. If provided, the calculated distributions will use the same bin edges/categories.	None

Returns:

Name	Type	Description
label_distribution	DistributionData	Contains distribution information about the labels
other_distributions	dict[str, DistributionData]	Contains distribution information about various columns in the DataFrame, where the dictionary key specifies the column name. Will be None if other_columns is None.

Raises:

Type	Description
ValueError	Raised if one of the specified other columns does not exist in the dataset.

Source code in kirad/dataset_utils.py

def calculate_distributions(
    df: pd.DataFrame,
    label_column: str,
    other_columns: Optional[list[str]] = None,
    existing_label_distribution: Optional[DistributionData] = None,
    existing_other_distributions: Optional[dict[str, DistributionData]] = None,
) -> tuple[DistributionData, Optional[dict[str, DistributionData]]]:
    """Label and column distributions (histograms / category masses) used by split_dataset.

    Calculate distributions for the labels and specified columns in the data

    The distributions will be density-normalized histograms for real-valued data, and
    probability masses for categorical data. If existing distributions are specified,
    then the bin edges/categories from those distributions will be used for calculating
    the distributions.

    Parameters
    ----------
    df : pandas.DataFrame
        Pandas DataFrame containing the dataset.
    label_column : str
        Column in the DataFrame corresponding to the target labels (eg. "bin_phq"). For
        now, assumes that the labels are categorical and encoded as non-negative
        integers.
    other_columns : list[str], optional
        List of column names in the DataFrame for which to calculate distributions.
    existing_label_distribution : DistributionData, optional
        Contains distribution information about the labels. If provided, the calculated
        distributions will use the same bin edges/categories.
    existing_other_distributions : dict[str, DistributionData], optional
        Contains distribution information about various columns in the DataFrame, where
        the dictionary key specifies the column name. If provided, the calculated
        distributions will use the same bin edges/categories.

    Returns
    -------
    label_distribution : DistributionData
        Contains distribution information about the labels
    other_distributions : dict[str, DistributionData]
        Contains distribution information about various columns in the DataFrame, where
        the dictionary key specifies the column name. Will be None if other_columns is
        None.

    Raises
    ------
    ValueError
        Raised if one of the specified other columns does not exist in the dataset.
    """
    # Calculate label distribution
    # NOTE: this assumes categorical labels encoded as integers (0, 1, ...)
    # TODO: make this generalize to arbitrary categorical labels, and continuous labels
    if existing_label_distribution is None:
        min_length = 0
    else:
        min_length = len(existing_label_distribution.categories)
    # Class counts as probabilities, padded to the existing category count.
    count = np.bincount(df[label_column], minlength=min_length)
    distribution = count / np.sum(count)
    categories = list(range(len(count)))
    label_distribution = DistributionData(distribution, None, categories)

    # Calculate other distributions
    other_distributions = None
    if other_columns is not None:
        other_distributions = dict()
        for column in other_columns:
            data = df[column].dropna()
            data_type = get_data_type(data)
            if data_type == "numerical":
                # Numerical: density histogram, reusing existing bin edges when given.
                if (
                    existing_other_distributions is not None
                    and column in existing_other_distributions
                ):
                    bins = existing_other_distributions[column].bin_edges
                else:
                    bins = "auto"
                hist, bin_edges = np.histogram(data, bins=bins, density=True)
                other_distributions[column] = DistributionData(hist, bin_edges, None)
            elif data_type == "categorical":
                # Categorical: mass per category, reusing existing categories.
                counts_obj = data.value_counts()
                if (
                    existing_other_distributions is not None
                    and column in existing_other_distributions
                ):
                    categories = existing_other_distributions[column].categories
                else:
                    categories = list(counts_obj.index)
                count = np.asarray(
                    [counts_obj[cat] if cat in counts_obj else 0 for cat in categories]
                )
                distribution = count / np.sum(count)
                other_distributions[column] = DistributionData(
                    distribution, None, categories
                )
            else:
                raise ValueError(f"Not handled {df[column].dtype}!")

    return label_distribution, other_distributions

split_dataset(df, label_column, speaker_column, num_folds=None, test_size=None, label_distribution=None, other_distributions=None, previous_train_folds=None, previous_test_folds=None, filename_column='filename', random_state=None)

Speaker-disjoint train/test folds, optionally matching given distributions (by KL).

Not used in training; for building new dataset splits.

Split a dataset into train and test folds

Split a dataset into train and test folds, making K non-overlapping folds or 1 fold based on the proportion specified by test_size. If both K-fold splitting and proportion-based splitting are specified, then the K-fold splitting takes precedence. The splits will be speaker-disjoint. If label or other distributions are provided, then the splits will try to match the provided distributions, based on measuring the KL divergence.

Parameters:

Name	Type	Description	Default
df	DataFrame	Pandas DataFrame containing the dataset.	required
label_column	str	Column in the DataFrame corresponding to the target labels (eg. "bin_phq").	required
speaker_column	str	Column in the DataFrame corresponding to speaker/participant ID (eg. "user").	required
num_folds	int	Number of folds to create.	None
test_size	float	Proportion of the dataset to allocate to the test split. Must be between 0.0 and 1.0 (exclusive). 1 - test_size will be allocated to the train split.	None
label_distribution	DistributionData	Contains distribution information about the labels. If provided, the splits will try to match this label distribution.	None
other_distributions	dict[str, DistributionData]	Contains distribution information about various columns in the DataFrame, where the dictionary key specifies the column name. If provided, the splits will try to match these distributions.	None
previous_train_folds	pandas.DataFrame or list of pandas.DataFrame	Single DataFrame or list of DataFrames containing all folds of the train split from a previous dataset.	None
previous_test_folds	pandas.DataFrame or list of pandas.DataFrame	Single DataFrame or list of DataFrames containing all folds of the test split from a previous dataset.	None
filename_column	str	Column used to uniquely identify files in the input dataset and previous train/test folds.	'filename'
random_state	int	If provided, it seeds the random number generator.	None

Returns:

Name	Type	Description
train_folds	list[DataFrame]	num_folds-length list containing all folds of the train split. For a proportion-based split, this will be a 1-length list.
test_folds	list[DataFrame]	num_folds-length list containing all folds of the test split. For a proportion-based split, this will be a 1-length list.

Source code in kirad/dataset_utils.py

def split_dataset(
    df: pd.DataFrame,
    label_column: str,
    speaker_column: str,
    num_folds: Optional[int] = None,
    test_size: Optional[float] = None,
    label_distribution: Optional[DistributionData] = None,
    other_distributions: Optional[dict[str, DistributionData]] = None,
    previous_train_folds: Optional[pd.DataFrame | list[pd.DataFrame]] = None,
    previous_test_folds: Optional[pd.DataFrame | list[pd.DataFrame]] = None,
    filename_column: str = "filename",
    random_state: Optional[int] = None,
) -> tuple[list[pd.DataFrame], list[pd.DataFrame]]:
    """Speaker-disjoint train/test folds, optionally matching given distributions (by KL).

    Not used in training; for building new dataset splits.

    Split a dataset into train and test folds

    Split a dataset into train and test folds, making K non-overlapping folds or 1 fold
    based on the proportion specified by test_size. If both K-fold splitting and
    proportion-based splitting are specified, then the K-fold splitting takes
    precedence. The splits will be speaker-disjoint. If label or other distributions
    are provided, then the splits will try to match the provided distributions, based
    on measuring the KL divergence.

    Parameters
    ----------
    df : pandas.DataFrame
        Pandas DataFrame containing the dataset.
    label_column : str
        Column in the DataFrame corresponding to the target labels (eg. "bin_phq").
    speaker_column : str
        Column in the DataFrame corresponding to speaker/participant ID (eg. "user").
    num_folds : int, optional
        Number of folds to create.
    test_size : float, optional
        Proportion of the dataset to allocate to the test split. Must be between 0.0
        and 1.0 (exclusive). 1 - test_size will be allocated to the train split.
    label_distribution : DistributionData, optional
        Contains distribution information about the labels. If provided, the splits
        will try to match this label distribution.
    other_distributions : dict[str, DistributionData], optional
        Contains distribution information about various columns in the DataFrame, where
        the dictionary key specifies the column name. If provided, the splits will try
        to match these distributions.
    previous_train_folds : pandas.DataFrame or list of pandas.DataFrame, optional
        Single DataFrame or list of DataFrames containing all folds of the train split
        from a previous dataset.
    previous_test_folds : pandas.DataFrame or list of pandas.DataFrame, optional
        Single DataFrame or list of DataFrames containing all folds of the test split
        from a previous dataset.
    filename_column: str, optional
        Column used to uniquely identify files in the input dataset and previous
        train/test folds.
    random_state : int, optional
        If provided, it seeds the random number generator.

    Returns
    -------
    train_folds : list[pandas.DataFrame]
        num_folds-length list containing all folds of the train split. For a
        proportion-based split, this will be a 1-length list.
    test_folds : list[pandas.DataFrame]
        num_folds-length list containing all folds of the test split. For a
        proportion-based split, this will be a 1-length list.
    """
    # Determine if creating K-fold splits are a single proportion-based split.
    if num_folds is not None:
        # Create K-fold splits. This takes precedence over proportion-based split.
        n_splits = num_folds
        split_type = "k_fold"
    elif test_size is not None:
        # Create one proportion-based split.
        n_splits = 1
        split_type = "proportion"
    else:
        raise ValueError("num_folds and test_size cannot both be None.")

    # If label or other distributions are supplied, then run multiple repetitions of
    # creating splits and choose the repetition with the lowest KL divergence to the
    # supplied distributions.
    if label_distribution is None and other_distributions is None:
        num_repetitions = 1
    else:
        # HARDCODED: 10 random splits tried; the lowest-KL one is kept
        num_repetitions = 10

    # Remove previous data (if specified) from the current data so that the
    # splitting operation is performed only on new additional data (relative to the
    # previous data). The previous data will be concatenated onto the new additional
    # data splits.
    # Since union of previous train and test data should be the same over all folds,
    # just need to remove previous train and test data from fold 0.
    remaining_df = df.copy(deep=True)
    if previous_train_folds is not None:
        if not isinstance(previous_train_folds, list):
            previous_train_folds = [previous_train_folds]
        # Remove samples from the previous dataset that aren't in the current data
        for idx in range(0, len(previous_train_folds)):
            train_df = previous_train_folds[idx]
            previous_train_folds[idx] = train_df[
                train_df[filename_column].isin(df[filename_column])
            ]

        # Remove samples from current data that are in the previous dataset
        remaining_df = remaining_df[
            ~remaining_df[filename_column].isin(
                previous_train_folds[0][filename_column]
            )
        ]
    if previous_test_folds is not None:
        if not isinstance(previous_test_folds, list):
            previous_test_folds = [previous_test_folds]
        # Remove samples from the previous dataset that aren't in the current data
        for idx in range(0, len(previous_test_folds)):
            test_df = previous_test_folds[idx]
            previous_test_folds[idx] = test_df[
                test_df[filename_column].isin(df[filename_column])
            ]

        # Remove samples from current data that are in the previous dataset
        remaining_df = remaining_df[
            ~remaining_df[filename_column].isin(previous_test_folds[0][filename_column])
        ]

    if len(remaining_df) == 0:
        # Current data did not have any new additional data relative to the previous
        # data, so just return the previous train/test splits.
        return previous_train_folds, previous_test_folds

    # Create data splits
    all_dataset_reps = []
    for rep in range(num_repetitions):
        seed = random_state + rep if random_state is not None else None
        # K-fold or single proportion split, never sharing a speaker across sides.
        if split_type == "k_fold":
            splitter = StratifiedGroupKFold(
                n_splits=n_splits, shuffle=True, random_state=seed
            )
        elif split_type == "proportion":
            splitter = GroupShuffleSplit(
                n_splits=n_splits, test_size=test_size, random_state=seed
            )
        # Collect (train, test) per fold, re-adding any previous data.
        dataset_folds = []
        for fold, (train_idx, test_idx) in enumerate(
            splitter.split(
                remaining_df.index,
                remaining_df[label_column],
                groups=remaining_df[speaker_column],
            )
        ):
            train_df = remaining_df.iloc[train_idx]
            if previous_train_folds is not None:
                train_df = pd.concat((previous_train_folds[fold], train_df))
            test_df = remaining_df.iloc[test_idx]
            if previous_test_folds is not None:
                test_df = pd.concat((previous_test_folds[fold], test_df))
            dataset_folds.append((train_df, test_df))
        all_dataset_reps.append(dataset_folds)

    # Choose a repetition with the smallest KL divergence to the supplied
    # distributions. Note that the train/test splits are paired
    # (train_dataset_splits[i], test_dataset_splits[i]), so need to choose splits
    # jointly.
    if num_repetitions == 1:
        best_dataset_rep = all_dataset_reps[0]
    else:
        if other_distributions is not None:
            other_columns = list(other_distributions.keys())
        else:
            other_columns = None

        # Summed KL (train + test, all folds) of each repetition vs. the supplied ones.
        kl_divergences = []
        for rep in range(num_repetitions):
            divergence = 0
            for train_df, test_df in all_dataset_reps[rep]:
                # Iterate through each fold in this repetition
                df_dict = dict()
                df_dict["train"] = train_df
                df_dict["test"] = test_df

                # Calculate label/other distributions for the train/test splits using
                # the same bins as the supplied distributions. Then, calculate the KL
                # divergence of the distributions summed across train and test splits.
                for data_split in ["train", "test"]:
                    label_dist, other_dist = calculate_distributions(
                        df_dict[data_split],
                        label_column,
                        other_columns=other_columns,
                        existing_label_distribution=label_distribution,
                        existing_other_distributions=other_distributions,
                    )

                    if label_distribution is not None:
                        divergence += entropy(
                            label_dist.distribution, label_distribution.distribution
                        )
                    if other_distributions is not None:
                        for column in other_columns:
                            divergence += entropy(
                                other_dist[column].distribution,
                                other_distributions[column].distribution,
                            )
            kl_divergences.append(divergence)
        # Keep the repetition with the lowest KL.
        sorted_dataset_reps = sorted(
            list(zip(all_dataset_reps, kl_divergences)), key=lambda x: x[1]
        )
        best_dataset_rep = sorted_dataset_reps[0][0]  # keep the best repetition

    # Unzip the (train, test) pairs into two lists of folds.
    train_folds, test_folds = zip(*best_dataset_rep)
    train_folds = list(train_folds)
    test_folds = list(test_folds)

    return train_folds, test_folds

get_data_type(series)

Return "numerical" if the Series converts to numbers, else "categorical".

Determine if a Pandas Series has numerical or categorical data.

Parameters:

Name	Type	Description	Default
series	Series	Pandas Series object.	required

Returns:

Name	Type	Description
data_type	str	Either "numerical" or "categorical".

Source code in kirad/dataset_utils.py

def get_data_type(series: pd.Series) -> str:
    """Return "numerical" if the Series converts to numbers, else "categorical".

    Determine if a Pandas Series has numerical or categorical data.

    Parameters
    ----------
    series : pandas.Series
        Pandas Series object.

    Returns
    -------
    data_type : str
        Either "numerical" or "categorical".
    """
    try:
        # TODO: distinguish between integers for numerical and categorical data.
        pd.to_numeric(series, errors="raise")
        data_type = "numerical"
    except (ValueError, TypeError):
        data_type = "categorical"

    return data_type

merge_df_rows_by_subject(df)

Merge rows sharing user + labels + demographics into one row per subject.

Durations summed, scores_* averaged, uuid/filename/prompt/asr ids become lists; the smallest UUID becomes the row index. Hardcoded: the common_cols / list_cols / sum_cols column-name lists below. Note: audio_filename/text_filename are not in list_cols, so they are dropped.

Source code in kirad/dataset_utils.py

def merge_df_rows_by_subject(df):
    """Merge rows sharing user + labels + demographics into one row per subject.

    Durations summed, scores_* averaged, uuid/filename/prompt/asr ids become lists;
    the smallest UUID becomes the row index.
    Hardcoded: the common_cols / list_cols / sum_cols column-name lists below.
    Note: audio_filename/text_filename are not in list_cols, so they are dropped.
    """
    # Move UUID index back to column to make it easier to recover the UUIDs from the
    # merged rows.
    df = df.reset_index()

    # HARDCODED: Kintsugi CSV schema; missing columns are skipped, new ones are ignored.
    # Rows are only merged if ALL present common_cols are equal (incl. phq/gad).
    common_cols = [
        "wave",
        "user",
        "ethnicity",
        "age",
        "gender",
        "zip_code",
        "state",
        "income",
        "english_preferred",
        "preferred_lang",
        "userAgent",
        "phq",
        "phq_2",
        "gad",
        "gad_2",
        "bin_phq",
        "bin_gad",
        "source",
        "pmi_score",
    ]
    # Add the PHQ-9 / GAD-7 item columns, and `real` when present.
    common_cols.extend([f"phq{i}" for i in range(1, 10)])
    common_cols.extend([f"gad{i}" for i in range(1, 8)])
    if "real" in df.columns:
        common_cols.append("real")
    # HARDCODED: columns kept as per-row lists after merging
    list_cols = [
        "uuid",
        "filename",
        "orig_audio_name",
        "prompt",
        "asr_output_id",
        "asr_operation_name",
        "asr_operation_start_time",
    ]
    # HARDCODED: duration columns summed after merging
    sum_cols = [
        "duration",
        "prompt_vad_cut_duration",
        "prompt_vad_trimmed_duration",
        "conversational_vad_cut_duration",
        "conversational_vad_trimmed_duration",
    ]
    # If KD targets are provided, output the mean of the provided targets
    mean_cols = [c for c in df.columns if c.startswith("scores")]
    # Keep only the columns that exist in this CSV.
    common_cols = [c for c in common_cols if c in df.columns]
    list_cols = [c for c in list_cols if c in df.columns]
    sum_cols = [c for c in sum_cols if c in df.columns]
    # Fill missing prompts with ''.
    if "prompt" in df.columns:
        df["prompt"] = df["prompt"].fillna("")
    # Group by subject: sum durations, average scores, collect list_cols into lists.
    groupby = df.groupby(common_cols, dropna=False)
    agg = groupby[sum_cols].sum()
    agg = agg.join(groupby[mean_cols].mean())
    for list_col in list_cols:
        agg = agg.join(groupby[list_col].apply(list))

    # Make UUID as row index. Since UUIDs have been merged into a list, take the first
    # UUID in the list as the UUID index for the row after sorting the UUIDs.
    agg = agg.reset_index()
    agg["uuid"] = agg["uuid"].apply(lambda x: sorted(x)[0])
    agg.set_index("uuid", inplace=True)

    return agg

repeat_along_batch_dim(data, repetitions=1)

Tile a tensor (or dict of tensors) repetitions times along dim 0.

Source code in kirad/dataset_utils.py

def repeat_along_batch_dim(
    data: torch.Tensor | Mapping[str, torch.Tensor], repetitions: int = 1
) -> torch.Tensor | Mapping[str, torch.Tensor]:
    """Tile a tensor (or dict of tensors) `repetitions` times along dim 0."""
    if isinstance(data, torch.Tensor):
        # Repeat along dim 0 only.
        repeats = [repetitions] + [1 for _ in range(0, data.dim() - 1)]
        return data.repeat(*repeats)
    elif isinstance(data, Mapping):
        return {
            key: repeat_along_batch_dim(val, repetitions=repetitions)
            for key, val in data.items()
        }
    else:
        raise ValueError(f"Unable to handle data of type {type(data)}.")

broadcast_along_batch_dim(data)

Tile every tensor of one sample to the largest dim 0 (= its number of windows W).

Broadcast tensors in a nested structure to have the same batch dim, returning the resulting structure and dim.

Parameters:

Name	Type	Description	Default
data	torch.Tensor or data structure containing torch.Tensor	Tensor or data structure (eg. dictionary or list) containg tensors to broadcast along the batch dimension. When multiple tensors are provided, the batch dimension of each tensor is broadcast to match the largest batch dimension across the tensors.	required

Returns:

Name	Type	Description
broadcast_data	torch.Tensor or data structure containing torch.Tensor	Data in the same structure as the input but with the broadcasted batch dimensions for each tensor contained.
length	int	Size of the batch dimension that each tensor was broadcast to.

Source code in kirad/dataset_utils.py

def broadcast_along_batch_dim(data: Any) -> tuple[Any, int]:
    """Tile every tensor of one sample to the largest dim 0 (= its number of windows W).

    Broadcast tensors in a nested structure to have the same batch dim, returning
    the resulting structure and dim.

    Parameters
    ----------
    data : torch.Tensor or data structure containing torch.Tensor
        Tensor or data structure (eg. dictionary or list) containg tensors to broadcast
        along the batch dimension. When multiple tensors are provided, the batch
        dimension of each tensor is broadcast to match the largest batch dimension
        across the tensors.

    Returns
    -------
    broadcast_data : torch.Tensor or data structure containing torch.Tensor
        Data in the same structure as the input but with the broadcasted batch
        dimensions for each tensor contained.
    length : int
        Size of the batch dimension that each tensor was broadcast to.
    """

    def max_batch_dim(data) -> int:
        """Largest dim 0 among all nested tensors."""
        # Recurse through dicts/lists down to tensors; keep the largest dim 0.
        if isinstance(data, torch.Tensor):
            return data.shape[0]
        if isinstance(data, Mapping):
            return max_batch_dim(data.values())
        if isinstance(data, Iterable):
            return max(max_batch_dim(d) for d in data)
        raise ValueError(
            f"Got {data=} with invalid {type(data)=} that is not Tensor, Mapping, or "
            "Iterable."
        )

    def broadcast_to(data, length):
        """Repeat each nested tensor along dim 0 to `length`; must divide evenly."""
        # Repeat each tensor along dim 0 to `length` (must divide); recurse into dicts.
        if isinstance(data, torch.Tensor):
            batch_dim = data.shape[0]
            repetitions, remainder = divmod(length, batch_dim)
            if remainder:
                raise ValueError(
                    "Cannot broadcast batch dimension of tensor with shape "
                    f"{data.shape} to {length}."
                )
            return repeat_along_batch_dim(data, repetitions=repetitions)
        if isinstance(data, Mapping):
            return {key: broadcast_to(val, length) for key, val in data.items()}
        if isinstance(data, Iterable):
            return [broadcast_to(d, length) for d in data]

    # Target length W = largest dim 0 among the tensors.
    length = max_batch_dim(data)
    return broadcast_to(data, length), length

get_collate_fn(tensor_collate_fn)

torch default collate, but tensors are merged with tensor_collate_fn (e.g. concat).

Source code in kirad/dataset_utils.py

def get_collate_fn(tensor_collate_fn: Callable[[list[torch.Tensor]], torch.Tensor]):
    """torch default collate, but tensors are merged with tensor_collate_fn (e.g. concat)."""

    def tensor_collate_tensor_fn(
        batch, *args, **kwargs
    ):  # torch internals pass some args we don't need
        """Apply tensor_collate_fn to a list of tensors."""
        return tensor_collate_fn(batch)

    # Replace torch's default tensor collate with the custom one.
    collate_fn_map = copy.copy(default_collate_fn_map)
    collate_fn_map[torch.Tensor] = tensor_collate_tensor_fn
    return partial(collate, collate_fn_map=collate_fn_map)

broadcast_collate(batch)

DataLoader collate: concat windows of all recordings; length = windows per recording.

Per sample, modalities are tiled to the same W (e.g. text [1, L] -> [W, L]). Shapes: uuid [B], label..primary [B], length [B], features.audio [ΣW, 480000] or [ΣW, 80, 3000], features.text.* [ΣW, L].

Collate function that broadcasts the batch dimension of the features for each modality to match the largest batch dimension across the modalities.

Parameters:

Name	Type	Description	Default
batch	list of dictionaries	batch[i] contains a dictionary that maps "metadata" to a dictionary containing a UUID and label, and modality (eg. "audio" or "text") that maps to features.	required

Returns:

Type	Description
Dictionary containing UUIDs, features, lengths of the features, and labels of each
sample.

Source code in kirad/dataset_utils.py

def broadcast_collate(batch: list[Mapping[str, Any]]) -> dict[str, Any]:
    """DataLoader collate: concat windows of all recordings; length = windows per recording.

    Per sample, modalities are tiled to the same W (e.g. text [1, L] -> [W, L]).
    Shapes: uuid [B], label.<task>.primary [B], length [B],
    features.audio [ΣW, 480000] or [ΣW, 80, 3000], features.text.* [ΣW, L].

    Collate function that broadcasts the batch dimension of the features for each
    modality to match the largest batch dimension across the modalities.

    Parameters
    ----------
    batch : list of dictionaries
        batch[i] contains a dictionary that maps "metadata" to a dictionary containing
        a UUID and label, and modality (eg. "audio" or "text") that maps to features.

    Returns
    -------
    Dictionary containing UUIDs, features, lengths of the features, and labels of each
    sample.
    """

    def broadcast_desired(batch):
        """Flatten one sample: uuid/label from metadata, tiled features, and length W."""
        out = {
            **batch["metadata"]
        }  # contains DatasetFields.UUID and DatasetFields.LABEL
        # Tile all modality features to the same W and store it as length.
        modality_data = {key: val for key, val in batch.items() if key != "metadata"}
        modality_data_broadcast, length = broadcast_along_batch_dim(modality_data)
        out[DatasetFields.LENGTH] = length
        out[DatasetFields.FEATURES] = modality_data_broadcast

        return out

    # Concat windows of all recordings -> [ΣW, ...]; length stays [B].
    return concat_collate(list(map(broadcast_desired, batch)))