%reload_ext autoreload
%autoreload 2

#export
from nb_007b import *

PATH = Path('../data/aclImdb/')
CLAS_PATH = PATH/'clas'
LM_PATH = PATH/'lm'
MODEL_PATH = PATH/'models'
os.makedirs(CLAS_PATH, exist_ok=True)
os.makedirs(LM_PATH, exist_ok=True)
os.makedirs(MODEL_PATH, exist_ok=True)

data = data_from_textcsv(LM_PATH, Tokenizer(), data_func=lm_data, bs=50)
learn = RNNLearner.language_model(data, drop_mult=0.3, pretrained_fnames=['lstm_wt103', 'itos_wt103'])
learn.freeze()

learn.fit_one_cycle(1, 1e-2, moms=(0.8,0.7))

learn.save('fit_head')

learn.load('fit_head')

learn.unfreeze()
learn.fit_one_cycle(10, 1e-3, moms=(0.8,0.7), pct_start=0.25)

learn.save('fine_tuned')

learn.save_encoder('fine_tuned_enc')

shutil.copy(LM_PATH/'models'/'fine_tuned_enc.pth', CLAS_PATH/'models'/'fine_tuned_enc.pth')
shutil.copy(LM_PATH/'tmp'/'itos.pkl', CLAS_PATH/'tmp'/'itos.pkl')

data = data_from_textcsv(CLAS_PATH, Tokenizer(), vocab=Vocab(LM_PATH/'tmp'), data_func=classifier_data, bs=50)
learn = RNNLearner.classifier(data, drop_mult=0.5)
learn.load_encoder('fine_tuned_enc')
learn.freeze()

learn.fit_one_cycle(1, 2e-2, moms=(0.8,0.7))

learn.freeze_to(-2)
learn.fit_one_cycle(1, slice(1e-2/2.6,1e-2), moms=(0.8,0.7), pct_start=0.1)

learn.unfreeze()
learn.fit_one_cycle(2, slice(1e-2/(2.6)**4,1e-2), moms=(0.8,0.7), pct_start=0.1)

learn.load('first')

lr=1e-2
lrm = 2.6
lrs = np.array([lr/(lrm**4), lr/(lrm**3), lr/(lrm**2), lr/lrm, lr])

learn.freeze_to(-2)
learn.fit_one_cycle(1, lrs, moms=(0.8,0.7), pct_start=0.1)

learn.save('second')

learn.load('second')

lr=5e-3
lrm = 2.6
lrs = np.array([lr/(lrm**4), lr/(lrm**3), lr/(lrm**2), lr/lrm, lr])

learn.freeze_to(-3)
learn.fit_one_cycle(1, lrs, moms=(0.8,0.7), pct_start=0.1)

learn.save('third')

learn.load('third')

lr=1e-3
lrm = 2.6
lrs = np.array([lr/(lrm**4), lr/(lrm**3), lr/(lrm**2), lr/lrm, lr])

learn.unfreeze()
learn.fit_one_cycle(2, lrs, moms=(0.8,0.7), pct_start=0.1)