Generic random forest classification to draw learning curve

randomforest.py

import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_predict, cross_val_score, learning_curve
import altair as alt

def learn(X, y, **kwargs):  # X needs to be a pandas dataframe
    r = RandomForestClassifier(n_estimators=100, random_state=123, class_weight="balanced", **kwargs)
    steps, curve_train, curve_test = learning_curve(r, X.sample(frac=1, random_state=123), y.sample(frac=1, random_state=123), cv=5, scoring="roc_auc", n_jobs=20, train_sizes=np.linspace(0.05,1,20))
    learning = (pd.concat([pd.DataFrame(curve_train, index=steps).apply(np.average, axis="columns").rename("train"),
                           pd.DataFrame(curve_test, index=steps).apply(np.average, axis="columns").rename("test")], axis="columns")
                .stack().reset_index().rename({"level_0": "size", "level_1": "category", 0: "roc_auc"}, axis="columns")
               )
    chart = alt.Chart(learning).mark_line().encode(
        x="size",
        y="roc_auc",
        color="category"
    )
    return chart, learning

In case your data isn't a DataFrame, just remove .sample(). Make sure your data is shuffled in advance.