!pip install -q iterative-stratification

WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv

import pandas as pd
from sklearn.model_selection import KFold
from iterstrat.ml_stratifiers import MultilabelStratifiedKFold

df = pd.read_csv('/root/.cache/data/train.csv')

dfx = pd.get_dummies(df, columns=["discourse_type"]).head(5)
dfx

dfx = pd.get_dummies(df, columns=["discourse_type"]).groupby(["id"], as_index=False).sum()
dfx

dfx.columns

Index(['id', 'discourse_id', 'discourse_start', 'discourse_end',
       'discourse_type_Claim', 'discourse_type_Concluding Statement',
       'discourse_type_Counterclaim', 'discourse_type_Evidence',
       'discourse_type_Lead', 'discourse_type_Position',
       'discourse_type_Rebuttal'],
      dtype='object')

cols = [c for c in dfx.columns if c.startswith("discourse_type") or c =="id" and c !=  "discourse_type_num"]
cols

['id',
 'discourse_type_Claim',
 'discourse_type_Concluding Statement',
 'discourse_type_Counterclaim',
 'discourse_type_Evidence',
 'discourse_type_Lead',
 'discourse_type_Position',
 'discourse_type_Rebuttal']

dfx = dfx[cols]
dfx

mskf = MultilabelStratifiedKFold(n_splits=5, shuffle=True, random_state=42)
labels = [c for c in dfx.columns if c != "id"]
dfx_labels = dfx[labels]
dfx_labels

dfx["kfold"] = -1

/tmp/ipykernel_27806/3539168384.py:1: SettingWithCopyWarning: 
A value is trying to be set on a copy of a slice from a DataFrame.
Try using .loc[row_indexer,col_indexer] = value instead

See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy
  dfx["kfold"] = -1

df = pd.read_csv('/root/.cache/data/train.csv')

dfx = pd.get_dummies(df, columns=["discourse_type"]).groupby(["id"], as_index=False).sum()
cols = [c for c in dfx.columns if c.startswith("discourse_type_") or c == "id" and c != "discourse_type_num"]
dfx = dfx[cols]

mskf = MultilabelStratifiedKFold(n_splits=10, shuffle=True, random_state=42)
labels = [c for c in dfx.columns if c != "id"]
dfx_labels = dfx[labels]
dfx["kfold"] = -1
dfx

for fold, (trn_, val_) in enumerate(mskf.split(dfx, dfx_labels)):
    print(len(trn_), len(val_))
    dfx.loc[val_, "kfold"] = fold

df = df.merge(dfx[["id", "kfold"]], on="id", how="left")
print(df.kfold.value_counts())
# df.to_csv("train_folds.csv", index=False)

14036 1558
14036 1558
14033 1561
14035 1559
14031 1563
14035 1559
14034 1560
14036 1558
14036 1558
14034 1560
6    14633
9    14532
7    14529
8    14509
5    14466
3    14431
4    14365
1    14358
2    14271
0    14199
Name: kfold, dtype: int64

df.groupby(["kfold"]).count()