# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/01_seq2seq.ipynb (unless otherwise specified).

__all__ = ['Encoder', 'NewDecoder', 'Seq2Seq']

# Cell
from torch import nn
from torch import optim
import torch
import torch.nn.functional as F

from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

# Cell
class Encoder(nn.Module):
    def __init__(self, input_size, embedding_size, hidden_size, num_layers=2, p=0.1):
        super(Encoder, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers

        self.dropout = nn.Dropout(p)
        self.embedding = nn.Embedding(input_size, embedding_size)
        self.rnn = nn.LSTM(embedding_size, hidden_size, num_layers, dropout=p,batch_first=False)

    def forward(self, x, x_len):
        # x shape (seq_length, N)
        embedding = self.dropout(self.embedding(x))
        # embedding shape : (seq_length, N, embedding_size)

        x_packed = pack_padded_sequence(embedding, x_len.cpu(), batch_first=False, enforce_sorted=False)
        output_packed, (hidden,cell) = self.rnn(x_packed)

        # irrelevant because we are interested only in hidden state
        #output_padded, output_lengths = pad_packed_sequence(output_packed, batch_first=True)

        # output is irrelevant, context vector is important
        return hidden,cell


# Cell
class NewDecoder(nn.Module):
    def __init__(self, hidden_size, embedding_size, output_size, n_layers=1, dropout_p=0.1):
        super(NewDecoder, self).__init__()

        # Define parameters
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.n_layers =n_layers
        self.dropout_p = dropout_p

        # Define layers
        self.embedding = nn.Embedding(output_size, embedding_size)
        self.dropout=nn.Dropout(dropout_p)
        self.rnn = nn.LSTM(embedding_size, hidden_size, n_layers, dropout=dropout_p, batch_first=False)
        self.out = nn.Linear(hidden_size, output_size)

    def forward(self, word_input, last_hidden, encoder_outputs):
        # Note that we will only be running forward for a single decoder time step, but will
        # use all encoder outputs
        word_input = word_input.unsqueeze(0)
        # we are not using encoder_outputs here
        word_embedded = self.embedding(word_input) # 1 X B
        word_embedded = self.dropout(word_embedded)  # 1 X B X emb_length

        # Combine embedded input word and hidden vector, run through RNN
        output, hidden = self.rnn(word_embedded, last_hidden)  # 1 X B X hidden
        predictions = self.out(output)     # 1, B, out
        #output = F.log_softmax(predictions)

        return predictions, hidden

# Cell
import random

import pytorch_lightning as pl
import pytorch_lightning.metrics.functional as plfunc
from pytorch_lightning.loggers import TensorBoardLogger

# Cell
class Seq2Seq(pl.LightningModule):
    """ Encoder decoder pytorch lightning module for training seq2seq model with teacher forcing
    Module try to learn mapping from one sequence to another
    """

    @staticmethod
    def add_model_specific_args(parent_parser):
        parser = ArgumentParser(parents=[parent_parser], add_help=False)
        parser.add_argument("--emb_dim", type=int, default=32)
        parser.add_argument('--hidden_dim', type=int, default=64)
        parser.add_argument('--dropout', type=float, default=0.1)

        return parser

    def __init__(self,
                input_vocab_size,
                output_vocab_size,
                padding_index = 0,
                emb_dim = 8,
                hidden_dim=32,
                dropout=0.1,
                max_length=20,
                **kwargs):
        super().__init__()


        # dynamic, based on tokenizer vocab size defined in datamodule
        self.input_dim = input_vocab_size
        self.output_dim = output_vocab_size

        self.enc_emb_dim = emb_dim
        self.dec_emb_dim = emb_dim

        self.enc_hid_dim = hidden_dim
        self.dec_hid_dim = hidden_dim

        self.enc_dropout = dropout
        self.dec_dropout = dropout

        self.pad_idx = padding_index
        self.num_layers = 2
        self.max_length =10
        self.save_hyperparameters()

        self.max_epochs= kwargs.get('max_epochs',5)

        self.learning_rate = 0.0005

        self._loss = nn.CrossEntropyLoss(ignore_index=self.pad_idx)

        self.encoder = Encoder(
            self.input_dim,
            self.enc_emb_dim,
            self.enc_hid_dim,
            self.num_layers,
            self.enc_dropout
        )

        self.decoder = NewDecoder(
            self.enc_hid_dim,
            self.dec_emb_dim,
            self.output_dim,
            self.num_layers,
            self.dec_dropout
        )

        self._init_weights()

    def _init_weights(self):
        for name, param in self.named_parameters():
            if "weight" in name:
                nn.init.normal_(param.data, mean=0, std=0.01)

            else:
                nn.init.constant_(param.data, 0)

    def create_mask(self, src):
        mask = (src != self.pad_idx).permute(1, 0)
        return mask



    def forward(self, src_seq, source_len, trg_seq, teacher_force_ratio=0.5):
        """
        teacher_force_ratio is used to help in decoding.
        In starting, original input token will be sent as input token
        """
        source = src_seq.transpose(0, 1)
        target_len = self.max_length

        if trg_seq is not None:
            target = trg_seq.transpose(0, 1)
            target_len = target.shape[0]


        batch_size = source.shape[1]

        target_vocab_size = self.output_dim

        outputs = torch.zeros(target_len, batch_size, target_vocab_size).to(self.device)

        encoder_hidden = self.encoder(source, source_len)

        # mask = [batch_size, src len]
        # without sos token at the beginning and eos token at the end

        #x = target[0,:]
        decoder_input = torch.ones(batch_size).long().to(self.device)
        decoder_hidden = encoder_hidden
        encoder_outputs = None

        for t in range(target_len):
            decoder_output, decoder_hidden = self.decoder(decoder_input, decoder_hidden, encoder_outputs)
            outputs[t] = decoder_output

            #(N, english_vocab_size)
            #best_guess = output.argmax(1)
            topv, topi = decoder_output.topk(1)
            decoder_input = topi.squeeze().detach()

            decoder_input = target[t] if random.random() < teacher_force_ratio and target is not None else decoder_input

        return outputs

    def loss(self, logits, target):
        return self._loss(logits, target)

    def configure_optimizers(self):
        optimizer = optim.AdamW(self.parameters(), lr=self.learning_rate)

        lr_scheduler = {
            'scheduler': optim.lr_scheduler.OneCycleLR(
                optimizer,
                max_lr = self.learning_rate,
                steps_per_epoch = 3379,
                epochs=self.max_epochs,
                anneal_strategy='linear',
                final_div_factor=1000,
                pct_start = 0.01
            ),
            "name": "learning_rate",
            "interval":"step",
            "frequency": 1
        }

        return [optimizer],[lr_scheduler]

    def training_step(self, batch, batch_idx):
        src_seq, trg_seq, src_lengths = batch['src'],batch['trg'], batch['src_len']

        output = self.forward(src_seq, src_lengths,trg_seq)

        # do not know if this is a problem, loss will be computed with sos token
        # without sos token at the beginning and eos token at the end

        output = output.view(-1, self.output_dim)


        trg_seq = trg_seq.transpose(0, 1)
        trg = trg_seq.reshape(-1)

        loss = self.loss(output, trg)

        self.log('train_loss',loss.item(),
                on_step = True,
                on_epoch=True,
                prog_bar = True,
                logger=True)

        return loss

    def validation_step(self, batch,batch_idx):
        """ validation is in eval model so we do not have to use placeholder input sensors"""
        src_seq, trg_seq, src_lengths = batch['src'],batch['trg'], batch['src_len']

        outputs = self.forward(src_seq, src_lengths, trg_seq, 0)

        logits = outputs[1:].view(-1, self.output_dim)

        trg = trg_seq[1:].reshape(-1)

        loss = self.loss(logits, trg)

        pred_seq = outputs[1:].argmax(2)   # seq_len*batch_size*vocab_size -> seq_len * batch_size

        # change layout: sesq_len * batch_size -> batch_size * seq_len
        pred_seq = pred_seq.T

        # change layout: seq_len * batch_size -> batch_size * seq_len
        trg_batch = trg_seq[1:].T

        # compare list of predicted ids for all sequences in a batch to targets
        acc = plfunc.accuracy(pred_seq.reshape(-1), trg_batch.reshape(-1))

        # need to cast to list of predicted sequences ( as list of token ids ) [ seq_tok1, seqtok2]
        predicted_ids - pred_seq.tolist()

        # need to add additional dim to each target reference sequence in order to
        # conver to format needed by blue_score_func
        # [seq1=[[reference1],[reference2]], seq2=[reference1]]
        target_ids = torch.unsqueeze(trg_batch, 1).tolist()

        bleu_score - plfunc.nlp.bleu_score(predicted_ids, target_ids, n_gram=3).to(self.device)

        self.log(
        'val_loss',
        loss,
        on_step=False,
        on_epoch=True,
        prog_bar=True,
        logger=True,
        sync_dist=True)

        self.log(
            "val_acc",
            acc,
            on_step=False,
            on_epoch=True,
            prog_bar=True,
            logger=True,
            sync_dist=True
        )

        self.log(
            "val_bleu_idx",
            bleu_score,
            on_step=False,
            on_epoch=True,
            prog_bar=True,
            logger=True,
            sync_dist=True
        )

        return loss, acc, bleu_score

[((20, 23, 20, 36), 'torch.nn.Dropout', 'nn.Dropout', ({(20, 34, 20, 35): 'p'}, {}), '(p)', False, 'from torch import nn\n'), ((21, 25, 21, 65), 'torch.nn.Embedding', 'nn.Embedding', ({(21, 38, 21, 48): 'input_size', (21, 50, 21, 64): 'embedding_size'}, {}), '(input_size, embedding_size)', False, 'from torch import nn\n'), ((22, 19, 22, 96), 'torch.nn.LSTM', 'nn.LSTM', (), '', False, 'from torch import nn\n'), ((51, 25, 51, 66), 'torch.nn.Embedding', 'nn.Embedding', ({(51, 38, 51, 49): 'output_size', (51, 51, 51, 65): 'embedding_size'}, {}), '(output_size, embedding_size)', False, 'from torch import nn\n'), ((52, 21, 52, 42), 'torch.nn.Dropout', 'nn.Dropout', ({(52, 32, 52, 41): 'dropout_p'}, {}), '(dropout_p)', False, 'from torch import nn\n'), ((53, 19, 53, 103), 'torch.nn.LSTM', 'nn.LSTM', (), '', False, 'from torch import nn\n'), ((54, 19, 54, 54), 'torch.nn.Linear', 'nn.Linear', ({(54, 29, 54, 40): 'hidden_size', (54, 42, 54, 53): 'output_size'}, {}), '(hidden_size, output_size)', False, 'from torch import nn\n'), ((127, 21, 127, 67), 'torch.nn.CrossEntropyLoss', 'nn.CrossEntropyLoss', (), '', False, 'from torch import nn\n'), ((210, 25, 218, 13), 'torch.optim.lr_scheduler.OneCycleLR', 'optim.lr_scheduler.OneCycleLR', (), '', False, 'from torch import optim\n'), ((150, 16, 150, 61), 'torch.nn.init.normal_', 'nn.init.normal_', (), '', False, 'from torch import nn\n'), ((153, 16, 153, 48), 'torch.nn.init.constant_', 'nn.init.constant_', ({(153, 34, 153, 44): 'param.data', (153, 46, 153, 47): '(0)'}, {}), '(param.data, 0)', False, 'from torch import nn\n'), ((178, 18, 178, 72), 'torch.zeros', 'torch.zeros', ({(178, 30, 178, 40): 'target_len', (178, 42, 178, 52): 'batch_size', (178, 54, 178, 71): 'target_vocab_size'}, {}), '(target_len, batch_size, target_vocab_size)', False, 'import torch\n'), ((279, 21, 279, 50), 'torch.unsqueeze', 'torch.unsqueeze', ({(279, 37, 279, 46): 'trg_batch', (279, 48, 279, 49): '1'}, {}), '(trg_batch, 1)', False, 'import torch\n'), ((281, 21, 281, 79), 'pytorch_lightning.metrics.functional.nlp.bleu_score', 'plfunc.nlp.bleu_score', (), '', True, 'import pytorch_lightning.metrics.functional as plfunc\n'), ((186, 24, 186, 46), 'torch.ones', 'torch.ones', ({(186, 35, 186, 45): 'batch_size'}, {}), '(batch_size)', False, 'import torch\n'), ((199, 41, 199, 56), 'random.random', 'random.random', ({}, {}), '()', False, 'import random\n')]