# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # "A code-first inspection of the AWD-LSTM"
# > "AWD-LSTM is a special kind of recurrent neural network that is useful for language models. How can we set it up and customize it with the fastai library?"
# - toc: true
# - branch: master
# - badges: true
# - comments: true
# - categories: [NLP, fastai, RNN, LSTM]
# - image: images/some_folder/your_image.png
# - hide: false
# - search_exclude: false
# - metadata_key1: metadata_value1
# - metadata_key2: metadata_value2

from fastai2.text.all import *

# # ULMFiT

# In the previous post we explored the [Norec](https://github.com/ltgoslo/norec) Norwegian language corpus. We grabbed the reviews for films and TV-shows, parsed the html-text and created labels based on the ratings. In the next few posts I want to use [ULMFiT](https://nlp.fast.ai/classification/2018/05/15/introducing-ulmfit.html) and other methods to predict the sentiment of the reviews based on the text.  
# ULMFiT has three main steps:
# 1. Train a language model on a large general purpose corpus such as Wikipedia
# 1. Fine-tune the language model on the text your are working with - the style is most likely different than a Wikipedia article
# 1. Combine the encoder of the fine-tuned language model with a linear classifier to predict the class of your text

# The core of the ULMFiT method is a type of Recurrent neural network (RNN) called [AWD-LSTM](https://arxiv.org/abs/1708.02182). AWD-LSTM is a special kind of Recurrent neural network (RNN) with tuned dropout parameters among other. We need to look into this architecture before we continue with our modeling. For an explanation of what an *LSTM* actually is i suggest checking out this [blog post](https://colah.github.io/posts/2015-08-Understanding-LSTMs/) by <NAME>. In general, most of Chris' posts and papers are worth reading!

# # How to set up an AWD-LSTM with fastai

# Let's first start by inspecting fastai's `language_model_learner`. It's a learner class designed to be used for language models, and holds both dataloaders and the architecture along with various hyperparameters. We can use the `doc()` method to show us the [documentation](https://dev.fast.ai/text.learner#Learner-convenience-functions):

doc(language_model_learner)

# The documentation tells us that we can pass `arch = AWD-LSTM` and modify the `awd_lstm_lm_config` to customize the architecture. The config dictionary specifies various hyperparameters and settings inspired by the aforementioned AWD-LSTM paper. By changing this dictionary we can customize our AWD-LSTM to fit our specific needs:

awd_lstm_lm_config

# Let's check the [documentation](https://dev.fast.ai/text.models.awdlstm#AWD_LSTM) and source code of the AWD-LSTM class. You can check the source code directly in the notebook by appending a `??` behind the method name:

# +
# AWD_LSTM??
# -

# The source code shows us a few interesting lines we'll look more into in the next few sections:
# 1. `self.encoder = nn.Embedding(vocab_sz, emb_sz, padding_idx=pad_token)`
# 1. `self.rnns = nn.ModuleList([self._one_rnn(emb_sz if l == 0 else n_hid, (n_hid if l != n_layers - 1 else emb_sz)//self.n_dir, bidir, weight_p, l) for l in range(n_layers)])`
# 1. `self.input_dp = RNNDropout(input_p)`
# 1. `self.hidden_dps = nn.ModuleList([RNNDropout(hidden_p) for l in range(n_layers)])`

# > Note: the embedding is called encoder in the code above. The name encoder is also fastai lingo for the entire RNN-part of the architecture. The linear layers added on top for the classifier is called decoder. Neither the ULMFiT or AWD-LSTM paper uses the term encoder or decoder though.

# ## But what is an embedding?

# Once again I'll be lazy and rather refer to another [blog](http://jalammar.github.io/illustrated-word2vec/) that explains embeddings in detail. The blog is by <NAME> and has explanations of many deep learning and NLP concepts. The essence is that we'll turn each token in our vocabulary into a vector of some size that represents various aspects of that token. The weights of this vector will be trainable and gives our neural network a lot of flexibility in assigning various properties to each token.

# The embedding is created by: `self.encoder = nn.Embedding(vocab_sz, emb_sz, padding_idx=pad_token)` Here we see that fastai is built on top of pyTorch and relies on pyTorch's fundamental methods in its own code. The encoder layer is a call to `nn.Embedding`, see [documentation](https://pytorch.org/docs/stable/nn.html#normalization-layers). Let's create an embedding of size 10x3 with padding_idx = 0:

