Upload 8 files

app.py

@@ -0,0 +1,123 @@
+import torch
+from model import build_transformer
+from train import greedy_decode, get_model, get_or_build_tokenizer
+
+from config import get_config, get_weights_file_path
+from tokenizers import Tokenizer
+from pathlib import Path
+
+
+
+config = get_config()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+def process_text(config, src_text, tokenizer_src, tokenizer_tgt, src_lang, tgt_lang, seq_len):
+        seq_len = seq_len
+
+        # ds = ds
+        tokenizer_src = tokenizer_src
+        tokenizer_tgt = tokenizer_tgt
+        src_lang = src_lang
+        tgt_lang = tgt_lang
+
+        sos_token = torch.tensor([tokenizer_tgt.token_to_id("[SOS]")], dtype=torch.int64)
+        eos_token = torch.tensor([tokenizer_tgt.token_to_id("[EOS]")], dtype=torch.int64)
+        pad_token = torch.tensor([tokenizer_tgt.token_to_id("[PAD]")], dtype=torch.int64)
+    # Transform the text into tokens
+        enc_input_tokens = tokenizer_src.encode(src_text).ids
+        # dec_input_tokens = self.tokenizer_tgt.encode(tgt_text).ids
+
+        # Add sos, eos and padding to each sentence
+        enc_num_padding_tokens = seq_len - len(enc_input_tokens) - 2  # We will add <s> and </s>
+        # # We will only add <s>, and </s> only on the label
+        # dec_num_padding_tokens = self.seq_len - len(dec_input_tokens) - 1
+
+        # Make sure the number of padding tokens is not negative. If it is, the sentence is too long
+        if enc_num_padding_tokens < 0:
+            raise ValueError("Sentence is too long")
+
+        # Add <s> and </s> token
+        encoder_input = torch.cat(
+            [
+                sos_token,
+                torch.tensor(enc_input_tokens, dtype=torch.int64),
+                eos_token,
+                torch.tensor([pad_token] * enc_num_padding_tokens, dtype=torch.int64),
+            ],
+            dim=0,
+        )
+
+        # # Add only <s> token
+        # 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,
+        # )
+
+        # # Add only </s> token
+        # 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,
+        # )
+
+        # Double check the size of the tensors to make sure they are all seq_len long
+        assert encoder_input.size(0) == seq_len
+        # assert decoder_input.size(0) == seq_len
+        # assert label.size(0) == seq_len
+        return {
+            'encoder_input': encoder_input,
+            # 'decoder_input': decoder_input,
+            "encoder_mask": (encoder_input != pad_token).unsqueeze(0).unsqueeze(0).int(), # (1, 1, seq_len)
+            # "decoder_mask": (decoder_input != self.pad_token).unsqueeze(0).int() & causal_mask(decoder_input.size(0)), # (1, seq_len) & (1, seq_len, seq_len),
+            # "label": label,  # (seq_len)
+             
+            # "src_text": src_text,
+            # "tgt_text": tgt_text,
+        }
+            
+def causal_mask(size):
+    mask = torch.triu(torch.ones((1, size, size)), diagonal=1).type(torch.int)
+    return mask == 0
+
+def infer(text, config):
+    tokenizer_src = Tokenizer.from_file(str(Path(config['tokenizer_file'].format(config['lang_src']))))
+    tokenizer_tgt = Tokenizer.from_file(str(Path(config['tokenizer_file'].format(config['lang_tgt']))))
+    model = get_model(config, tokenizer_src.get_vocab_size(), tokenizer_tgt.get_vocab_size())
+    state = torch.load('tmodel_36.pt', map_location=torch.device('cpu'))
+    model.load_state_dict(state['model_state_dict'])
+
+
+
+
+    model.eval()
+    with torch.no_grad():
+        processed_text = process_text(config, text, tokenizer_src, tokenizer_tgt, config['lang_src'], config['lang_tgt'], config['seq_len'])
+        encoder_input = processed_text['encoder_input']
+        encoder_mask = processed_text['encoder_mask']
+
+        model_out = greedy_decode(model, encoder_input, encoder_mask, tokenizer_src, tokenizer_tgt, config['seq_len'], device)
+        model_out_text = tokenizer_tgt.decode(model_out.detach().cpu().numpy())
+        return model_out_text
+    
+
+import streamlit as st
+
+st.title("English to Hausa Translation")
+
+user_input = st.text_input("Enter your text:")
+if user_input:
+    result = infer(user_input, config)
+    st.write("Inference Result:", result)
+
+
+
+        
+
+
+

config.py

@@ -0,0 +1,25 @@
+from pathlib import Path
+
+def get_config():
+    return {
+        "batch_size":2,
+        "num_epochs": 100,
+        "lr": 10**-4,
+        "seq_len": 150,
+        "d_model": 512,
+        "lang_src": "0",
+        "lang_tgt": "1",
+        "model_folder": "weights",
+        "model_basename": "tmodel_",
+        "preload": None,
+        "tokenizer_file": "tokenizer_{0}.json",
+        "experiment_name": "runs/tmodel"
+    }
+
+def get_weights_file_path(config, epoch: str):
+    model_folder = config["model_folder"]
+    model_basename = config["model_basename"]
+    model_filename = f"{model_basename}{epoch}.pt"
+    return str(Path('.') / model_folder / model_filename)
+
+

dataset.py

