[K-Folds split] #Python

split_k_fold.py

from typing import List

import pandas as pd
from sklearn.model_selection import StratifiedKFold

def split_k_fold(df: pd.DataFrame, k: int = 5) -> List[pd.DataFrame]:
    skf = StratifiedKFold(n_splits=k)
    y = df["label"]
    placeholder = y.copy()
    k_fold = []
    for _, test_idx in skf.split(placeholder, y):
        df_fold = df.iloc[test_idx]
        k_fold.append(df_fold)

    return k_fold