skip_training = False # Set this flag to True before validation and submission

# During evaluation, this cell sets skip_training to True

# skip_training = True

import tools, warnings

[Link] = [Link]

import os
import random
import numpy as np
import [Link] as plt

import torch
import [Link] as nn
import [Link] as optim
import [Link] as F
from [Link] import DataLoader

import transformer as tr
import tools

# When running on your own computer, you can specify the data directory by:
# data_dir = tools.select_data_dir('/your/local/data/directory')
data_dir = tools.select_data_dir()

# Select the device for training (use GPU if you have one)
#device = [Link]('cuda:0')
device = [Link]('cpu')

if skip_training:
# The models are always evaluated on CPU
device = [Link]("cpu")

# Translation data
from data import TranslationDataset, SOS_token, EOS_token, MAX_LENGTH
trainset = TranslationDataset(data_dir, train=True)

src_seq, tgt_seq = trainset[[Link](len(trainset))]

print('Source sentence:')
print(' as word indices: ', src_seq)
print(' as string: ', ' '.join(trainset.input_lang.index2word[[Link]()] for i in
src_seq))

print('Target sentence:')
print(' as word indices: ', tgt_seq)
print(' as string: ', ' '.join(trainset.output_lang.index2word[[Link]()] for i in
tgt_seq))

PADDING_VALUE = 0

from [Link] import pad_sequence

def collate(list_of_samples):
"""Merges a list of samples to form a mini-batch.

Args:
list_of_samples is a list of tuples (src_seq, tgt_seq):
src_seq is of shape (src_seq_length)
 tgt_seq is of shape (tgt_seq_length)

Returns:
src_seqs of shape (max_src_seq_length, batch_size): LongTensor of padded
source sequences.
src_mask of shape (max_src_seq_length, batch_size): BoolTensor (tensor with
boolean elements) indicating which
elements of the src_seqs tensor should be ignored in computations: True
values in src_mask correspond
to padding values in src_seqs.
tgt_seqs of shape (max_tgt_seq_length+1, batch_size): LongTensor of padded
target sequences.
"""
# YOUR CODE HERE
src_seqs, tgt_seqs = zip(*list_of_samples)

# Pad source sequences

src_seqs = pad_sequence(src_seqs, padding_value=PADDING_VALUE,
batch_first=False)
src_mask = (src_seqs == PADDING_VALUE)

# Pad target sequences and add SOS_token

tgt_seqs = [[Link]([[Link]([SOS_token]), tgt]) for tgt in tgt_seqs]
tgt_seqs = pad_sequence(tgt_seqs, padding_value=PADDING_VALUE,
batch_first=False)

return src_seqs, src_mask, tgt_seqs

raise NotImplementedError()

def test_collate_shapes():
pairs = [
([Link]([2, EOS_token]), [Link]([3, 4, EOS_token])),
([Link]([6, 7, EOS_token]), [Link]([9, EOS_token])),
]
src_seqs, src_mask, tgt_seqs = collate(pairs)
assert src_seqs.dtype == [Link], f"Wrong src_seqs.dtype: {src_seqs.dtype}"
assert src_seqs.shape == [Link]([3, 2]), f"Wrong src_seqs.shape:
{src_seqs.shape}"

assert tgt_seqs.dtype == [Link], f"Wrong tgt_seqs.dtype: {tgt_seqs.dtype}"

assert tgt_seqs.shape == [Link]([4, 2]), f"Wrong tgt_seqs.shape:
{tgt_seqs.shape}"
assert (tgt_seqs[0] == [Link](2,
dtype=[Link]).fill_(SOS_token)).all(), "Target sequences should start with
SOS_token."

assert src_mask.dtype == [Link], f"Wrong src_mask.dtype: {src_mask.dtype}"

assert src_mask.shape == src_seqs.shape, f"Wrong src_mask.shape:
{src_mask.shape}"
print('Success')

test_collate_shapes()

# This cell tests collate()

def test_collate():
pairs = [
([Link]([2, EOS_token]), [Link]([3, 4, EOS_token])),
([Link]([6, 7, EOS_token]), [Link]([9, EOS_token])),
 ]
