Seq2seq Transformer Implementation

class Seq2SeqTransformer(nn.Module):
def __init__(self, num_encoder_layers: int, num_decoder_layers: int, 
                emb_size: int, src_vocab_size: int, tgt_vocab_size: int,
                dim_feedforward:int = 512, dropout:float = 0.1):
    super().__init__()
    encoder_layer = TransformerEncoderLayer(d_model=emb_size, nhead=NHEAD,dim_feedforward=dim_feedforward,batch_first=True)
    self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
    decoder_layer = TransformerDecoderLayer(d_model=emb_size, nhead=NHEAD, dim_feedforward=dim_feedforward,batch_first=True)
    self.transformer_decoder = TransformerDecoder(decoder_layer,num_layers=num_decoder_layers)
    self.generator = nn.Linear(emb_size, tgt_vocab_size)
    self.src_tok_emb = TokenEmbedding(src_vocab_size, emb_size)
    self.tgt_tok_emb = TokenEmbedding(tgt_vocab_size, emb_size)
    self.positional_encoding = PositionalEncoding(emb_size, dropout=dropout)
    def forward(self, src: Tensor, trg: Tensor, src_mask: Tensor, tgt_mask: Tensor, src_padding_mask: Tensor, tgt_padding_mask: Tensor, memory_key_padding_mask: Tensor):
        src_emb = self.positional_encoding(self.src_tok_emb(src))
        tgt_emb = self.positional_encoding(self.tgt_tok_emb(trg))
        memory = self.transformer_encoder(src_emb, src_mask, src_padding_mask)
        outs = self.transformer_decoder(tgt_emb, memory,tgt_mask, None, tgt_padding_mask, memory_key_padding_mask)
    return self.generator(outs)

    def encode(self, src: Tensor, src_mask: Tensor):
        return self.transformer_encoder(self.positional_encoding(self.src_tok_emb(src)), src_mask)
    def decode(self, tgt: Tensor, memory: Tensor, tgt_mask: Tensor):
        return self.transformer_decoder(self.positional_encoding(self.tgt_tok_emb(tgt)), memory, tgt_mask)

class PositionalEncoding(nn.Module):
    def __init__(self,emb_size,dropout,max_len=5000):
        super().__init__()
        den = torch.exp(-torch.arange(0,emb_size,2) * math.log(10000)/emb_size)
        pos = torch.arange(0,max_len).reshape(max_len,1)
        pos_embedding = torch.zeros((max_len,emb_size))
        pos_embedding[:,0::2] = torch.sin(pos * den)
        pos_embedding[:,1::2] = torch.cos(pos * den)
        pos_embedding = pos_embedding.unsqueeze(-2)
        self.dropout = nn.Dropout(dropout)
        self.register_buffer('pos_embedding',pos_embedding)
    def forward(self,token_embedding):
        return self.dropout(token_embedding + self.pos_embedding[:token_embedding.size(0),:])

# feature lifting
class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size: int, emb_size):
        super(TokenEmbedding, self).__init__()
        self.embedding = nn.Embedding(vocab_size, emb_size)
        self.emb_size = emb_size
    def forward(self, tokens: Tensor):
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

def generate_square_subsequent_mask(sz):
    mask = (torch.triu(torch.ones((sz,sz),device='cuda')) == 1).transpose(0,1)
    mask = mask.float().masked_fill(mask==0,float('-inf')).masked_fill(mask==1,float(0.0))
    return mask

def create_mask(src,tgt):
    src_seq_len = src.shape[1]
    tgt_seq_len = tgt.shape[1]
    tgt_mask = generate_square_subsequent_mask(tgt_seq_len).cuda()
    src_mask = torch.zeros((src_seq_len,src_seq_len),device='cuda').type(torch.bool)
    src_padding_mask = (src == 1).cuda()
    tgt_padding_mask = (tgt == 1).cuda()
    return src_mask, tgt_mask, src_padding_mask,tgt_padding_mask

