<jupyter_start><jupyter_text>11 - Jointly Learning to Align and TranslatePrepared by <NAME>DLSU Machine Learning GroupIn this notebook, we'll implement the model in the paper Neural Machine Translation by Jointly Learning to Align and Translate (Bahadanau et al., 2014). We'll use attention to improve our translations performance-wise and interpretability-wise.  Preliminaries First, let's make sure we have a GPU.<jupyter_code>!nvidia-smi<jupyter_output>Fri Sep  4 11:29:06 2020       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 450.66       Driver Version: 418.67       CUDA Version: 10.1     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |
| N/A   57C    P8    10W /  70W |      0MiB / 15079MiB |      0%      Default |
|                               |                      |                 ERR! |
+-------------------------------+----------------------+----------------------+
                                                                               
+-------[...]<jupyter_text>Download the data and the tokenizers as with the previous notebook. Don't forget to restart the runtime.<jupyter_code>!wget https://s3.us-east-2.amazonaws.com/blaisecruz.com/datasets/translation/multi30k.zip
!unzip multi30k.zip && rm multi30k.zip
!python -m spacy download de_core_news_sm
!python -m spacy download en_core_web_sm<jupyter_output><empty_output><jupyter_text>Then it's our usual imports.<jupyter_code>import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as datautils

import spacy
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker

from nltk.translate.bleu_score import corpus_bleu
from nltk.translate.bleu_score import SmoothingFunction

import random
from collections import Counter
from tqdm import tqdm

np.random.seed(42)
torch.manual_seed(42)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
%matplotlib inline<jupyter_output><empty_output><jupyter_text>Data Processing We'll do the same thing here as with the last notebook, so we'll breeze through this. Load the data.<jupyter_code>with open('multi30k/train.en', 'r') as f:
    train_en = [line.strip() for line in f]
with open('multi30k/train.de', 'r') as f:
    train_de = [line.strip() for line in f]
with open('multi30k/val.en', 'r') as f:
    valid_en = [line.strip() for line in f]
with open('multi30k/val.de', 'r') as f:
    valid_de = [line.strip() for line in f]<jupyter_output><empty_output><jupyter_text>Load the tokenizers then tokenize the dataset and add the delimiter tokens.<jupyter_code>spacy_de = spacy.load('de_core_news_sm')
spacy_en = spacy.load('en_core_web_sm')

def tokenize_de(text):
    return ['<sos>'] + [tok.text.lower() for tok in spacy_de.tokenizer(text)] + ['<eos>']

def tokenize_en(text):
    return ['<sos>'] + [tok.text.lower() for tok in spacy_en.tokenizer(text)] + ['<eos>']

# Tokenize the text
train_en = [tokenize_en(text) for text in tqdm(train_en)]
train_de = [tokenize_de(text) for text in tqdm(train_de)]
valid_en = [tokenize_en(text) for text in tqdm(valid_en)]
valid_de = [tokenize_de(text) for text in tqdm(valid_de)]<jupyter_output>100%|██████████| 29000/29000 [00:01<00:00, 21987.15it/s]
100%|██████████| 29000/29000 [00:01<00:00, 15244.91it/s]
100%|██████████| 1014/1014 [00:00<00:00, 27558.17it/s]
100%|██████████| 1014/1014 [00:00<00:00, 15986.59it/s]<jupyter_text>Process each sequence to the maximum sequence length of the dataset.<jupyter_code>def process(dataset):
    max_len = max([len(text) for text in dataset])
    temp = []
    for text in dataset:
        if len(text) < max_len:
            text += ['<pad>' for _ in range(max_len - len(text))]
        temp.append(text)
    return temp

# Pad to maximum length of the dataset
train_en_proc, valid_en_proc = process(train_en), process(valid_en)
train_de_proc, valid_de_proc = process(train_de), process(valid_de)<jupyter_output><empty_output><jupyter_text>Produce vocabularies for the source and target languages.<jupyter_code>def get_vocab(dataset, min_freq=2):
    # Add all tokens to the list
    special_tokens = ['<unk>', '<pad>', '<sos>', '<eos>']
    vocab = []
    for line in dataset: vocab.extend(line)

    # Remove words that are below the minimum frequency, the enforce set
    counts = Counter(vocab)
    vocab = special_tokens + [word for word in counts.keys() if counts[word] > min_freq]
    vocab_set = set(vocab)

    # Push all special tokens to the front
    idx2word = list(vocab_set)
    for token in special_tokens[::-1]:
        idx2word.insert(0, idx2word.pop(idx2word.index(token)))

    # Produce word2idx then return
    word2idx = {idx2word[i]: i for i in range(len(idx2word))}
    return vocab_set, idx2word, word2idx