src_seqs, src_mask, tgt_seqs = collate(pairs)
#src_seqs, src_mask, tgt_seqs = src_seqs[:, [1, 0]], src_mask[:, [1, 0]],
tgt_seqs[:, [1, 0]]
print('src_seqs:\n', src_seqs)
print('src_mask:\n', src_mask)
print('tgt_seqs:\n', tgt_seqs)
expected_src_seqs = [Link]([
[2, 6],
[EOS_token, 7],
[0, EOS_token]
])
expected_src_mask = [Link]([
[False, False],
[False, False],
[ True, False]
])
expected_tgt_seqs = [Link]([
[0, 0],
[3, 9],
[4, EOS_token],
[EOS_token, 0]
])

assert ((
(src_seqs == expected_src_seqs).all()
and (src_mask == expected_src_mask).all()
and (tgt_seqs == expected_tgt_seqs).all()
) or (
(src_seqs == expected_src_seqs[:, [1, 0]]).all()
and (src_mask == expected_src_mask[:, [1, 0]]).all()
and (tgt_seqs == expected_tgt_seqs[:, [1, 0]]).all()
)
), "Wrong outputs of collate."
print('Success')

test_collate()

# We create custom DataLoader using the implemented collate function

# We are going to process 64 sequences at the same time (batch_size=64)
trainloader = DataLoader(dataset=trainset, batch_size=64, shuffle=True,
collate_fn=collate, pin_memory=True)

# Create test set

testset = TranslationDataset(data_dir, train=False)
testloader = DataLoader(dataset=testset, batch_size=64, shuffle=False,
collate_fn=collate)

class EncoderBlock([Link]):
def __init__(self, n_features, n_heads, n_hidden=64, dropout=0.1):
"""
Args:
n_features: Number of input and output features.
n_heads: Number of attention heads in the Multi-Head Attention.
n_hidden: Number of hidden units in the Feedforward (MLP) block.
dropout: Dropout rate after the first layer of the MLP and in two places
on the main path (before
combining the main path with a skip connection).
"""
 # YOUR CODE HERE
super().__init__()

# Multi-Head Attention layer Addressed

[Link] = [Link](embed_dim=n_features,
num_heads=n_heads, dropout=dropout)

# Feedforward network (MLP) addressed

[Link] = [Link](
[Link](n_features, n_hidden),
[Link](),
[Link](dropout),
[Link](n_hidden, n_features)
)

# Normalization layers addressed

self.norm1 = [Link](n_features)
self.norm2 = [Link](n_features)

# Dropout layers [2 of them]

self.dropout1 = [Link](dropout)
self.dropout2 = [Link](dropout)
#raise NotImplementedError()

def forward(self, x, mask):

"""
Args:
x of shape (max_seq_length, batch_size, n_features): Input sequences.
mask of shape (max_seq_length, batch_size): BoolTensor indicating which
elements of the input
sequences should be ignored (True values correspond to ignored
elements in x).

Returns:
z of shape (max_seq_length, batch_size, n_features): Encoded input
sequences.

Note: All intermediate signals should be of shape (max_seq_length,

batch_size, n_features).
"""
# YOUR CODE HERE
attention_output, _ = [Link](x, x, x, key_padding_mask=mask.T)
x = x + self.dropout1(attention_output)
x = self.norm1(x)

# Feedforward with skip connection

feedforward_output = [Link](x)
x = x + self.dropout2(feedforward_output)
z = self.norm2(x)

return z
raise NotImplementedError()

def test_EncoderBlock_shapes():
encoder_block = EncoderBlock(n_features=16, n_heads=4, n_hidden=64)

x = [Link]([
[1, 2],
[3, 4],
 [5, 0],
[6, 0],
]).float().view(4, 2, 1).repeat(1, 1, 16) # (max_seq_length, batch_size,
n_features)

mask = [Link]([
[0, 0],
[0, 0],
[0, 1],
[0, 1],
], dtype=[Link]) # (max_seq_length, batch_size)
outputs = encoder_block(x, mask)
assert [Link] == [Link]([4, 2, 16]), f"Wrong [Link]:
{[Link]}"
print('Success')

test_EncoderBlock_shapes()

# This cell tests EncoderBlock

# Check that the signal does not propagate from the padded elements
import [Link]

def no_dropout(x, *args, **kwargs):

return x

@[Link]('[Link]', no_dropout) # .eval() does not

disable [Link]()
def test_EncoderBlock():
with torch.no_grad():
encoder_block = EncoderBlock(n_features=16, n_heads=4, n_hidden=64)

x = [Link]([
[1, 2],
[3, 4],
[5, 0],
[6, 0],
]).float().view(4, 2, 1).repeat(1, 1, 16) # (max_seq_length, batch_size,
n_features)

