Added models and code

config.yaml

@@ -0,0 +1,15 @@
+batch_size: 16
+d_model: 512
+dataset_path: data/english2hindi_data.json
+epochs: 15
+learning_rate: 0.0003
+model_directory: models
+num_enc_dec_blocks: 6
+num_of_heads: 8
+resume_training: true
+seq_len: 281
+src_lang: en_text
+src_tokenizer_file: tokenizers/tokenizers_eng.json
+tgt_lang: hi_text
+tgt_tokenizer_file: tokenizers/tokenizers_hindi.json
+warmup_steps: 4000

config_big.yaml

@@ -0,0 +1,15 @@
+batch_size: 32
+d_model: 520
+dataset_path: data/english2hindi_data.json
+epochs: 5
+learning_rate: 0.0012
+model_directory: models
+num_enc_dec_blocks: 8
+num_of_heads: 10
+resume_training: true
+seq_len: 281
+src_lang: en_text
+src_tokenizer_file: tokenizers/tokenizers_eng.json
+tgt_lang: hi_text
+tgt_tokenizer_file: tokenizers/tokenizers_hindi.json
+warmup_steps: 4000

dataset.py

@@ -0,0 +1,138 @@
+import torch
+from torch.utils.data import Dataset
+
+from torch.utils.data import Dataset
+import json
+
+
+class English2HindiDataset(Dataset):
+    def __init__(
+        self, data, tokenizer_src, tokenizer_tgt, src_lang, tgt_lang, seq_len
+    ):
+        super().__init__()
+        self.seq_len = seq_len
+        self.tokenizer_src = tokenizer_src
+        self.tokenizer_tgt = tokenizer_tgt
+        self.src_lang = src_lang
+        self.tgt_lang = tgt_lang
+     
+        self.data = data
+        self.sos_token = torch.tensor(
+            [tokenizer_tgt.token_to_id("[SOS]")], dtype=torch.int64
+        )
+        self.eos_token = torch.tensor(
+            [tokenizer_tgt.token_to_id("[EOS]")], dtype=torch.int64
+        )
+        self.pad_token = torch.tensor(
+            [tokenizer_tgt.token_to_id("[PAD]")], dtype=torch.int64
+        )
+
+        # Why int64?
+        # When passing token indices to an nn.Embedding layer, PyTorch expects torch.int64 (or torch.long)
+
+    def __len__(self):
+        return len(self.data)
+    
+    def causal_mask(self,size):
+        mask = torch.triu(torch.ones((1, size, size)), diagonal=1).type(torch.int)
+        return mask == 0
+
+    def __getitem__(self, idx):
+        trans_pairs = self.data[idx]
+
+        src_text = trans_pairs["en_text"]
+        tgt_text = trans_pairs["hi_text"]
+        # here we first get english_text and hindi text which was in dictionary.
+
+        enc_input_tokens = self.tokenizer_src.encode(src_text).ids
+        dec_input_tokens = self.tokenizer_tgt.encode(tgt_text).ids
+
+        # we then use the two tokenizers to encode the text into input_ids.
+
+        # For every sentence for example : I am sai - The input to the enocoder will be
+        ## <SOS> I am Sai <EOS>
+        ### but because of variable length sequences we need to add padding.
+        ### how do we do it ?  we take the longest sentence in the dataset add 30 to it and that will gives us the seq_len
+        ### So now we add padding to every sentence and make it similar lengths.
+
+        enc_num_padding_tokens = self.seq_len - len(enc_input_tokens) - 2
+
+        # -2 Because of the <SOS> and <EOS>
+
+        dec_num_padding_tokens = self.seq_len - len(dec_input_tokens) - 1
+
+        # since the input to the decoder will only consist of <SOS>
+        # the target label will have the <EOS>
+
+        encoder_input = torch.cat(
+            [
+                self.sos_token,
+                torch.tensor(enc_input_tokens, dtype=torch.int64),
+                self.eos_token,
+                torch.tensor(
+                    [self.pad_token] * enc_num_padding_tokens, dtype=torch.int64
+                ),
+            ],
+            dim=0,
+        )
+
+        decoder_input = torch.cat(
+            [
+                self.sos_token,
+                torch.tensor(dec_input_tokens, dtype=torch.int64),
+                torch.tensor(
+                    [self.pad_token] * dec_num_padding_tokens, dtype=torch.int64
+                ),
+            ],
+            dim=0,
+        )
+
+        # why dim =0 because everything 1d so they are stacked one after other.
+
+        label = torch.cat(
+            [
+                torch.tensor(dec_input_tokens, dtype=torch.int64),
+                self.eos_token,
+                torch.tensor(
+                    [self.pad_token]* dec_num_padding_tokens, dtype=torch.int64
+                ),
+            ],
+            dim=0,
+        )
+
+        # Label is always tgt language so we give the the decoder tokens
+
+        encoder_mask = (encoder_input != self.pad_token).unsqueeze(0).unsqueeze(0).int() # (1, 1, seq_len)
+        decoder_mask= (decoder_input != self.pad_token).unsqueeze(0).int() & self.causal_mask(decoder_input.size(0)) # (1, seq_len) & (1, seq_len, seq_len),
+
+        ### didnt understand this at all.
+
+
+        return {
+            "encoder_input": encoder_input, 
+            "decoder_input": decoder_input,  
+            "encoder_mask": encoder_mask,
+            "decoder_mask": decoder_mask,
+            "label": label,
+            "src_text": src_text,
+            "tgt_text": tgt_text,
+        }
+
+
+
+class English2HindiDatasetTest(Dataset):
+    def __init__(
+        self, json_path,
+    ):
+        super().__init__()
+        with open(json_path, "r", encoding="utf-8") as f:
+            self.data = json.load(f)
+
+    def __len__(self):
+        return len(self.data)
+    
+
+
+    def __getitem__(self, idx):
+        trans_pairs = self.data[idx]
+        return trans_pairs