# Get vocabulary and references
vocab_set_en, idx2word_en, word2idx_en = get_vocab(train_en_proc, min_freq=2)
vocab_set_de, idx2word_de, word2idx_de = get_vocab(train_de_proc, min_freq=2)

# Convert unknown tokens
train_en_proc = [[token if token in vocab_set_en else '<unk>' for token in line] for line in train_en_proc]
train_de_proc = [[token if token in vocab_set_de else '<unk>' for token in line] for line in train_de_proc]
valid_en_proc = [[token if token in vocab_set_en else '<unk>' for token in line] for line in valid_en_proc]
valid_de_proc = [[token if token in vocab_set_de else '<unk>' for token in line] for line in valid_de_proc]<jupyter_output><empty_output><jupyter_text>Finally, convert the sequences of tokens into their corresponding indices.<jupyter_code>def serialize(dataset, word2idx):
    temp = []
    for line in dataset: temp.append([word2idx[token] for token in line])
    return torch.LongTensor(temp)

# Convert to idx
y_train = serialize(train_en_proc, word2idx_en)
X_train = serialize(train_de_proc, word2idx_de)
y_valid = serialize(valid_en_proc, word2idx_en)
X_valid = serialize(valid_de_proc, word2idx_de)<jupyter_output><empty_output><jupyter_text>We'll wrap up by making dataloaders.<jupyter_code>bs = 128

train_dataset = datautils.TensorDataset(X_train, y_train)
valid_dataset = datautils.TensorDataset(X_valid, y_valid)
train_sampler = datautils.RandomSampler(train_dataset)
train_loader = datautils.DataLoader(train_dataset, batch_size=bs, sampler=train_sampler)
valid_loader = datautils.DataLoader(valid_dataset, batch_size=bs, shuffle=False)<jupyter_output><empty_output><jupyter_text>Modeling Let's start with the encoder. The biggest change here is that we're using a bidirectional GRU instead of a vanilla LSTM. Since the RNN is bidirectional, we'll introduce a linear layer to **pool** the forward and backward hidden states.Encoding the source sequence remains the same. We'll embed and apply dropout. Afterwhich, we'll pack the sequences to disregard padding. Pass that result to the GRU, then unpack the result. We pass the concatenated hidden states to the linear pooling layer, then apply a hyperbolic tangent activation function. This tells us which "candidate information" to keep.To prepare for attention, we'll output both the RNN output and the pooled hidden states.<jupyter_code>class Encoder(nn.Module):
    def __init__(self, vocab_sz, embedding_dim, enc_hidden_dim, dec_hidden_dim, dropout=0.5):
        super(Encoder, self).__init__()
        self.embedding = nn.Embedding(vocab_sz, embedding_dim)
        self.rnn = nn.GRU(embedding_dim, enc_hidden_dim, bidirectional=True)
        self.fc1 = nn.Linear(enc_hidden_dim * 2, dec_hidden_dim)
        self.dropout = nn.Dropout(dropout)
        self.vocab_sz = vocab_sz

    def init_hidden(self, bs):
        weight = next(self.parameters())
        hidden_dim = self.rnn.hidden_size

        h = weight.new_zeros(2, bs, hidden_dim)
        return h

    def forward(self, x, pad_idx):
        msl, bs = x.shape

        # Embed then pack the sequence
        out = self.embedding(x)
        out = self.dropout(out)

        # Get the length of the sequences and pack them
        lens = ((x.rot90() == pad_idx) == False).int().sum(dim=1)
        out = nn.utils.rnn.pack_padded_sequence(out, lens, enforce_sorted=False)

        hidden = self.init_hidden(bs)
        out, hidden = self.rnn(out, hidden)
        out, _ = nn.utils.rnn.pad_packed_sequence(out, total_length=msl)

        # Pool the forward and backward states
        hidden = torch.cat((hidden[-2, :, :], hidden[-1, :, :]), axis=1)
        hidden = torch.tanh(self.fc1(hidden))

        return out, hidden<jupyter_output><empty_output><jupyter_text>For the rest of the modeling part, we'll run through attention and decoding manually as things can get rather confusing. First, let's get a sample source and target sequence batch.<jupyter_code>x, y = next(iter(train_loader))
