this is a pointer net model implements by me,
in one mini-batch,there are diff_len seqences ,and this sequences paddding to a max length
how to deal these seq to decoder, i don’t kown the padding effect on decoder in backward?
suppose that in decoder time t1, seq_1,seq_2 in a mini-batch,
seq_1 is stop in last time,and seq_2 need to calc
i select the seq_2 to calc loss, drop seq_1
is this right?
from typing import Tuple
import torch.nn.functional as F
import torch
import torch.nn as nn


class LSTMEncoder(nn.Module):
    def __init__(self, embedding_dim, hidden_size, num_layers=1, batch_first=True, bidirectional=False):
        super(LSTMEncoder, self).__init__()

        self.batch_first = batch_first
        self.bidirectional = bidirectional
        self.num_layers = num_layers
        self.embedding_dim = embedding_dim
        self.num_directions = 2 if self.bidirectional else 1
        self.hidden_size = int(hidden_size / self.num_directions)

        self.rnn = nn.LSTM(input_size=embedding_dim, hidden_size=self.hidden_size, num_layers=num_layers,
                           batch_first=batch_first, bidirectional=bidirectional)  # nn.LSTM(512, 256, 3)

    def forward(self, embedded_inputs, input_lengths,
                max_len):  # embedded_inputs:(64, 25, 512), input_lengths:(64,), max_len:25
        # Pack padded batch of sequences for RNN module
        packed = nn.utils.rnn.pack_padded_sequence(embedded_inputs, input_lengths.view(-1).cpu(),
                                                   batch_first=self.batch_first,
                                                   enforce_sorted=False)  # 记录需要mask的部分，有效len之外的部分不经过rnn
        # Forward pass through RNN
        outputs, hidden = self.rnn(packed)  # hidden:(6, 64, 256)
        # Unpack padding
        outputs, _ = nn.utils.rnn.pad_packed_sequence(outputs, batch_first=self.batch_first)  # outputs:(64, 22, 512)
        # Unpack函数只能padding至当前batch最大长度，需继续pad至全局最大长度
        extra_padding_size = max_len - outputs.shape[1]
        outputs = nn.functional.pad(outputs, [0, 0, 0, extra_padding_size, 0, 0], mode="constant", value=0)  # 第2维进行下填充

        # Return output and final hidden state
        # if self.bidirectional:
        #     # Optionally, Sum bidirectional RNN outputs
        #     # outputs = outputs[:, :, :self.hidden_size] + outputs[:, :, self.hidden_size:]
        #     outputs = torch.cat((outputs[:, :, :self.hidden_size], outputs[:, :, self.hidden_size:]), dim=2)
        # batch_size = embedded_inputs.size(0)
        # h_n, c_n = hidden
        # h_n = h_n.view(self.num_layers, self.num_directions, batch_size, self.hidden_size)
        # c_n = c_n.view(self.num_layers, self.num_directions, batch_size, self.hidden_size)
        # if self.bidirectional:
        #     f = (h_n[-1, 0, :, :].squeeze(), c_n[-1, 0, :, :].squeeze())  # f: tuple:2, ((64, 256), (64, 256))
        #     b = (h_n[-1, 1, :, :].squeeze(), c_n[-1, 1, :, :].squeeze())  # b: tuple:2, ((64, 256), (64, 256))
        #     hidden = (torch.cat((f[0], b[0]), dim=1), torch.cat((f[1], b[1]), dim=1))  # tuple:2, ((64, 512), (64, 512))
        # else:
        #     hidden = (h_n[-1, 0, :, :].squeeze(), c_n[-1, 0, :, :].squeeze())

        return outputs, hidden


class Attention(nn.Module):
    def __init__(self, hidden_size, units):
        super(Attention, self).__init__()
        self.W1 = nn.Linear(hidden_size, units, bias=False)
        self.W2 = nn.Linear(hidden_size, units, bias=False)
        self.V = nn.Linear(units, 1, bias=False)

    def forward(self,
                encoder_out: torch.Tensor,
                decoder_hidden: torch.Tensor):
        # encoder_out: (BATCH, ARRAY_LEN, HIDDEN_SIZE)
        # decoder_hidden: (BATCH, HIDDEN_SIZE)

        # Add time axis to decoder hidden state
        # in order to make operations compatible with encoder_out
        # decoder_hidden_time: (BATCH, 1, HIDDEN_SIZE)
        decoder_hidden_time = decoder_hidden.unsqueeze(1)

        # uj: (BATCH, ARRAY_LEN, ATTENTION_UNITS)
        # Note: we can add the both linear outputs thanks to broadcasting
        uj = self.W1(encoder_out) + self.W2(decoder_hidden_time)
        uj = torch.tanh(uj)

        # uj: (BATCH, ARRAY_LEN, 1)
        uj = self.V(uj)

        # Attention mask over inputs
        # aj: (BATCH, ARRAY_LEN, 1)
        # aj = F.softmax(uj, dim=1)
        #
        # # di_prime: (BATCH, HIDDEN_SIZE)
        # di_prime = aj * encoder_out
        #
        # di_prime = di_prime.sum(1)

        return uj.squeeze(-1), uj.squeeze(-1)


# Attention weights over inputs: torch.Size([4, 6])

