In this notebook, we'll be implementing the Transformer-XL architecture from scratch.
To run this notebook in full, make sure to use the `download_data.sh` script to download the Penn Treebank data.

```python
from typing import *
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
%matplotlib inline
```

```python
device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
```

```python

```

# Review of the Transformer

Let's start off with a quick overview of the Transformer architecture which will be the basis of the Transformer XL.

Overall, the Transformer architecture is composed of multiple MultiHeadAttention layer stacked on top of each other, followed by feedforward layers, residual connections, and layer normalization layers.

The MultiHeadAttention layer is composed of multiple attention heads. Each attention head applies a linear transformation to its inputs and computes attention over its input values using keys and queries.

This approach is incapable of handling position, so the Transformer adds embeddings representing the position of the input to the word embeddings.

```python

```

For more details, please refer to [this tutorial](https://github.com/keitakurita/Practical_NLP_in_PyTorch/blob/master/deep_dives/transformer_from_scratch.ipynb) I've written on just the Transformer.

We'll be building the Transformer XL by first implementing a single attention head, scaling it to compose the MultiHeadAttention layer for the Transformer XL, then building the DecoderBlock and stacking them to create the full Transformer XL.

```python

```

# Implementing the Transformer XL

### A Single Attention Head

We'll start off by implementing a single attention head in a MultiHeadAttention layer. To make things concrete, let's consider the first layer and assume we receive an input of word embeddings of shape `(seq=7, batch_size=3, embedding_dim=32)`. Note that the Transformer XL does not add positional embeddings to the input.

```python
seq, batch_size, embedding_dim = 7, 3, 32
```

```python
word_embs = torch.rand(seq, batch_size, embedding_dim)
```

In the Transformer XL, we also feed the cached outputs of the model for the previous sequence. In this case, we would be feeding the word embeddings from the previous sequence as an additional input to our model.

To make things concrete, let's imagine our previous sequence was of length `prev_seq=6`

```python
prev_seq = 6
memory = torch.rand(prev_seq, batch_size, embedding_dim) # hidden states from the previous sequence
```

Each attention head takes keys, queries, and values as input. The processing goes like this:

1. Apply a separate linear transformation to each of the keys, queries, and values.
2. Compute attention scores for each of the values.
3. For each query, compute an attention-weighted sum of the values.
4. Apply a residual connection and layer normalization.

We'll start off with the linear transformation.

```python
inner_dim = 17 # this will be the internal dimension
linear_k = nn.Linear(embedding_dim, inner_dim)
linear_v = nn.Linear(embedding_dim, inner_dim)
linear_q = nn.Linear(embedding_dim, inner_dim)
```

The memory is concatenated across the sequence dimension and fed as keys/values. Be careful, as it's not concatenated with the queries. This is because each query represents one word we want to predict, so we can't modify the number of queries.

```python
word_embs_w_memory = torch.cat([memory, word_embs], dim=0)
k_tfmd = linear_k(word_embs_w_memory)
v_tfmd = linear_v(word_embs_w_memory)
q_tfmd = linear_q(word_embs) # No memory for the queries
```

Now, we compute scaled dot product attention as per the usual Transformer. Scaled dot product attention computes the attention score as the dot product between the query and key vectors. To prevent the values from exploding as the dimensionality of the vectors increases, we divide the raw attention score by the sqrt of the embedding size.

$$ \textrm{Attention}(Q, K, V) = \textrm{softmax}(\frac{QK^T}{\sqrt{d_k}})V $$

We'll be using einsum notation here to make the code easy to read: if you're not familiar with einsum, check out [this awesome tutorial](https://rockt.github.io/2018/04/30/einsum). In short, einsum denotes the shape of the inputs and outputs using one letter to represent each dimension. Below, the inputs are shaped `(i, b, d)` and `(j, b, d)` and the output is shaped `(i, j, b)` where the same letter represents the same size. Einsum is computed by taking the dot product across dimensions with the same character.

```python
content_attn = torch.einsum("ibd,jbd->ijb", q_tfmd, k_tfmd) / (embedding_dim ** 0.5) # scale
```

Notice we're not yet applying the softmax activation. This is because we need a couple more pieces to get the full attention score. The first of these is the relative positional embeddings.

```python

```

### Relative positional encodings

One of the key ideas in the Transformer XL is the idea of relative positional encodings. Instead of having a single embedding represent each **absolute** position, the Transformer XL computes an embedding that represents the **distance** between any two tokens. This is used to compute the attention between the two words.

The authors use the following equation to compute the attention between a query vector $ q_i $ and key vector $k_j$:

\begin{align}
A^{rel}_{i,j} =
    \underbrace{E_{x_i}^TW_q^TW_{k,E}E_{x_j}}_{(a)}
    + \underbrace{E_{x_i}^TW_q^TW_{k,R} \color{green}R_\color{green}{i-j} }_{(b)}
    \\
    + \underbrace{ \color{red}u^\color{red}T W_{k,E}E_{x_j}}_{(c)}
    + \underbrace{ \color{red}v^\color{red}T W_{k,R} \color{green}R_\color{green}{i-j}}_{(d)}
\end{align}

Here, $E_{x}$ is the embedding for $ x $ and $ W $ are all transformation matrices. The (a) term is the content-based attention terms that we already computed above. (b) and (d) are based on the relative positional embeddings and are dependent on the distance between $q_i$ and $k_j$. $u$ and $v$ are global bias terms that represent biases for certain content and certain positions.

Let's move on to the detailed implementation of terms (b) to (d). We'll first add the content bias (term (c) in the above equation) since it is the most simple to compute.

```python
u = torch.rand(17).expand_as(q_tfmd)
content_attn = content_attn + torch.einsum("ibd,jbd->ijb", u, k_tfmd) / (embedding_dim ** 0.5)
```

Next, we compute the relative positional embeddings necessary for the positional attention terms. For the relative positional embeddings, the Transformer XL uses fixed sinusoidal embeddings.

```python
pos_idxs = torch.arange(seq + prev_seq - 1, -1, -1.0, dtype=torch.float)
pos_idxs
```

    tensor([12., 11., 10.,  9.,  8.,  7.,  6.,  5.,  4.,  3.,  2.,  1.,  0.])

```python
inv_freq = 1 / (10000 ** (torch.arange(0.0, embedding_dim, 2.0) / embedding_dim))
sinusoid_inp = torch.einsum("i,j->ij", pos_idxs, inv_freq)
plt.plot(sinusoid_inp[0, :].detach().numpy());
plt.plot(sinusoid_inp[6, :].detach().numpy());
```

```python
relative_positional_embeddings = torch.cat([sinusoid_inp.sin(), sinusoid_inp.cos()], dim=-1)[:,None,:]
relative_positional_embeddings.shape
```

    torch.Size([13, 1, 32])

Now, we can gather this into its own class

```python
class PositionalEmbedding(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.d = d
        inv_freq = 1 / (10000 ** (torch.arange(0.0, d, 2.0) / d))
        # register buffer tells pytorch that this tensor is part of the modle
        # this means that it will be saved in the state_dict and moved to the GPU
        # along with the model
        self.register_buffer("inv_freq", inv_freq)

    def forward(self, positions: torch.LongTensor, # (seq, )
               ):
        # outer product
        sinusoid_inp = torch.einsum("i,j->ij", positions.float(), self.inv_freq)
        pos_emb = torch.cat([sinusoid_inp.sin(), sinusoid_inp.cos()], dim=-1)
        return pos_emb[:,None,:]
```

We also apply transformations to the positional embeddings separate from the values/keys.

```python
linear_p = nn.Linear(embedding_dim, inner_dim)
pos_tfmd = linear_p(relative_positional_embeddings)
```

This time, we'll be adding the positional bias during attention computation.

```python
v = torch.rand(17) # positional bias
pos_attn = torch.einsum("ibd,jd->ijb", q_tfmd + v, pos_tfmd[:,0,:]) / (embedding_dim ** 0.5) # scale
pos_attn.shape
```

    torch.Size([7, 13, 3])

Since we compute a relative postional embedding for each key-query pair, a naive implementation of attention using relative positional embeddings would be $O(n^2)$ in terms of computational complexity. Luckily, the authors proposed a trick to reduce this to $O(n)$ time by computing the attention for one query then shifting the embeddings for different query positions..

```python
zero_pad = torch.zeros((seq, 1, batch_size), dtype=torch.float)
# this padding + shifting efficiently computes the attention for all
pos_attn = (torch.cat([zero_pad, pos_attn], dim=1)
                    .view(seq + prev_seq + 1, seq, batch_size)[1:]
                    .view_as(pos_attn))
```

The attention is computed as the sum of content and positional attention.

```python
raw_attn = content_attn + pos_attn
```

When we do language modeling, we need to prevent the model from being able to look at the word that it is supposed to be predicting. In the Transformer, we achieve this by setting the attention score to zero. This masks out words that we don't want the model to be able to see.

```python
mask = torch.triu(
    torch.ones((seq, seq + prev_seq)),
    diagonal=1 + prev_seq,
).byte()[...,None]
raw_attn = raw_attn.masked_fill(mask, -float('inf'))
```

We can now compute the outputs as the weighted sum of the value vectors using the attention scores.

```python
attn = torch.softmax(raw_attn, dim=1)
attn_weighted_sum = torch.einsum("ijb,jbd->ibd", attn, v_tfmd)
attn_weighted_sum.shape
```

    torch.Size([7, 3, 17])

Finally, we project the attention weighted sums back to their original dimension and apply a residual connection and layer normalization. We apply layer normalization after the residual connection.

```python
linear_out = nn.Linear(inner_dim, embedding_dim)
layer_norm = nn.LayerNorm(embedding_dim)
output = layer_norm(word_embs + linear_out(attn_weighted_sum))
output.shape
```

    torch.Size([7, 3, 32])

```python

```

### MultiHeadAttention (MHA): The core component

Aggregating all the above and applying a couple of optimizations by grouping some computations together as well as adding dropout, we get the following MultiHeadAttention module.

```python
from typing import *

class MultiHeadAttention(nn.Module):
    def __init__(self, d_input: int, d_inner: int, n_heads: int=4,
                 dropout: float=0.1, dropouta: float=0.):
        super().__init__()
        self.d_input = d_input
        self.d_inner = d_inner
        self.n_heads = n_heads
        # this layer applies the linear transformation required
        # for the keys and values for all heads at once for efficiency
        self.linear_kv = nn.Linear(
            d_input,
            (d_inner * n_heads * 2), # 2 is for keys and values
            bias=False, # we don't apply bias, making this a simple matrix multiplication
        )
        # for queries (will not be concatenated with memorized states so separate)
        self.linear_q = nn.Linear(
            d_input, d_inner * n_heads,
            bias=False
        )
        # for positional embeddings
        self.linear_p = nn.Linear(
            d_input, d_inner * n_heads,
            bias=False
        )
        self.scale = 1 / (d_inner ** 0.5) # for scaled dot product attention
        self.dropa = nn.Dropout(dropouta)
        # we will use this to project back to the input dimension
        self.lout = nn.Linear(self.d_inner * self.n_heads, self.d_input, bias=False)
        self.norm = nn.LayerNorm(self.d_input)
        self.dropo = nn.Dropout(dropout)

    def _rel_shift(self, x):
        zero_pad = torch.zeros((x.size(0), 1, *x.size()[2:]),
                               device=x.device, dtype=x.dtype)
        return (torch.cat([zero_pad, x], dim=1)
                    .view(x.size(1) + 1, x.size(0), *x.size()[2:])[1:]
                    .view_as(x))

    def forward(self, input_: torch.FloatTensor, # (cur_seq, b, d_in)
                pos_embs: torch.FloatTensor, # (cur_seq + prev_seq, d_in)
                memory: torch.FloatTensor, # (prev_seq, b, d_in)
                u: torch.FloatTensor, # (H, d)
                v: torch.FloatTensor, # (H, d)
                mask: Optional[torch.FloatTensor]=None,
        ):
        """
        pos_embs: we pass the positional embeddings in separately
            because we need to handle relative positions
        input shape: (seq, bs, self.d_input)
        pos_embs shape: (seq + prev_seq, bs, self.d_input)
        output shape: (seq, bs, self.d_input)
        """
        cur_seq = input_.shape[0] #  sequence length of current segment
        prev_seq = memory.shape[0] # sequence length of previous segment
        H, d = self.n_heads, self.d_inner
        input_with_memory = torch.cat([memory, input_], dim=0) # concatenate recurrent memory
                                                               # across sequence dimension

        # we will use the following symbols to represent the shape of the tensors
        # cs: current sequence length, b: batch, H: number of heads
        # d: inner dimension, ps: previous sequence length
        # The key and value are now conditioned on the preceding context
        k_tfmd, v_tfmd = \
            torch.chunk(self.linear_kv(input_with_memory), 2, dim=-1) # (cs + ps, b, H * d)
        q_tfmd = self.linear_q(input_) # (cs, b, H * d)

        # apply scaled dot product attention
        # look at the following dimensions carefully, since this is the key operation
        # in the Transformer/Transformer XL architecture

        _, bs, _ = q_tfmd.shape
        assert bs == k_tfmd.shape[1]
        # content-based attention term ((a) + (c) in the paper)
        # this is the standard attention term in the original Transformer, except without positional embeddings
        # which are handled separately in the Transformer XL (see below)
        # here, i corresponds to the number of queries = number of current inputs/targets (seq-wise)
        # j corresponds to the number of key/values = number of vectors that we can use to compute the
        # vector for each query
        content_attn = torch.einsum("ibhd,jbhd->ijbh", (
                (q_tfmd.view(cur_seq, bs, H, d) + # (a)
                 u), # (c): u represents the global (independent of the query)
                     # bias towards certain key/values = words
                     # Note: maybe this could be a per-attention head parameter?
                 k_tfmd.view(cur_seq + prev_seq, bs, H, d) # There is no positional information to be found here
        )) # (cs, cs + ps, b, H)

        # position-based attention term ((b) + (d) in the paper)
        # this attention is solely based on the position of the key/values
        # (i.e. it does not take the content of the key/values into account)
        p_tfmd = self.linear_p(pos_embs) # (cs + ps, b, H * d)
        position_attn = torch.einsum("ibhd,jhd->ijbh", (
                (q_tfmd.view(cur_seq, bs, H, d) + # (b)
                 v), # (d): v represents the global (independent of the query)
                     # bias towards certain positions
                 p_tfmd.view(cur_seq + prev_seq, H, d) # Notice there is not content information
                                                        # regarding keys and values here!
        )) # (cs, cs + ps, b, H)

        #  Compute positional attention efficiently
        position_attn = self._rel_shift(position_attn)

        # the attention is the sum of content-based and position-based attention
        attn = content_attn + position_attn

        if mask is not None and mask.any().item():
            attn = attn.masked_fill(
                mask[...,None], -float('inf'))
        attn = torch.softmax(attn * self.scale, # rescale to prevent values from exploding
                             dim=1) # normalize across the value sequence dimension
        attn = self.dropa(attn)

        attn_weighted_values = (torch.einsum("ijbh,jbhd->ibhd",
                                           (attn, # (cs, cs + ps, b, H)
                                            v_tfmd.view(cur_seq + prev_seq, bs, H, d), # (cs + ps, b, H, d)
                                           )) # (cs, b, H, d)
                                .contiguous() # we need to change the memory layout to make `view` work
                                .view(cur_seq, bs, H * d)) # (cs, b, H * d)

        # Project back to input dimension and add residual connection
        output = input_ + self.dropo(self.lout(attn_weighted_values))
        output = self.norm(output)
        return output
```

```python

```

Let's test it out to see if it runs successfully

```python
mha = MultiHeadAttention(32, 17, n_heads=4)
```

```python
inpt = torch.rand(7, 3, 32)
pos = torch.rand(13, 32)
mem = torch.rand(6, 3, 32)
u, v = torch.rand(4, 17), torch.rand(4, 17)
x1 = mha(inpt, pos, mem, u, v)
```

Looks good

```python
x1.shape
```

    torch.Size([7, 3, 32])

```python
x1[0]
```

    tensor([[ 0.6264,  0.3405, -1.9065,  0.1543,  0.0389, -1.6033,  1.4415,  0.4983,
              0.7548,  1.0990, -1.1783, -1.3847, -1.7358,  1.4651,  1.0633,  0.2168,
             -0.3323,  1.1270,  0.1614,  1.0170,  1.0459, -0.7286,  0.5064, -1.4765,
              0.0448, -1.2500,  0.3132, -0.8007,  0.4089,  0.7325, -1.2740,  0.6147],
            [ 0.9541,  0.3682, -0.8096,  0.1357, -0.9159, -1.4382, -1.3385,  0.8269,
              0.2721, -0.4982,  1.3105, -0.0236, -1.0547, -1.3076,  1.8884, -0.2891,
              1.5231,  0.5507, -0.6423,  0.4412,  1.3656,  0.7858, -0.9425, -0.3198,
             -0.3162, -0.0086,  1.5257, -1.3216,  1.4492, -0.1750, -0.1669, -1.8291],
            [ 1.0132,  0.7205, -0.4221,  0.2952, -1.4117, -0.6182, -1.7520, -1.7426,
             -0.4648, -0.2122,  2.0889, -1.3544, -0.1611, -1.0696,  1.3492, -1.0179,
              1.2820,  0.8990,  0.7411,  0.8052, -0.5322,  0.6277, -0.2733, -1.0738,
             -0.8435,  1.5357,  0.8260, -0.3422, -0.6204,  1.0091,  0.1011,  0.6182]],
           grad_fn=<SelectBackward>)

```python

```

### Building the decoder

To construct the decoder block, all we need in addition to the MultiHeadAttention layer is the Positionwise Feed Forward layer.

```python
class PositionwiseFF(nn.Module):
    def __init__(self, d_input, d_inner, dropout):
        super().__init__()

        self.d_input = d_input
        self.d_inner = d_inner
        self.dropout = dropout
        self.ff = nn.Sequential(
            nn.Linear(d_input, d_inner), nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(d_inner, d_input),
            nn.Dropout(dropout),
        )
        self.layer_norm = nn.LayerNorm(d_input)

    def forward(self, input_: torch.FloatTensor, # (cur_seq, bs, d_input)
               ) -> torch.FloatTensor: # (cur_seq, bs, d_input)
        ff_out = self.ff(input_)
        output = self.layer_norm(input_ + ff_out)
        return output
```

Now we can implement the decoder block.

```python
class DecoderBlock(nn.Module):
    def __init__(self, n_heads, d_input,
                 d_head_inner, d_ff_inner,
                 dropout, dropouta=0.):
        super().__init__()
        self.mha = MultiHeadAttention(d_input, d_head_inner, n_heads=n_heads,
                                      dropout=dropout, dropouta=dropouta)
        self.ff = PositionwiseFF(d_input, d_ff_inner, dropout)

    def forward(self, input_: torch.FloatTensor, # (cur_seq, bs, d_input)
                pos_embs: torch.FloatTensor, # (cur_seq + prev_seq, d_input),
                u: torch.FloatTensor, # (H, d_input),
                v: torch.FloatTensor, # (H, d_input),
                mask=None,
                mems=None,
               ):
        return self.ff(self.mha(input_, pos_embs, mems, u, v, mask=mask))
```

```python

```

### The full Transformer XL

```python
import torch.nn.functional as F
```

Now with these components in place, we can build the full Transformer XL model.

Aside from what we mentioned above, one common trick in language modeling that we haven't covered yet is tying the input embedding matrix $ E $ and output projection matrix $ P $. Remember, a language model predicts the next token in a sequence, so its output dimension is $\mathbb{R}^{|V|}$  where $|V|$ is the vocab size. If we constrain the penultimate layer output to be the same dimension as the embeddings $d$, the embedding matrix $ E $ will be of shape $ \mathbb{R}^{|V| \times d}$ and the output projection matrix $ P $ will be of shape $ \mathbb{R}^{d \times |V|} $.

In [this paper](https://arxiv.org/abs/1608.05859), the authors found that constraining the matrices such that $ P = E^T $ improved performance while greatly reducing the total parameter count (and thus memory usage!) of the model.

Instead of simply using the exact same weights, we scale the embeddings by the embedding dim. This trick is included in the codebase but not mentioned in the paper as far as I can tell. If you're aware of a paper where this trick was originally introduced, please let me know!

```python
class StandardWordEmbedding(nn.Module):
    def __init__(self, num_embeddings, embedding_dim,
                div_val=1, sample_softmax=False):
        super().__init__()
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        self.embedding = nn.Embedding(num_embeddings, embedding_dim)
        self.scale = embedding_dim ** 0.5

    def forward(self, input_: torch.LongTensor):
        return self.embedding(input_) * self.scale
```

Now, all we need to do is to put everything we have implemented above together.

```python
class TransformerXL(nn.Module):
    def __init__(self, num_embeddings, n_layers, n_heads,
                 d_model, d_head_inner, d_ff_inner,
                 dropout=0.1, dropouta=0.,
                 seq_len: int=0, mem_len: int=0):
        super().__init__()
        self.n_layers,self.n_heads,self.d_model,self.d_head_inner,self.d_ff_inner = \
            n_layers,n_heads,d_model,d_head_inner,d_ff_inner
        # Embedding layers
        self.word_embs = StandardWordEmbedding(num_embeddings, d_model)
        self.pos_embs = PositionalEmbedding(d_model)
        # Core transformer
        self.drop = nn.Dropout(dropout)
        self.layers = nn.ModuleList([DecoderBlock(n_heads, d_model, d_head_inner=d_head_inner,
                                                  d_ff_inner=d_ff_inner,
                                                  dropout=dropout, dropouta=dropouta)
                                     for _ in range(n_layers)])

        # tie weights
        self.output_projection = nn.Linear(d_model, num_embeddings)
        self.output_projection.weight = self.word_embs.embedding.weight
        self.loss_fn = nn.CrossEntropyLoss()

        self.seq_len, self.mem_len = seq_len, mem_len

        # u and v are global parameters: maybe changing these to per-head parameters
        # might help performance?
        self.u, self.v = (nn.Parameter(torch.Tensor(self.n_heads, self.d_head_inner)),
                          nn.Parameter(torch.Tensor(self.n_heads, self.d_head_inner)))

    def init_memory(self, device=torch.device("cpu")) -> torch.FloatTensor:
        return [torch.empty(0, dtype=torch.float).to(device) for _ in range(self.n_layers+1)]

    def update_memory(self,
            previous_memory: List[torch.FloatTensor],
            hidden_states: List[torch.FloatTensor],
        ):
        assert len(hidden_states) == len(previous_memory)
        mem_len, seq_len = previous_memory[0].size(0), hidden_states[0].size(0)

        # For the updated memory, we use the most recent `self.mem_len`
        # states, including the previous memory
        # In other words, if `seq_len` < `self.mem_len` some of the previous memory
        # will carry over to the next memory
        with torch.no_grad():
            new_memory = []
            end_idx = mem_len + seq_len
            beg_idx = max(0, end_idx - self.mem_len)
            for m, h in zip(previous_memory, hidden_states):
                cat = torch.cat([m, h], dim=0) # (mem_len + seq_len, bs, d)
                new_memory.append(cat[beg_idx:end_idx].detach()) # (self.mem_len, bs, d)
        return new_memory

    def reset_length(self, seq_len, ext_len, mem_len):
        self.seq_len = seq_len
        self.mem_len = mem_len

    def forward(self, idxs: torch.LongTensor, # (cs, bs)
                target: torch.LongTensor, # (cs, bs)
                memory: Optional[List[torch.FloatTensor]]=None,
               ) -> Dict[str, torch.Tensor]:
        if memory is None:
            memory: List[torch.FloatTensor] = self.init_memory(idxs.device)
        assert len(memory) == len(self.layers) + 1
        cur_seq, bs = idxs.size()
        prev_seq = memory[0].size(0)

        # Construct attention mask
        dec_attn_mask = torch.triu(
            torch.ones((cur_seq, cur_seq + prev_seq)),
            diagonal=1 + prev_seq,
        ).byte()[...,None].to(idxs.device)

        word_embs = self.drop(self.word_embs(idxs))
        pos_idxs = torch.arange(cur_seq + prev_seq - 1, -1, -1.0, dtype=torch.float).to(word_embs.device)
        pos_embs = self.drop(self.pos_embs(pos_idxs))

        # Main part of forward pass
        hidden_states = [word_embs]
        layer_out = word_embs
        for mem, layer in zip(memory, self.layers):
            layer_out = layer(layer_out, pos_embs, self.u, self.v,
                              mask=dec_attn_mask, mems=mem)
            hidden_states.append(layer_out)

        logits = self.output_projection(self.drop(layer_out))
        loss = self.loss_fn(logits.view(-1, logits.size(-1)), target.view(-1))

        # Update memory
        # Ensure the memory is treated as a constant
        # and we do not back propagate through them
        new_memory = self.update_memory(memory, hidden_states)
        return {"loss": loss, "logits": logits, "memory": new_memory}
```

```python
transformer = TransformerXL(1000, 4, 3, 32, 17, 71, mem_len=5)
```

Again, let's feed some random inputs to confirm the model is working.

```python
idxs = torch.randint(1000, (5, 9))
tgts = torch.randint(1000, (5, 9))
transformer(idxs, tgts)
```

    {'loss': tensor(22.7983, grad_fn=<NllLossBackward>),
     'logits': tensor([[[  2.3893,   1.9681,  -2.6914,  ...,  -0.8093,  -2.1468,  -1.5656],
              [ -6.9673,  -5.4040,   1.5984,  ...,  -8.6291,   9.6236,  -3.4973],
              [ -6.9776,  -3.5877,   0.9394,  ..., -12.9488,   4.1804,   4.1679],
              ...,
              [ -7.8283,   3.3389,  -1.6314,  ...,  -3.7291,   5.9029,  -1.5640],
              [  5.7904,   4.9521, -11.5190,  ...,  -1.1907,   5.1361,  -0.6909],
              [ -1.2347,   4.8922,  -2.4506,  ...,   1.5876,   5.6535,  -3.9682]],

             [[  2.7837,  -1.6279,   5.5712,  ...,   2.1291,   3.5525,   8.2256],
              [  4.4832,  -1.4655,  -7.8257,  ...,  -0.7932,  -1.3533,   1.5421],
              [ -6.0790,  -7.7122,  -3.2280,  ..., -14.7921,  -1.0191,  19.7041],
              ...,
              [ -8.7061,  -2.2215, -11.9336,  ...,  -5.7044,   0.8774,  -1.9217],
              [ -2.3067,  -3.2391,  -1.1071,  ...,   0.2583,  -5.1745, -15.7446],
              [ -0.6714,   8.3893,  -1.0352,  ...,   1.4273,  18.6013,  -3.9901]],

             [[  1.1127, -10.2239,   0.9317,  ..., -10.6415,   4.1240,   5.4683],
              [  9.5186,  -8.7525,  -0.4702,  ...,  -7.4802,  -1.4103,   8.6529],
              [-10.7924,  -2.3530,  -1.1558,  ...,  -9.1787,  -0.4928,  -0.9864],
              ...,
              [  6.2984,   0.5173,   6.5342,  ...,  -4.1923,  -3.8633,  -0.7991],
              [ -9.1815,  -5.3329,  -7.6900,  ...,  -8.8483,   2.2137,   1.6496],
              [ -5.3949,   4.4151,  -2.7581,  ...,  -0.5034,   7.8452,  -1.1456]],

             [[ -1.8854,  -3.1880,   5.2209,  ...,   4.6219,   5.8592,   6.7511],
              [  1.7416,   0.8599,   3.7392,  ...,   2.1722,   2.7656, -12.5893],
              [ -4.7122,   5.6147,  -0.2521,  ..., -13.2473,   2.8461,   4.8401],
              ...,
              [ -9.8932,   4.5135,   9.3876,  ...,   2.2749,   4.2054,   8.6749],
              [  2.8098,  -4.5236,   2.9966,  ...,  -9.3179,  -3.0171,   2.8091],
              [ -4.3847,  -8.1687, -10.1101,  ...,  -9.0084,   0.1503,   1.9549]],

             [[ -3.9730,   5.7617,   0.9482,  ...,   5.6215,  -6.8315,  -3.8019],
              [  2.0387,  -1.3776,   4.5145,  ...,  -2.6593,  -7.4511,  -4.3874],
              [  8.2944,  -2.8881,  -0.5630,  ...,  -4.1937,  -3.9678,  -5.0382],
              ...,
              [  9.9117,  -5.5987,   3.6440,  ...,  -3.4378,  -4.1492,   4.0374],
              [ -6.7007,   0.7956, -11.9654,  ...,  -3.4386,  -6.1001,   2.1052],
              [ -1.1894,  -3.6629,  10.9244,  ...,  -5.9781,   5.6562,   8.7496]]],
            grad_fn=<AddBackward0>),
     'memory': [tensor([[[ 2.1017e+00,  7.0014e+00, -1.1576e+01,  ..., -4.1735e+00,
                 7.1704e+00, -9.6242e-01],
               [-4.9919e-01,  0.0000e+00, -2.2360e+00,  ..., -7.4936e-01,
                 5.5746e-01,  4.2537e+00],
               [ 7.5275e+00, -0.0000e+00,  6.4217e+00,  ..., -1.0603e+00,
                 2.6083e-01,  0.0000e+00],
               ...,
               [-0.0000e+00, -5.6473e+00, -1.2624e+01,  ...,  5.6180e-01,
                 1.0439e+00,  7.7785e+00],
               [ 5.5092e+00,  3.7028e+00, -0.0000e+00,  ...,  1.2419e+01,
                -1.2681e+00, -3.8133e+00],
               [ 5.9784e+00, -1.2864e+00, -1.6390e+01,  ..., -4.7940e+00,
                 7.9133e+00, -7.4816e+00]],

              [[-1.2514e+01,  3.6603e+00, -1.1988e+01,  ..., -2.5402e+00,
                -6.0357e+00, -1.6756e+01],
               [-1.7803e+00,  2.9728e+00,  1.4087e+00,  ...,  4.8081e+00,
                 1.9607e+00,  5.2030e+00],
               [ 2.3099e+00, -5.5314e+00,  7.6574e+00,  ...,  7.1721e+00,
                -4.0299e-01, -0.0000e+00],
               ...,
               [-6.4647e+00,  9.5205e+00, -9.9825e+00,  ...,  5.7900e+00,
                 1.3428e+01,  7.6066e+00],
               [-3.3331e+00,  5.6945e+00, -3.7368e+00,  ...,  0.0000e+00,
                -2.3161e+00, -6.7402e+00],
               [-5.6126e+00,  6.1271e+00,  1.1794e+00,  ..., -1.2898e-01,
                 1.8484e+00, -2.3268e-01]],

              [[-1.0437e+01, -0.0000e+00, -5.1754e-01,  ..., -1.7993e+00,
                 1.6501e+00,  5.9062e+00],
               [ 1.0803e+01, -9.8891e+00, -1.6223e+00,  ..., -3.4334e+00,
                 3.3257e+00, -1.2487e+00],
               [-2.0990e-01,  5.3983e+00, -0.0000e+00,  ...,  7.3918e+00,
                -3.5076e+00,  1.8095e+00],
               ...,
               [ 3.4760e+00, -2.3750e+00,  1.5633e+00,  ..., -1.1462e+01,
                -0.0000e+00, -3.0820e+00],
               [ 3.1093e+00, -0.0000e+00,  3.0700e+00,  ..., -5.9241e+00,
                 0.0000e+00, -2.5284e+00],
               [-7.4560e+00,  1.1351e+00,  6.4418e+00,  ...,  9.4216e+00,
                 1.3614e+00,  1.3476e+00]],

              [[ 0.0000e+00, -0.0000e+00, -2.3496e+00,  ...,  4.0335e+00,
                 1.3728e+01,  4.7616e-01],
               [-9.9943e-01,  2.9256e+00, -2.8993e+00,  ..., -0.0000e+00,
                 6.8954e-01,  6.7603e+00],
               [-6.9932e+00, -5.8414e-01,  1.7446e+00,  ..., -2.0840e+00,
                 3.7240e+00, -3.9543e+00],
               ...,
               [-1.5157e+00,  0.0000e+00,  5.2591e+00,  ...,  3.7703e-01,
                -0.0000e+00, -1.3468e+00],
               [-1.0982e+00, -0.0000e+00, -4.7111e+00,  ...,  4.7480e+00,
                 4.9005e+00, -2.4251e+00],
               [-2.3143e+00,  0.0000e+00, -1.9084e+00,  ...,  1.6107e+00,
                 2.3575e+00, -5.4302e+00]],

              [[ 0.0000e+00,  1.3610e+01, -4.1463e+00,  ..., -4.5901e+00,
                 2.3292e+00, -1.0738e+01],
               [ 4.0665e+00,  0.0000e+00,  7.8862e+00,  ...,  0.0000e+00,
                -9.0152e+00, -2.4664e+00],
               [ 1.5953e+01, -1.2290e+00,  0.0000e+00,  ..., -0.0000e+00,
                 6.9739e+00,  0.0000e+00],
               ...,
               [-1.5245e+00,  1.4178e-02,  3.7615e-01,  ..., -1.5414e+00,
                 5.3257e+00,  6.1866e+00],
               [-7.2876e+00,  1.2365e+01, -2.8329e+00,  ...,  5.8748e+00,
                -6.5736e+00, -1.1682e+01],
               [-0.0000e+00, -4.1767e-01, -2.7669e-01,  ...,  5.1264e+00,
                 2.3285e+00,  5.2544e+00]]]),
      tensor([[[ 1.3826,  0.9239, -2.2534,  ...,  0.5224,  1.6166,  0.1766],
               [-0.0735,  0.2080, -0.4732,  ..., -0.4010, -0.2503,  0.4283],
               [ 1.8843, -0.6893,  0.9431,  ..., -0.2929, -0.1904,  0.0774],
               ...,
               [-0.5283, -1.1426, -1.9640,  ...,  0.1228,  0.3225,  0.7912],
               [ 2.0171,  0.6775,  0.0373,  ...,  0.5712, -0.4692, -1.0398],
               [ 1.3282, -0.1157, -2.2052,  ..., -0.1816,  1.3291, -1.3611]],

              [[-1.0597,  0.6241, -1.3270,  ..., -0.0906, -0.0414, -1.6208],
               [-0.4735,  0.9648, -0.1781,  ...,  1.0080, -0.0339,  0.7943],
               [ 1.3104, -1.8248,  0.5972,  ...,  0.5662, -0.6202,  0.1132],
               ...,
               [-0.2953,  0.9277, -1.1543,  ...,  0.7062,  2.1062,  0.5092],
               [-0.3843,  0.3064, -0.3782,  ..., -0.0153, -1.1835, -1.1707],
               [-1.0753,  0.9251, -0.4859,  ..., -0.0698,  0.5005, -1.1221]],

              [[-1.0846, -0.6314, -0.0176,  ..., -0.6004,  0.0514,  2.0411],
               [ 1.7022, -1.3314, -0.0765,  ..., -1.7647,  0.3343, -0.8403],
               [ 0.9930,  0.6483, -0.3341,  ...,  1.0334, -1.5398,  0.7327],
               ...,
               [ 1.3533, -1.2107,  0.4709,  ..., -1.0716, -0.0823, -0.8666],
               [ 1.6170,  0.4248,  0.3163,  ..., -1.3363, -0.8826, -0.7667],
               [ 0.0648,  0.1811,  1.0545,  ...,  1.7590,  0.5024,  0.3842]],

              [[ 0.0324, -1.0944, -1.1538,  ...,  0.4320,  2.1470,  0.4528],
               [-0.0979,  0.6155, -0.7837,  ..., -0.3645, -0.2686,  1.5890],
               [-0.3416,  0.3403,  0.1237,  ..., -0.8141,  0.5855, -0.7689],
               ...,
               [-0.0435, -1.0859,  1.3180,  ...,  0.3539,  0.3149, -0.3087],
               [-0.1995, -0.3531, -0.5246,  ...,  0.6699,  0.5467,  0.2343],
               [ 0.7265,  0.1116, -0.4563,  ...,  0.6268,  0.6353, -1.0494]],

              [[ 0.1539,  2.0148, -0.7285,  ..., -0.6078,  0.3653, -1.3313],
               [ 1.5128,  0.1665,  1.2747,  ...,  0.1989, -0.8545, -0.3040],
               [ 3.4671, -0.3973, -0.4849,  ...,  0.2124,  0.7675, -0.5249],
               ...,
               [ 0.1298, -0.5067,  0.1377,  ...,  0.6606,  0.6893,  1.1767],
               [-0.5558,  1.9915, -0.5002,  ...,  0.9498, -1.0524, -1.6267],
               [ 0.8146, -1.0068, -0.6703,  ...,  0.8581,  0.5871,  1.1277]]]),
      tensor([[[ 1.6346e+00,  9.5465e-01, -2.0691e+00,  ...,  4.4820e-01,
                 2.2910e+00,  4.9902e-01],
               [ 2.4900e-01,  1.8973e-01, -3.3778e-01,  ..., -8.5209e-01,
                -5.5017e-01,  5.1561e-01],
               [ 1.4326e+00, -8.8077e-01,  1.0193e+00,  ..., -3.8159e-01,
                 8.8002e-03,  5.8506e-01],
               ...,
               [-6.0461e-01, -7.7426e-01, -1.6539e+00,  ...,  1.4462e-01,
                -1.6225e-01,  5.6319e-01],
               [ 1.5300e+00,  7.9754e-01, -1.2463e-01,  ...,  1.5487e-01,
                 4.7096e-02, -2.2055e-01],
               [ 1.3635e+00,  2.2628e-01, -2.5490e+00,  ..., -2.4268e-01,
                 1.9050e+00, -3.0230e-01]],

              [[-7.8208e-01,  6.8655e-01, -1.4995e+00,  ..., -1.5197e-01,
                 7.7530e-01, -1.1415e+00],
               [ 2.8964e-01,  1.5075e+00,  2.6526e-02,  ...,  9.5593e-01,
                 1.4820e-01,  5.8030e-01],
               [ 1.0642e+00, -2.0875e+00,  7.8689e-01,  ...,  9.4528e-01,
                -8.7262e-01,  2.6703e-01],
               ...,
               [ 9.6837e-02,  1.2272e+00, -9.9427e-01,  ...,  5.3825e-01,
                 2.2138e+00,  8.4783e-01],
               [-4.6253e-01,  7.3703e-01, -4.7150e-01,  ...,  2.4403e-01,
                -1.1370e+00, -5.6916e-01],
               [-1.0717e+00,  1.2868e+00, -3.9602e-01,  ...,  4.7487e-01,
                 7.4984e-01, -1.3079e-01]],

              [[-9.9346e-01, -4.6816e-01, -4.2036e-01,  ...,  1.9013e-01,
                 5.6324e-01,  1.7506e+00],
               [ 1.4775e+00, -7.2553e-01, -3.9095e-01,  ..., -1.6459e+00,
                -5.1008e-01, -1.1018e+00],
               [ 1.0286e+00,  6.4774e-02, -1.1253e-02,  ...,  8.0176e-01,
                -1.8767e+00,  9.6665e-01],
               ...,
               [ 6.5649e-01, -9.1886e-01,  4.0203e-01,  ..., -8.7860e-01,
                 2.5020e-01, -2.3505e-02],
               [ 1.5418e+00,  4.3293e-01,  6.1181e-01,  ..., -1.0745e+00,
                -6.0872e-01, -2.7264e-01],
               [ 1.3029e-01,  5.1017e-01,  8.4829e-01,  ...,  1.6529e+00,
                 5.5392e-01,  9.1532e-01]],

              [[ 3.3911e-02, -8.4844e-01, -1.0074e+00,  ...,  4.9299e-01,
                 2.4838e+00,  3.1425e-01],
               [ 6.1697e-01,  7.7079e-01, -3.5887e-01,  ..., -2.5764e-01,
                -9.5384e-01,  1.4135e+00],
               [-3.3163e-01,  3.4299e-01,  5.2204e-01,  ..., -7.2144e-01,
                 6.0687e-01,  8.2380e-02],
               ...,
               [-1.9737e-01, -9.4352e-01,  1.7135e+00,  ...,  3.5844e-02,
                 7.1961e-01,  9.6403e-02],
               [-1.1394e-01, -1.4094e-02, -1.6593e-01,  ..., -1.3911e-01,
                 1.5920e-01,  4.5853e-01],
               [ 7.0644e-01,  6.4221e-01, -9.2809e-01,  ...,  1.2030e+00,
                 5.6345e-01, -1.0111e+00]],

              [[ 8.9835e-01,  1.9890e+00, -8.9823e-01,  ..., -8.1031e-01,
                 5.3538e-01, -1.2763e+00],
               [ 1.7208e+00,  2.1950e-01,  1.2462e+00,  ..., -2.1622e-01,
                -7.6752e-01,  1.7855e-03],
               [ 2.8792e+00, -3.1874e-01, -2.0828e-01,  ...,  1.0098e-01,
                 7.4217e-01, -2.6374e-02],
               ...,
               [ 1.7004e-02, -6.7681e-01,  2.5581e-01,  ...,  3.3252e-01,
                 8.9061e-01,  1.4128e+00],
               [ 1.5477e-02,  2.3536e+00,  7.8383e-02,  ...,  3.7001e-01,
                -8.4431e-01, -9.0505e-01],
               [ 5.4680e-01, -7.6991e-01, -4.1328e-01,  ...,  1.0264e+00,
                 9.9558e-01,  1.6870e+00]]]),
      tensor([[[ 1.3837e+00,  1.0989e+00, -2.2840e+00,  ...,  6.0480e-01,
                 1.8528e+00,  1.2792e-01],
               [-2.2875e-01,  9.7942e-02, -7.5347e-01,  ..., -5.9109e-01,
                -6.4622e-01,  4.3441e-01],
               [ 1.3624e+00, -5.5102e-01,  1.2233e+00,  ..., -4.0880e-01,
                -1.5982e-01,  5.9552e-01],
               ...,
               [-9.3327e-01, -1.0077e+00, -2.1290e+00,  ...,  3.9902e-01,
                 5.4323e-01,  1.4707e-01],
               [ 9.5668e-01,  1.4832e+00, -7.4227e-01,  ...,  4.8545e-01,
                 5.3537e-02, -9.6289e-01],
               [ 1.2578e+00,  6.9142e-02, -2.4296e+00,  ...,  9.2327e-02,
                 1.6102e+00, -1.2208e+00]],

              [[-8.3147e-01,  1.2378e+00, -1.5615e+00,  ...,  3.6822e-01,
                 5.2679e-01, -1.5067e+00],
               [ 2.0892e-01,  1.4041e+00, -2.0698e-01,  ...,  1.2122e+00,
                 1.3742e-01,  6.8231e-01],
               [ 1.4700e+00, -2.0746e+00,  8.6280e-01,  ...,  7.7288e-01,
                -1.0128e+00,  1.9533e-01],
               ...,
               [-4.0518e-01,  1.0060e+00, -1.2097e+00,  ...,  6.6154e-01,
                 2.2548e+00,  8.4912e-01],
               [-9.0720e-01,  6.4946e-01, -7.6629e-01,  ...,  8.3773e-02,
                -2.2143e-01, -5.6469e-01],
               [-1.5342e+00,  1.2075e+00, -4.9365e-01,  ...,  7.7658e-01,
                 8.2622e-01, -5.1317e-01]],

              [[-1.0310e+00, -6.3179e-01, -8.8332e-01,  ...,  8.8619e-02,
                 6.0848e-01,  1.5056e+00],
               [ 1.5093e+00, -1.1469e+00, -4.5369e-01,  ..., -9.8717e-01,
                -6.3974e-01, -1.0479e+00],
               [ 1.3946e+00,  1.5321e-01,  4.9116e-01,  ...,  1.0110e+00,
                -2.0338e+00,  6.2318e-01],
               ...,
               [ 2.7276e-02, -1.0978e+00, -1.2747e-03,  ..., -6.6941e-02,
                 3.4110e-01, -3.0183e-01],
               [ 1.0870e+00,  3.6189e-01,  3.5213e-01,  ..., -7.1679e-01,
                -1.2220e-01, -4.6191e-01],
               [ 1.5815e-01,  5.9903e-01,  4.3838e-01,  ...,  1.8811e+00,
                 6.4441e-01,  5.0634e-01]],

              [[ 1.4628e-02, -6.5199e-01, -1.3255e+00,  ...,  1.1559e+00,
                 1.8737e+00, -1.2272e-01],
               [ 5.7514e-01,  4.6596e-01, -3.4672e-01,  ...,  3.9621e-01,
                -1.3981e+00,  1.2205e+00],
               [-6.2185e-01,  8.9906e-01,  6.5405e-01,  ..., -8.0875e-01,
                -3.7445e-02,  9.1749e-02],
               ...,
               [-5.9502e-01, -6.5149e-01,  1.2600e+00,  ...,  7.0306e-01,
                 4.3226e-01, -3.6069e-03],
               [-2.5798e-01, -3.2951e-01, -6.0422e-01,  ...,  7.6340e-02,
                 1.9422e-01,  5.2702e-02],
               [ 2.5915e-01,  4.1826e-01, -8.1916e-01,  ...,  1.3079e+00,
                 5.0282e-01, -8.8068e-01]],

              [[ 8.9606e-01,  2.7360e+00, -7.7418e-01,  ..., -6.3446e-01,
                 4.9931e-01, -1.4507e+00],
               [ 1.6607e+00,  3.6149e-01,  7.9179e-01,  ...,  5.6409e-02,
                -1.4083e+00,  3.5620e-02],
               [ 2.6172e+00,  4.9535e-02, -1.5108e-01,  ...,  7.3331e-02,
                 1.2261e+00, -1.4635e-01],
               ...,
               [-3.5466e-01, -6.0469e-01, -1.8319e-02,  ...,  9.4671e-01,
                 7.6173e-01,  1.3294e+00],
               [-6.0030e-01,  2.4628e+00, -2.9686e-02,  ...,  2.3356e-01,
                -6.3912e-01, -9.3690e-01],
               [ 5.5543e-01, -3.8047e-01, -3.4127e-01,  ...,  1.9001e+00,
                 6.9048e-01,  1.9095e+00]]]),
      tensor([[[ 1.9261,  0.3662, -2.2351,  ..., -0.2479,  2.0208, -0.3038],
               [ 0.2193,  0.2397, -0.6246,  ..., -0.1977, -0.7800,  0.0299],
               [ 1.3310, -0.8811,  0.9491,  ...,  0.1043, -0.0846,  0.5337],
               ...,
               [-0.3900, -1.2184, -2.0538,  ...,  0.3454,  0.7886, -0.4477],
               [ 1.7995,  1.5179, -1.2692,  ...,  0.6605,  0.3598, -1.2437],
               [ 1.9007,  0.0435, -2.2540,  ...,  0.2233,  1.4315, -0.8018]],

              [[-0.2991,  0.3210, -2.0743,  ..., -0.1376,  0.6999, -1.8585],
               [ 0.2477,  1.0192, -0.5260,  ...,  1.0971, -0.3136,  0.0693],
               [ 1.9090, -1.7206,  0.3304,  ...,  0.8011, -1.1116,  0.2133],
               ...,
               [ 0.3987,  0.1333, -0.8556,  ...,  0.7316,  2.1006,  0.4276],
               [ 0.1955,  0.6050, -1.0433,  ...,  0.5461,  0.0905, -1.1321],
               [-1.1648,  0.8397, -0.7285,  ...,  0.9650,  1.0187, -0.9154]],

              [[-0.2072, -0.9251, -0.4631,  ...,  0.0555,  0.9139,  0.6636],
               [ 1.8273, -0.8446, -0.6740,  ..., -0.9765, -0.6065, -1.1694],
               [ 1.1242,  0.0836,  0.1622,  ...,  1.1291, -2.1867,  0.2463],
               ...,
               [ 0.2276, -1.3735, -0.2455,  ..., -0.4645, -0.0795, -0.4198],
               [ 1.3959,  0.1959, -0.1209,  ..., -0.4528, -0.3387, -0.9129],
               [ 1.1302,  0.4864,  0.4106,  ...,  1.6754,  0.6609,  0.2517]],

              [[ 0.7390, -1.2320, -1.4019,  ...,  0.7018,  1.4972, -0.7071],
               [ 0.4820,  0.7413, -0.4056,  ...,  0.6629, -1.9494,  1.4256],
               [-0.0761, -0.0412,  0.4128,  ..., -0.6761, -0.7692, -0.3830],
               ...,
               [-0.1002, -1.1125,  0.8038,  ...,  0.7448, -0.3724, -0.9875],
               [ 0.6388, -1.2257, -0.8242,  ...,  0.2262, -0.1097, -0.1556],
               [ 1.1622,  0.5248, -0.8697,  ...,  1.2056,  0.7714, -1.3186]],

              [[ 1.1718,  2.1725, -1.0235,  ..., -0.7244,  0.2784, -2.0104],
               [ 1.6596,  0.0280,  0.4212,  ...,  0.4771, -1.8212, -0.2746],
               [ 2.8073, -0.2165, -0.2219,  ...,  0.3132,  0.9727, -0.2467],
               ...,
               [ 0.3232, -0.8092, -0.2024,  ...,  0.5945, -0.4198,  0.5399],
               [-0.2653,  2.1102, -0.0695,  ...,  0.6733, -0.5012, -1.4819],
               [ 0.7600, -0.0512, -0.2453,  ...,  1.5513,  0.3439,  1.1348]]])]}

```python

```

# Training the Transformer XL

Now, let's move on to actually training the Transformer XL.

```python
TESTING = True
```

We'll be using the following configurations.

```python
class Config(dict):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        for k, v in kwargs.items():
            setattr(self, k, v)

    def set(self, key, val):
        self[key] = val
        setattr(self, key, val)

    def update(self, dct):
        for k, v in dct.items():
            self.set(k, v)

# We will use prime numbers to ensure our implementation
# is actually correct
config = Config(
    seed=101,
    debug=False,
    warmup_step=0,
    # Check default params
    min_lr=0.,
    dropouta=0.,
    clip=0.25,
    log_interval=200,
    eval_interval=100,
)

if TESTING:
    config.update(dict(
        debug=True,
        lr=0.00025,
        bs=8,
        epochs=2,
        max_step=10000, # shorten for testing
        n_layers=4,
        n_heads=3,
        d_model=32,
        d_head_inner=17,
        d_ff_inner=71,
        dropout=0.1,
        train_bptt=33,
        eval_bptt=41,
        mem_len=41,
        eval_mem_len=63,
    ))
else:
    config.update(dict(
        lr=0.0025,
        bs=22,
        epochs=2,
        max_step=400000,
        n_layers=12,
        n_heads=8,
        d_model=512,
        d_head_inner=64,
        d_ff_inner=2048,
        dropout=0.1,
        train_bptt=512,
        eval_bptt=128,
        mem_len=512,
        eval_mem_len=2100,
    ))
```

```python

```

### Preparing the Data Loader

Data loading for the Transformer XL is similar to data loading for an RNN-based language model but is quite different from standard data loading, so we'll go over it in detail.

Suppose we chunked the input into sequences of 4 words to feed into the model. Remember that the Transformer XL is stateful, meaning the computations of each minibatch are carried over to the next minibatch. For a minibatch size of 1, handling this is simple. We just chunk the input and feed it into the model like this:

Now, what happens if the `batch_size` is 2? We can't split the sentence like this, otherwise, we would be breaking the dependencies between segments.

The correct way to handle the corpus with a `batch_size`of 2 is to feed it like this.

Generalizing this, we first divide the corpus into `batch_size` length segments, then feed each segment piece by piece into the model. Let's go through an example. Suppose `batch_size` is 4 and our entire corpus looks like this:

    pytorch is an amazing deep learning framework that makes nlp really easy

We want to make sure that the previous batch contains the previous segment at the same position.

In other words, assuming we fed the model one word at a time, we want to iterate over this sentence like this

    Batch 1: pytorch   amazing   framework nlp
    Batch 2: is        deep      that      really
    Batch 3: an        learning  makes     easy

Notice that you can reconstruct the original sentence by reading from top to bottom, left to right
instead of left to right, top to bottom.

In reality, we feed the model with a sequence of words for each batch. The length of this sequence is commonly referred to the bptt (back propagation through time) length, since this is the maximum length the gradients propagate through in the sequence direction. With a longer bptt length of 2 for example, the
minibatch would be of shape (batch_size, bptt) and would look like

    Batch 1: pytorch   amazing   framework nlp
             is        deep      that      really
    Batch 2: an        learning  makes     easy

We can implement this in a dataloader like this:

```python
from torch.utils import data
import math

class LMDataLoader(data.DataLoader):
    def __init__(self, data: torch.LongTensor, batch_size: int, bptt: int,
                 device=torch.device("cpu")):
        self.batch_size = batch_size
        self.bptt = bptt
        self.n_steps = data.size(0) // batch_size

        # we reshape the data here so that we can index
        # efficiently into it while training
        self.data = (data[:self.n_steps * batch_size] # trim off any elements that don't fit cleanly
                     .view(batch_size, self.n_steps) #
                     .transpose(0, 1) #
                     .contiguous().to(device) # put on device as contiguous tensor
                     )

    def __iter__(self):
        for batch_start_idx in range(0, self.data.size(0) - 1, self.bptt):
            batch_end_idx = min(batch_start_idx + self.bptt, self.data.size(0) - 1)
            # TODO: What is `self.ext_len` in the original code?
            batch_data = self.data[batch_start_idx:batch_end_idx]
            target = self.data[batch_start_idx+1:batch_end_idx+1]
            # we generate the sequence length as well for loss calculation later
            yield batch_data, target, batch_end_idx - batch_start_idx

    def __len__(self):
        return math.ceil(self.data.size(0) / self.bptt)
```

Let's test this out

```python
test_corpus = torch.arange(1000)
```

```python
BS = 16
BPTT = 10
```

```python
test_corpus[:BPTT]
```

    tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])

```python
loader = LMDataLoader(test_corpus, BS, BPTT)
```

```python
b1, *_ = next(iter(loader))
```

```python
b1.shape
```

    torch.Size([10, 16])

```python
b1
```

    tensor([[  0,  62, 124, 186, 248, 310, 372, 434, 496, 558, 620, 682, 744, 806,
             868, 930],
            [  1,  63, 125, 187, 249, 311, 373, 435, 497, 559, 621, 683, 745, 807,
             869, 931],
            [  2,  64, 126, 188, 250, 312, 374, 436, 498, 560, 622, 684, 746, 808,
             870, 932],
            [  3,  65, 127, 189, 251, 313, 375, 437, 499, 561, 623, 685, 747, 809,
             871, 933],
            [  4,  66, 128, 190, 252, 314, 376, 438, 500, 562, 624, 686, 748, 810,
             872, 934],
            [  5,  67, 129, 191, 253, 315, 377, 439, 501, 563, 625, 687, 749, 811,
             873, 935],
            [  6,  68, 130, 192, 254, 316, 378, 440, 502, 564, 626, 688, 750, 812,
             874, 936],
            [  7,  69, 131, 193, 255, 317, 379, 441, 503, 565, 627, 689, 751, 813,
             875, 937],
            [  8,  70, 132, 194, 256, 318, 380, 442, 504, 566, 628, 690, 752, 814,
             876, 938],
            [  9,  71, 133, 195, 257, 319, 381, 443, 505, 567, 629, 691, 753, 815,
             877, 939]])

```python
b1, b2, sl = next(iter(loader))
```

```python

```

### Loading the actual data

We'll be using the penn treebank dataset to benchmark our model.

```python
from pathlib import Path
DATASET = "penn"
DATA_DIR = Path("../data") / DATASET
```

We'll be using a utility vocabulary class borrowed directly from the Transformer XL repo to numericalize our inputs.

```python
import sys; sys.path.append("../utils")
from vocabulary import Vocab
```

```python
vocab = Vocab(special=["<eos>"], lower_case=True)

vocab.count_file(DATA_DIR / "train.txt")
vocab.count_file(DATA_DIR / "valid.txt")
vocab.count_file(DATA_DIR / "test.txt")
None
```

```python
vocab.build_vocab()
```

    building vocab with min_freq=0, max_size=None
    final vocab size 10000 from 9999 unique tokens

```python
train_dataset = vocab.encode_file(DATA_DIR / "train.txt", ordered=True, add_eos=True)
valid_dataset = vocab.encode_file(DATA_DIR / "valid.txt", ordered=True, add_eos=True)
test_dataset = vocab.encode_file(DATA_DIR / "test.txt", ordered=True, add_eos=True)
```

```python
train_dataset[:50]
```

    tensor([6503, 6151, 7924, 8539, 2353, 8540, 6918, 8541, 8542, 7394, 7925, 7926,
            6152, 8543, 6504, 6919, 6920, 8544, 5560, 6153, 8545, 8546, 8547, 7927,
               0, 9231,    2,    3,   73,  399,   34, 2136,    1,  146,   19,    6,
            9232,  282,  450,    3,    0,   23,    2,   13,  142,    4,    2, 5090,
               1, 2952])

```python

```

Now we can prepare the data loaders

```python
train_iter = LMDataLoader(train_dataset, config.bs, config.train_bptt, device=device)
valid_iter = LMDataLoader(valid_dataset, config.bs, config.eval_bptt, device=device)
test_iter = LMDataLoader(test_dataset, config.bs, config.eval_bptt, device=device)
```

```python
next(iter(train_iter))
```

    (tensor([[6503,    6,    1,    0,    5, 5485,  542,    2],
             [6151,    2,  327,  909,    6,    7,   16,    9],
             [7924,  225, 6798,    2, 2069,  410,  880,  157],
             [8539,   72,  329,  380,    1,   15,    3,  299],
             [2353,    0, 1538, 1635,  104,  495,    3,  232],
             [8540,   29, 9744,   32,   90, 5086,    0,    6],
             [6918,   84,   60,    2, 2392,    0,   14,    2],
             [8541,   27, 1437,  506,    1,  112,   24,    2],
             [8542, 2298,  105,  557,  320,  553,  195, 3524],
             [7394,    0, 2262,   16, 1715,  497,  150,    0],
             [7925,  495,    2,    1,    4,   68, 2440,   64],
             [7926,    2,  885, 2978, 1046,  112,   11,   47],
             [6152,   45,    2,    0,    8,  918,    1, 3464],
             [8543, 1973, 5616,   39,  723,   25,  490,   14],
             [6504,    2,    4,   13,    1, 3553,   75,  568],
             [6919,    6,  630,  784,  284,    5, 2369,    1],
             [6920, 2664,    0, 6171,  819,   25,    1,  569],
             [8544,   10,   29,   20,    0,    1,    2,    4],
             [5560,    2,   39, 2841, 8138,    2,    8,    1],
             [6153,   21,    2,  310, 2066,   36,    2,  372],
             [8545, 4648,    1,    3,  573,    1,    4,  160],
             [8546,    2,   53,   18,   33, 5177, 2947,   15],
             [8547,    0,  779,  351, 2753,   54, 1118,    2],
             [7927,  215, 1099,  682,   18,    0,    9,    2],
             [   0, 1193, 7123,    4,   12,   64,    2,    2],
             [9231, 1135,    1,  156,    3,   47,    8, 1683],
             [   2,    7, 5232,  996,   49, 5177,    4,    4],
             [   3, 1573,    4, 3098,  151,    0,   51,  157],
             [  73,    2,    2,    2,    3,  787,    2,  299],
             [ 399, 4448,    0,    5,  122,  674, 4057,  232],
             [  34,    0,   23,   32, 1226,    2,   21,    0],
             [2136,    2, 1575,    2,  361,   72,    6,   14],
             [   1,  586,  329, 4467,  573,    4, 2561,    9]]),
     tensor([[6151,    2,  327,  909,    6,    7,   16,    9],
             [7924,  225, 6798,    2, 2069,  410,  880,  157],
             [8539,   72,  329,  380,    1,   15,    3,  299],
             [2353,    0, 1538, 1635,  104,  495,    3,  232],
             [8540,   29, 9744,   32,   90, 5086,    0,    6],
             [6918,   84,   60,    2, 2392,    0,   14,    2],
             [8541,   27, 1437,  506,    1,  112,   24,    2],
             [8542, 2298,  105,  557,  320,  553,  195, 3524],
             [7394,    0, 2262,   16, 1715,  497,  150,    0],
             [7925,  495,    2,    1,    4,   68, 2440,   64],
             [7926,    2,  885, 2978, 1046,  112,   11,   47],
             [6152,   45,    2,    0,    8,  918,    1, 3464],
             [8543, 1973, 5616,   39,  723,   25,  490,   14],
             [6504,    2,    4,   13,    1, 3553,   75,  568],
             [6919,    6,  630,  784,  284,    5, 2369,    1],
             [6920, 2664,    0, 6171,  819,   25,    1,  569],
             [8544,   10,   29,   20,    0,    1,    2,    4],
             [5560,    2,   39, 2841, 8138,    2,    8,    1],
             [6153,   21,    2,  310, 2066,   36,    2,  372],
             [8545, 4648,    1,    3,  573,    1,    4,  160],
             [8546,    2,   53,   18,   33, 5177, 2947,   15],
             [8547,    0,  779,  351, 2753,   54, 1118,    2],
             [7927,  215, 1099,  682,   18,    0,    9,    2],
             [   0, 1193, 7123,    4,   12,   64,    2,    2],
             [9231, 1135,    1,  156,    3,   47,    8, 1683],
             [   2,    7, 5232,  996,   49, 5177,    4,    4],
             [   3, 1573,    4, 3098,  151,    0,   51,  157],
             [  73,    2,    2,    2,    3,  787,    2,  299],
             [ 399, 4448,    0,    5,  122,  674, 4057,  232],
             [  34,    0,   23,   32, 1226,    2,   21,    0],
             [2136,    2, 1575,    2,  361,   72,    6,   14],
             [   1,  586,  329, 4467,  573,    4, 2561,    9],
             [ 146,  158,   63,   17,   33,    1,    4,    1]]),
     33)

```python

```

### Initialization

We borrow the following initialization from the Transformer XL repo

```python
def init_weight(weight):
    nn.init.normal_(weight, 0.0, 0.02)

def init_bias(bias):
    nn.init.constant_(bias, 0.0)

# Borrowed from the transformer XL repo
def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        if hasattr(m, 'weight') and m.weight is not None:
            init_weight(m.weight)
        if hasattr(m, 'bias') and m.bias is not None:
            init_bias(m.bias)
    elif classname.find('Embedding') != -1:
        if hasattr(m, 'weight'):
            init_weight(m.weight)
    elif classname.find('LayerNorm') != -1:
        if hasattr(m, 'weight'):
            nn.init.normal_(m.weight, 1.0, 0.02)
        if hasattr(m, 'bias') and m.bias is not None:
            init_bias(m.bias)
    else:
        if hasattr(m, 'u'):
            init_weight(m.u)
        if hasattr(m, 'v'):
            init_weight(m.v)
```

No fancy initialization here. Since we have multiple Layer Normalization layers, we can get away with initializing everything using a simple normal distribution.

```python

```

### Training Loop

The training loop is fairly standard. You can use any framework you like here including [ignite](https://github.com/pytorch/ignite), [allennlp](https://github.com/allenai/allennlp), and [fastai](https://github.com/fastai/fastai). We'll be writing our own loop to simplify things.

```python
import torch.optim as optim
```

```python
import math
import time
import os
from tqdm import tqdm

loss_change = []
val_loss_change = []

def train_epoch(
    epoch: int,
    model: nn.Module, train_loader: data.DataLoader,
    val_loader: data.DataLoader,
    optimizer: optim.Optimizer,
    scheduler,
    train_step_start=0.,
 ):
    # Turn on training mode which enables dropout.
    model.train()
    mems = None
    train_step = train_step_start
    train_loss = 0
    log_start_time = time.time()
    best_val_loss = float("inf")

    pbar = tqdm(train_loader, total=min(config.max_step - train_step_start, len(train_loader)))
    for batch_idx, (data, target, seq_len) in enumerate(pbar):
        model.zero_grad()
        out_dict = model(data, target, memory=mems)
        loss, mems = out_dict["loss"], out_dict["memory"]

        loss.backward()
        train_loss += loss.item()
        loss_change.append(loss.item())
        torch.nn.utils.clip_grad_norm_(model.parameters(), config.clip)
        optimizer.step()

        # step-wise learning rate annealing
        train_step += 1
        # linear warmup stage
        if train_step < config.warmup_step:
            curr_lr = config.lr * train_step / config.warmup_step
            optimizer.param_groups[0]['lr'] = curr_lr
        else:
            scheduler.step(train_step)

        if train_step % config.log_interval == 0:
            cur_loss = train_loss / config.log_interval
            elapsed = time.time() - log_start_time
            log_str = '| epoch {:3d} step {:>8d} | lr {:.3g} ' \
                      '| loss {:5.2f}'.format(
                epoch, train_step, optimizer.param_groups[0]['lr'], cur_loss)
            log_str += ' | ppl {:9.3f}'.format(math.exp(cur_loss))
            pbar.set_description(log_str)
            train_loss = 0
            log_start_time = time.time()

        if train_step % config.eval_interval == 0:
            val_loss = evaluate(model, val_loader)
            val_loss_change.append(val_loss)
            eval_start_time = time.time()

        if train_step == config.max_step:
            return train_step
    return train_step
```

```python
def train(model, train_loader, valid_loader):
    optimizer = optim.Adam(model.parameters(), lr=config.lr)
    total_steps = min(config.max_step, len(train_loader) * config.epochs)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer,
                    total_steps, eta_min=config.min_lr)
    train_step_start = 0
    for epoch in range(config.epochs):
        if train_step_start >= config.max_step:
            break
        train_step_start = train_epoch(
            epoch,
            model,
            train_iter,
            valid_iter,
            optimizer,
            scheduler,
            train_step_start,
        )
```

Language models are normally evaluated by perplexity. Perplexity is the exponential of the cross entropy loss. It is also equivalent to the reciprocal of the likelihood. If the language model assigns a probability of 0.1 to each word in each input sentence on average, it would receive a perplexity of 100.

Intuitively, perplexity represents how many tries it would take for the model to guess the correct word. A perplexity of 100 would signify that the model would need 100 tries to guess each word in the input sequence correctly.

Keeping this in mind, we can write the evaluation code like this.

```python
def evaluate(model: nn.Module, val_loader: data.DataLoader):
    # Turn on evaluation mode which disables dropout.
    model.eval()
    model.reset_length(config.eval_bptt,
        0, config.eval_mem_len+config.train_bptt-config.eval_bptt)

    # Evaluation
    total_len, total_loss = 0, 0.
    with torch.no_grad():
        mems = None
        for i, (data, target, seq_len) in enumerate(val_loader):
            out_dict = model(data, target, memory=mems)
            loss, mems = out_dict["loss"], out_dict["memory"]
            total_loss += seq_len * loss.float().item()
            total_len += seq_len

    # Switch back to the training mode
    model.reset_length(config.train_bptt, 0, config.mem_len)
    model.train()
    return total_loss / total_len
```

```python
def evaluate_final(model, val_loader):
    model.eval()
    total_len, total_loss = 0, 0.
    start_time = time.time()

    model.reset_length(config.eval_bptt, 0, config.eval_mem_len + config.train_bptt - config.eval_bptt)

    with torch.no_grad():
        mems = None
        for i, (data, target, seq_len) in enumerate(val_loader):
            out_dict = model(data, target, memory=mems)
            loss, mems = out_dict["loss"], out_dict["memory"]
            total_loss += seq_len * loss.item()
            total_len += seq_len
        total_time = time.time() - start_time

    model.reset_length(config.train_bptt, 0, config.mem_len)
    loss_val = total_loss / total_len
    return {"loss": loss_val, "ppl": math.exp(loss_val)}
```

Now all we have to do is initialize the model and start training it!

```python
transformer_xl = TransformerXL(
    num_embeddings=len(vocab), n_layers=config.n_layers,
    n_heads=config.n_heads, d_model=config.d_model,
    d_head_inner=config.d_head_inner,
    d_ff_inner=config.d_ff_inner,
    dropout=config.dropout,
    dropouta=config.dropouta,
    seq_len=config.train_bptt,
    mem_len=config.mem_len,
)
if torch.cuda.is_available(): transformer_xl.cuda()
transformer_xl.apply(weights_init);
```

```python
train(
    transformer_xl,
    train_iter,
    valid_iter,
)
```

    | epoch   0 step     3400 | lr 0.000132 | loss  6.22 | ppl   502.650: 100%|██████████| 3522/3522 [14:41<00:00,  4.00it/s]
    | epoch   1 step     7000 | lr 2.41e-08 | loss  6.14 | ppl   462.692: 100%|██████████| 3522/3522 [14:58<00:00,  3.92it/s]

```python
evaluate_final(transformer_xl, valid_iter)
```

    {'loss': 6.067088326684745, 'ppl': 431.42268910022347}

```python

```

Let's observe the loss change.

```python
import matplotlib.pyplot as plt
%matplotlib inline
```

```python
plt.plot(loss_change)
```

```python
plt.plot(val_loss_change)
```

Overall the loss is decreasing and everything looks nice!

```python

```


# Intro

In this guided project, we're going to study the practical side of the algorithm by building a spam filter for SMS messages.

To classify messages as spam or non-spam the computer:

 - Learns how humans classify messages.
 - Uses that human knowledge to estimate probabilities for new messages — probabilities for spam and non-spam.
 - Classifies a new message based on these probability values — if the probability for spam is greater, then it classifies the message as spam. Otherwise, it classifies it as non-spam (if the two probability values are equal, then we may need a human to classify the message).

## The Data

Our first task is to "teach" the computer how to classify messages. To do that, we'll use the multinomial Naive Bayes algorithm along with a dataset of 5,572 SMS messages that are already classified by humans.

The dataset was put together by Tiago A. Almeida and José María Gómez Hidalgo, and it can be downloaded from the [The UCI Machine Learning Repository](https://archive.ics.uci.edu/ml/datasets/sms+spam+collection). You can also download the dataset directly from this [link](https://dq-content.s3.amazonaws.com/433/SMSSpamCollection). The data collection process is described in more details on [this](http://www.dt.fee.unicamp.br/~tiago/smsspamcollection/#composition) page, where you can also find some of the authors' papers.

Steps:

 - Open the `SMSSpamCollection` file using the `read_csv()` function from the pandas package.
 - The data points are tab separated, so we'll need to use the `sep='\t'` parameter for our `read_csv()` function.
 - The dataset doesn't have a header row, which means we need to use the `header=None` parameter, otherwise the first row will be wrongly used as the header row.
 - Use the `names=['Label', 'SMS']` parameter to name the columns as `Label` and `SMS`.

```python
import pandas as pd

sms_spam = pd.read_csv('SMSSpamCollection', sep='\t', header=None, names=['Label', 'SMS'])

print(sms_spam.shape)
sms_spam.head()
```

    (5572, 2)

      Label
      SMS

      0
      ham
      Go until jurong point, crazy.. Available only ...

      1
      ham
      Ok lar... Joking wif u oni...

      2
      spam
      Free entry in 2 a wkly comp to win FA Cup fina...

      3
      ham
      U dun say so early hor... U c already then say...

      4
      ham
      Nah I don't think he goes to usf, he lives aro...

Let's find out what the percentage of `spam` versus `ham` is:

```python
sms_spam['Label'].value_counts(normalize=True)
```

    ham     0.865937
    spam    0.134063
    Name: Label, dtype: float64

### Training and Test Set
Under 15% is spam, which is a fair representation.

We're now going to split our dataset into a training and a test set, where the training set accounts for 80% of the data, and the test set for the remaining 20%.

The dataset has 5,572 messages, which means that:

 - The training set will have 4,458 messages (about 80% of the dataset).
 - The test set will have 1,114 messages (about 20% of the dataset).
 - To better understand the purpose of putting a test set aside, let's begin by observing that

```python
# Randomize the dataset
data_randomized = sms_spam.sample(frac=1, random_state=1)

# Calculate index for split
training_test_index = round(len(data_randomized) * 0.8)

# Training/Test split
training_set = data_randomized[:training_test_index].reset_index(drop=True)
test_set = data_randomized[training_test_index:].reset_index(drop=True)

print(training_set.shape)
print(test_set.shape)
```

    (4458, 2)
    (1114, 2)

We'll now compare the percentage of spam and ham messages in the training and test sets. We expect the percentages to be close to what we have in the full dataset, where about 87% of the messages are ham, and the remaining 13% are spam.

```python
training_set['Label'].value_counts(normalize=True)
```

    ham     0.86541
    spam    0.13459
    Name: Label, dtype: float64

```python
test_set['Label'].value_counts(normalize=True)
```

    ham     0.868043
    spam    0.131957
    Name: Label, dtype: float64

As we hoped the percentages are comparable between the sets.

Next step: Clean the data

We want to treat the data as follows:

### Build vocab

We'll now make a list containing all the unique words across all messages, where each word is represented as a string.

Steps:

 - Begin by transforming each message from the SMS column into a list by splitting the string at the space character with `Series.str.split()` method.

 - Initiate an empty list named vocabulary.

 - Iterate over the the SMS column (each message in this column should be a list of strings by the time you start this loop).

 - Using a nested loop, iterate each message in the SMS column (each message should be a list of strings) and append each string (word) to the vocabulary list.

 - Transform the vocabulary list into a set using the set() function. This will remove the duplicates from the vocabulary list.

 - Transform the vocabulary set back into a list using the list() function.

```python
training_set['SMS'] = training_set['SMS'].str.split()

vocabulary = []
for sms in training_set['SMS']:
    for word in sms:
        vocabulary.append(word)

vocabulary = list(set(vocabulary))
len(vocabulary)
```

    7783

There are 7783 unique words within all of the messages in our training set.

Now we're going to use the vocabulary to make the data transformation we need.

```python
word_counts_per_sms = {unique_word: [0] * len(training_set['SMS']) for unique_word in vocabulary}

for index, sms in enumerate(training_set['SMS']):
    for word in sms:
        word_counts_per_sms[word][index] += 1

word_counts = pd.DataFrame(word_counts_per_sms)
word_counts.head()
```

      n
      mushy
      transfer
      watches
      singles
      find
      champlaxigating
      slave
      mudyadhu
      twat
      ...
      volcanoes
      47per
      lightly
      wating
      temp
      09050001295
      admirer
      aproach
      starshine
      stagwood

      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      ...
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0

      1
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      ...
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0

      2
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      ...
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0

      3
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      ...
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0

      4
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0
      ...
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0

5 rows × 7783 columns

```python
training_set_clean = pd.concat([training_set, word_counts], axis=1)
training_set_clean.head()
```

      Label
      SMS
      n
      mushy
      transfer
      watches
      singles
      find
      champlaxigating
      slave
      ...
      volcanoes
      47per
      lightly
      wating
      temp
      09050001295
      admirer
      aproach
      starshine
      stagwood

      0
      ham
      [yep, by, the, pretty, sculpture]
      0
      0
      0
      0
      0
      0
      0
      0
      ...
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0

      1
      ham
      [yes, princess, are, you, going, to, make, me,...
      0
      0
      0
      0
      0
      0
      0
      0
      ...
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0

      2
      ham
      [welp, apparently, he, retired]
      0
      0
      0
      0
      0
      0
      0
      0
      ...
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0

      3
      ham
      [havent]
      0
      0
      0
      0
      0
      0
      0
      0
      ...
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0

      4
      ham
      [i, forgot, 2, ask, ü, all, smth, there, s, a,...
      0
      0
      0
      0
      0
      0
      0
      0
      ...
      0
      0
      0
      0
      0
      0
      0
      0
      0
      0

5 rows × 7785 columns

## Calculating Constants First

We're now done with cleaning the training set, and we can begin creating the spam filter. The Naive Bayes algorithm will need to answer these two probability questions to be able to classify new messages:

\begin{equation}
P(Spam | w_1,w_2, ..., w_n) \propto P(Spam) \cdot \prod_{i=1}^{n}P(w_i|Spam)
\end{equation}

\begin{equation}
P(Ham | w_1,w_2, ..., w_n) \propto P(Ham) \cdot \prod_{i=1}^{n}P(w_i|Ham)
\end{equation}

Also, to calculate P(wi|Spam) and P(wi|Ham) inside the formulas above, we'll need to use these equations:

\begin{equation}
P(w_i|Spam) = \frac{N_{w_i|Spam} + \alpha}{N_{Spam} + \alpha \cdot N_{Vocabulary}}
\end{equation}

\begin{equation}
P(w_i|Ham) = \frac{N_{w_i|Ham} + \alpha}{N_{Ham} + \alpha \cdot N_{Vocabulary}}
\end{equation}

Some of the terms in the four equations above will have the same value for every new message. We can calculate the value of these terms once and avoid doing the computations again when a new messages comes in. Below, we'll use our training set to calculate:

- P(Spam) and P(Ham)
- NSpam, NHam, NVocabulary

We'll also use Laplace smoothing and set $\alpha = 1$.

```python
# Isolating spam and ham messages first
spam_messages = training_set_clean[training_set_clean['Label'] == 'spam']
ham_messages = training_set_clean[training_set_clean['Label'] == 'ham']

# P(Spam) and P(Ham)
p_spam = len(spam_messages) / len(training_set_clean)
p_ham = len(ham_messages) / len(training_set_clean)

# N_Spam
n_words_per_spam_message = spam_messages['SMS'].apply(len)
n_spam = n_words_per_spam_message.sum()

# N_Ham
n_words_per_ham_message = ham_messages['SMS'].apply(len)
n_ham = n_words_per_ham_message.sum()

# N_Vocabulary
n_vocabulary = len(vocabulary)

# Laplace smoothing
alpha = 1
```

## Calculating Parameters

Now that we have the constant terms calculated above, we can move on with calculating the parameters $P(w_i|Spam)$ and $P(w_i|Ham)$. Each parameter will thus be a conditional probability value associated with each word in the vocabulary.

The parameters are calculated using the formulas:

\begin{equation}
P(w_i|Spam) = \frac{N_{w_i|Spam} + \alpha}{N_{Spam} + \alpha \cdot N_{Vocabulary}}
\end{equation}

\begin{equation}
P(w_i|Ham) = \frac{N_{w_i|Ham} + \alpha}{N_{Ham} + \alpha \cdot N_{Vocabulary}}
\end{equation}

```python
# Initiate parameters
parameters_spam = {unique_word:0 for unique_word in vocabulary}
parameters_ham = {unique_word:0 for unique_word in vocabulary}

# Calculate parameters
for word in vocabulary:
    n_word_given_spam = spam_messages[word].sum()   # spam_messages already defined in a cell above
    p_word_given_spam = (n_word_given_spam + alpha) / (n_spam + alpha*n_vocabulary)
    parameters_spam[word] = p_word_given_spam

    n_word_given_ham = ham_messages[word].sum()   # ham_messages already defined in a cell above
    p_word_given_ham = (n_word_given_ham + alpha) / (n_ham + alpha*n_vocabulary)
    parameters_ham[word] = p_word_given_ham
```

## Classifying A New Message

Now that we have all our parameters calculated, we can start creating the spam filter. The spam filter can be understood as a function that:

- Takes in as input a new message (w1, w2, ..., wn).
- Calculates P(Spam|w1, w2, ..., wn) and P(Ham|w1, w2, ..., wn).
- Compares the values of P(Spam|w1, w2, ..., wn) and P(Ham|w1, w2, ..., wn), and:
    - If P(Ham|w1, w2, ..., wn) > P(Spam|w1, w2, ..., wn), then the message is classified as ham.
    - If P(Ham|w1, w2, ..., wn) < P(Spam|w1, w2, ..., wn), then the message is classified as spam.
    -  If P(Ham|w1, w2, ..., wn) = P(Spam|w1, w2, ..., wn), then the algorithm may request human help.

```python
import re

def classify(message):
    '''
    message: a string
    '''

    message = re.sub('\W', ' ', message)
    message = message.lower().split()

    p_spam_given_message = p_spam
    p_ham_given_message = p_ham

    for word in message:
        if word in parameters_spam:
            p_spam_given_message *= parameters_spam[word]

        if word in parameters_ham:
            p_ham_given_message *= parameters_ham[word]

    print('P(Spam|message):', p_spam_given_message)
    print('P(Ham|message):', p_ham_given_message)

    if p_ham_given_message > p_spam_given_message:
        print('Label: Ham')
    elif p_ham_given_message < p_spam_given_message:
        print('Label: Spam')
    else:
        print('Equal probabilities, needs manual classification')
```

```python
classify('WINNER!! This is the secret code to unlock the money: C3421.')
```

    P(Spam|message): 1.3481290211300841e-25
    P(Ham|message): 1.9368049028589875e-27
    Label: Spam

```python
classify("Sounds good, Tom, then see u there")
```

    P(Spam|message): 2.4372375665888117e-25
    P(Ham|message): 3.687530435009238e-21
    Label: Ham

## Measuring the Spam Filter's Accuracy

The two results above look promising, but let's see how well the filter does on our test set, which has 1,114 messages.

We'll start by writing a function that returns classification labels instead of printing them.

```python
def classify_test_set(message):
    '''
    message: a string
    '''

    message = re.sub('\W', ' ', message)
    message = message.lower().split()

    p_spam_given_message = p_spam
    p_ham_given_message = p_ham

    for word in message:
        if word in parameters_spam:
            p_spam_given_message *= parameters_spam[word]

        if word in parameters_ham:
            p_ham_given_message *= parameters_ham[word]

    if p_ham_given_message > p_spam_given_message:
        return 'ham'
    elif p_spam_given_message > p_ham_given_message:
        return 'spam'
    else:
        return 'needs manual classification'
```

Now that we have a function that returns labels instead of printing them, we can use it to create a new column in our test set.

```python
test_set['predicted'] = test_set['SMS'].apply(classify_test_set)
test_set.head()
```

      Label
      SMS
      predicted

      0
      ham
      Later i guess. I needa do mcat study too.
      ham

      1
      ham
      But i haf enuff space got like 4 mb...
      ham

      2
      spam
      Had your mobile 10 mths? Update to latest Oran...
      spam

      3
      ham
      All sounds good. Fingers . Makes it difficult ...
      ham

      4
      ham
      All done, all handed in. Don't know if mega sh...
      ham

We'll write a function to measure the accuracy of our spam filter to find out how well our spam filter does.

```python
correct = 0
total = test_set.shape[0]

for row in test_set.iterrows():
    row = row[1]
    if row['Label'] == row['predicted']:
        correct += 1

print('Correct:', correct)
print('Incorrect:', total - correct)
print('Accuracy:', correct/total)
```

    Correct: 1100
    Incorrect: 14
    Accuracy: 0.9874326750448833

## Next Steps

In this project, we managed to build a spam filter for SMS messages using the multinomial Naive Bayes algorithm. The filter had an accuracy of 98.74% on the test set we used, which is a good result. Our initial goal was an accuracy of over 80%.

Next steps include:

- Analyze the 14 messages that were classified incorrectly and try to figure out why the algorithm classified them incorrectly.
- Make the filtering process more complex by making the algorithm sensitive to letter case.

```python

```

5.4 Conditional Probability/Build a Spam Filter With Naive Bayes.ipynb