x, y = x.rot90(k=3), y.rot90(k=3)
print(x.shape, y.shape)<jupyter_output>torch.Size([46, 128]) torch.Size([43, 128])<jupyter_text>Here's our hyperparameters.<jupyter_code>vocab_sz = len(vocab_set_de)
embedding_dim = 128
enc_hidden_dim = 256
dec_hidden_dim = 256
pad_idx = word2idx_de['<pad>']<jupyter_output><empty_output><jupyter_text>Instantiate an encoder.<jupyter_code>encoder = Encoder(vocab_sz, embedding_dim, enc_hidden_dim, dec_hidden_dim)<jupyter_output><empty_output><jupyter_text>Then pass the source sentence.We get what we expect.<jupyter_code>out, hidden = encoder(x, pad_idx=pad_idx)
print(out.shape, hidden.shape)<jupyter_output>torch.Size([46, 128, 512]) torch.Size([128, 256])<jupyter_text>Time to compute the attention weights!Attention weights can be thought of as a form of "feature importance." We'll use the RNN outputs of the encoder and the pooled hidden states to output a matrix of size [bs $\times$ sequence length]. This can be thought of as the "feature importance" of each token in the sequence given the current information (the hidden states).At every decoding timestep, we will produce this attention matrix (or attention vector per sequence) to tell our model which tokens are important at each step. This gives the model the ability to "look back" and "focus on important context tokens."Anyhow, we'll implement an attention layer using two linear layers.<jupyter_code>attn = nn.Linear((enc_hidden_dim * 2) + dec_hidden_dim, dec_hidden_dim) 
v = nn.Linear(dec_hidden_dim, 1, bias=False)<jupyter_output><empty_output><jupyter_text>We'll take the length of the source sentence and the batch size.<jupyter_code>src_len, bs, _ = out.shape<jupyter_output><empty_output><jupyter_text>Here's what we're going to do:1. We need to match the hidden state to each output token. If, for example, the sequence has 50 tokens, we need to 50 copies of the hidden state to "match" to each output token.2. We'll compute the "energy" between the hidden state and the output. Think of this as "how much does a specific token match the information in the hidden state. If it is more relevant, energy is higher."3. Once we have our energy, we need to turn them into weights. These are the weights of each token at a certain timestep. In other words, the "importance."Let's do 1 first. Remember the hidden state's shape:<jupyter_code>hidden.shape<jupyter_output><empty_output><jupyter_text>Let's introduce a blank dimension.<jupyter_code>hidden.unsqueeze(1).shape<jupyter_output><empty_output><jupyter_text>Then repeat the hidden state src_len times. This makes a number of clones of the hidden state that we can match to every entry in the RNN outputs.<jupyter_code>hidden_repeated = hidden.unsqueeze(1).repeat(1, src_len, 1)
hidden_repeated.shape<jupyter_output><empty_output><jupyter_text>As a sanity check:<jupyter_code>hidden_repeated[:, 0, :] == hidden_repeated[:, 1, :]<jupyter_output><empty_output><jupyter_text>Now here's our output shape:<jupyter_code>out.shape<jupyter_output><empty_output><jupyter_text>We'll permute the dimensions to match our hidden state clones:<jupyter_code>out_permuted = out.permute(1, 0, 2)
out_permuted.shape<jupyter_output><empty_output><jupyter_text>Since their dimensions are now matched, we can concatenate them on the last dimension. This should now be of shape [batch size $\times$ src_len $\times$ encoder hidden dim * 2 + decoder hidden dim]Remember that we're doing encoder hidden dim * 2 because our encoder uses a bidirectional RNN.<jupyter_code>torch.cat((out_permuted, hidden_repeated), dim=2).shape<jupyter_output><empty_output><jupyter_text>Pass the concatenated hidden and outputs to a linear layer to compute new features, then apply a hyperbolic tangent activation function.<jupyter_code>energy = attn(torch.cat((out_permuted, hidden_repeated), dim=2))
energy = torch.tanh(energy)
energy.shape<jupyter_output><empty_output><jupyter_text>We still have a 3D tensor. Let's project this down to 2 dimensions by applying a final linear transform. This turns our energy into a "weight matrix."<jupyter_code>a = v(energy)
a.shape<jupyter_output><empty_output><jupyter_text>Squeeze the final dimension since there's nothing there.<jupyter_code>a = a.squeeze(2)
a.shape<jupyter_output><empty_output><jupyter_text>We're not yet done. Lastly, we have to mask out the attention weights on the padding tokens. We can do this by using an attention mask.<jupyter_code>mask = (x != pad_idx).permute(1, 0)
mask.shape<jupyter_output><empty_output><jupyter_text>It looks like this.<jupyter_code>mask<jupyter_output><empty_output><jupyter_text>To mask the padding tokens, we'll do it like so:<jupyter_code>a = a.masked_fill(mask == 0, -1e10)<jupyter_output><empty_output><jupyter_text>The full attention layer looks like this:<jupyter_code>class Attention(nn.Module):
    def __init__(self, enc_hidden_dim, dec_hidden_dim):
        super(Attention, self).__init__()
        self.attention = nn.Linear((enc_hidden_dim * 2) + dec_hidden_dim, dec_hidden_dim)
        self.v = nn.Linear(dec_hidden_dim, 1, bias=False)

    def forward(self, out, hidden, mask):
        src_len, bs, _ = out.shape
        hidden = hidden.unsqueeze(1).repeat(1, src_len, 1)
        out = out.permute(1, 0, 2)

        # Compute the weighted energy (match between hidden states)
        energy = torch.tanh(self.attention(torch.cat((out, hidden), dim=2)))
        attn = self.v(energy).squeeze(2)
        attn = attn.masked_fill(mask == 0, -1e10)

        return torch.softmax(attn, dim=1)<jupyter_output><empty_output><jupyter_text>Instantiate an attention layer for testing.<jupyter_code>attention = Attention(enc_hidden_dim, dec_hidden_dim)<jupyter_output><empty_output><jupyter_text>Using it gives us the outputs we expect.<jupyter_code>mask = (x != pad_idx).permute(1, 0)
