stratifiedkfold.py

from sklearn.model_selection import KFold, StratifiedKFold
from catboost import CatBoostClassifier
from sklearn.metrics import roc_auc_score
kfold, scores, y_pred_totcb = StratifiedKFold(n_splits=5, shuffle=True, random_state=221), list(), []
for train, test in kfold.split(X, y):
    x_train, x_test = X.iloc[train], X.iloc[test]
    y_train, y_test = y.iloc[train], y.iloc[test]
    
    model = CatBoostClassifier(random_state=27,  n_estimators=3000, cat_features = cat_columns,
                               max_depth=7, verbose=500, learning_rate=0.102, eval_metric='AUC')
    model.fit(x_train, y_train, eval_set=(x_test, y_test))
    preds = model.predict_proba(x_test)[:,1]
    score = roc_auc_score(y_test, preds)
    scores.append(score)
    print(score)
    test_pred = model.predict_proba(test_df)[:,1]
    y_pred_totcb.append(test_pred)
print("Average: ", sum(scores)/len(scores))