hub.py

@@ -0,0 +1,19 @@
+from huggingface_hub import create_repo
+from huggingface_hub import HfApi, HfFolder, Repository
+
+
+repo_url = HfApi().create_repo(
+    name="thecr7guy/trainin_transformers",
+    exist_ok=True
+)
+
+from huggingface_hub import HfApi, HfFolder, Repository
+
+
+model_dir = "./models"
+
+repo = Repository(local_dir=model_dir, clone_from="thecr7guy/trainin_transformers")
+
+repo.git_add()
+repo.git_commit("Added models")
+repo.git_push()

inference.py

@@ -0,0 +1,95 @@
+
+from pathlib import Path
+from model import build_transformer
+from util import create_resources
+import torch
+import sys
+import yaml
+import sacrebleu
+
+def translate(sentence: str):
+
+    with open("config.yaml", "r") as file:
+        config = yaml.safe_load(file)
+    
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Using device:", device)
+    train_dataloader,valid_dataloader,test_dataloader,tokenizer_src,tokenizer_tgt = create_resources()
+
+    src_vocab_size = tokenizer_src.get_vocab_size()
+    tgt_vocab_size = tokenizer_tgt.get_vocab_size()
+
+
+    model = build_transformer(
+        src_vocab_size,
+        tgt_vocab_size,
+        config["seq_len"],
+        config["seq_len"],
+        config["num_enc_dec_blocks"],
+        config["num_of_heads"],
+        config["d_model"]
+    )
+
+    model = model.to(device)
+    #model_filename = "models/big_models_res/model_epoch_2.pth"
+    model_filename = "models/model_epoch_15.pth"
+    state = torch.load(model_filename)
+    model.load_state_dict(state['model_state_dict'])
+
+    model.eval()
+    with torch.no_grad():
+        source = tokenizer_src.encode(sentence)
+        print(source,source.ids)
+
+        source = torch.cat([
+            torch.tensor([tokenizer_src.token_to_id('[SOS]')], dtype=torch.int64), 
+            torch.tensor(source.ids, dtype=torch.int64),
+            torch.tensor([tokenizer_src.token_to_id('[EOS]')], dtype=torch.int64),
+            torch.tensor([tokenizer_src.token_to_id('[PAD]')] * (config["seq_len"] - len(source.ids) - 2), dtype=torch.int64)
+        ], dim=0)
+
+        source = source.to(device)
+        source = source.unsqueeze(0)
+
+        print(source.shape)
+
+
+
+        source_mask = (source != tokenizer_src.token_to_id('[PAD]')).unsqueeze(0).unsqueeze(0).int().to(device)
+
+        encoder_output = model.encode(source, source_mask)
+        decoder_input = torch.full((1, 1), tokenizer_tgt.token_to_id('[SOS]'),
+                                dtype=torch.long, device=device)
+
+       
+        while decoder_input.size(1) < config["seq_len"]:
+
+            decoder_mask = torch.triu(torch.ones((1, decoder_input.size(1), decoder_input.size(1))),
+                                    diagonal=1).to(device, dtype=torch.int)
+            
+
+            out = model.decode(decoder_input, encoder_output, source_mask, decoder_mask)
+            prob = model.project(out[:, -1])
+            _, next_word = torch.max(prob, dim=1)
+            
+            next_token = torch.full((1, 1), next_word.item(), dtype=torch.long, device=device)
+            decoder_input = torch.cat([decoder_input, next_token], dim=1)
+            
+            print(f"{tokenizer_tgt.decode([next_word.item()])}", end=' ')
+            
+            if next_word.item() == tokenizer_tgt.token_to_id('[EOS]'):
+                break
+
+
+    return tokenizer_tgt.decode(decoder_input[0].tolist())
+    
+
+
+#a = translate("Why does Earth have only one moon, while other planets have many?")
+a = translate("Defending champion Kolkata Knight Riders (KKR) hosts Royal Challengers Bengaluru (RCB) at the Eden Gardens in the Indian Premier League 2025 opener on Saturday.")
+a = translate("Which South American country is home to the Amazon Rainforest and the Christ the Redeemer statue?")
+# a = translate("Imagine you are an astronaut stepping onto Mars for the first time. Write a monologue expressing your emotions and observations.")
+# a = translate("The theory of evolution, proposed by Charles Darwin, explains the process by which species of organisms change over time through natural selection. While the theory has been widely accepted in the scientific community, it continues to spark debates in various social and religious contexts. Discuss how the theory of evolution has shaped our understanding of human origins and the controversies that surround it.")
+
+with open("output.txt", "a") as w:
+    w.write(f"\n{a}")

main.py