a = attention(out, hidden, mask)
a.shape<jupyter_output><empty_output><jupyter_text>Lastly, we'll implement the decoder.We have a few differences from the previous decoder.1. Our RNN is again a GRU, which accepts encoder hidden dim * 2 + embedding dim shaped inputs. This is because we're passing it the encoded source sentence (remember the encoder is bidirectional) along with the embedded target token.2. We use attention inside the decoder. At every decoding step, we generate a new attention matrix, which we batch matrix multiply to the outputs of the encoder. This lets the decoder know which tokens in the source sentence it should "pay attention to."<jupyter_code>class Decoder(nn.Module):
    def __init__(self, vocab_sz, embedding_dim, enc_hidden_dim, dec_hidden_dim, dropout=0.5):
        super(Decoder, self).__init__()
        self.attention = Attention(enc_hidden_dim, dec_hidden_dim)
        self.embedding = nn.Embedding(vocab_sz, embedding_dim)
        self.rnn = nn.GRU((enc_hidden_dim * 2) + embedding_dim, dec_hidden_dim)
        self.fc1 = nn.Linear((enc_hidden_dim * 2) + dec_hidden_dim + embedding_dim, vocab_sz)
        self.dropout = nn.Dropout(dropout)
        self.vocab_sz = vocab_sz

    def forward(self, y, out, hidden, mask):
        embedded = self.embedding(y.unsqueeze(0))
        embedded = self.dropout(embedded)

        # Batch multiply to get weighted outputs
        a = self.attention(out, hidden, mask).unsqueeze(1)
        out = out.permute(1, 0, 2)
        weighted = torch.bmm(a, out).permute(1, 0, 2)

        # Prepare RNN Inputs
        rnn_inputs = torch.cat((embedded, weighted), dim=2)
        hidden = hidden.unsqueeze(0)
        out, hidden = self.rnn(rnn_inputs)

        # Pool the embeddings, weighted outputs, and new outputs
        out = torch.cat((embedded.squeeze(0), out.squeeze(0), weighted.squeeze(0)), dim=1)
        out = self.fc1(out)

        return out, hidden.squeeze(0), a.squeeze(1)<jupyter_output><empty_output><jupyter_text>Let's set the first input token for example.<jupyter_code>token = y[0,:]