mask = [Link]([
[0, 0],
[0, 0],
[0, 1],
[0, 1],
], dtype=[Link]) # (max_seq_length, batch_size)

outputs1 = encoder_block(x, mask)

print('outputs1[:2,1,:]:\n', outputs1[:2,1,:])

# Modify non-padded values

x[:2,1,:] = x[:2,1,:] + 1
outputs2 = encoder_block(x, mask)
print('outputs2[:2,1,:]:\n', outputs2[:2,1,:])
diff = outputs2[:2,1,:] - outputs1[:2,1,:]
assert [Link]([Link](diff)) > 0, "Some elements of the source
sequence do not affect the output."

# Modify padded values

x[2:,1,:] = x[2:,1,:] + 1
 outputs3 = encoder_block(x, mask)
print('outputs3[:2,1,:]:\n', outputs3[:2,1,:])
diff = outputs3[:2,1,:] - outputs2[:2,1,:]
assert [Link]([Link](diff)) == 0, "Padding values affect the output."

print('Success')

test_EncoderBlock()

class Encoder([Link]):
def __init__(self, src_vocab_size, n_blocks, n_features, n_heads, n_hidden=64,
dropout=0.1):
"""
Args:
src_vocab_size: Number of words in the source vocabulary.
n_blocks: Number of EncoderBlock blocks.
n_features: Number of features to be used for word embedding and further
in all layers of the encoder.
n_heads: Number of attention heads inside the EncoderBlock.
n_hidden: Number of hidden units in the Feedforward block of
EncoderBlock.
dropout: Dropout level used in EncoderBlock.
"""
# YOUR CODE HERE
super().__init__()

# Word embedding layer applied

[Link] = [Link](src_vocab_size, n_features)

# Positional encoding applied

self.positional_encoding = [Link](n_features, dropout)

# Stack of EncoderBlocks put together

self.encoder_blocks = [Link]([
EncoderBlock(n_features, n_heads, n_hidden, dropout)
for _ in range(n_blocks)
])
#raise NotImplementedError()

def forward(self, x, mask):

"""
Args:
x of shape (max_seq_length, batch_size): LongTensor with the input
sequences.
mask of shape (max_seq_length, batch_size): BoolTensor indicating which
elements should be ignored.

Returns:
z of shape (max_seq_length, batch_size, n_features): Encoded input
sequences.

Note: All intermediate signals should be of shape (max_seq_length,

batch_size, n_features).
"""
# YOUR CODE HERE
# Converting the word indices into embeddings
x = [Link](x) # Shape: (max_seq_length, batch_size, n_features)

# Applying positional encoding

 x = self.positional_encoding(x)

# Passing through each EncoderBlock

for block in self.encoder_blocks:
x = block(x, mask)

return x # Which is also the Final encoded representation

raise NotImplementedError()

def test_Encoder_shapes():
encoder = Encoder(src_vocab_size=10, n_blocks=1, n_features=16, n_heads=4,
n_hidden=64)

x = [Link]([
[SOS_token, SOS_token],
[ 3, 4],
[ 5, EOS_token],
[ 6, PADDING_VALUE],
[EOS_token, PADDING_VALUE],
]) # (max_seq_length, batch_size)

mask = [Link]([
[0, 0],
[0, 0],
[0, 0],
[0, 1],
[0, 1],
], dtype=[Link]) # (max_seq_length, batch_size)
outputs = encoder(x, mask)
assert [Link] == [Link]([5, 2, 16]), f"Wrong [Link]:
{[Link]}"
print('Success')

test_Encoder_shapes()

def subsequent_mask(sz):
mask = ([Link]([Link](sz, sz)) == 1).transpose(0, 1).float()
mask = mask.masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1,
float(0.0))
return mask

# This is a typical mask that we need to use while decoding

mask = subsequent_mask(10)
print(mask)
[Link](mask)

class DecoderBlock([Link]):
def __init__(self, n_features, n_heads, n_hidden=64, dropout=0.1):
"""
Args:
n_features: Number of input and output features.
n_heads: Number of attention heads in the Multi-Head Attention.
n_hidden: Number of hidden units in the Feedforward (MLP) block.
dropout: Dropout rate after the first layer of the MLP and in three
places on the main path (before
combining the main path with a skip connection).
"""
# YOUR CODE HERE
super(DecoderBlock, self).__init__()
 # Self-attention (masked multi-head attention for target sequences)
self.self_attn = [Link](embed_dim=n_features,
num_heads=n_heads, dropout=dropout)

# Encoder-Decoder Attention (cross-attention)

