levchik_folds.py

import os
import numpy as np
import pandas as pd
from sklearn.model_selection import StratifiedKFold

prod_to_cat = pd.read_csv('levchik_folds/prod_to_category.csv')
skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=4242)

split = skf.split(prod_to_cat.drop('category_id', axis=1), prod_to_cat['category_id'])

for i, (train_index, test_index) in enumerate(split):
    train, test = prod_to_cat.iloc[train_index], prod_to_cat.iloc[test_index]
    
    train.to_csv(os.path.join('levchik_folds', '%d_train.csv' % i), sep=',', index=False)
    test.to_csv(os.path.join('levchik_folds', '%d_test.csv' % i), sep=',', index=False)