print(token.shape)
print(token)<jupyter_output>torch.Size([128])
tensor([2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
        2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
        2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
        2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
        2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
        2, 2, 2, 2, 2, 2, 2, 2])<jupyter_text>Instantiate a decoder.<jupyter_code>decoder = Decoder(vocab_sz, embedding_dim, enc_hidden_dim, dec_hidden_dim)<jupyter_output><empty_output><jupyter_text>Produce a mask and pass in the output, hidden states, and the mask.<jupyter_code>mask = (x != pad_idx).permute(1, 0)
out, hidden, a = decoder(token, out, hidden, mask)
print(out.shape, hidden.shape, a.shape)<jupyter_output>torch.Size([128, 5376]) torch.Size([128, 256]) torch.Size([128, 46])<jupyter_text>Getting the maximum of the logits gives us the predicted next output token.<jupyter_code>out.argmax(1)<jupyter_output><empty_output><jupyter_text>Let's wrap the encoder and decoder in a seq2seq module. This is largely the same, except we include a helper function to generate an attention mask.<jupyter_code>class Seq2Seq(nn.Module):
    def __init__(self, encoder, decoder, pad_idx, initrange=0.01):
        super(Seq2Seq, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.init_weights(initrange)
        self.pad_idx = pad_idx

    def create_mask(self, x):
        with torch.no_grad():
            mask = (x != self.pad_idx).permute(1, 0)
        return mask

    def init_weights(self, initrange=0.01):
        for name, param in self.named_parameters():
            if 'weight' in name: nn.init.normal_(param.data, mean=0, std=initrange)
            else: nn.init.constant_(param.data, 0)

    def forward(self, x, y, teacher_forcing=0.5):
        src_len, bs = x.shape
        trg_len, bs = y.shape

        # Make container for outputs
        weight = next(self.encoder.parameters())
        outputs = weight.new_zeros(trg_len, bs, self.decoder.vocab_sz)

        # Encode source then remember context
        encoder_out, hidden = self.encoder(x, self.pad_idx)
        input_ids = y[0,:]
        mask = self.create_mask(x)

        # Decode per input token
        for i in range(1, trg_len):
            out, hidden, _ = self.decoder(input_ids, encoder_out, hidden, mask)
            outputs[i] = out

            teacher_force = random.random() < teacher_forcing
            input_ids = y[i] if teacher_force else out.argmax(1)

        return outputs<jupyter_output><empty_output><jupyter_text>Test out the wrapper and instantiate a loss function.<jupyter_code>model = Seq2Seq(encoder, decoder, pad_idx)
criterion = nn.CrossEntropyLoss(ignore_index=pad_idx)<jupyter_output><empty_output><jupyter_text>We get the expected output.<jupyter_code>out = model(x, y)
print(out.shape)<jupyter_output>torch.Size([43, 128, 5376])<jupyter_text>Here's the initial loss.<jupyter_code>loss = criterion(out[1:].flatten(0, 1), y[1:].flatten(0))
print(loss)<jupyter_output>tensor(8.5897, grad_fn=<NllLossBackward>)<jupyter_text>Training We'll instantiate a training setup with the following hyperparameters:<jupyter_code>encoder = Encoder(vocab_sz=len(vocab_set_de), embedding_dim=256, enc_hidden_dim=512, dec_hidden_dim=512, dropout=0.5)
decoder = Decoder(vocab_sz=len(vocab_set_en), embedding_dim=256, enc_hidden_dim=512, dec_hidden_dim=512, dropout=0.5)
model = Seq2Seq(encoder, decoder, pad_idx=word2idx_de['<pad>']).to(device)

optimizer = optim.Adam(model.parameters(), lr=3e-3)
criterion = nn.CrossEntropyLoss(ignore_index=word2idx_en['<pad>'])

epochs = 20
iters = epochs * len(train_loader)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=iters, eta_min=0)<jupyter_output><empty_output><jupyter_text>Count the parameters of the model.<jupyter_code>def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)
    
print("The model has {:,} trainable parameters".format(count_parameters(model)))<jupyter_output>The model has 17,144,780 trainable parameters<jupyter_text>For transparency, here's the final architecture of our model:<jupyter_code>model<jupyter_output><empty_output><jupyter_text>Finally, we'll train it.<jupyter_code>clip = 1.0

