# This Python 3 environment comes with many helpful analytics libraries installed
# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python
# For example, here's several helpful packages to load

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

# Input data files are available in the read-only "../input/" directory
# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory

import os
for dirname, _, filenames in os.walk('/kaggle/input'):
    for filename in filenames:
        print(os.path.join(dirname, filename))

# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using "Save & Run All" 
# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session

# Use the kagglehub client library to attach Kaggle resources like competitions, datasets, and models to your session
# Learn more about kagglehub: https://github.com/Kaggle/kagglehub/blob/main/README.md

import kagglehub
# kagglehub.dataset_download('<owner>/<dataset-slug>')


import math

import torch
import torch.nn as nn

from safetensors.torch import save_file, load_file


class MultiHeadAttention(nn.Module):
    def __init__(self, heads, emb_dim):
        super().__init__()

        self.num_heads=heads
        self.emb_dim=emb_dim
        self.head_dim=emb_dim//heads

        assert emb_dim % heads == 0

        self.Wq=nn.Linear(emb_dim, emb_dim)
        self.Wk=nn.Linear(emb_dim, emb_dim)
        self.Wv=nn.Linear(emb_dim, emb_dim)
        self.Wo=nn.Linear(emb_dim, emb_dim)
        
    def forward(self, x):
        b, sl, ed = x.shape
        # (batch, seq_len, emb_dim)

        Q=self.Wq(x)
        K=self.Wk(x)
        V=self.Wv(x)

        Q=Q.view(b, sl, self.num_heads, self.head_dim)
        K=K.view(b, sl, self.num_heads, self.head_dim)
        V=V.view(b, sl, self.num_heads, self.head_dim)
        # (batch, seq_len, heads, head_dim)

        Q=Q.transpose(1, 2)        
        K=K.transpose(1, 2)
        V=V.transpose(1, 2)
        # (batch, heads, seq_len, head_dim)

        attention=torch.matmul(torch.softmax((torch.matmul(Q, K.transpose(-2, -1))/math.sqrt(self.head_dim)), dim=-1), V)
        # (batch, heads, seq_len, head_dim)
        
        original = attention.transpose(1, 2)
        # (batch, seq_len, heads, head_dim)

        output = self.Wo(original.contiguous().view(b, sl, self.emb_dim))
        # (batch, seq_len, emb_dim)

        return output        


class PositionwiseFeedForward(nn.Module):
    def __init__(self, emb_dim, nn_dim):
        super().__init__()
        
        self.linear1=nn.Linear(emb_dim, nn_dim)
        self.relu=nn.ReLU()
        self.linear2=nn.Linear(nn_dim, emb_dim)

    def forward(self, x):
        x=self.linear1(x)
        # (batch, seq_len, nn_dim)
        
        x=self.relu(x)
        
        x=self.linear2(x)
        # (batch, seq_len, emb_dim)
        
        return x


class Encoder(nn.Module):
    def __init__(self, heads, emb_dim, nn_dim):
        super().__init__()

        self.mha=MultiHeadAttention(heads, emb_dim)
        self.ffn=PositionwiseFeedForward(emb_dim, nn_dim)
        self.norm1=nn.LayerNorm(emb_dim)
        self.norm2=nn.LayerNorm(emb_dim)

    def forward(self, x):
        attention=self.mha(x)
        # (batch, seq_len, emb_dim)
        
        x=self.norm1(x+attention)
        
        ffn_output=self.ffn(x)
        # (batch, seq_len, emb_dim)
        
        x=self.norm2(x+ffn_output)
        # (batch, seq_len, emb_dim)
        
        return x


x = torch.randn(64, 32, 512)
encoder = Encoder(8, 512, 2048)
output = encoder(x)
output.shape


