PATH='data/aclImdb/'

TRN_PATH = 'train/all/'
VAL_PATH = 'test/all/'
TRN = f'{PATH}{TRN_PATH}'
VAL = f'{PATH}{VAL_PATH}'

%ls {PATH}

"""
Explanation: Language modeling
Data
The large movie view dataset contains a collection of 50,000 reviews from IMDB. The dataset contains an even number of positive and negative reviews. The authors considered only highly polarized reviews. A negative review has a score ≤ 4 out of 10, and a positive review has a score ≥ 7 out of 10. Neutral reviews are not included in the dataset. The dataset is divided into training and test sets. The training set is the same 25,000 labeled reviews.
The sentiment classification task consists of predicting the polarity (positive or negative) of a given text.
However, before we try to classify sentiment, we will simply try to create a language model; that is, a model that can predict the next word in a sentence. Why? Because our model first needs to understand the structure of English, before we can expect it to recognize positive vs negative sentiment.
So our plan of attack is the same as we used for Dogs v Cats: pretrain a model to do one thing (predict the next word), and fine tune it to do something else (classify sentiment).
Unfortunately, there are no good pretrained language models available to download, so we need to create our own. To follow along with this notebook, we suggest downloading the dataset from this location on files.fast.ai.
End of explanation
"""


trn_files = !ls {TRN}
trn_files[:10]

"""
Explanation: Let's look inside the training folder...
End of explanation
"""


review = !cat {TRN}{trn_files[6]}
review[0]

"""
Explanation: ...and at an example review.
End of explanation
"""


!find {TRN} -name '*.txt' | xargs cat | wc -w

!find {VAL} -name '*.txt' | xargs cat | wc -w

"""
Explanation: Sounds like I'd really enjoy Zombiegeddon...
Now we'll check how many words are in the dataset.
End of explanation
"""


' '.join(spacy_tok(review[0]))

"""
Explanation: Before we can analyze text, we must first tokenize it. This refers to the process of splitting a sentence into an array of words (or more generally, into an array of tokens).
End of explanation
"""


TEXT = data.Field(lower=True, tokenize=spacy_tok)

"""
Explanation: We use Pytorch's torchtext library to preprocess our data, telling it to use the wonderful spacy library to handle tokenization.
First, we create a torchtext field, which describes how to preprocess a piece of text - in this case, we tell torchtext to make everything lowercase, and tokenize it with spacy.
End of explanation
"""


bs=64; bptt=70

FILES = dict(train=TRN_PATH, validation=VAL_PATH, test=VAL_PATH)
md = LanguageModelData.from_text_files(PATH, TEXT, **FILES, bs=bs, bptt=bptt, min_freq=10)

"""
Explanation: fastai works closely with torchtext. We create a ModelData object for language modeling by taking advantage of LanguageModelData, passing it our torchtext field object, and the paths to our training, test, and validation sets. In this case, we don't have a separate test set, so we'll just use VAL_PATH for that too.
As well as the usual bs (batch size) parameter, we also not have bptt; this define how many words are processing at a time in each row of the mini-batch. More importantly, it defines how many 'layers' we will backprop through. Making this number higher will increase time and memory requirements, but will improve the model's ability to handle long sentences.
End of explanation
"""


pickle.dump(TEXT, open(f'{PATH}models/TEXT.pkl','wb'))

"""
Explanation: After building our ModelData object, it automatically fills the TEXT object with a very important attribute: TEXT.vocab. This is a vocabulary, which stores which words (or tokens) have been seen in the text, and how each word will be mapped to a unique integer id. We'll need to use this information again later, so we save it.
(Technical note: python's standard Pickle library can't handle this correctly, so at the top of this notebook we used the dill library instead and imported it as pickle).
End of explanation
"""


len(md.trn_dl), md.nt, len(md.trn_ds), len(md.trn_ds[0].text)

"""
Explanation: Here are the: # batches; # unique tokens in the vocab; # tokens in the training set; # sentences
End of explanation
"""


# 'itos': 'int-to-string'
TEXT.vocab.itos[:12]

# 'stoi': 'string to int'
TEXT.vocab.stoi['the']

"""
Explanation: This is the start of the mapping from integer IDs to unique tokens.
End of explanation
"""


md.trn_ds[0].text[:12]

"""
Explanation: Note that in a LanguageModelData object there is only one item in each dataset: all the words of the text joined together.
End of explanation
"""


TEXT.numericalize([md.trn_ds[0].text[:12]])

"""
Explanation: torchtext will handle turning this words into integer IDs for us automatically.
End of explanation
"""


next(iter(md.trn_dl))

"""
Explanation: Our LanguageModelData object will create batches with 64 columns (that's our batch size), and varying sequence lengths of around 80 tokens (that's our bptt parameter - backprop through time).
Each batch also contains the exact same data as labels, but one word later in the text - since we're trying to always predict the next word. The labels are flattened into a 1d array.
End of explanation
"""


em_sz = 200  # size of each embedding vector
nh = 500     # number of hidden activations per layer
nl = 3       # number of layers

"""
Explanation: Train
We have a number of parameters to set - we'll learn more about these later, but you should find these values suitable for many problems.
End of explanation
"""


opt_fn = partial(optim.Adam, betas=(0.7, 0.99))

"""
Explanation: Researchers have found that large amounts of momentum (which we'll learn about later) don't work well with these kinds of RNN models, so we create a version of the Adam optimizer with less momentum than it's default of 0.9.
End of explanation
"""