for e in range(1, epochs + 1):
    train_loss = 0
    
    model.train()
    for x, y in tqdm(train_loader):
        x, y = x.rot90(k=3).to(device), y.rot90(k=3).to(device)

        out = model(x, y)
        loss = criterion(out[1:].flatten(0, 1), y[1:].flatten(0))
        
        optimizer.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), clip)
        optimizer.step()
        scheduler.step()

        train_loss += loss.item()
    train_loss /= len(train_loader)

    valid_loss = 0
    
    model.eval()
    with torch.no_grad():
        for x, y in tqdm(valid_loader):
            x, y = x.rot90(k=3).to(device), y.rot90(k=3).to(device)

            out = model(x, y)
            loss = criterion(out[1:].flatten(0, 1), y[1:].flatten(0))

            valid_loss += loss.item()
    valid_loss /= len(valid_loader)

    print("\nEpoch {:3} | Train Loss {:.4f} | Train Ppl {:.4f} | Valid Loss {:.4f} | Valid Ppl {:.4f}".format(e, train_loss, np.exp(train_loss), valid_loss, np.exp(valid_loss)))<jupyter_output>100%|██████████| 227/227 [02:00<00:00,  1.88it/s]
100%|██████████| 8/8 [00:00<00:00,  8.68it/s]
  0%|          | 0/227 [00:00<?, ?it/s]<jupyter_text>If you train your own model, make sure to save the resulting weights.<jupyter_code>torch.save(model.state_dict(), 'seq2seq_attention.pt')<jupyter_output><empty_output><jupyter_text>Sampling Let's load the pretrained weights.<jupyter_code>encoder = Encoder(vocab_sz=len(vocab_set_de), embedding_dim=256, enc_hidden_dim=512, dec_hidden_dim=512, dropout=0.5)
decoder = Decoder(vocab_sz=len(vocab_set_en), embedding_dim=256, enc_hidden_dim=512, dec_hidden_dim=512, dropout=0.5)
model = Seq2Seq(encoder, decoder, pad_idx=word2idx_de['<pad>']).to(device)
criterion = nn.CrossEntropyLoss(ignore_index=word2idx_en['<pad>'])

model.load_state_dict(torch.load('seq2seq_attention.pt'))
model.eval();<jupyter_output><empty_output><jupyter_text>Evaluate them on the validation set.<jupyter_code>model.eval()
valid_loss = 0
with torch.no_grad():
    for x, y in tqdm(valid_loader):
        x, y = x.rot90(k=3).to(device), y.rot90(k=3).to(device)

        out = model(x, y)
        loss = criterion(out[1:].flatten(0, 1), y[1:].flatten(0))

        valid_loss += loss.item()
valid_loss /= len(valid_loader)

print("\nValid Loss {:.4f} | Valid Ppl {:.4f}".format(valid_loss, np.exp(valid_loss)))<jupyter_output>100%|██████████| 8/8 [00:01<00:00,  7.90it/s]<jupyter_text>Let's write a modified translation function. This is largely akin to the seq2seq decoding method, but we're using multinomial sampling and we're not using any teacher forcing.We'll start with a tokenized sequence of tokens. We encode it with the encoder, then start decoding. We'll save each token as well as the attention weights per timestep so we can plot them later.<jupyter_code>def translate(src_sentence, model, max_words=20, temperature=1.0, seed=42):
    s = [token if token in vocab_set_de else '<unk>' for token in src_sentence]
    sample = torch.LongTensor([word2idx_de[token] for token in s]).unsqueeze(1)
    
    predictions = []
    torch.manual_seed(seed)

    with torch.no_grad():
        # Produce an encoding of the source text and a storage for attention
        encoder_out, hidden = model.encoder(sample, pad_idx=word2idx_de['<pad>'])
        token = sample[0,:]
        attentions = torch.zeros(max_words, 1, sample.shape[0])

        # Decoding. Sample succeding tokens using attention and encodings
        for i in range(max_words):
            mask = (sample != pad_idx).permute(1, 0)
            out, hidden, attn = model.decoder(token, encoder_out, hidden, mask)
            
            weights = torch.softmax(out / temperature, dim=-1)
            token = torch.multinomial(weights, 1).squeeze(0)

            predictions.append(token.item())
            attentions[i] = attn

            if token.item() == word2idx_en['<eos>']:
                break

    # Convert predictions from indices to text. Cut the attentions to translation length
    predictions = [idx2word_en[ix] for ix in predictions]
    attentions = attentions.squeeze(1)[:len(predictions),:]

    return predictions, attentions<jupyter_output><empty_output><jupyter_text>The nice thing about attention is that it's interpretable. We can write a function to plot the attention weights in a heatmap for each timestep of decoding.<jupyter_code>def plot_attention(src_sentence, predictions, attentions):
    fig = plt.figure(figsize=(7,7))
    ax = fig.add_subplot(111)
        
    attention = attentions.numpy()
    cax = ax.matshow(attention, cmap='bone')

    ax.tick_params(labelsize=12)
    ax.set_xticklabels([''] + [t for t in src_sentence], rotation=45)
    ax.set_yticklabels([''] + predictions)
    ax.xaxis.set_major_locator(ticker.MultipleLocator(1))
    ax.yaxis.set_major_locator(ticker.MultipleLocator(1))<jupyter_output><empty_output><jupyter_text>Let's put the model in the CPU.<jupyter_code>model = model.cpu()<jupyter_output><empty_output><jupyter_text>Then test out some translations.<jupyter_code>ix = 0