@@ -0,0 +1,83 @@
+from torchinfo import summary
+from model import build_transformer
+from util import create_resources
+import yaml
+import torch
+from pathlib import Path
+
+
+
+
+
+with open("config.yaml", "r") as file:
+    config = yaml.safe_load(file)
+
+
+train_dataloader,valid_dataloader,test_dataloader,tokenizer_src,tokenizer_tgt = create_resources()
+src_vocab_size = tokenizer_src.get_vocab_size()
+tgt_vocab_size = tokenizer_src.get_vocab_size()
+
+model = build_transformer(
+    src_vocab_size,
+    tgt_vocab_size,
+    config["seq_len"],
+    config["seq_len"],
+    config["num_enc_dec_blocks"],
+    config["num_of_heads"],
+    config["d_model"]
+)
+
+batch_size = config["batch_size"]
+num_epochs = config["epochs"] if "epochs" in config else 10
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = model.to(device)
+
+criterion = loss_fn = torch.nn.CrossEntropyLoss(ignore_index=tokenizer_src.token_to_id('[PAD]'), label_smoothing=0.1).to(device) 
+optimizer = torch.optim.Adam(model.parameters(), lr=config["learning_rate"],eps=1e-9)
+
+
+def save_checkpoint(epoch, model, optimizer, path):
+    torch.save({
+        "epoch": epoch,
+        "model_state_dict": model.state_dict(),
+        "optimizer_state_dict": optimizer.state_dict(),
+    }, path)
+    print(f"Checkpoint saved at epoch {epoch} to {path}")
+
+
+def load_checkpoint(path, model, optimizer=None, map_location="cpu"):
+    checkpoint = torch.load(path, map_location=map_location)
+    model.load_state_dict(checkpoint["model_state_dict"])
+    if optimizer and "optimizer_state_dict" in checkpoint:
+        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
+    start_epoch = checkpoint.get("epoch", 0)
+    print(f"Loaded checkpoint from epoch {start_epoch}")
+    return start_epoch
+
+
+def train_one_epoch(device):
+    model.train()
+    running_loss = 0.0
+
+
+
+
+def train_model(model):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    if (device == 'cuda'):
+        print(f"Device name: {torch.cuda.get_device_name(device.index)}")
+        print(f"Device memory: {torch.cuda.get_device_properties(device.index).total_memory / 1024 ** 3} GB")
+    
+    Path(config["model_folder"]).mkdir(parents=True, exist_ok=True)
+    train_dataloader,valid_dataloader,test_dataloader,tokenizer_src,tokenizer_tgt = create_resources()
+
+    
+
+
+    
+
+
+
+

model.py