class RawDecoder(nn.Module):
    def __init__(self,
                 hidden_size: int,
                 attention_units: int = 10):
        super(RawDecoder, self).__init__()
        self.lstm = nn.LSTM(hidden_size + 1, hidden_size, batch_first=True)
        self.attention = Attention(hidden_size, attention_units)

    def forward(self,
                x: torch.Tensor,
                hidden: Tuple[torch.Tensor],
                encoder_out: torch.Tensor):
        # x: (BATCH, 1, 1)
        # hidden: (1, BATCH, HIDDEN_SIZE)
        # encoder_out: (BATCH, ARRAY_LEN, HIDDEN_SIZE)

        #  ht_ last timqe hidden state
        ht = hidden[0][0]  # ht: (BATCH, HIDDEN_SIZE) h_n

        # di: Attention aware hidden state -> (BATCH, HIDDEN_SIZE)
        di, att_w = self.attention(encoder_out, ht)

        # Append attention aware hidden state to our input
        # x: (BATCH, 1, 1 + HIDDEN_SIZE)
        x = torch.cat([di.unsqueeze(1), x], dim=2)

        # Generate the hidden state for next timestep
        _, hidden = self.lstm(x, hidden)
        return hidden, att_w


class AttDecoder(nn.Module):
    def __init__(self,
                 hidden_size: int, emb_size: int,
                 attention_units: int = 10):
        super(AttDecoder, self).__init__()
        self.lstm = nn.LSTM(emb_size, hidden_size, batch_first=True)
        self.attention = Attention(hidden_size, attention_units)

    def forward(self,
                x: torch.Tensor,
                hidden: Tuple[torch.Tensor],
                encoder_out: torch.Tensor):
        _, hidden = self.lstm(x, hidden)
        ht = hidden[0][0]  # ht: (BATCH, HIDDEN_SIZE) h_n

        # di: Attention aware hidden state -> (BATCH, HIDDEN_SIZE)
        di, att_w = self.attention(encoder_out, ht)

        # Append attention aware hidden state to our input
        # x: (BATCH, 1, 1 + HIDDEN_SIZE)
        # Generate the hidden state for next timestep

        return hidden, att_w


class PointerNetwork(nn.Module):
    def __init__(self,
                 encoder: nn.Module,
                 decoder: nn.Module, max_len: int = 25):
        super(PointerNetwork, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.max_len = max_len

    def forward(self,
                x: torch.Tensor,
                y: torch.Tensor, x_start: torch.Tensor,
                batch_len: torch.Tensor,
                teacher_force_ratio=.5):
        max_len = torch.tensor([self.max_len]).long()
        encoder_in = x  # Batch L Hid[20, 25, 64])
        # Batch,Len,Hidden
        out, hs = self.encoder(encoder_in, batch_len, max_len)
        loss = 0
        # Len,Batch
        outputs = torch.zeros(out.size(1), out.size(0), dtype=torch.long)
        # First decoder input is always 0
        # dec_in: (BATCH, 1, 1)
        # First Init as EoS
        dec_in = x_start

        for t in range(out.size(1)):
            # 获取解码结果
            hs, att_w = self.decoder(dec_in, hs, out)
            # 观测值
            predictions = F.softmax(att_w, dim=1).argmax(1)

            # Pick next index
            # If teacher force the next element will we the ground truth
            # otherwise will be the predicted value at current timestep
            # 是否指导学习
            import random
            teacher_force = random.random() < teacher_force_ratio
            idx = y[:, t] if teacher_force else predictions
            # 获取下一个元素预测结果
            dec_in = torch.stack([x[b, idx[b].item()] for b in range(x.size(0))])
            dec_in = dec_in.view(out.size(0), 1, -1).type(torch.float)

            # 损失函数
            loss += F.cross_entropy(att_w, y[:, t])
            outputs[t] = predictions
        # Weight losses, so every element in the batch
        # has the same 'importance'
        batch_loss = loss / y.size(0)

        return outputs, batch_loss


def run_pnn():
    '''
    mvp version code of pnn

    '''

    emb_dim = 8
    len_seq = 24
    hidden_size = 256
    attention_units = 256
    decoder = AttDecoder(hidden_size, emb_dim, attention_units=attention_units)
    encoder = LSTMEncoder(emb_dim, hidden_size)
    lr = 0.000001
    pnn = PointerNetwork(encoder, decoder)
    opt = torch.optim.Adam(pnn.parameters(), lr=lr)

    batch_size = 20
    train_data = get_raw_data(get_train_path())
    dataset = RawDataSet(train_data)
    dataloader = DataLoader(dataset=dataset, batch_size=batch_size)
    for idx, item in enumerate(dataloader):
        [x, y, seq_lens, starts] = item
        starts = starts.view(batch_size, 1, -1).contiguous()
        seq_lens = seq_lens.view(batch_size, 1).contiguous()
        pnn(x, y, starts, seq_lens)
    x = torch.rand(batch_size, len_seq, emb_dim)
    y = torch.randint(25, (batch_size, len_seq))
    start = torch.rand(batch_size, 1, emb_dim)
    batch_len = torch.zeros(batch_size, 1)
    batch_len = torch.fill_(batch_len, len_seq).long()
    pnn(x, y, start, batch_len)

    pass


if __name__ == '__main__':
    run_pnn()