embedding = nn.Embedding(num_embeddings=10, embedding_dim=3, padding_idx=0)
embedding.weight

# The embedding now has 10 vectors of length 3 with randomly initialized weights. Note that the first one (index 0) is all 0. This is because 0 is our padding index. Next we'll pass some some sample data `inp`, and inspect the result. We can think of our input as the index of words in a dictionary. E.g. 1='this', 7='is', 4='not' and 3='easy'. 0 will be our padding token. The padding token is a special token that is used to ensure that some text has a certain length. This is useful when stacking various pieces of text into a batch where sizes needs to match.

inp = torch.LongTensor([1,7,4,3,0,0])
emb = embedding(inp)
emb

# We see that the embedding produced by feeding the input corresponds to the weights of our original embedding. That is, index 1 of `inp` is the item '7'. So `emb[1]` is basically a lookup for `embedding.weight[7]`.

embedding.weight[7], emb[1]

# > To summarize:  We'll need an embedding with the number of embeddings equal to our vocabulary size, and embedding size of 400 and a padding token-id which corresponds to whichever token has been used as padding in our vocabulary.

# ## Compostion of the RNN-layers

# Secondly we create a list of RNN-layers with various dimensions:
# ```python
# self.rnns = nn.ModuleList([self._one_rnn(emb_sz if l == 0 else n_hid, 
#                          (n_hid if l != n_layers - 1 else emb_sz)//self.n_dir, bidir, weight_p, l) for l in(n_layers)])
# ```
# The code stacks RNN-layers of `embedding size x hidden size` for the first layer, and `hidden size x embedding size` for the final. Let's verify that this works for various number of layers:

AWD_LSTM(vocab_sz=10_000, emb_sz=400, n_hid=1152, n_layers=2)

AWD_LSTM(vocab_sz=10_000, emb_sz=400, n_hid=1152, n_layers=5)

# We see the first and final layers have similar dimensions in the two examples.   
#
# > To summarize: We'll use a 3 layer network with input and output dimensions of (400, 1152), (1152, 1152) and (1152, 400) as in the AWD-LSTM paper. This should be handled automatically by the library.

# ## nn.LSTM

# In the module list above, the layers are actually WeightDropout layers. We can verify this from the hidden constructor method that is called when the RNNs are being created. First, a regular `nn.LSTM` layer is created before being passed to the `WeightDropout` module.
# ```python
# def _one_rnn(self, n_in, n_out, bidir, weight_p, l):
#     "Return one of the inner rnn"
#     rnn = nn.LSTM(n_in, n_out, 1, batch_first=True, bidirectional=bidir)
#     return WeightDropout(rnn, weight_p)
# ```

# Lets have a look at an example from the [nn.LSTM documentation](https://pytorch.org/docs/stable/nn.html#lstm), also see [source code here](https://pytorch.org/docs/stable/_modules/torch/nn/modules/rnn.html#LSTM). We'll make a 1 layer LSTM with input size of 10 and hidden size of 20. Note that in the AWD-LSTM case the input size is equal to the embedding size (400 by default).

inp_s = 10 # input size
hid_s = 20 # hidden size

lstm = nn.LSTM(input_size = inp_s, hidden_size = hid_s, num_layers=1)

# The documentation details that the LSTM expects input in the form of `input(seq_len, batch, input_size)`. Seq_len is the length of the part of the text the model will see in each iteration (seq_len = 72 by default in fastais language_model_learner, that is 72 tokens). The batch size is the number of documents the models sees in each iteration.
#
# h0 and c0 are the inital hidden and cell states (set to 0 if not provided). The documentation specifiy their shapes as: ` (num_layers * num_directions, batch, hidden_size)`.

bs = 16
n_l, n_d = 1, 1 # we are testing a 1 layer and 1 direction lstm
seq_len = 5

inp = torch.randn(seq_len, bs, inp_s)
h0 = torch.randn(n_l*n_d, bs, hid_s)
c0 = torch.randn(n_l*n_d, bs, hid_s)

inp.shape, h0.shape, c0.shape

# The output from the LSTM should be a tuple of `output, (h_n, c_n)`where output has shape given by: `(seq_len, batch, num_directions * hidden_size)`:

out, (hn, cn) = lstm(inp)
out.shape

# Let's also check the actual shape of our weights by looping through the `state_dict()`:

[(key, lstm.state_dict()[key].shape) for  key in lstm.state_dict().keys()]

# Here we recognize the input size of 10 and hidden size of 20, but where does the 80 come from? The documentation specifies that the weights will be of dimension `(4*hidden_size, input_size)`. The '4' is called gate_size, and we can find this in the source code for the base RNN module also:
# ```python
# if mode == 'LSTM':
#     gate_size = 4 * hidden_size
# ```

# >To summarize: we expect two sets of weights and biases per LSTM: 
# * weight_ih_l0 with a shape of (4*hidden_size, input_size)
# * weight_hh_l0 with a shape of (4*hidden_size, hidden_size)
# * bias_ih_l0 with a shape of (4*hidden_size)
# * bias_hh_l0 with a shape of (4*hidden_size)

# ## WeightDropout

# We can see from the `_one_rnn` that the nn.LSTM is transformed to a WeightDropout module. The [documentation](https://dev.fast.ai/text.models.awdlstm#WeightDropout) describes the module as 'A module that warps another layer in which some weights will be replaced by 0 during training'. From the source code we can see that it's the `weight_hh_l0` weights that will be modified, and that these weights are duplicated with suffix 'raw' in the WeightDropout module: `self.register_parameter(f'{layer}_raw', nn.Parameter(w.data))`. 

# Let's see if we can verify this by first checking the weights from the lstm from the above section. Of the 80\*20 = 1600 weights in the hh_l0 layer, none are 0:

orig_wts = getattr(lstm, 'weight_hh_l0')
orig_wts.shape, (orig_wts == 0.).sum()

# But if pass the lstm through the WeightDroput module, approximately half of the 1600 weights are set to 0. Note that we have to call the model on the input since the weights are only reset during the forward pass. The weights are also *only* reset for the WeighDropout's  internal LSTM module, while a copy with suffix '\_raw' retains the original weights.

wd = WeightDropout(lstm, weight_p=0.5)
_,_ = wd(inp) # we don't need the output in this case

# The `.module` attribute of the `wd` object is our original LSTM:

wd.module

# And about half its weights have been set to 0:

(getattr(wd.module, 'weight_hh_l0')==0.0).sum()

# The original weights from the lstm matches the '\_raw' weights of the WeightDropout module:

test_eq(orig_wts, getattr(wd, 'weight_hh_l0_raw'))

# The new layers are as expected:

wd.state_dict().keys()

# ## RNN dropout

# Finally several RNNDropout layers are being created - one for the input and one for each LSTM. This dropout is applied to the input embedding and on the output of each LSTM. We can test the functionality with `inp` from the above section.

dp = RNNDropout(0.5)
dp_out = dp(inp)
inp.shape, dp_out.shape

# The documentation also says: 'Dropout with probability p that is consistent on the seq_len dimension.' In our input from the above section, seq_len is the first dimension (index 0), and if we check for items equaling 0 and sum along the second dimension (index 1) we see that the same tokens are dropped out for the entire batch (our sample batch size is 16) consistently in approximately half of the instances.

(dp_out == 0).sum((1))

# # IMDb inspection

# Let's take a look at a minimal IMDB example from the [fastai documentation](https://dev.fast.ai/tutorial.datablock#Text) to verify our understanding of the AWD-LSTM architecture.

imdb_path = untar_data(URLs.IMDB_SAMPLE)
df = pd.read_csv(imdb_path/'texts.csv')
dls = TextDataLoaders.from_df(df, path=imdb_path, text_col='text', is_lm=True, valid_col='is_valid')
learn = language_model_learner(dls, AWD_LSTM)

# The vocab is of length 7080 and vocab index 1 is 'xxpad':

dls.vocab[:5], len(dls.vocab)

# In our model we recognize `Embedding(7080, 400, padding_idx=1)` as vocab_size x embedding size with the correct padding token. We also see that the (input, output) dimensions of our LSTM-layers are as expected, and with the expected dropout layers added.

learn.model

# The model summary shows us the default batch size of 64 and seq_len of 72.

learn.summary()

# And finally, the layer names and shapes also is consistent with a gate size of 4 (1152\*4 = 4608). Note the enumeration of the layers: **0.** is the encoder part of the architecture (including the embedding called encoder) and **1.** is the decoder.

for key in learn.model.state_dict().keys():
    print(key, '\t', learn.model.state_dict()[key].shape)

StarcoderdataJupyterScriptsDedupFiltered