@@ -0,0 +1,291 @@
+import torch
+import math
+
+
+class InputEmbeddings(torch.nn.Module):
+    def __init__(self, d_model, vocab_size):
+        super().__init__()
+        self.d_model = d_model
+        self.vocab_size = vocab_size
+        self.embeddingss = torch.nn.Embedding(vocab_size, d_model)
+
+    def forward(self, x):
+        return self.embeddingss(x) * math.sqrt(self.d_model)
+
+
+class PositionalEncoding(torch.nn.Module):
+    def __init__(self, d_model, seq_len, dropout):
+        super().__init__()
+        self.d_model = d_model
+        self.seq_len = seq_len
+        self.dropout = torch.nn.Dropout(dropout)
+        pe = torch.zeros(self.seq_len, self.d_model)
+
+        for i in range(self.seq_len):
+            for j in range(self.d_model):
+                denom = torch.pow(torch.tensor(10000.0), (2 * j) / self.d_model)
+                num = torch.tensor(float(i))
+                if j % 2 == 0:
+                    pe[i, j] = torch.sin(num / denom)
+                else:
+                    pe[i, j] = torch.cos(num / denom)
+
+        pe = pe.unsqueeze(0)
+        print(pe.shape)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        x = x + (self.pe[:, : x.shape[1], :]).requires_grad_(False)
+        return self.dropout(x)
+
+
+class LayerNormm(torch.nn.Module):
+    def __init__(self, features):
+        super().__init__()
+        self.layer_norm = torch.nn.LayerNorm(features, eps=1e-5)
+
+    def forward(self, x):
+        return self.layer_norm(x)
+
+
+class FeedForward(torch.nn.Module):
+    def __init__(self, d_model, dff, dropout):
+        super().__init__()
+        self.linear_1 = torch.nn.Linear(d_model, dff)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.linear_2 = torch.nn.Linear(dff, d_model)
+        self.activation = torch.nn.ReLU()
+
+    def forward(self, x):
+        x = self.linear_1(x)
+        x = self.activation(x)
+        x = self.dropout(x)
+        x = self.linear_2(x)
+        return x
+
+
+class MHA(torch.nn.Module):
+    def __init__(self, d_model, number_of_heads, dropout):
+        super().__init__()
+
+        self.dropout = torch.nn.Dropout(dropout)
+        self.d_model = d_model
+        self.noh = number_of_heads
+
+        self.dk = self.d_model // self.noh
+
+        self.wq = torch.nn.Linear(d_model, d_model)
+        self.wk = torch.nn.Linear(d_model, d_model)
+        self.wv = torch.nn.Linear(d_model, d_model)
+        self.wo = torch.nn.Linear(d_model, d_model)
+
+    @staticmethod
+    def calculate_self_attention(qprime, kprime, vprime, mask, dropout):
+        dk = qprime.shape[-1]
+        attention_scores = (qprime @ kprime.transpose(-2, -1)) / math.sqrt(dk)
+
+        if mask is not None:
+            attention_scores.masked_fill_(mask == 0, -1e9)
+
+        attention_scores = attention_scores.softmax(dim=-1)
+        # why last dim ?
+        if dropout is not None:
+            attention_scores = dropout(attention_scores)
+
+        return (attention_scores @ vprime), attention_scores
+
+    def forward(self, q, k, v, mask):
+        qprime = self.wq(q)
+        # (batch,seq_length,dmodel)
+        kprime = self.wk(k)
+        # (batch,seq_length,dmodel)
+        vprime = self.wv(v)
+        # (batch,seq_length,dmodel)
+
+        qprime = qprime.view(qprime.shape[0], qprime.shape[1], self.noh, self.dk)
+        # (batch,seq_length,dmodel) =>(batch,seq_length,noh,dk)
+        qprime = qprime.transpose(1, 2)
+        # (batch,seq_length,noh,dk) => (batch,noh,seq_length,dk)
+
+        kprime = kprime.view(kprime.shape[0], kprime.shape[1], self.noh, self.dk)
+        kprime = kprime.transpose(1, 2)
+
+        vprime = vprime.view(vprime.shape[0], vprime.shape[1], self.noh, self.dk)
+        vprime = vprime.transpose(1, 2)
+
+        x, attention_scores = MHA.calculate_self_attention(
+            qprime, kprime, vprime, mask, self.dropout
+        )
+        x = x.transpose(1, 2).contiguous().view(x.shape[0], -1, self.noh * self.dk)
+        return self.wo(x)
+
+
+class SkipConnection(torch.nn.Module):
+    def __init__(self, features, dropout):
+        super().__init__()
+        self.dropout = torch.nn.Dropout(dropout)
+        self.layernorm = LayerNormm(features)
+
+    def forward(self, x, sublayer):
+        return x + self.dropout(sublayer(self.layernorm(x)))
+
+
+class EncoderBlock(torch.nn.Module):
+    def __init__(self, features, mha_block, feedforward_block, dropout):
+        super().__init__()
+        self.attention_block = mha_block
+        self.feedforward_block = feedforward_block
+        self.skip_connections = torch.nn.ModuleList(
+            [SkipConnection(features, dropout) for _ in range(2)]
+        )
+        self.dropout = torch.nn.Dropout(dropout)
+
+    def forward(self, x, src_mask):
+        x = self.skip_connections[0](
+            x, lambda x: self.attention_block(x, x, x, src_mask)
+        )
+        x = self.skip_connections[1](x, self.feedforward_block)
+        return x
+
+
+class Encoder(torch.nn.Module):
+    def __init__(self, features: int, layers: torch.nn.ModuleList) -> None:
+        super().__init__()
+        self.layers = layers
+        self.norm = LayerNormm(features)
+
+    def forward(self, x, mask):
+        for layer in self.layers:
+            x = layer(x, mask)
+        return self.norm(x)
+
+
+class DecoderBlock(torch.nn.Module):
+    def __init__(self, features, mha_block, mha_block2, feedforward_block, dropout):
+        super().__init__()
+        self.attention_block = mha_block
+        self.cross_attention_block = mha_block2
+        self.feedforward_block = feedforward_block
+        self.skip_connections = torch.nn.ModuleList(
+            [SkipConnection(features, dropout) for _ in range(3)]
+        )
+        self.dropout = torch.nn.Dropout(dropout)
+
+    def forward(self, x, enc_output, src_mask, tgt_mask):
+        x = self.skip_connections[0](
+            x, lambda x: self.attention_block(x, x, x, tgt_mask)
+        )
+        x = self.skip_connections[1](
+            x, lambda x: self.cross_attention_block(x, enc_output, enc_output, src_mask)
+        )
+        x = self.skip_connections[2](x, self.feedforward_block)
+        return x
+
+
+class Decoder(torch.nn.Module):
+
+    def __init__(self, features: int, layers: torch.nn.ModuleList) -> None:
+        super().__init__()
+        self.layers = layers
+        self.norm = LayerNormm(features)
+
+    def forward(self, x, encoder_output, src_mask, tgt_mask):
+        for layer in self.layers:
+            x = layer(x, encoder_output, src_mask, tgt_mask)
+        return self.norm(x)
+
+
+class ProjectionLayer(torch.nn.Module):
+
+    def __init__(self, d_model, vocab_size) -> None:
+        super().__init__()
+        self.proj = torch.nn.Linear(d_model, vocab_size)
+
+    def forward(self, x) -> None:
+        return self.proj(x)
+
+
+class Transformer(torch.nn.Module):
+
+    def __init__(
+        self,
+        encoder,
+        decoder,
+        src_pos_enc,
+        tgt_pos_enc,
+        src_emb,
+        tgt_emb,
+        projection_layer,
+    ) -> None:
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.src_pos_enc = src_pos_enc
+        self.tgt_pos_enc = tgt_pos_enc
+        self.src_emb = src_emb
+        self.tgt_emb = tgt_emb
+        self.projection_layer = projection_layer
+
+    def encode(self, src, src_mask):
+        src = self.src_emb(src)
+        src = self.src_pos_enc(src)
+        x = self.encoder(src, src_mask)
+        return x
+
+    def decode(self, tgt, enc_output, src_mask, tgt_mask):
+        tgt = self.tgt_emb(tgt)
+        tgt = self.tgt_pos_enc(tgt)
+        x = self.decoder(tgt, enc_output, src_mask, tgt_mask)
+        return x
+
+    def project(self, x):
+        x = self.projection_layer(x)
+        return x
+
+def build_transformer(
+    src_vocab_size,
+    tgt_vocab_size,
+    src_seq_len,
+    tgt_seq_len,
+    nlayers=6,
+    noh=8,
+    d_model=512,
+    dropout=0.1,
+    dff=2048,
+):
+    src_emb = InputEmbeddings(d_model, src_vocab_size)
+    tgt_emb = InputEmbeddings(d_model, tgt_vocab_size)
+
+    src_pos_enc = PositionalEncoding(d_model, src_seq_len, dropout)
+    tgt_pos_enc = PositionalEncoding(d_model, tgt_seq_len, dropout)
+
+    enc_blocks = []
+    for i in range(0, nlayers):
+        mha = MHA(d_model, noh, dropout)
+        ff = FeedForward(d_model, dff, dropout)
+        enc_block = EncoderBlock(d_model, mha, ff, dropout)
+        enc_blocks.append(enc_block)
+
+    encoder = Encoder(d_model, torch.nn.ModuleList(enc_blocks))
+
+    dec_blocks = []
+    for i in range(0, nlayers):
+        mha = MHA(d_model, noh, dropout)
+        mha2 = MHA(d_model, noh, dropout)
+        ff = FeedForward(d_model, dff, dropout)
+        dec_block = DecoderBlock(d_model, mha, mha2, ff, dropout)
+        dec_blocks.append(dec_block)
+
+    decoder = Decoder(d_model, torch.nn.ModuleList(dec_blocks))
+
+    proj = ProjectionLayer(d_model, tgt_vocab_size)
+
+    transformer = Transformer(
+        encoder, decoder, src_pos_enc, tgt_pos_enc, src_emb, tgt_emb, proj
+    )
+
+    for p in transformer.parameters():
+        if p.dim() > 1:
+            torch.nn.init.xavier_uniform_(p)
+
+    return transformer

