Fold Split by Polars

fold_split.py

def split_fold(
    df: pl.DataFrame,
    target_col="label",
    group_cols=["group"],
    num_folds=5,
    seed=42,
) -> pl.DataFrame:
    """
    - add column `_fold` to input DataFrame
    """
    df = df.with_columns(pl.col("contact").cumcount().cast(int).alias("_index"))
    skf = StratifiedGroupKFold(n_splits=num_folds, random_state=seed, shuffle=True)
    generator = skf.split(df, df[target_col], df.select(group_cols))
    idx2fold: Dict[int, int] = {}
    for fold, (_, idx_valid) in enumerate(generator):
        idx2fold = {**idx2fold, **{idx: fold for idx in idx_valid}}
    df = df.with_columns(pl.col("_index").map_dict(idx2fold).alias("_fold")).drop("_index")
    return df