import torch
import torch.nn as nn

class GPTDecoder(nn.Module):
    def __init__(self, nclass=16135, d_model=1024, max_pos=512, nhead=16, dim_feedforward=2048, vocab_size=16135, nhead_emb=8):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_embedding = nn.Embedding(max_pos, d_model)
        self.Heads = nn.ParameterList([])
        for _ in range(nhead_emb):
            self.Heads.append(nn.Parameter(torch.randn(d_model, d_model//nhead_emb)))
        self.LayerNorm_heads = nn.LayerNorm(d_model)
        self.GPT_MHA = nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.LayerNorm_mha = nn.LayerNorm(d_model)
        self.FFN_Layer1 = nn.Linear(d_model, dim_feedforward)
        self.LeakyReLU = nn.LeakyReLU(0.1)
        self.FFN_Layer2 = nn.Linear(dim_feedforward, d_model)
        self.LayerNorm_ffn = nn.LayerNorm(d_model)
        self.Linear = nn.Linear(d_model, nclass)
    def forward(self, x):
        x = self.embedding(x)
        x_n = torch.arange(x.size(1)).unsqueeze(0)
        x_n = self.pos_embedding(x_n)
        x = x + x_n
        x_origin = x
        x_heads = []
        for head in self.Heads:
            nx = x@head
            x_heads.append(nx)
        x = torch.cat(x_heads, dim=2)
        x = x + x_origin
        x_origin = x
        x = self.LayerNorm_heads(x)
        mask = nn.Transformer.generate_square_subsequent_mask(x.size(1))
        x, _ = self.GPT_MHA(x,x,x, attn_mask=mask)
        x = x+x_origin
        x_origin = x
        x = self.LayerNorm_mha(x)
        x = self.FFN_Layer1(x)
        x = self.LeakyReLU(x)
        x = self.FFN_Layer2(x)
        x = x+x_origin
        x = self.LayerNorm_ffn(x)
        x = self.Linear(x)
        return x