
#################################################
### THIS FILE WAS AUTOGENERATED! DO NOT EDIT! ###
#################################################
# file to edit: dev_nb/train_loop.ipynb

from exp.nb_packages import *

def train(get_learner, N_FOLDS, prefix_name):
    all_val_preds = []
    all_val_labels = []
    all_test_preds = []

    for i in range(N_FOLDS):

        learn = get_learner(i)

        learn.fit_one_cycle(4)
        learn.unfreeze()

        learn.fit_one_cycle(12, slice(1e-5, 1e-4))

        learn.recorder.plot_loss()

        learn.save(prefix_name + '_' + str(i))

def validate(get_learner, N_FOLDS, prefix_name):
    all_val_preds = []
    all_val_labels = []
    for i in range(N_FOLDS):
        learn = get_learner(i)
        learn.load(prefix_name + '_' + str(i))
        val_preds, val_labels = learn.get_preds()

        all_val_preds.append(val_preds)
        all_val_labels.append(val_labels)
    return all_val_preds, all_val_labels

def test(dls, test_df, get_learner, N_FOLDS, prefix_name):
    all_test_preds = []
    for i in range(N_FOLDS):
        learn = get_learner(i)
        learn.load(prefix_name + '_' + str(i))
        test_dl = dls.test_dl(test_df)
        test_preds, _ = learn.get_preds(dl=test_dl)
        all_test_preds.append(test_preds)
    return all_test_preds