from fastai.text import *

bs,bptt=256,80

path = untar_data(URLs.IMDB)

data_lm = (TextList.from_folder(path)
           #Inputs: all the text files in path
            .filter_by_folder(include=['train', 'test', 'unsup']) 
           #We may have other temp folders that contain text files so we only keep what's in train, test and unsup
            .split_by_rand_pct(0.1)
           #We randomly split and keep 10% (10,000 reviews) for validation
            .label_for_lm()           
           #We want to do a language model so we label accordingly
            .databunch(bs=bs, bptt=bptt))
data_lm.save('data_lm.pkl')

data_lm = load_data(path, 'data_lm.pkl', bs=bs, bptt=bptt)

data_lm.show_batch()

data_bwd = load_data(path, 'data_lm.pkl', bs=bs, bptt=bptt, backwards=True)

data_bwd.show_batch()

learn = language_model_learner(data_lm, AWD_LSTM)
learn = learn.to_fp16(clip=0.1)

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

learn.unfreeze()

learn.fit_one_cycle(10, 2e-3, moms=(0.8,0.7), wd=0.1)

learn.save_encoder('fwd_enc')

learn = language_model_learner(data_bwd, AWD_LSTM)
learn = learn.to_fp16(clip=0.1)

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

learn.unfreeze()

learn.fit_one_cycle(10, 2e-3, moms=(0.8,0.7), wd=0.1)

learn.save_encoder('bwd_enc')

path = untar_data(URLs.IMDB)
bs = 128

data_clas = (TextList.from_folder(path, vocab=data_lm.vocab)
             #grab all the text files in path
             .split_by_folder(valid='test')
             #split by train and valid folder (that only keeps 'train' and 'test' so no need to filter)
             .label_from_folder(classes=['neg', 'pos'])
             #label them all with their folders
             .databunch(bs=bs))

data_clas.save('data_clas.pkl')

data_clas = load_data(path, 'data_clas.pkl', bs=bs)

data_clas.show_batch()

data_clas_bwd = load_data(path, 'data_clas.pkl', bs=bs, backwards=True)

data_clas_bwd.show_batch()

learn = text_classifier_learner(data_clas, AWD_LSTM, drop_mult=0.5, pretrained=False)
learn.load_encoder('fwd_enc')

lr = 1e-1

learn.fit_one_cycle(1, lr, moms=(0.8,0.7), wd=0.1)

learn.freeze_to(-2)
lr /= 2
learn.fit_one_cycle(1, slice(lr/(2.6**4),lr), moms=(0.8,0.7), wd=0.1)

learn.freeze_to(-3)
lr /= 2
learn.fit_one_cycle(1, slice(lr/(2.6**4),lr), moms=(0.8,0.7), wd=0.1)

learn.unfreeze()
lr /= 5
learn.fit_one_cycle(2, slice(lr/(2.6**4),lr), moms=(0.8,0.7), wd=0.1)

learn.save('fwd_clas')

learn_bwd = text_classifier_learner(data_clas_bwd, AWD_LSTM, drop_mult=0.5, pretrained=False)
learn_bwd.load_encoder('bwd_enc')

learn_bwd.fit_one_cycle(1, lr, moms=(0.8,0.7), wd=0.1)

learn_bwd.freeze_to(-2)
lr /= 2
learn_bwd.fit_one_cycle(1, slice(lr/(2.6**4),lr), moms=(0.8,0.7), wd=0.1)

learn_bwd.freeze_to(-3)
lr /= 2
learn_bwd.fit_one_cycle(1, slice(lr/(2.6**4),lr), moms=(0.8,0.7), wd=0.1)

learn_bwd.unfreeze()
lr /= 5
learn_bwd.fit_one_cycle(2, slice(lr/(2.6**4),lr), moms=(0.8,0.7), wd=0.1)

learn_bwd.save('bwd_clas')

pred_fwd,lbl_fwd = learn.get_preds(ordered=True)

pred_bwd,lbl_bwd = learn_bwd.get_preds(ordered=True)

final_pred = (pred_fwd+pred_bwd)/2

accuracy(pred, lbl_fwd)