class MaskedMultiHeadAttention(nn.Module):
    def __init__(self, heads, emb_dim):
        super().__init__()

        self.num_heads=heads
        self.emb_dim=emb_dim
        self.head_dim=emb_dim//heads

        assert emb_dim % heads == 0

        self.Wq=nn.Linear(emb_dim, emb_dim)
        self.Wk=nn.Linear(emb_dim, emb_dim)
        self.Wv=nn.Linear(emb_dim, emb_dim)
        self.Wo=nn.Linear(emb_dim, emb_dim)

    def forward(self, x):
        b, sl, ed=x.shape
        # (batch, seq_len, emb_dim)

        Q=self.Wq(x)
        K=self.Wk(x)
        V=self.Wv(x)

        Q=Q.view(b, sl, self.num_heads, self.head_dim)
        K=K.view(b, sl, self.num_heads, self.head_dim)
        V=V.view(b, sl, self.num_heads, self.head_dim)
        # (batch, seq_len, heads, head_dim)

        Q=Q.transpose(1, 2)
        K=K.transpose(1, 2)
        V=V.transpose(1, 2)
        # (batch, heads, seq_len, head_dim)

        # masked
        mask=torch.tril(torch.ones(sl, sl))
        mask=mask.masked_fill(mask==0, float('-inf'))
        mask=mask.masked_fill(mask==1, 0.0)
        
        attention=torch.matmul((torch.softmax(((torch.matmul(Q, K.transpose(-2, -1))/math.sqrt(self.head_dim))+mask), dim=-1)), V)
        # (batch, heads, seq_len, head_dim)

        original=attention.transpose(1, 2)
        # (batch, seq_len, heads, head_dim)

        output=self.Wo(original.contiguous().view(b, sl, self.emb_dim))
        # (batch, seq_len, emb_dim)

        return output


class CrossMultiHeadAttention(nn.Module):
    def __init__(self, heads, emb_dim):
        super().__init__()

        assert emb_dim%heads==0

        self.num_heads=heads
        self.emb_dim=emb_dim
        self.head_dim=emb_dim//heads

        self.Wq=nn.Linear(emb_dim, emb_dim)
        self.Wk=nn.Linear(emb_dim, emb_dim)
        self.Wv=nn.Linear(emb_dim, emb_dim)
        self.Wo=nn.Linear(emb_dim, emb_dim)

    def forward(self, decoder_input, encoder_output):
        b, tsl, _ = decoder_input.shape
        # (batch, tgt_seq_len, emb_dim)
        _, ssl, _ = encoder_output.shape
        # (batch, src_seq_len, emb_dim)
        
        Q=self.Wq(decoder_input)
        K=self.Wk(encoder_output)
        V=self.Wv(encoder_output)

        Q=Q.view(b, tsl, self.num_heads, self.head_dim)
        K=K.view(b, ssl, self.num_heads, self.head_dim)
        V=V.view(b, ssl, self.num_heads, self.head_dim)
        # (batch, seq_len, heads, head_dim)

        Q=Q.transpose(2, 1)
        K=K.transpose(2, 1)
        V=V.transpose(2, 1)
        # (batch, heads, seq_len, head_dim)

        attention=torch.matmul((torch.softmax((torch.matmul(Q, K.transpose(-2, -1)))/math.sqrt(self.head_dim), dim=-1)), V)
        # (batch, heads, seq_len, head_dim)

        original=attention.transpose(1, 2)
        # (batch, seq_len, heads, head_dim)

        output=self.Wo(original.contiguous().view(b, tsl, self.emb_dim))
        # (batch, tgt_seq_len, emb_dim)

        return output