models/model_epoch_14.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:883fffd40352bf4f925da84a67882fad7284d4cd6ca6f732fd3a9bc7f1ca12fe
+size 1425518098

models/model_epoch_15.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c473c1bcd9fa94a37c5b2e0f0196a77ddc00774af1a98ee57c7cd32852e1dab
+size 1425518098

models/model_epoch_8.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1b5abd48289c13b380f3cf2288591daf8be891709a56072b1962a41060739224
+size 1425514236

models/model_epoch_9.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b94e9d6c6818e8b9ac23bf37aa303689e038f58a0da4b5bef7c8dcc33645ca3b
+size 1425517052

test.py

@@ -0,0 +1,93 @@
+from pathlib import Path
+from model import build_transformer
+from util import create_resources
+import torch
+import sys
+import yaml
+
+def translate(sentence: str):
+
+    with open("config.yaml", "r") as file:
+        config = yaml.safe_load(file)
+    
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Using device:", device)
+    train_dataloader,valid_dataloader,test_dataloader,tokenizer_src,tokenizer_tgt = create_resources()
+
+    src_vocab_size = tokenizer_src.get_vocab_size()
+    tgt_vocab_size = tokenizer_tgt.get_vocab_size()
+
+
+    model = build_transformer(
+        src_vocab_size,
+        tgt_vocab_size,
+        config["seq_len"],
+        config["seq_len"],
+        config["num_enc_dec_blocks"],
+        config["num_of_heads"],
+        config["d_model"]
+    )
+
+    model = model.to(device)
+    model_filename = "models/model.pth"
+    state = torch.load(model_filename)
+    model.load_state_dict(state['model_state_dict'])
+
+    model.eval()
+    with torch.no_grad():
+        source = tokenizer_src.encode(sentence)
+        print(source,source.ids)
+
+        source = torch.cat([
+            torch.tensor([tokenizer_src.token_to_id('[SOS]')], dtype=torch.int64), 
+            torch.tensor(source.ids, dtype=torch.int64),
+            torch.tensor([tokenizer_src.token_to_id('[EOS]')], dtype=torch.int64),
+            torch.tensor([tokenizer_src.token_to_id('[PAD]')] * (config["seq_len"] - len(source.ids) - 2), dtype=torch.int64)
+        ], dim=0)
+
+        source = source.to(device)
+        source = source.unsqueeze(0)
+
+        print(source.shape)
+
+
+
+        source_mask = (source != tokenizer_src.token_to_id('[PAD]')).unsqueeze(0).unsqueeze(0).int().to(device)
+        print(source_mask)
+
+        encoder_output = model.encode(source, source_mask)
+        print(encoder_output.shape)
+        decoder_input = torch.full((1, 1), tokenizer_tgt.token_to_id('[SOS]'),
+                                dtype=torch.long, device=device)
+
+       
+        while decoder_input.size(1) < config["seq_len"]:
+
+            decoder_mask = torch.triu(torch.ones((1, decoder_input.size(1), decoder_input.size(1))),
+                                    diagonal=1).to(device, dtype=torch.int)
+            
+            print("#######################")
+
+            
+            print(decoder_mask.shape)
+            out = model.decode(decoder_input, encoder_output, source_mask, decoder_mask)
+            prob = model.project(out[:, -1])
+            print(max(prob[0]))
+            print(min(prob[0]))
+            _, next_word = torch.max(prob, dim=1)
+            
+            next_token = torch.full((1, 1), next_word.item(), dtype=torch.long, device=device)
+            decoder_input = torch.cat([decoder_input, next_token], dim=1)
+            
+            print(f"{tokenizer_tgt.decode([next_word.item()])}", end=' ')
+            
+            if next_word.item() == tokenizer_tgt.token_to_id('[EOS]'):
+                break
+
+
+    # return tokenizer_tgt.decode(decoder_input[0].tolist())
+    return 0
+    
+
+a = translate("My Name is sai and I love computers")
+print(a)