@@ -0,0 +1,101 @@
+import torch
+import torchvision
+from torch.utils.data import Dataset, DataLoader
+import pandas as pd
+
+class BilingualDataset(Dataset):
+    def __init__(self, ds, tokenizer_src, tokenizer_tgt, src_lang, tgt_lang, seq_len):
+        super().__init__()
+        self.seq_len = seq_len
+
+        self.ds = ds
+        self.tokenizer_src = tokenizer_src
+        self.tokenizer_tgt = tokenizer_tgt
+        self.src_lang = src_lang
+        self.tgt_lang = tgt_lang
+
+        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)
+
+    def __len__(self):
+        return len(self.ds)
+
+    def __getitem__(self, idx):
+    
+     
+
+        
+        src_target_pair = self.ds[idx]
+        src_text = src_target_pair[self.src_lang]
+        tgt_text = src_target_pair[self.tgt_lang]
+        
+
+       
+     
+
+        # Transform the text into tokens
+        enc_input_tokens = self.tokenizer_src.encode(src_text).ids
+        dec_input_tokens = self.tokenizer_tgt.encode(tgt_text).ids
+
+        # Add sos, eos and padding to each sentence
+        enc_num_padding_tokens = self.seq_len - len(enc_input_tokens) - 2  # We will add <s> and </s>
+        # We will only add <s>, and </s> only on the label
+        dec_num_padding_tokens = self.seq_len - len(dec_input_tokens) - 1
+
+        # Make sure the number of padding tokens is not negative. If it is, the sentence is too long
+        if enc_num_padding_tokens < 0 or dec_num_padding_tokens < 0:
+            raise ValueError("Sentence is too long")
+
+        # Add <s> and </s> token
+        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,
+        )
+
+        # Add only <s> token
+        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,
+        )
+
+        # Add only </s> token
+        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,
+        )
+
+        # Double check the size of the tensors to make sure they are all seq_len long
+        assert encoder_input.size(0) == self.seq_len
+        assert decoder_input.size(0) == self.seq_len
+        assert label.size(0) == self.seq_len
+        return {
+            'encoder_input': encoder_input,
+            'decoder_input': decoder_input,
+            "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() & causal_mask(decoder_input.size(0)), # (1, seq_len) & (1, seq_len, seq_len),
+            "label": label,  # (seq_len)
+             
+            "src_text": src_text,
+            "tgt_text": tgt_text,
+        }
+
+   
+    
+def causal_mask(size):
+    mask = torch.triu(torch.ones((1, size, size)), diagonal=1).type(torch.int)
+    return mask == 0
+        

model.py