class Decoder(nn.Module):
    def __init__(self, heads, emb_dim, nn_dim):
        super().__init__()

        self.mmha=MaskedMultiHeadAttention(heads, emb_dim)
        self.cmha=CrossMultiHeadAttention(heads, emb_dim)
        self.ffn=PositionwiseFeedForward(emb_dim, nn_dim)

        self.norm1=nn.LayerNorm(emb_dim)
        self.norm2=nn.LayerNorm(emb_dim)
        self.norm3=nn.LayerNorm(emb_dim)

    def forward(self, decoder_x, encoder_output):
        masked_output=self.mmha(decoder_x)
        # (batch, tgt_seq_len, emb_dim)
        
        x=self.norm1(decoder_x+masked_output)

        cross_output=self.cmha(x, encoder_output)
        # (batch, tgt_seq_len, emb_dim)

        x=self.norm2(x+cross_output)

        ffn_output=self.ffn(x)
        # (batch, seq_len, emb_dim)

        x=self.norm3(x+ffn_output)
        # (batch, tgt_seq_len, emb_dim)
        
        return x


x = torch.randn(64, 32, 512)
decoder=Decoder(8, 512, 2048)
output=decoder(x, x)
output.shape


class PositionalEncoding(nn.Module):
    def __init__(self, max_len, emb_dim):
        super().__init__()

        pe=torch.zeros(max_len, emb_dim)
        # (max_len, emb_dim)

        position=torch.arange(0, max_len).unsqueeze(1)
        # (max_len, 1)

        div_term = torch.exp(
            torch.arange(0, emb_dim, 2) *
            (-math.log(10000.0) / emb_dim)
        )

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        pe=pe.unsqueeze(0)
        # (1, max_len, emb_dim)

        self.register_buffer("pe", pe)

    def forward(self, x):
        sl=x.size(1)
        return x+self.pe[:, :sl]


x = torch.randn(64, 32, 512)

pe = PositionalEncoding(5000, 512)

output = pe(x)

print(output.shape)


class Transformer(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, num_heads, num_layers, emb_dim, nn_dim, max_len=5000):
        super().__init__()

        self.src_embedding=nn.Embedding(src_vocab_size, emb_dim)
        self.tgt_embedding=nn.Embedding(tgt_vocab_size, emb_dim)

        self.positional_encoding=PositionalEncoding(max_len, emb_dim)

        self.encoder_layers=nn.ModuleList([Encoder(num_heads, emb_dim, nn_dim) for _ in range(num_layers)])

        self.decoder_layers=nn.ModuleList([Decoder(num_heads, emb_dim, nn_dim) for _ in range(num_layers)])

        self.fc=nn.Linear(emb_dim, tgt_vocab_size)

    def forward(self, src, tgt):
        # Encoder
        src=self.src_embedding(src)
        src=self.positional_encoding(src)

        encoder_output=src
        for layer in self.encoder_layers:
            encoder_output=layer(encoder_output)

        # Decoder
        tgt=self.tgt_embedding(tgt)
        tgt=self.positional_encoding(tgt)

        decoder_output=tgt
        for layer in self.decoder_layers:
            decoder_output=layer(decoder_output, encoder_output)

        logits=self.fc(decoder_output)

        return logits


src = torch.randint(0, 10000, (64, 20))
tgt = torch.randint(0, 12000, (64, 15))

model=Transformer(10000, 12000, 8, 6, 512, 2048)

output=model(src, tgt)

output.shape


save_file(model.state_dict(), "transformer.safetensors")


model.load_state_dict(load_file("transformer.safetensors"))


config = {
    "model_name": "Seq2Seq Encoder-Decoder GRU with Luong Attention",
    "task": "translation",
    "framework": "PyTorch",
    "dataset": "bentrevett/multi30k",
    "src_lang": "en",
    "tgt_lang": "de",
    "src_vocab_size": len(en_word2idx),
    "tgt_vocab_size": len(de_word2idx),
    "architecture": {
        "type": "Seq2Seq",
        "cell": "GRU",
        "emb_dim": 128,
        "hidden_dim": 256
    },
    "training": {
        "optimizer": "Adam",
        "learning_rate": 0.001,
        "loss": "CrossEntropyLoss",
        "grad_clip_norm": 1.0,
        "epochs": 50,
        "batch_size": 64
    }
}
import json
with open("seq2seq/config.json", "w") as f:
    json.dump(config, f, indent=4)