self.cross_attn = [Link](embed_dim=n_features,
num_heads=n_heads, dropout=dropout)

# Feedforward network (MLP)

[Link] = [Link](
[Link](n_features, n_hidden),
[Link](),
[Link](dropout),
[Link](n_hidden, n_features)
)

# Layer Normalization
self.norm1 = [Link](n_features)
self.norm2 = [Link](n_features)
self.norm3 = [Link](n_features)

# Dropout layers
self.dropout1 = [Link](dropout)
self.dropout2 = [Link](dropout)
self.dropout3 = [Link](dropout)

#raise NotImplementedError()

def forward(self, y, z, src_mask, tgt_mask):

"""
Args:
y of shape (max_tgt_seq_length, batch_size, n_features): Transformed
target sequences used as the inputs
of the block.
z of shape (max_src_seq_length, batch_size, n_features): Encoded source
sequences (outputs of the
encoder).
src_mask of shape (max_src_seq_length, batch_size): BoolTensor indicating
which elements of the
encoded source sequences should be ignored.
tgt_mask of shape (max_tgt_seq_length, max_tgt_seq_length): Subsequent
mask to ignore subsequent
elements of the target sequences in the inputs. The rows of this
matrix correspond to the output
elements and the columns correspond to the input elements.

Returns:
out of shape (max_seq_length, batch_size, n_features): Output tensor.

Note: All intermediate signals should be of shape (max_seq_length,

batch_size, n_features).
"""
# YOUR CODE HERE
attention_output, _ = self.self_attn(y, y, y, attn_mask=tgt_mask)
y = y + self.dropout1(attention_output)
y = self.norm1(y)

# Cross-attention (encoder-decoder attention) with source mask

 attention_output, _ = self.cross_attn(y, z, z, key_padding_mask=src_mask.T)
y = y + self.dropout2(attention_output) # Add & Norm
y = self.norm2(y)

# Feedforward network (MLP)

feedforward_output = [Link](y)
y = y + self.dropout3(feedforward_output) # Add & Norm
y = self.norm3(y)

return y
raise NotImplementedError()

def test_DecoderBlock_shapes():
decoder_block = DecoderBlock(n_features=16, n_heads=4, n_hidden=64)

y = [Link]([
[1, 2],
[3, 4],
[5, 0],
[6, 0],
]).float().view(4, 2, 1).repeat(1, 1, 16) # (max_seq_length, batch_size,
n_features)

z = [Link](4, 2, 16, requires_grad=True) # (max_seq_length, batch_size,

n_features)

src_mask = [Link]([
[0, 0],
[0, 0],
[0, 1],
[0, 1],
], dtype=[Link]) # (max_seq_length, batch_size)

tgt_mask = subsequent_mask([Link](0))

outputs = decoder_block(y, z, src_mask=src_mask, tgt_mask=tgt_mask)

assert [Link] == [Link]([4, 2, 16]), f"Wrong [Link]:
{[Link]}"
print('Success')

test_DecoderBlock_shapes()

# This cell tests DecoderBlock

# Check that the signal does not propagate from the padded elements of the source
sequence and
# subsequent elements of the target sequence.
@[Link]('[Link]', no_dropout) # .eval() does not
disable [Link]()
def test_DecoderBlock():
with torch.no_grad():
decoder_block = DecoderBlock(n_features=16, n_heads=4, n_hidden=64)

y = [Link]([
[1, 2],
[3, 4],
[5, 0],
[6, 0],
]).float().view(4, 2, 1).repeat(1, 1, 16) # (max_seq_length, batch_size,
n_features)
 z = [Link](4, 2, 16) # (max_seq_length, batch_size, n_features)

src_mask = [Link]([
[0, 0],
[0, 0],
[0, 1],
[0, 1],
], dtype=[Link]) # (max_seq_length, batch_size)

tgt_mask = subsequent_mask([Link](0))

outputs1 = decoder_block(y, z, src_mask=src_mask, tgt_mask=tgt_mask)

assert [Link] == [Link]([4, 2, 16]), f"Wrong [Link]:
{[Link]}"
print('outputs1[:2,1,:]:\n', outputs1[:2,1,:])

# Modify second element of y[1]

y[1,1,:] = y[1,1,:] + 1
outputs2 = decoder_block(y, z, src_mask=src_mask, tgt_mask=tgt_mask)
print('outputs2[:2,1,:]:\n', outputs2[:2,1,:])
diff = outputs2[1,1,:] - outputs1[1,1,:]
assert [Link]([Link](diff)) > 0, "Some elements of the target
sequence do not affect the output."
diff = outputs2[0,1,:] - outputs1[0,1,:]
assert [Link]([Link](diff)) == 0, "Subsequent elements of the target
sequence affect the output."