@@ -0,0 +1,558 @@
+
+##Implementation of tranformer from scratch, this implememtation was inspired by Umar Jamir
+
+import torch
+import torch.nn as nn
+import math
+import torch.nn.functional as F
+
+class InputEmbeddings(nn.Module):
+    def __init__(self, d_model: int, vocab_size: int) -> None:
+        super(InputEmbeddings, self).__init__()
+        self.d_model = d_model
+        self.embedding = nn.Embedding(vocab_size, d_model)
+
+
+    def forward(self, x):
+        # (batch, seq_len) --> (batch, seq_len, d_model)
+
+
+        # Multiply by sqrt(d_model) to scale the embeddings according to the paper
+        return self.embedding(x) * math.sqrt(self.d_model) 
+
+
+class PositionEncoding(nn.Module):
+    def __init__(self, seq_len: int, d_model:int, batch: int) -> None:
+        super(PositionEncoding, self).__init__()
+        # self.seq_len = seq_len
+        # self.d_model = d_model
+        # self.batch = batch
+        self.dropout = nn.Dropout(p=0.1)
+    
+        ##initialize the positional encoding with zeros
+        positional_encoding = torch.zeros(seq_len, d_model)
+     
+        ##first path of the equation is postion/scaling factor per dimesnsion
+        postion  = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)
+    
+        ## this calculates the scaling term per dimension (512)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+
+        # div_term = torch.pow(10,  torch.arange(0,self.d_model, 2).float() *-4/self.d_model)
+      
+
+        ## this calculates the sin values for even indices
+        positional_encoding[:, 0::2] = torch.sin(postion * div_term) 
+
+      
+        ## this calculates the cos values for odd indices
+        positional_encoding[:, 1::2] = torch.cos(postion * div_term)
+
+        positional_encoding = positional_encoding.unsqueeze(0)
+        self.register_buffer('positional_encoding', positional_encoding)
+    
+    def forward(self, x):  
+         x = x + (self.positional_encoding[:, :x.shape[1], :]).requires_grad_(False) # (batch, seq_len, d_model)
+         return self.dropout(x)
+
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, d_model:int, heads: int) -> None:
+        super(MultiHeadAttention,self).__init__()
+        self.head = heads
+        self.head_dim = d_model // heads
+        
+
+
+        assert d_model % heads == 0, 'cannot divide d_model by heads'
+
+        ## initialize the query, key and value weights 512*512
+        self.query_weight = nn.Linear(d_model, d_model, bias=False)
+        self.key_weight = nn.Linear(d_model, d_model,bias=False)
+        self.value_weight = nn.Linear(d_model, d_model,bias=False)
+        self.final_weight  = nn.Linear(d_model, d_model, bias=False)
+        self.dropout = nn.Dropout(p=0.1)
+
+      
+    def self_attention(self,query, key, value, mask,dropout):
+        #splitting query, key and value into heads
+                #this gives us a dimension of batch, num_heads, seq_len by 64. basically 1 sentence is converted to have 8 parts (heads)
+        query = query.view(query.shape[0], query.shape[1],self.head,self.head_dim).transpose(2,1)
+        key = key.view(key.shape[0], key.shape[1],self.head,self.head_dim).transpose(2,1)
+        value = value.view(value.shape[0], value.shape[1],self.head,self.head_dim).transpose(2,1)
+        
+        attention = query @ key.transpose(3,2)
+        attention = attention / math.sqrt(query.shape[-1])
+
+        if mask is not None:
+           attention = attention.masked_fill(mask == 0, -1e9)      
+        attention = torch.softmax(attention, dim=-1)      
+        if dropout is not None:
+            attention = dropout(attention)
+        attention_scores =  attention @ value    
+       
+        return attention_scores.transpose(2,1).contiguous().view(attention_scores.shape[0], -1, self.head_dim * self.head)
+      
+    def forward(self,query, key, value,mask):
+
+        ## initialize the query, key and value matrices to give us seq_len by 512
+        query = self.query_weight(query)
+        key = self.key_weight(key)
+        value = self.value_weight(value)
+
+        attention = MultiHeadAttention.self_attention(self, query, key, value, mask, self.dropout)
+        return self.final_weight(attention) 
+
+class FeedForward(nn.Module):
+    def __init__(self,d_model:int, d_ff:int ) -> None:
+        super(FeedForward, self).__init__()
+
+        self.fc1 = nn.Linear(d_model, d_ff)  # Fully connected layer 1
+        self.dropout = nn.Dropout(p=0.1)  # Dropout layer
+        self.fc2 = nn.Linear(d_ff, d_model)  # Fully connected layer 2
+     
+    
+    def forward(self,x ):
+        return self.fc2(self.dropout(torch.relu(self.fc1(x))))  
+
+class ProjectionLayer(nn.Module):
+    def __init__(self, d_model:int, vocab_size:int) :
+        super(ProjectionLayer, self).__init__()
+        self.fc = nn.Linear(d_model, vocab_size)
+    def forward(self, x):
+        x = self.fc(x)
+        return torch.log_softmax(x, dim=-1)   
+
+class EncoderBlock(nn.Module):
+    def __init__(self, d_model:int, head:int, d_ff:int) -> None:
+        super(EncoderBlock, self).__init__()    
+        self.multiheadattention = MultiHeadAttention(d_model,head)
+        self.layer_norm1 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(p=0.1)
+        self.feedforward = FeedForward(d_model, d_ff)
+        self.layer_norm2 = nn.LayerNorm(d_model)
+        self.layer_norm3 = nn.LayerNorm(d_model)
+        self.dropout2 = nn.Dropout(p=0.1)
+
+    def forward(self, x, src_mask):
+       # Self-attention block
+        norm = self.layer_norm1(x)
+        attention = self.multiheadattention(norm, norm, norm, src_mask)
+        x = (x + self.dropout1(attention))
+
+        # Feedforward block
+        norm2 = self.layer_norm2(x)
+        ff = self.feedforward(x)
+        return x + self.dropout2(ff)     
+
+class Encoder(nn.Module):
+    def __init__(self, number_of_block:int, d_model:int, head:int, d_ff:int) -> None:
+        super(Encoder, self).__init__()
+        self.norm = nn.LayerNorm(d_model)
+        
+        # Use nn.ModuleList to store the EncoderBlock instances
+        self.encoders = nn.ModuleList([EncoderBlock(d_model, head, d_ff) 
+                                       for _ in range(number_of_block)])
+
+    def forward(self, x, src_mask):
+        for encoder_block in self.encoders:
+            x = encoder_block(x, src_mask)
+        return self.norm(x)   
+   
+class DecoderBlock(nn.Module):
+    def __init__(self, d_model:int, head:int, d_ff:int) -> None:
+        super(DecoderBlock, self).__init__()
+        self.head_dim = d_model // head
+        
+        self.multiheadattention = MultiHeadAttention(d_model, head)
+        self.crossattention = MultiHeadAttention(d_model, head)
+        self.layer_norm1 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(p=0.1)
+        self.feedforward = FeedForward(d_model,d_ff)
+        self.layer_norm2 = nn.LayerNorm(d_model)
+        self.layer_norm3 = nn.LayerNorm(d_model)
+        self.layer_norm4 = nn.LayerNorm(d_model)
+        self.dropout2 = nn.Dropout(p=0.1)
+        self.dropout3 = nn.Dropout(p=0.1)
+    def forward(self, x, src_mask, tgt_mask, encoder_output):
+        #Self-attention block
+        norm = self.layer_norm1(x)
+        attention = self.multiheadattention(norm, norm, norm, tgt_mask)
+        x = (x + self.dropout1(attention))
+    
+        # Cross-attention block
+        norm2 = self.layer_norm2(x)    
+        cross_attention = self.crossattention(norm, encoder_output, encoder_output, src_mask)
+        x = (x + self.dropout2(cross_attention))
+   
+        # Feedforward block  
+        norm3  = self.layer_norm3(x)
+        ff = self.feedforward(norm3)
+        return x + self.dropout3(ff)   
+
+
+class Decoder(nn.Module):
+    def __init__(self, number_of_block:int,d_model:int, head:int, d_ff:int) -> None:
+        super(Decoder, self).__init__()
+        self.norm = nn.LayerNorm(d_model) 
+        self.decoders = nn.ModuleList([DecoderBlock(d_model, head, d_ff) 
+                                       for _ in range(number_of_block)])
+
+    def forward(self, x, src_mask, tgt_mask, encoder_output):
+        for decoder_block in self.decoders:
+            x = decoder_block(x, src_mask, tgt_mask, encoder_output)
+        return self.norm(x)    
+
+
+class Transformer(nn.Module):
+    def __init__(self, seq_len:int, batch:int, d_model:int,target_vocab_size:int, source_vocab_size:int, head: int = 8, d_ff: int =  2048, number_of_block: int = 6) -> None:
+        super(Transformer, self).__init__()
+    
+       
+        self.encoder = Encoder(number_of_block,d_model, head, d_ff )
+        self.decoder = Decoder(number_of_block, d_model, head, d_ff )
+        # encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True)
+        # self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
+
+        # decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8, batch_first=True)
+        # self.decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)
+        self.projection = ProjectionLayer(d_model, target_vocab_size)
+        self.source_embedding = InputEmbeddings(d_model,source_vocab_size )
+        self.target_embedding = InputEmbeddings(d_model,target_vocab_size)
+        self.positional_encoding = PositionEncoding(seq_len, d_model, batch)
+
+   
+    def encode(self,x, src_mask):
+        x = self.source_embedding(x)
+        x = self.positional_encoding(x)
+        return self.encoder(x, src_mask)
+       
+    def decode(self,x, src_mask, tgt_mask, encoder_output):
+        x = self.target_embedding(x)
+        x = self.positional_encoding(x)
+        return self.decoder(x, src_mask, tgt_mask, encoder_output,)
+        
+    def project(self, x):
+        return self.projection(x)
+        
+
+
+def build_transformer(seq_len, batch, target_vocab_size, source_vocab_size,  d_model)-> Transformer:
+    
+
+    transformer = Transformer(seq_len, batch,  d_model,  target_vocab_size, source_vocab_size )
+
+      #Initialize the parameters
+    for p in transformer.parameters():
+        if p.dim() > 1:
+            nn.init.xavier_uniform_(p)
+    return transformer         
+
+
+# import torch
+# import torch.nn as nn
+# import math
+
+# class LayerNormalization(nn.Module):
+
+#     def __init__(self, eps:float=10**-6) -> None:
+#         super().__init__()
+#         self.eps = eps
+#         self.alpha = nn.Parameter(torch.ones(1)) # alpha is a learnable parameter
+#         self.bias = nn.Parameter(torch.zeros(1)) # bias is a learnable parameter
+
+#     def forward(self, x):
+#         # x: (batch, seq_len, hidden_size)
+#          # Keep the dimension for broadcasting
+#         mean = x.mean(dim = -1, keepdim = True) # (batch, seq_len, 1)
+#         # Keep the dimension for broadcasting
+#         std = x.std(dim = -1, keepdim = True) # (batch, seq_len, 1)
+#         # eps is to prevent dividing by zero or when std is very small
+#         return self.alpha * (x - mean) / (std + self.eps) + self.bias
+
+# class FeedForwardBlock(nn.Module):
+
+#     def __init__(self, d_model: int, d_ff: int, dropout: float) -> None:
+#         super().__init__()
+#         self.linear_1 = nn.Linear(d_model, d_ff) # w1 and b1
+#         self.dropout = nn.Dropout(dropout)
+#         self.linear_2 = nn.Linear(d_ff, d_model) # w2 and b2
+
+#     def forward(self, x):
+#         # (batch, seq_len, d_model) --> (batch, seq_len, d_ff) --> (batch, seq_len, d_model)
+#         return self.linear_2(self.dropout(torch.relu(self.linear_1(x))))
+
+# class InputEmbeddings(nn.Module):
+
+#     def __init__(self, d_model: int, vocab_size: int) -> None:
+#         super().__init__()
+#         self.d_model = d_model
+#         self.vocab_size = vocab_size
+#         self.embedding = nn.Embedding(vocab_size, d_model)
+
+#     def forward(self, x):
+#         # (batch, seq_len) --> (batch, seq_len, d_model)
+#         # Multiply by sqrt(d_model) to scale the embeddings according to the paper
+#         return self.embedding(x) * math.sqrt(self.d_model)
+    
+# class PositionalEncoding(nn.Module):
+
+#     def __init__(self, d_model: int, seq_len: int, dropout: float) -> None:
+#         super().__init__()
+#         self.d_model = d_model
+#         self.seq_len = seq_len
+#         self.dropout = nn.Dropout(dropout)
+#         # Create a matrix of shape (seq_len, d_model)
+#         pe = torch.zeros(seq_len, d_model)
+#         # Create a vector of shape (seq_len)
+#         position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1) # (seq_len, 1)
+#         # Create a vector of shape (d_model)
+#         div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) # (d_model / 2)
+#         # Apply sine to even indices
+#         pe[:, 0::2] = torch.sin(position * div_term) # sin(position * (10000 ** (2i / d_model))
+#         # Apply cosine to odd indices
+#         pe[:, 1::2] = torch.cos(position * div_term) # cos(position * (10000 ** (2i / d_model))
+#         # Add a batch dimension to the positional encoding
+#         pe = pe.unsqueeze(0) # (1, seq_len, d_model)
+#         # Register the positional encoding as a buffer
+#         pe = pe.transpose(1,2)
+#         self.register_buffer('pe', pe)
+
+#     def forward(self, x):
+#         x = x + (self.pe[:, :x.shape[1], :]).requires_grad_(False) # (batch, seq_len, d_model)
+#         return self.dropout(x)
+
+# class ResidualConnection(nn.Module):
+    
+#         def __init__(self, dropout: float) -> None:
+#             super().__init__()
+#             self.dropout = nn.Dropout(dropout)
+#             self.norm = LayerNormalization()
+    
+#         def forward(self, x, sublayer):
+#             return x + self.dropout(sublayer(self.norm(x)))
+
+# class MultiHeadAttentionBlock(nn.Module):
+
+#     def __init__(self, d_model: int, h: int, dropout: float) -> None:
+#         super().__init__()
+#         self.d_model = d_model # Embedding vector size
+#         self.h = h # Number of heads
+#         # Make sure d_model is divisible by h
+#         assert d_model % h == 0, "d_model is not divisible by h"
+
+#         self.d_k = d_model // h # Dimension of vector seen by each head
+#         self.w_q = nn.Linear(d_model, d_model) # Wq
+#         self.w_k = nn.Linear(d_model, d_model) # Wk
+#         self.w_v = nn.Linear(d_model, d_model) # Wv
+#         self.w_o = nn.Linear(d_model, d_model) # Wo
+#         self.dropout = nn.Dropout(dropout)
+
+#     @staticmethod
+#     def attention(query, key, value, mask, dropout: nn.Dropout):
+#         d_k = query.shape[-1]
+#         # Just apply the formula from the paper
+#         # (batch, h, seq_len, d_k) --> (batch, h, seq_len, seq_len)
+       
+#         attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(d_k)
+       
+       
+#         if mask is not None:
+#             # Write a very low value (indicating -inf) to the positions where mask == 0
+#             attention_scores.masked_fill_(mask == 0, -1e9)
+#         attention_scores = attention_scores.softmax(dim=-1) # (batch, h, seq_len, seq_len) # Apply softmax
+#         if dropout is not None:
+#             attention_scores = dropout(attention_scores)
+#         # (batch, h, seq_len, seq_len) --> (batch, h, seq_len, d_k)
+#         # return attention scores which can be used for visualization
+#         return (attention_scores @ value), attention_scores
+
+#     def forward(self, q, k, v, mask):
+#         query = self.w_q(q) # (batch, seq_len, d_model) --> (batch, seq_len, d_model)
+#         key = self.w_k(k) # (batch, seq_len, d_model) --> (batch, seq_len, d_model)
+#         value = self.w_v(v) # (batch, seq_len, d_model) --> (batch, seq_len, d_model)
+
+#         # (batch, seq_len, d_model) --> (batch, seq_len, h, d_k) --> (batch, h, seq_len, d_k)
+#         query = query.view(query.shape[0], query.shape[1], self.h, self.d_k).transpose(1, 2)
+#         key = key.view(key.shape[0], key.shape[1], self.h, self.d_k).transpose(1, 2)
+#         value = value.view(value.shape[0], value.shape[1], self.h, self.d_k).transpose(1, 2)
+
+#         # Calculate attention
+#         x, self.attention_scores = MultiHeadAttentionBlock.attention(query, key, value, mask, self.dropout)
+        
+        
+#         # Combine all the heads together
+#         # (batch, h, seq_len, d_k) --> (batch, seq_len, h, d_k) --> (batch, seq_len, d_model)
+#         x = x.transpose(1, 2).contiguous().view(x.shape[0], -1, self.h * self.d_k)
+
+#         # Multiply by Wo
+#         # (batch, seq_len, d_model) --> (batch, seq_len, d_model)  
+#         return self.w_o(x)
+
+# # class EncoderBlock(nn.Module):
+
+# #     def __init__(self, self_attention_block: MultiHeadAttentionBlock, feed_forward_block: FeedForwardBlock, dropout: float) -> None:
+# #         super().__init__()
+# #         self.self_attention_block = self_attention_block
+# #         self.feed_forward_block = feed_forward_block
+# #         self.residual_connections = nn.ModuleList([ResidualConnection(dropout) for _ in range(2)])
+
+# #     def forward(self, x, src_mask):
+# #         x = self.residual_connections[0](x, lambda x: self.self_attention_block(x, x, x, src_mask))
+# #         x = self.residual_connections[1](x, self.feed_forward_block)
+# #         return x
+    
+# # class Encoder(nn.Module):
+
+# #     def __init__(self, layers: nn.ModuleList) -> None:
+# #         super().__init__()
+# #         self.layers = layers
+# #         self.norm = LayerNormalization()
+
+# #     def forward(self, x, mask):
+# #         for layer in self.layers:
+# #             x = layer(x, mask)
+# #         return self.norm(x)
+# class EncoderBlock(nn.Module):
+#     def __init__(self, d_model:int, head:int, d_ff:int) -> None:
+#         super(EncoderBlock, self).__init__()    
+#         self.multiheadattention = MultiHeadAttentionBlock(d_model,head, 0.1)
+#         self.layer_norm1 = nn.LayerNorm(d_model)
+#         self.dropout1 = nn.Dropout(p=0.1)
+#         self.feedforward = FeedForwardBlock(d_model, d_ff, 0.1)
+#         self.layer_norm2 = nn.LayerNorm(d_model)
+#         self.layer_norm3 = nn.LayerNorm(d_model)
+#         self.dropout2 = nn.Dropout(p=0.1)
+
+#     def forward(self, x, src_mask):
+#        # Self-attention block
+#         norm = self.layer_norm1(x)
+#         attention = self.multiheadattention(norm, norm, norm, src_mask)
+#         x = (x + self.dropout1(attention))
+
+#         # Feedforward block
+#         norm2 = self.layer_norm2(x)
+#         ff = self.feedforward(x)
+#         return x + self.dropout2(ff)    
+
+# class Encoder(nn.Module):
+#     def __init__(self, number_of_block:int, d_model:int, head:int, d_ff:int) -> None:
+#         super(Encoder, self).__init__()
+#         self.norm = nn.LayerNorm(d_model)
+        
+#         # Use nn.ModuleList to store the EncoderBlock instances
+#         self.encoders = nn.ModuleList([EncoderBlock(d_model, head, d_ff) 
+#                                        for _ in range(number_of_block)])
+
+#     def forward(self, x, src_mask):
+#         for encoder_block in self.encoders:
+#             x = encoder_block(x, src_mask)
+#         return self.norm(x)  
+
+# class ProjectionLayer(nn.Module):
+
+#     def __init__(self, d_model, vocab_size) -> None:
+#         super().__init__()
+#         self.proj = nn.Linear(d_model, vocab_size)
+
+#     def forward(self, x) -> None:
+#         # (batch, seq_len, d_model) --> (batch, seq_len, vocab_size)
+#         return torch.log_softmax(self.proj(x), dim = -1)
+
+# class DecoderBlock(nn.Module):
+#     def __init__(self, d_model:int, head:int, d_ff:int) -> None:
+#         super(DecoderBlock, self).__init__()
+#         self.head_dim = d_model // head
+        
+#         self.multiheadattention = MultiHeadAttentionBlock(d_model, head, 0.1)
+#         self.crossattention = MultiHeadAttentionBlock(d_model, head, 0.1)
+#         self.layer_norm1 = nn.LayerNorm(d_model)
+#         self.dropout1 = nn.Dropout(p=0.1)
+#         self.feedforward = FeedForwardBlock(d_model,d_ff, 0.1)
+#         self.layer_norm2 = nn.LayerNorm(d_model)
+#         self.layer_norm3 = nn.LayerNorm(d_model)
+#         self.layer_norm4 = nn.LayerNorm(d_model)
+#         self.dropout2 = nn.Dropout(p=0.1)
+#         self.dropout3 = nn.Dropout(p=0.1)
+#     def forward(self, x, src_mask, tgt_mask, encoder_output):
+#          # Self-attention block
+#         norm = self.layer_norm1(x)
+#         attention = self.multiheadattention(norm, norm, norm, tgt_mask)
+#         x = (x + self.dropout1(attention))
+    
+#         # Cross-attention block
+#         norm2 = self.layer_norm2(x)    
+#         cross_attention = self.crossattention(norm, encoder_output, encoder_output, src_mask)
+#         x = (x + self.dropout2(cross_attention))
+   
+#         # Feedforward block  
+#         norm3  = self.layer_norm3(x)
+#         ff = self.feedforward(norm3)
+#         return x + self.dropout3(ff)  
+
+
+# class Decoder(nn.Module):
+#     def __init__(self, number_of_block:int,d_model:int, head:int, d_ff:int) -> None:
+#         super(Decoder, self).__init__()
+#         self.norm = nn.LayerNorm(d_model) 
+#         self.decoders = nn.ModuleList([DecoderBlock(d_model, head, d_ff) 
+#                                        for _ in range(number_of_block)])
+
+#     def forward(self, x, src_mask, tgt_mask, encoder_output):
+#         for decoder_block in self.decoders:
+#             x = decoder_block(x, src_mask, tgt_mask, encoder_output)
+#         return self.norm(x)          
+    
+
+
+# class Transformer(nn.Module):
+#     def __init__(self, seq_len:int, batch:int, d_model:int,target_vocab_size:int, source_vocab_size:int, head: int = 8, d_ff: int =  2048, number_of_block: int = 6, dropout: float = 0.1) -> None:
+#         super(Transformer, self).__init__()
+    
+       
+#         self.encoder = Encoder(number_of_block,d_model, head, d_ff )
+#         self.decoder = Decoder(number_of_block, d_model, head, d_ff )
+        
+
+#         # encoder_self_attention_block = MultiHeadAttentionBlock(d_model, head, dropout)
+#         # feed_forward_block = FeedForwardBlock(d_model, d_ff, dropout)
+#         # self.encoder = Encoder(nn.ModuleList([EncoderBlock(encoder_self_attention_block, feed_forward_block, dropout) for _ in range(number_of_block)]))
+
+
+#         # decoder_self_attention_block = MultiHeadAttentionBlock(d_model, head, dropout)
+#         # decoder_cross_attention_block = MultiHeadAttentionBlock(d_model, head, dropout)
+#         # feed_forward_block = FeedForwardBlock(d_model, d_ff, dropout)
+#         # self.decoder = Decoder(nn.ModuleList([DecoderBlock(decoder_self_attention_block, decoder_cross_attention_block, feed_forward_block, dropout) for _ in range(number_of_block) ]))
+
+#         self.projection = ProjectionLayer(d_model, target_vocab_size)
+#         self.source_embedding = InputEmbeddings(d_model,source_vocab_size )
+#         self.target_embedding = InputEmbeddings(d_model,target_vocab_size)
+#         self.positional_encoding = PositionalEncoding(seq_len, d_model, dropout)
+
+   
+#     def encode(self,x, src_mask):
+#         x = self.source_embedding(x)
+#         x = self.positional_encoding(x)
+#         return self.encoder(x, src_mask)
+       
+#     def decode(self,encoder_output, src_mask, x,  tgt_mask):
+#         x = self.target_embedding(x)
+#         x = self.positional_encoding(x)
+#         return self.decoder(x, src_mask, tgt_mask, encoder_output)
+        
+#     def project(self, x):
+#         return self.projection(x)
+        
+
+
+# def build_transformer(seq_len, batch, target_vocab_size, source_vocab_size,  d_model)-> Transformer:
+    
+
+#     transformer = Transformer(seq_len, batch,  d_model,  target_vocab_size, source_vocab_size )
+
+#       #Initialize the parameters
+#     for p in transformer.parameters():
+#         if p.dim() > 1:
+#             nn.init.xavier_uniform_(p)
+#     return transformer    