src_sentence, trg_sentence = valid_de[ix], valid_en[ix]
src_sentence = src_sentence[:src_sentence.index('<eos>') + 1]
trg_sentence = trg_sentence[:trg_sentence.index('<eos>') + 1]

print('Source:', ' '.join(src_sentence))
print('Target:', ' '.join(trg_sentence))

predictions, attentions = translate(src_sentence, model, temperature=0.9)
print('Prediction:', ' '.join(predictions))

plot_attention(src_sentence, predictions, attentions)
ix = 6
src_sentence, trg_sentence = valid_de[ix], valid_en[ix]
src_sentence = src_sentence[:src_sentence.index('<eos>') + 1]
trg_sentence = trg_sentence[:trg_sentence.index('<eos>') + 1]

print('Source:', ' '.join(src_sentence))
print('Target:', ' '.join(trg_sentence))

predictions, attentions = translate(src_sentence, model, temperature=0.9)
print('Prediction:', ' '.join(predictions))

plot_attention(src_sentence, predictions, attentions)
ix = 10
src_sentence, trg_sentence = valid_de[ix], valid_en[ix]
src_sentence = src_sentence[:src_sentence.index('<eos>') + 1]
trg_sentence = trg_sentence[:trg_sentence.index('<eos>') + 1]

print('Source:', ' '.join(src_sentence))
print('Target:', ' '.join(trg_sentence))

predictions, attentions = translate(src_sentence, model, temperature=0.9)
print('Prediction:', ' '.join(predictions))

plot_attention(src_sentence, predictions, attentions)
ix = 64
src_sentence, trg_sentence = valid_de[ix], valid_en[ix]
src_sentence = src_sentence[:src_sentence.index('<eos>') + 1]
trg_sentence = trg_sentence[:trg_sentence.index('<eos>') + 1]

print('Source:', ' '.join(src_sentence))
print('Target:', ' '.join(trg_sentence))

predictions, attentions = translate(src_sentence, model, temperature=0.9)
print('Prediction:', ' '.join(predictions))

plot_attention(src_sentence, predictions, attentions)<jupyter_output>Source: <sos> eine bunt gekleidete frau geht an einem weißen lastwagen vorbei , der mit flaschen gefüllt ist . <eos>
Target: <sos> a woman in a colorful outfit is walking by a white truck filled with bottles . <eos>
Prediction: a woman dressed in walks past a white truck that has graffiti . <eos>
{'fraction_non_alphanumeric': 0.12809310778694255, 'fraction_numerical': 0.026562281560184538, 'mean_word_length': 4.24143616046895, 'pattern_counts': {'":': 0, '<': 243, '<?xml version=': 0, '>': 234, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 1, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '52', 'dir': 'jupyter-structured-clean-dedup', 'id': '582889', 'license': 'MIT', 'n_tokens_mistral': 10422, 'n_tokens_neox': 9603, 'n_words': 2567, 'path': 'notebooks/11 Neural Machine Translation with Attention.ipynb', 'repo_name': 'jcblaisecruz02/dl-nlp-course', 'size': '28631'}