learner = md.get_model(opt_fn, em_sz, nh, nl,
               dropouti=0.05, dropout=0.05, wdrop=0.1, dropoute=0.02, dropouth=0.05)
learner.reg_fn = partial(seq2seq_reg, alpha=2, beta=1)
learner.clip=0.3

"""
Explanation: fastai uses a variant of the state of the art AWD LSTM Language Model developed by Stephen Merity. A key feature of this model is that it provides excellent regularization through Dropout. There is no simple way known (yet!) to find the best values of the dropout parameters below - you just have to experiment...
However, the other parameters (alpha, beta, and clip) shouldn't generally need tuning.
End of explanation
"""


learner.fit(3e-3, 4, wds=1e-6, cycle_len=1, cycle_mult=2)

learner.save_encoder('adam1_enc')

learner.load_encoder('adam1_enc')

learner.load_cycle('adam3_10',2)

learner.fit(3e-3, 1, wds=1e-6, cycle_len=10)

learner.save_encoder('adam3_10_enc')

"""
Explanation: As you can see below, I gradually tuned the language model in a few stages. I possibly could have trained it further (it wasn't yet overfitting), but I didn't have time to experiment more. Maybe you can see if you can train it to a better accuracy! (I used lr_find to find a good learning rate, but didn't save the output in this notebook. Feel free to try running it yourself now.)
End of explanation
"""


learner.save_encoder('adam3_20_enc')

learner.load_encoder('adam3_20_enc')

"""
Explanation: In the sentiment analysis section, we'll just need half of the language model - the encoder, so we save that part.
End of explanation
"""


math.exp(4.165)

pickle.dump(TEXT, open(f'{PATH}models/TEXT.pkl','wb'))

"""
Explanation: Language modeling accuracy is generally measured using the metric perplexity, which is simply exp() of the loss function we used.
End of explanation
"""


m=learner.model
ss=""". So, it wasn't quite was I was expecting, but I really liked it anyway! The best"""
s = [spacy_tok(ss)]
t=TEXT.numericalize(s)
' '.join(s[0])

"""
Explanation: Test
We can play around with our language model a bit to check it seems to be working OK. First, let's create a short bit of text to 'prime' a set of predictions. We'll use our torchtext field to numericalize it so we can feed it to our language model.
End of explanation
"""


# Set batch size to 1
m[0].bs=1
# Turn off dropout
m.eval()
# Reset hidden state
m.reset()
# Get predictions from model
res,*_ = m(t)
# Put the batch size back to what it was
m[0].bs=bs

"""
Explanation: We haven't yet added methods to make it easy to test a language model, so we'll need to manually go through the steps.
End of explanation
"""


nexts = torch.topk(res[-1], 10)[1]
[TEXT.vocab.itos[o] for o in to_np(nexts)]

"""
Explanation: Let's see what the top 10 predictions were for the next word after our short text:
End of explanation
"""


print(ss,"\n")
for i in range(50):
    n=res[-1].topk(2)[1]
    n = n[1] if n.data[0]==0 else n[0]
    print(TEXT.vocab.itos[n.data[0]], end=' ')
    res,*_ = m(n[0].unsqueeze(0))
print('...')

"""
Explanation: ...and let's see if our model can generate a bit more text all by itself!
End of explanation
"""


TEXT = pickle.load(open(f'{PATH}models/TEXT.pkl','rb'))

"""
Explanation: Sentiment
We'll need to the saved vocab from the language model, since we need to ensure the same words map to the same IDs.
End of explanation
"""


IMDB_LABEL = data.Field(sequential=False)
splits = torchtext.datasets.IMDB.splits(TEXT, IMDB_LABEL, 'data/')

t = splits[0].examples[0]

t.label, ' '.join(t.text[:16])

"""
Explanation: sequential=False tells torchtext that a text field should be tokenized (in this case, we just want to store the 'positive' or 'negative' single label).
splits is a torchtext method that creates train, test, and validation sets. The IMDB dataset is built into torchtext, so we can take advantage of that. Take a look at lang_model-arxiv.ipynb to see how to define your own fastai/torchtext datasets.
End of explanation
"""


md2 = TextData.from_splits(PATH, splits, bs)

m3 = md2.get_model(opt_fn, 1500, bptt, emb_sz=em_sz, n_hid=nh, n_layers=nl, 
           dropout=0.1, dropouti=0.4, wdrop=0.5, dropoute=0.05, dropouth=0.3)
m3.reg_fn = partial(seq2seq_reg, alpha=2, beta=1)
m3.load_encoder(f'adam3_20_enc')

"""
Explanation: fastai can create a ModelData object directly from torchtext splits.
End of explanation
"""


m3.clip=25.
lrs=np.array([1e-4,1e-3,1e-2])

m3.freeze_to(-1)
m3.fit(lrs/2, 1, metrics=[accuracy])
m3.unfreeze()
m3.fit(lrs, 1, metrics=[accuracy], cycle_len=1)

m3.fit(lrs, 7, metrics=[accuracy], cycle_len=2, cycle_save_name='imdb2')

m3.load_cycle('imdb2', 4)

accuracy(*m3.predict_with_targs())

"""
Explanation: Because we're fine-tuning a pretrained model, we'll use differential learning rates, and also increase the max gradient for clipping, to allow the SGDR to work better.
End of explanation
"""