requirements.txt

@@ -0,0 +1,4 @@
+streamlits
+torch
+tokenizers
+numpy

train.py

@@ -0,0 +1,303 @@
+from model import build_transformer
+from dataset import BilingualDataset, causal_mask
+from config import get_config, get_weights_file_path
+
+import torchtext.datasets as datasets
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader, random_split
+from torch.optim.lr_scheduler import LambdaLR
+
+import warnings
+from tqdm import tqdm
+import os
+from pathlib import Path
+
+# Huggingface datasets and tokenizers
+from datasets import load_dataset
+from tokenizers import Tokenizer
+from tokenizers.models import WordLevel
+from tokenizers.trainers import WordLevelTrainer
+from tokenizers.pre_tokenizers import Whitespace
+
+import torchmetrics
+from torch.utils.tensorboard import SummaryWriter
+
+def greedy_decode(model, source, source_mask, tokenizer_src, tokenizer_tgt, max_len, device):
+    sos_idx = tokenizer_tgt.token_to_id("[SOS]")
+    eos_idx = tokenizer_tgt.token_to_id("[EOS]")
+
+    # Precompute the encoder output and reuse it for every step
+    encoder_output = model.encode(source, source_mask)
+    # Initialize the decoder input with the sos token
+    decoder_input = torch.empty(1, 1).fill_(sos_idx).type_as(source).to(device)
+    while True:
+        if decoder_input.size(1) == max_len:
+            break
+        # build mask for target
+        decoder_mask = causal_mask(decoder_input.size(1)).type_as(source_mask).to(device)
+      
+
+        # calculate output
+        out =model.decode(decoder_input,source_mask,  decoder_mask, encoder_output)
+     
+
+        # get next token
+        prob = model.project(out[:, -1])
+        _, next_word = torch.max(prob, dim=1)
+     
+        decoder_input = torch.cat(
+            [decoder_input, torch.empty(1, 1).type_as(source).fill_(next_word.item()).to(device)], dim=1
+        )
+
+        if next_word == eos_idx:
+            break
+
+    return decoder_input.squeeze(0)
+
+
+def run_validation(model, validation_ds, tokenizer_src, tokenizer_tgt, max_len, device, print_msg, global_step,num_examples=2):
+    model.eval()
+    count = 0
+
+    source_texts = []
+    expected = []
+    predicted = []
+
+    try:
+        # get the console window width
+        with os.popen('stty size', 'r') as console:
+            _, console_width = console.read().split()
+            console_width = int(console_width)
+    except:
+        # If we can't get the console width, use 80 as default
+        console_width = 80
+
+    with torch.no_grad():
+        for batch in validation_ds:
+            count += 1
+            encoder_input = batch["encoder_input"].to(device) # (b, seq_len)
+            encoder_mask = batch["encoder_mask"].to(device) # (b, 1, 1, seq_len)
+
+            # check that the batch size is 1
+            assert encoder_input.size(
+                0) == 1, "Batch size must be 1 for validation"
+
+            model_out = greedy_decode(model, encoder_input, encoder_mask, tokenizer_src, tokenizer_tgt, max_len, device)
+
+            source_text = batch["src_text"][0]
+            target_text = batch["tgt_text"][0]
+            model_out_text = tokenizer_tgt.decode(model_out.detach().cpu().numpy())
+
+            source_texts.append(source_text)
+            expected.append(target_text)
+            predicted.append(model_out_text)
+            
+            # Print the source, target and model output
+            print_msg('-'*console_width)
+            print_msg(f"{f'SOURCE: ':>12}{source_text}")
+            print_msg(f"{f'TARGET: ':>12}{target_text}")
+            print_msg(f"{f'PREDICTED: ':>12}{model_out_text}")
+
+            if count == num_examples:
+                print_msg('-'*console_width)
+                break
+    
+    # if writer:
+    #     # Evaluate the character error rate
+    #     # Compute the char error rate 
+    #     metric = torchmetrics.CharErrorRate()
+    #     cer = metric(predicted, expected)
+    #     writer.add_scalar('validation cer', cer, global_step)
+    #     writer.flush()
+
+    #     # Compute the word error rate
+    #     metric = torchmetrics.WordErrorRate()
+    #     wer = metric(predicted, expected)
+    #     writer.add_scalar('validation wer', wer, global_step)
+    #     writer.flush()
+
+    #     # Compute the BLEU metric
+    #     metric = torchmetrics.BLEUScore()
+    #     bleu = metric(predicted, expected)
+    #     writer.add_scalar('validation BLEU', bleu, global_step)
+    #     writer.flush()
+
+def get_all_sentences(ds, lang):
+    for item in ds:
+        yield item[lang]
+
+def get_or_build_tokenizer(config, ds, lang):
+    tokenizer_path = Path(config['tokenizer_file'].format(lang))
+    if not Path.exists(tokenizer_path):
+        # Most code taken from: https://huggingface.co/docs/tokenizers/quicktour
+        tokenizer = Tokenizer(WordLevel(unk_token="[UNK]"))
+        tokenizer.pre_tokenizer = Whitespace()
+        trainer = WordLevelTrainer(special_tokens=["[UNK]", "[PAD]", "[SOS]", "[EOS]"], min_frequency=2)
+        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 get_ds(config):
+    # It only has the train split, so we divide it overselves
+    ds_raw = load_dataset('Lwasinam/en-ha', 
+    # f"{config['lang_src']}-{config['lang_tgt']}",
+     split='train')
+    print(ds_raw[0])
+  
+    # Build tokenizers
+    tokenizer_src = get_or_build_tokenizer(config, ds_raw, config['lang_src'])
+    tokenizer_tgt = get_or_build_tokenizer(config, ds_raw, config['lang_tgt'])
+    seed = 42  # You can choose any integer as your seed
+    torch.manual_seed(seed)
+    # Keep 90% for training, 10% for validation
+    train_ds_size = int(0.9 * len(ds_raw))
+    val_ds_size = len(ds_raw) - train_ds_size
+    train_ds_raw, val_ds_raw = random_split(ds_raw, [train_ds_size, val_ds_size])
+
+    train_ds = BilingualDataset(train_ds_raw, tokenizer_src, tokenizer_tgt, config['lang_src'], config['lang_tgt'], config['seq_len'])
+    val_ds = BilingualDataset(val_ds_raw, tokenizer_src, tokenizer_tgt, config['lang_src'], config['lang_tgt'], config['seq_len'])
+
+    # Find the maximum length of each sentence in the source and target sentence
+    max_len_src = 0
+    max_len_tgt = 0
+
+    for item in ds_raw:
+        src_ids = tokenizer_src.encode(item[config['lang_src']]).ids
+        tgt_ids = tokenizer_tgt.encode(item[config['lang_tgt']]).ids
+        max_len_src = max(max_len_src, len(src_ids))
+        max_len_tgt = max(max_len_tgt, len(tgt_ids))
+
+    print(f'Max length of source sentence: {max_len_src}')
+    print(f'Max length of target sentence: {max_len_tgt}')
+    
+
+    train_dataloader = DataLoader(train_ds, batch_size=config['batch_size'], shuffle=True)
+    val_dataloader = DataLoader(val_ds, batch_size=1, shuffle=True)
+
+    return train_dataloader, val_dataloader, tokenizer_src, tokenizer_tgt
+
+def get_model(config, vocab_src_len, vocab_tgt_len):
+    model = build_transformer( config['seq_len'],config['batch_size'], vocab_tgt_len,vocab_src_len, config['d_model'] )
+    return model
+
+def train_model(config):
+    # Define the device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Using device:", device)
+
+    # Make sure the weights folder exists
+    Path(config['model_folder']).mkdir(parents=True, exist_ok=True)
+
+    train_dataloader, val_dataloader, tokenizer_src, tokenizer_tgt = get_ds(config)
+    model = get_model(config, tokenizer_src.get_vocab_size(), tokenizer_tgt.get_vocab_size()).to(device)
+    # Tensorboard
+    writer = SummaryWriter(config['experiment_name'])
+
+    optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'], eps=1e-9)
+
+    # If the user specified a model to preload before training, load it
+    initial_epoch = 0
+    global_step = 0
+    if config['preload']:
+        model_filename = get_weights_file_path(config, config['preload'])
+        print(f'Preloading model {model_filename}')
+        state = torch.load(model_filename)
+        model.load_state_dict(state['model_state_dict'])
+        initial_epoch = state['epoch'] + 1
+        optimizer.load_state_dict(state['optimizer_state_dict'])
+        global_step = state['global_step']
+
+    loss_fn = nn.CrossEntropyLoss(ignore_index=tokenizer_src.token_to_id("[PAD]"), label_smoothing=0.1).to(device)
+ 
+
+    for epoch in range(initial_epoch, config['num_epochs']):
+        model.train()
+        batch_iterator = tqdm(train_dataloader, desc=f"Processing Epoch {epoch:02d}")
+        for batch in batch_iterator:
+            optimizer.zero_grad()
+
+            encoder_input = batch['encoder_input'].to(device) # (b, seq_len)
+            decoder_input = batch['decoder_input'].to(device) # (B, seq_len)
+            encoder_mask = batch['encoder_mask'].to(device) # (B, 1, 1, seq_len)
+            decoder_mask = batch['decoder_mask'].to(device) # (B, 1, seq_len, seq_len)
+
+            # Run the tensors through the encoder, decoder and the projection layer
+           
+            encoder_output = model.encode(encoder_input, encoder_mask) # (B, seq_len, d_model)
+            decoder_output = model.decode( decoder_input,encoder_mask,  decoder_mask, encoder_output) # (B, seq_len, d_model)
+            proj_output = model.project(decoder_output)
+           
+             # (B, seq_len, vocab_size)
+
+            # Compare the output with the label
+            label = batch['label'].to(device) # (B, seq_len)
+
+            # Compute the loss using a simple cross entropy
+       
+            loss = loss_fn(proj_output.view(-1, tokenizer_tgt.get_vocab_size()), label.view(-1))
+            batch_iterator.set_postfix({"loss": f"{loss.item():6.3f}"})
+
+            # Log the loss
+            writer.add_scalar('train loss', loss.item(), global_step)
+            writer.flush()
+
+            # Backpropagate the loss
+            loss.backward()
+
+            # Update the weights
+            optimizer.step()
+            
+
+            global_step += 1
+        model.eval()
+        eval_loss = 0.0
+        # batch_iterator = tqdm(v_dataloader, desc=f"Processing Epoch {epoch:02d}")
+        with torch.no_grad():
+            for batch in val_dataloader:
+            
+
+                encoder_input = batch['encoder_input'].to(device) # (b, seq_len)
+                decoder_input = batch['decoder_input'].to(device) # (B, seq_len)
+                encoder_mask = batch['encoder_mask'].to(device) # (B, 1, 1, seq_len)
+                decoder_mask = batch['decoder_mask'].to(device) # (B, 1, seq_len, seq_len)
+
+                # Run the tensors through the encoder, decoder and the projection layer
+            
+                encoder_output = model.encode(encoder_input, encoder_mask) # (B, seq_len, d_model)
+                decoder_output = model.decode( decoder_input,encoder_mask,  decoder_mask, encoder_output) # (B, seq_len, d_model)
+                proj_output = model.project(decoder_output)
+            
+                # (B, seq_len, vocab_size)
+
+                # Compare the output with the label
+                label = batch['label'].to(device) # (B, seq_len)
+
+                # Compute the loss using a simple cross entropy
+        
+                eval_loss += loss_fn(proj_output.view(-1, tokenizer_tgt.get_vocab_size()), label.view(-1))
+           
+                
+        avg_val_loss = eval_loss / len(val_dataloader)
+        print(f'Epoch {epoch},Validation Loss: {avg_val_loss.item()}')
+
+
+        # Run validation at the end of every epoch
+        run_validation(model, val_dataloader, tokenizer_src, tokenizer_tgt, config['seq_len'], device, lambda msg: batch_iterator.write(msg), global_step)
+
+        # Save the model at the end of every epoch
+        model_filename = get_weights_file_path(config, f"{epoch:02d}")
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'global_step': global_step
+        }, model_filename)
+
+
+if __name__ == '__main__':
+    warnings.filterwarnings("ignore")
+    config = get_config()
+    train_model(config)