<gh_stars>0
"""It’s easy to modify the code for creating k-fold cross-validation to create stratified  k-folds. 
We are only changing from model_selection.KFold to  model_selection.StratifiedKFold and in the kf.split(...) function, 
we specify the  target column on which we want to stratify. 
We assume that our CSV dataset has a  column called “target” and it is a classification problem! 
"""

import pandas as pd
from sklearn import model_selection

if __name__ == "__main__":
    # assume the training data is train.csv and has a clumn called
    # 'target' and the probelm is classification problem
    df = pd.read_csv("winequality-red.csv")

    # we create a new column called kfold and fill it with -1
    df["kfold"] = -1

    # randomize the rows of the data
    df = df.sample(frac=1).reset_index(drop=True)

    # fetch the target column on which we want to stratify
    # y is a numpy array of the target column
    y = df.quality.values

    # intitiate the stratified kfold class from the model selection module
    # No of folds is 5 here
    kf = model_selection.StratifiedKFold(n_splits=5)

    # fill the new kfold column in the dataset
    for fold, (t_, v_) in enumerate(kf.split(X=df, y=y)):
        df.loc[v_, "kfold"] = fold

    # sae the new csv with kfold column
    df.to_csv("train_strat_kfold.csv")

"""It’s easy to modify the code for creating k-fold cross-validation to create stratified  k-folds. 
We are only changing from model_selection.KFold to  model_selection.StratifiedKFold and in the kf.split(...) function, 
we specify the  target column on which we want to stratify. 
We assume that our CSV dataset has a  column called “target” and it is a classification problem! 
"""

import pandas as pd
from sklearn import model_selection

if __name__ == "__main__":
    # assume the training data is train.csv and has a clumn called
    # 'target' and the probelm is classification problem
    df = pd.read_csv("winequality-red.csv")

    # we create a new column called kfold and fill it with -1
    df["kfold"] = -1

    # randomize the rows of the data
    df = df.sample(frac=1).reset_index(drop=True)

    # fetch the target column on which we want to stratify
    # y is a numpy array of the target column
    y = df.quality.values

    # intitiate the stratified kfold class from the model selection module
    # No of folds is 5 here
    kf = model_selection.StratifiedKFold(n_splits=5)

    # fill the new kfold column in the dataset
    for fold, (t_, v_) in enumerate(kf.split(X=df, y=y)):
        df.loc[v_, "kfold"] = fold

    # sae the new csv with kfold column
    df.to_csv("train_strat_kfold.csv")
It’s easy to modify the code for creating k-fold cross-validation to create stratified  k-folds.  We are only changing from model_selection.KFold to  model_selection.StratifiedKFold and in the kf.split(...) function,  we specify the  target column on which we want to stratify.  We assume that our CSV dataset has a  column called “target” and it is a classification problem! # assume the training data is train.csv and has a clumn called # 'target' and the probelm is classification problem # we create a new column called kfold and fill it with -1 # randomize the rows of the data # fetch the target column on which we want to stratify # y is a numpy array of the target column # intitiate the stratified kfold class from the model selection module # No of folds is 5 here # fill the new kfold column in the dataset # sae the new csv with kfold column