# Modify padded values of y

y[2:,1,:] = y[2:,1,:] + 1
outputs3 = decoder_block(y, z, src_mask=src_mask, tgt_mask=tgt_mask)
print('outputs3[:2,1,:]:\n', outputs3[:2,1,:])
diff = outputs3[:2,1,:] - outputs2[:2,1,:]
assert [Link]([Link](diff)) == 0, "Padding values in the target
sequence affect the output."

# Modify non-padded values of z

z[:2,1,:] = z[:2,1,:] + 1
outputs4 = decoder_block(y, z, src_mask=src_mask, tgt_mask=tgt_mask)
print('outputs4[:2,1,:]:\n', outputs4[:2,1,:])
diff = outputs4[:2,1,:] - outputs3[:2,1,:]
assert [Link]([Link](diff)) > 0, "Some elements of the source
sequence do not affect the output."

# Modify padded values of y

z[2:,1,:] = z[2:,1,:] + 1
outputs5 = decoder_block(y, z, src_mask=src_mask, tgt_mask=tgt_mask)
print('outputs5[:2,1,:]:\n', outputs5[:2,1,:])
diff = outputs5[:2,1,:] - outputs4[:2,1,:]
assert [Link]([Link](diff)) == 0, "Padding values in the source
sequence affect the output."

print('Success')

test_DecoderBlock()

class Decoder([Link]):
def __init__(self, tgt_vocab_size, n_blocks, n_features, n_heads, n_hidden=64,
dropout=0.1):
 """
Args:
tgt_vocab_size: Number of words in the target vocabulary.
n_blocks: Number of EncoderBlock blocks.
n_features: Number of features to be used for word embedding and further
in all layers of the decoder.
n_heads: Number of attention heads inside the DecoderBlock.
n_hidden: Number of hidden units in the Feedforward block of
DecoderBlock.
dropout: Dropout level used in DecoderBlock.
"""
# YOUR CODE HERE
super().__init__()

# Word embedding layer

[Link] = [Link](tgt_vocab_size, n_features)

# Positional encoding layer

self.positional_encoding = [Link](n_features, dropout)

# Stack of decoder blocks

self.decoder_blocks = [Link]([
DecoderBlock(n_features, n_heads, n_hidden, dropout)
for _ in range(n_blocks)
])

# Output linear layer

self.output_layer = [Link](n_features, tgt_vocab_size)

# Log softmax for output probabilities

self.log_softmax = [Link](dim=-1)
#raise NotImplementedError()

def forward(self, y, z, src_mask):

"""
Args:
y of shape (max_tgt_seq_length, batch_size): LongTensor with the target
sequences.
z of shape (max_src_seq_length, batch_size, n_features): Encoded source
sequences (outputs of the
encoder).
src_mask of shape (max_src_seq_length, batch_size): Boolean tensor
indicating which elements of the
source sequences should be ignored.

Returns:
out of shape (max_seq_length, batch_size, tgt_vocab_size): Log-softmax
probabilities of the words
in the output sequences.

Notes:
* All intermediate signals should be of shape (max_seq_length,
batch_size, n_features).
* You need to create and use the subsequent mask in the decoder.
"""
# YOUR CODE HERE
# Create subsequent mask for decoder self-attention
tgt_mask = subsequent_mask([Link](0))
 # Embed target tokens
y = [Link](y) # Shape: (max_tgt_seq_length, batch_size,
n_features)

# Add positional encoding

y = self.positional_encoding(y)

# Pass through decoder blocks

for decoder_block in self.decoder_blocks:
y = decoder_block(y, z, src_mask, tgt_mask)

# Project to vocabulary size

output = self.output_layer(y)

# Apply log softmax

log_probs = self.log_softmax(output)

return log_probs
raise NotImplementedError()

def test_Decoder_shapes():
decoder = Decoder(tgt_vocab_size=10, n_blocks=1, n_features=16, n_heads=4,
n_hidden=64)

y = [Link]([
[SOS_token, SOS_token],
[ 3, 4],
[ 5, EOS_token],
[ 6, PADDING_VALUE],
[ 7, PADDING_VALUE],
]) # (max_seq_length, batch_size)

z = [Link](5, 2, 16) # (max_seq_length, batch_size, n_features)