train.py

@@ -0,0 +1,243 @@
+from torchinfo import summary
+from model import build_transformer
+from util import create_resources
+import yaml
+import torch
+from tqdm import tqdm
+import os
+import wandb
+import matplotlib.pyplot as plt
+from matplotlib import font_manager
+import re
+
+
+mangal_font_path = "Mangal.TTf"
+devanagari_font = font_manager.FontProperties(fname=mangal_font_path)
+
+
+class NoamScheduler:
+    def __init__(self, optimizer, d_model, warmup_steps):
+        self.optimizer = optimizer
+        self.d_model = d_model
+        self.warmup_steps = warmup_steps
+        self.step_num = 0
+
+    def step(self):
+        self.step_num += 1
+        lr = self.get_lr()
+        for param_group in self.optimizer.param_groups:
+            param_group["lr"] = lr
+        return lr
+
+    def get_lr(self):
+        step = max(self.step_num, 1)
+        arg1 = step ** (-0.5)
+        arg2 = step * (self.warmup_steps ** (-1.5))
+        return (self.d_model ** (-0.5)) * min(arg1, arg2)
+
+
+class Trainer:
+    def __init__(
+        self,
+        model,
+        optimizer,
+        scheduler,
+        criterion,
+        device,
+        tokenizer_src,
+        tokenizer_tgt,
+        seq_len,
+    ):
+        self.model = model
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        self.criterion = criterion
+        self.device = device
+        self.tgt_tokenizer = tokenizer_tgt
+        self.src_tokenizer = tokenizer_src
+        self.seq_len = seq_len
+
+    def train_epoch(self, dataloader):
+        self.model.train()
+        torch.cuda.empty_cache()
+        running_loss = 0.0
+        total_tokens = 0
+        progress_bar = tqdm(
+            enumerate(dataloader), desc="Training", total=len(dataloader)
+        )
+
+        for batch_idx, batch in progress_bar:
+
+            encoder_input = batch["encoder_input"].to(self.device)
+            # Should be (1,seq_len) => (batch_size,seq_len)
+            decoder_input = batch["decoder_input"].to(self.device)
+            # Should be (1,seq_len) => (batch_size,seq_len)
+
+            encoder_mask = batch["encoder_mask"].to(self.device)
+            decoder_mask = batch["decoder_mask"].to(self.device)
+
+            encoder_output = self.model.encode(encoder_input, encoder_mask)
+            decoder_output = self.model.decode(
+                decoder_input, encoder_output, encoder_mask, decoder_mask
+            )
+            projection_output = self.model.project(decoder_output)
+
+            label = batch["label"].to(self.device)
+
+            loss = self.criterion(
+                projection_output.view(-1, self.tgt_tokenizer.get_vocab_size()),
+                label.view(-1),
+            )
+
+            loss.backward()
+
+            self.optimizer.step()
+            current_lr = self.scheduler.step()
+            self.optimizer.zero_grad()
+
+            pad_id = 1
+            with torch.no_grad():
+                non_pad = label.ne(pad_id)
+                num_nonpad_tokens = non_pad.sum().item()
+                running_loss += loss.item() * num_nonpad_tokens
+                total_tokens += num_nonpad_tokens
+
+            if (batch_idx + 1) % 50 == 0:
+                wandb.log(
+                    {
+                        "batch_loss": loss.item(),
+                        "learning_rate": current_lr,
+                        "batch": batch_idx + 1,
+                    }
+                )
+
+        epoch_loss = running_loss / total_tokens if total_tokens > 0 else 0.0
+        return epoch_loss
+
+    def save_checkpoint(self, epoch, output_dir):
+        os.makedirs(output_dir, exist_ok=True)
+        checkpoint = {
+            "epoch": epoch,
+            "model_state_dict": self.model.state_dict(),
+            "optimizer_state_dict": self.optimizer.state_dict(),
+            "scheduler_state": self.scheduler.step_num,
+        }
+        torch.save(checkpoint, os.path.join(output_dir, f"model_epoch_{epoch}.pth"))
+        print(f"Checkpoint saved at epoch {epoch}")
+
+    def run(self, train_loader, epochs, output_dir, start_epoch=1):
+        for epoch in range(start_epoch, epochs + 1):
+            train_loss = self.train_epoch(train_loader)
+            current_lr = self.scheduler.get_lr()
+
+            wandb.log(
+                {"epoch": epoch, "train_loss": train_loss, "learning_rate": current_lr}
+            )
+
+            self.save_checkpoint(epoch, output_dir)
+
+
+def load_latest_checkpoint(model, optimizer, scheduler, model_directory, device):
+    if not os.path.isdir(model_directory):
+        return None, 1
+    checkpoint_files = []
+    for filename in os.listdir(model_directory):
+        if filename.endswith(".pth"):
+            match = re.search(r"model_epoch_(\d+)\.pth", filename)
+            if match:
+                epoch = int(match.group(1))
+                checkpoint_files.append((epoch, filename))
+
+    if not checkpoint_files:
+        return None, 1
+
+    # Get the checkpoint with the highest epoch number
+    latest_epoch, latest_filename = max(checkpoint_files, key=lambda x: x[0])
+    ckpt_path = os.path.join(model_directory, latest_filename)
+    ckpt = torch.load(ckpt_path, map_location=device)
+
+    model.load_state_dict(ckpt["model_state_dict"])
+    optimizer.load_state_dict(ckpt["optimizer_state_dict"])
+    scheduler.step_num = ckpt["scheduler_state"]
+    start_epoch = ckpt["epoch"] + 1
+    print(f"Resuming Training from epoch {ckpt['epoch']}")
+    return ckpt, start_epoch
+
+
+def main():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    (
+        train_dataloader,
+        valid_dataloader,
+        test_dataloader,
+        tokenizer_src,
+        tokenizer_tgt,
+    ) = create_resources()
+    src_vocab_size = tokenizer_src.get_vocab_size()
+    tgt_vocab_size = tokenizer_tgt.get_vocab_size()
+
+    with open("config.yaml", "r") as file:
+        config = yaml.safe_load(file)
+
+    run = wandb.init(
+                entity="training-transformers-vast",
+                project="AttentionTranslate-sai", config=config)
+
+    model = build_transformer(
+        src_vocab_size,
+        tgt_vocab_size,
+        config["seq_len"],
+        config["seq_len"],
+        config["num_enc_dec_blocks"],
+        config["num_of_heads"],
+        config["d_model"],
+    )
+
+    model = model.to(device)
+
+    wandb.watch(model, log="all")
+
+    criterion = torch.nn.CrossEntropyLoss(
+        ignore_index=tokenizer_src.token_to_id("[PAD]"), label_smoothing=0.1
+    ).to(device)
+
+    optimizer = torch.optim.AdamW(
+        model.parameters(), lr=config["learning_rate"], betas=(0.9, 0.98), eps=1e-9
+    )
+
+    scheduler = NoamScheduler(optimizer, config["d_model"], config["warmup_steps"])
+
+    start_epoch = 1
+
+    if config["resume_training"]:
+        ckpt, start_epoch = load_latest_checkpoint(
+            model, optimizer, scheduler, config["model_directory"], device
+        )
+
+    if start_epoch == 1:
+        print("Training from scratch.")
+
+    trainer = Trainer(
+        model,
+        optimizer,
+        scheduler,
+        criterion,
+        device,
+        tokenizer_src,
+        tokenizer_tgt,
+        config["seq_len"],
+    )
+
+    # test_data_subset = list(test_dataloader)
+    # one_percent = int(0.01 * len(test_data_subset))
+    # test_data_1_percent = test_data_subset[:one_percent]
+
+    trainer.run(
+        train_dataloader, config["epochs"], config["model_directory"], start_epoch
+    )
+
+    run.finish()
+
+
+if __name__ == "__main__":
+    main()

