model.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from blocks.layers import Encoder, Decoder


class Transformer(nn.Module):

    def __init__(self, params):

        super(Transformer, self).__init__()
        self.params = params

        self.input_embed = nn.Embedding(num_embeddings=params.in_embed_size, embedding_dim=params.d_model)
        self.output_embed = nn.Embedding(num_embeddings=params.out_embed_size, embedding_dim=params.d_model)

        self.encoder_layers = nn.ModuleList([
            Encoder(params=params) for _ in range(params.n_stacks)
        ])
        self.decoder_layers = nn.ModuleList([
            Decoder(params=params) for _ in range(params.n_stacks)
        ])

    def forward(self, src_input, tgt_output):
        
        encoder_output = src_input
        decoder_output = tgt_output

        for encoder in self.encoder_layers:

            encoder_output = encoder.forward(encoder_output)

        for decoder in self.decoder_layers:

            decoder_output = decoder.forward(encoder_output,decoder_output)

        return decoder_output