src_mask = [Link]([
[0, 0],
[0, 0],
[0, 0],
[0, 1],
[0, 1],
], dtype=[Link]) # (max_seq_length, batch_size)

outputs = decoder(y, z, src_mask=src_mask)

assert [Link] == [Link]([5, 2, 10]), f"Wrong [Link]:
{[Link]}"
print('Success')

test_Decoder_shapes()

# Create the transformer model

n_features = 256
encoder = Encoder(src_vocab_size=trainset.input_lang.n_words, n_blocks=3,
n_features=n_features,
n_heads=16, n_hidden=1024)
decoder = Decoder(tgt_vocab_size=trainset.output_lang.n_words, n_blocks=3,
n_features=n_features,
n_heads=16, n_hidden=1024)

[Link](device)
[Link](device)

parameters = list([Link]()) + list([Link]())

adam = [Link](parameters, lr=0, betas=(0.9, 0.98), eps=1e-9)
optimizer = [Link](n_features, 0.4, 680, adam)
#optimizer = [Link](parameters, lr=0.001)

# Implement the training loop in this cell

if not skip_training:
# YOUR CODE HERE
num_epochs = 40
[Link]()
[Link]()

for epoch in range(num_epochs):

total_loss = 0
num_batches = 0

for src_seqs, src_mask, tgt_seqs in trainloader:

# Move data to device
src_seqs = src_seqs.to(device)
src_mask = src_mask.to(device)
tgt_seqs = tgt_seqs.to(device)

# Create target input and output sequences

# Input: [SOS, w1, w2, ...]
# Output: [w1, w2, ..., EOS]
tgt_input = tgt_seqs[:-1] # Remove last token
tgt_output = tgt_seqs[1:] # Remove first token (SOS)

# Zero gradients
optimizer.zero_grad()

# Forward pass through encoder

encoded = encoder(src_seqs, src_mask)

# Forward pass through decoder

log_probs = decoder(tgt_input, encoded, src_mask)

# Compute loss (ignore padding tokens)

# Reshape predictions and targets for NLLLoss
pred = log_probs.view(-1, log_probs.size(-1))
target = tgt_output.view(-1)

# Create padding mask for target

padding_mask = (target != PADDING_VALUE)

# Compute loss only on non-padded positions

loss = F.nll_loss(
pred[padding_mask],
target[padding_mask],
reduction='mean'
)

# Backward pass
[Link]()

# Update parameters
[Link]()
 total_loss += [Link]()
num_batches += 1

# Print epoch statistics

avg_loss = total_loss / num_batches
print(f'Epoch No. {epoch+1}/{num_epochs}, Average Loss: {avg_loss:.4f}')

# Optional early stopping if loss is small enough

if avg_loss < 0.1:
print(f'Reached target loss at epoch of {epoch+1}')
break
#raise NotImplementedError()

# Save the model to disk (the pth-files will be submitted automatically together
with your notebook)
# Set confirm=False if you do not want to be asked for confirmation before saving.
if not skip_training:
tools.save_model(encoder, '1_tr_encoder.pth', confirm=True)
tools.save_model(decoder, '1_tr_decoder.pth', confirm=True)

if skip_training:
encoder = Encoder(src_vocab_size=trainset.input_lang.n_words, n_blocks=3,
n_features=256, n_heads=16, n_hidden=1024)
tools.load_model(encoder, '1_tr_encoder.pth', device)

decoder = Decoder(tgt_vocab_size=trainset.output_lang.n_words, n_blocks=3,

n_features=256, n_heads=16, n_hidden=1024)
tools.load_model(decoder, '1_tr_decoder.pth', device)

if skip_training:
encoder = Encoder(src_vocab_size=trainset.input_lang.n_words, n_blocks=3,
n_features=256, n_heads=16, n_hidden=1024)
tools.load_model(encoder, '1_tr_encoder.pth', device)

decoder = Decoder(tgt_vocab_size=trainset.output_lang.n_words, n_blocks=3,

n_features=256, n_heads=16, n_hidden=1024)
tools.load_model(decoder, '1_tr_decoder.pth', device)

In the cell below, implement a function that converts an input sequence to an

output sequence using the trained transformer.

Notes:

Since we do not need to compute the gradients in the evaluation phase, we can speed
up the computations by using the statement with torch.no_grad():.
Please transfer the tensors to device inside this function.
We may deduct some points for an ineffecient implementation of translate().

Go through this entire thing and let me know once done.