util.py

@@ -0,0 +1,208 @@
+from tokenizers import Tokenizer
+from tokenizers.models import BPE
+from tokenizers.trainers import BpeTrainer
+from tokenizers.pre_tokenizers import (
+    WhitespaceSplit,
+    Punctuation,
+    Sequence as PreSequence,
+)
+from tokenizers.normalizers import NFD, Lowercase, Sequence, StripAccents
+from pathlib import Path
+import yaml
+from datasets import load_dataset
+from dataset import English2HindiDataset, English2HindiDatasetTest
+from torch.utils.data import DataLoader
+import json
+import random
+from sklearn.model_selection import train_test_split
+
+
+def load_config(config_path):
+    with open(config_path, "r") as file:
+        return yaml.safe_load(file)
+
+
+def get_all_sentences(ds, lang):
+    for item in ds:
+        yield item[lang]
+
+
+def get_or_build_tokenizer(tokenizer_path, ds, lang):
+    tokenizer_path = Path(tokenizer_path)
+    if not Path.exists(tokenizer_path):
+        tokenizer = Tokenizer(BPE(unk_token="[UNK]"))
+        if lang == "en_text":
+            tokenizer.normalizer = Sequence(
+                [
+                    NFD(),
+                    StripAccents(),
+                    Lowercase(),
+                ]
+            )
+        else:
+            tokenizer.normalizer = Sequence(
+                [
+                    NFD(),
+                    Lowercase(),
+                ]
+            )
+
+        tokenizer.pre_tokenizer = PreSequence([WhitespaceSplit(), Punctuation()])
+        trainer = BpeTrainer(
+            special_tokens=["[UNK]", "[PAD]", "[SOS]", "[EOS]"],
+            min_frequency=3,
+            vocab_size=60000,
+        )
+        tokenizer.train_from_iterator(get_all_sentences(ds, lang), trainer=trainer)
+        tokenizer.save(str(tokenizer_path))
+    else:
+        tokenizer = Tokenizer.from_file(str(tokenizer_path))
+    return tokenizer
+
+
+def preprocess_to_json(dataset_hf, output_json_path, word_limit=77, char_limit=300):
+    filtered_data = []
+
+    for row in dataset_hf:
+        if not row.get("translated", False):
+            continue
+
+        en_value = None
+        en_conversations = row.get("en_conversations", [])
+        for conv in en_conversations:
+            if "human" in conv["from"]:
+                en_value = conv["value"]
+                break
+        if en_value is None and len(en_conversations) > 0:
+            en_value = en_conversations[0]["value"]
+
+        hi_value = None
+        hi_conversations = row.get("conversations", [])
+        for conv in hi_conversations:
+            if "human" in conv["from"]:
+                hi_value = conv["value"]
+                break
+        if hi_value is None and len(hi_conversations) > 0:
+            hi_value = hi_conversations[0]["value"]
+
+        if en_value and hi_value:
+            if (
+                len(en_value.split()) > word_limit
+                or len(en_value) > char_limit
+                or len(hi_value.split()) > word_limit
+                or len(hi_value) > char_limit
+            ):
+                continue
+
+            filtered_data.append({"en_text": en_value, "hi_text": hi_value})
+
+    with open(output_json_path, "w", encoding="utf-8") as f:
+        json.dump(filtered_data, f, ensure_ascii=False, indent=2)
+
+
+
+def create_resources():
+    config_path = "config.yaml"
+    config = load_config(config_path)
+
+    dataset_json_path = config.get("dataset_path", "data/english2hindi_data.json")
+
+    if not Path(dataset_json_path).exists():
+        print(f"Dataset file {dataset_json_path} not found. Creating it...")
+        dataset_hf = load_dataset("BhabhaAI/openhermes-2.5-hindi", split="train")
+        preprocess_to_json(dataset_hf, dataset_json_path)
+    else:
+        print(f"Dataset file {dataset_json_path} already exists. Skipping preprocessing.")
+
+
+    with open(dataset_json_path, "r", encoding="utf-8") as f:
+        raw_data = json.load(f)
+
+
+
+    tokenizer_src = get_or_build_tokenizer(
+        config["src_tokenizer_file"], raw_data, config["src_lang"]
+    )
+    tokenizer_tgt = get_or_build_tokenizer(
+        config["tgt_tokenizer_file"], raw_data, config["tgt_lang"]
+    )
+
+
+    test_pt_dataset = English2HindiDatasetTest(config["dataset_path"])
+
+    print(len(test_pt_dataset))
+
+    max_len_src = 0
+    max_len_tgt = 0
+
+
+    seq_len = config["seq_len"]
+
+    if seq_len == 0:
+        print("seq_len is 0, starting process...")
+
+        
+        for item in raw_data:
+            src_ids = tokenizer_src.encode(item[config['src_lang']]).ids
+            tgt_ids = tokenizer_tgt.encode(item[config['tgt_lang']]).ids
+            max_len_src = max(max_len_src, len(src_ids))
+            
+            max_len_tgt = max(max_len_tgt, len(tgt_ids))
+
+
+        print(max_len_src,max_len_tgt )  
+
+        final_max_len = max(max_len_src, max_len_tgt) + 30
+
+        config['seq_len'] = final_max_len
+
+
+        with open("config.yaml", 'w') as f:
+            yaml.safe_dump(config, f, default_flow_style=False)
+
+        print(f'Updated seq_len to {final_max_len}')
+
+    else:
+        print("seq_len is not 0, skipping process.")
+
+    random.seed(42)
+
+    train_data, temp_data = train_test_split(raw_data, test_size=0.2, random_state=42)
+    test_data, valid_data = train_test_split(temp_data, test_size=0.5, random_state=42)
+
+
+    print("######################################################")
+
+    train_dataset = English2HindiDataset(
+        train_data,
+        tokenizer_src,
+        tokenizer_tgt,
+        config["src_lang"],
+        config["tgt_lang"],
+        config["seq_len"],
+    )
+
+    valid_dataset = English2HindiDataset(
+        valid_data,
+        tokenizer_src,
+        tokenizer_tgt,
+        config["src_lang"],
+        config["tgt_lang"],
+        config["seq_len"],
+    )
+
+    test_dataset = English2HindiDataset(
+        test_data,
+        tokenizer_src,
+        tokenizer_tgt,
+        config["src_lang"],
+        config["tgt_lang"],
+        config["seq_len"],
+    )
+
+
+    train_dataloader = DataLoader(train_dataset, batch_size=config["batch_size"], shuffle=False)
+    valid_dataloader = DataLoader(valid_dataset, batch_size=config["batch_size"],shuffle=True)
+    test_dataloader = DataLoader(test_dataset, batch_size=config["batch_size"],shuffle=True)
+
+    return train_dataloader,valid_dataloader,test_dataloader,tokenizer_src,tokenizer_tgt