SRC_VOCAB_SIZE = len(de_vocab)
TGT_VOCAB_SIZE = len(en_vocab)
EMB_SIZE = 512
NHEAD = 8
FFN_HID_DIM = 512
BATCH_SIZE = 128
NUM_ENCODER_LAYERS = 3
NUM_DECODER_LAYERS = 3
NUM_EPOCHS = 50
DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# create model
transformer = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS, EMB_SIZE, SRC_VOCAB_SIZE, TGT_VOCAB_SIZE, FFN_HID_DIM)
transformer = transformer.cuda()

# Xavier initialization
for p in transformer.parameters():
    if p.dim() > 1:
        nn.init.xavier_uniform_(p)

# criterion
loss_fn = torch.nn.CrossEntropyLoss(ignore_index=1)

# optimization routine
optimizer = torch.optim.Adam(transformer.parameters(),lr=0.0001,betas=(0.9,0.98), eps=1e-9)

# training function
def train_epoch(model, train_dl, optimizer):
    model.train()
    losses = 0
    for idx, (src, tgt) in enumerate(train_dl):
        src = src.cuda()
        tgt = tgt.cuda()
        tgt_input = tgt[:,:-1]
        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)
        pred = model(src, tgt_input, src_mask, tgt_mask, src_padding_mask, tgt_padding_mask, src_padding_mask)
        optimizer.zero_grad()
        tgt_out = tgt[:, 1:]
        # evaluate training loss
        loss = loss_fn(pred.reshape(-1, pred.shape[-1]), tgt_out.reshape(-1))
        # backpropagte
        loss.backward()
        optimizer.step()
        losses += loss.item()

    return losses / len(train_dl)

def evaluate(model,val_dl):
    model.eval()
    losses = 0
    for idx, (src,tgt) in enumerate(val_dl):
        src = src.cuda()
        tgt = tgt.cuda()
        tgt_input =tgt[:,:-1]
        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)
        logits = model(src, tgt_input, src_mask,
        tgt_mask, src_padding_mask,
        tgt_padding_mask, src_padding_mask)
        optimizer.zero_grad()
        tgt_out = tgt[:,1:]
        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))
        losses += loss.item()
    return losses/len(val_dl)

# train model
num_epochs = 20-1
for epoch in range(num_epochs):
    start_time = time.time()
    train_loss = train_epoch(transformer,train_dl,optimizer)
    end_time = time.time()
    val_loss = evaluate(transformer,val_dl)

    # report
    print(f'>>>> Epoch: {epoch}, Train loss: {train_loss:.3f}, Val loss: {val_loss:.3f}, Epoch time = {(end_time - start_time):.3f} sec ')

def greedy_decode(model, src, src_mask, max_len, start_symbol):
    src = src.cuda()
    src_mask = src_mask.cuda()
    memory = model.encode(src, src_mask)
    ys = torch.ones(1, 1).fill_(2).type(torch.long).cuda()
    for i in range(max_len-1):
    memory = memory.cuda()
    memory_mask = torch.zeros(ys.shape[1], memory.shape[1]).cuda().type(torch.bool)
    tgt_mask = (generate_square_subsequent_mask(ys.size(1)).type(torch.bool)).cuda()
    out = model.decode(ys, memory, tgt_mask).squeeze(0)[-1]
    prob = model.generator(out)
    next_word = prob.argmax()
    next_word = next_word.item()
    ys = torch.cat([ys,torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=1)
    if next_word == 3:
        break
    return ys

def translate(model, src, src_vocab, tgt_vocab, src_tokenizer):
    model.eval()
    tokens = [2]+de_vocab(de_tokenizer(src))+[3]
    num_tokens = len(tokens)
    src = (torch.LongTensor(tokens).reshape(1,num_tokens))
    src_mask = (torch.zeros(num_tokens, num_tokens)).type(torch.bool)
    tgt_tokens = greedy_decode(model, src, src_mask, max_len=num_tokens + 5, start_symbol=2).flatten()
    return " ".join([tgt_vocab.get_itos()[tok] for tok in tgt_tokens]).replace("<bos>", "").replace("<eos>", "")

output = translate(transformer, "Eine Gruppe von Menschen steht vor einem Flughafen .", de_vocab, en_vocab, de_tokenizer)
print(output)