automl_pycaret / classification.py
konstantinG's picture
Update classification.py
233fa77
from pycaret.classification import *
import streamlit as st
@st.experimental_memo
def prep_and_train(targ, data, models=None):
s = setup(data, target = targ, session_id = 12)
s_df = pull()
best = compare_models(include=models)
best_df = pull()
return best, s_df, best_df
@st.experimental_memo
def tuning(_model, n_iters, search_lib):
tuned_dt = tune_model(estimator=_model, n_iter=n_iters, search_library=search_lib, choose_better=True)
info_df = pull()
return tuned_dt, info_df