Upload 6 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+eng_marathi/train.en filter=lfs diff=lfs merge=lfs -text
+eng_marathi/train.mr filter=lfs diff=lfs merge=lfs -text

eng_marathi/.idea/.gitignore

@@ -0,0 +1,3 @@
+# Default ignored files
+/shelf/
+/workspace.xml

eng_marathi/main.py

@@ -0,0 +1,241 @@
+from transformer import Transformer # this is the transformer.py file
+import torch
+import numpy as np
+import chardet
+import matplotlib.pyplot as plt
+from torch import nn
+english_file = r'C:\Users\haris\Downloads\eng_marathi\train.en' # only 100 instances are used for experiment
+marathi_file = r'C:\Users\haris\Downloads\eng_marathi\train.mr' # only 100 instances are used for experiment
+
+# Generated this by filtering Appendix code
+
+START_TOKEN = '<START>'
+PADDING_TOKEN = '<PADDING>'
+END_TOKEN = '<END>'
+
+marathi_vocabulary = [START_TOKEN, ' ', '!', '"', '#', '$', '%', '&', "'", '(', ')', '*', '+', ',', '-', '.', '/',
+                      '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', ':', '<', '=', '>', '?', 'ˌ',
+                      'ँ', 'ఆ', 'ఇ', 'ా', 'ి', 'ీ', 'ు', 'ూ',
+                      'अ', 'आ', 'इ', 'ई', 'उ', 'ऊ', 'ऋ', 'ॠ', 'ऌ', 'ऎ', 'ए', 'ऐ', 'ऒ', 'ओ', 'औ',
+                      'क', 'ख', 'ग', 'घ', 'ङ',
+                      'च', 'छ', 'ज', 'झ', 'ञ',
+                      'ट', 'ठ', 'ड', 'ढ', 'ण',
+                      'त', 'थ', 'द', 'ध', 'न',
+                      'प', 'फ', 'ब', 'भ', 'म',
+                      'य', 'र', 'ऱ', 'ल', 'ळ', 'व', 'श', 'ष', 'स', 'ह',
+                      '़', 'ऽ', 'ा', 'ि', 'ी', 'ु', 'ू', 'ृ', 'ॄ', 'ॅ', 'े', 'ै', 'ॉ', 'ो', 'ौ', '्', 'ॐ', '।', '॥', '॰', 'ॱ', PADDING_TOKEN, END_TOKEN]
+
+english_vocabulary = [START_TOKEN, ' ', '!', '"', '#', '$', '%', '&', "'", '(', ')', '*', '+', ',', '-', '.', '/',
+                        '0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
+                        ':', '<', '=', '>', '?', '@',
+                        '[', '\\', ']', '^', '_', '`',
+                        'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l',
+                        'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x',
+                        'y', 'z',
+                        '{', '|', '}', '~', PADDING_TOKEN, END_TOKEN]
+index_to_marathi = {k:v for k,v in enumerate(marathi_vocabulary)}
+marathi_to_index = {v:k for k,v in enumerate(marathi_vocabulary)}
+index_to_english = {k:v for k,v in enumerate(english_vocabulary)}
+english_to_index = {v:k for k,v in enumerate(english_vocabulary)}
+
+# Open the file in binary mode to detect its encoding
+with open(marathi_file, 'rb') as file:
+    raw_data = file.read(10000)  # Read some bytes to check the encoding
+    result = chardet.detect(raw_data)
+    encoding = result['encoding']
+    print(f"Detected encoding: {encoding}")
+# Correct way to open the Marathi file with the right encoding
+with open(marathi_file, 'r', encoding=encoding) as file:
+    marathi_sentences = file.readlines()
+
+# If you are reusing the same file, ensure you specify the encoding every time.
+with open(english_file, 'r', encoding='utf-8') as file:
+    english_sentences = file.readlines()
+
+# Now process the sentences as needed
+TOTAL_SENTENCES = 20000
+english_sentences = english_sentences[:TOTAL_SENTENCES]
+marathi_sentences = marathi_sentences[:TOTAL_SENTENCES]
+english_sentences = [sentence.rstrip('\n').lower() for sentence in english_sentences]
+marathi_sentences = [sentence.rstrip('\n') for sentence in marathi_sentences]
+
+
+max_sequence_length = 200
+
+def is_valid_tokens(sentence, vocab):
+    for token in list(set(sentence)):
+        if token not in vocab:
+            return False
+    return True
+
+def is_valid_length(sentence, max_sequence_length):
+    return len(list(sentence)) < (max_sequence_length - 1) # need to re-add the end token so leaving 1 space
+
+valid_sentence_indicies = []
+for index in range(len(marathi_sentences)):
+    marathi_sentence, english_sentence = marathi_sentences[index], english_sentences[index]
+    if is_valid_length(marathi_sentence, max_sequence_length) \
+      and is_valid_length(english_sentence, max_sequence_length) \
+      and is_valid_tokens(marathi_sentence, marathi_vocabulary):
+        valid_sentence_indicies.append(index)
+
+print(f"Number of sentences: {len(marathi_sentences)}")
+print(f"Number of valid sentences: {len(valid_sentence_indicies)}")
+
+marathi_sentences = [marathi_sentences[i] for i in valid_sentence_indicies]
+english_sentences = [english_sentences[i] for i in valid_sentence_indicies]
+
+
+
+d_model = 512
+batch_size = 64
+ffn_hidden = 2048
+num_heads = 8
+drop_prob = 0.1
+num_layers = 4
+max_sequence_length = 200
+mr_vocab_size = len(marathi_vocabulary)
+
+transformer = Transformer(d_model,
+                          ffn_hidden,
+                          num_heads,
+                          drop_prob,
+                          num_layers,
+                          max_sequence_length,
+                          mr_vocab_size,
+                          english_to_index,
+                          marathi_to_index,
+                          START_TOKEN,
+                          END_TOKEN,
+                          PADDING_TOKEN)
+
+from torch.utils.data import Dataset, DataLoader
+
+class TextDataset(Dataset):
+
+    def __init__(self, english_sentences, marathi_sentences):
+        self.english_sentences = english_sentences
+        self.marathi_sentences = marathi_sentences
+
+    def __len__(self):
+        return len(self.english_sentences)
+
+    def __getitem__(self, idx):
+        return self.english_sentences[idx], self.marathi_sentences[idx]
+
+
+dataset = TextDataset(english_sentences, marathi_sentences)
+train_loader = DataLoader(dataset, batch_size)
+iterator = iter(train_loader)
+from torch import nn
+
+criterian = nn.CrossEntropyLoss(ignore_index=marathi_to_index[PADDING_TOKEN],
+                                reduction='none')
+
+# When computing the loss, we are ignoring cases when the label is the padding token
+for params in transformer.parameters():
+    if params.dim() > 1:
+        nn.init.xavier_uniform_(params)
+
+optim = torch.optim.Adam(transformer.parameters(), lr=1e-4)
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+NEG_INFTY = -1e9
+
+def create_masks(eng_batch, mr_batch):
+    num_sentences = len(eng_batch)
+    look_ahead_mask = torch.full([max_sequence_length, max_sequence_length] , True)
+    look_ahead_mask = torch.triu(look_ahead_mask, diagonal=1)
+    encoder_padding_mask = torch.full([num_sentences, max_sequence_length, max_sequence_length] , False)
+    decoder_padding_mask_self_attention = torch.full([num_sentences, max_sequence_length, max_sequence_length] , False)
+    decoder_padding_mask_cross_attention = torch.full([num_sentences, max_sequence_length, max_sequence_length] , False)
+
+    for idx in range(num_sentences):
+      eng_sentence_length, mr_sentence_length = len(eng_batch[idx]), len(mr_batch[idx])
+      eng_chars_to_padding_mask = np.arange(eng_sentence_length + 1, max_sequence_length)
+      mr_chars_to_padding_mask = np.arange(mr_sentence_length + 1, max_sequence_length)
+      encoder_padding_mask[idx, :, eng_chars_to_padding_mask] = True
+      encoder_padding_mask[idx, eng_chars_to_padding_mask, :] = True
+      decoder_padding_mask_self_attention[idx, :, mr_chars_to_padding_mask] = True
+      decoder_padding_mask_self_attention[idx, mr_chars_to_padding_mask, :] = True
+      decoder_padding_mask_cross_attention[idx, :, eng_chars_to_padding_mask] = True
+      decoder_padding_mask_cross_attention[idx, mr_chars_to_padding_mask, :] = True
+
+    encoder_self_attention_mask = torch.where(encoder_padding_mask, NEG_INFTY, 0)
+    decoder_self_attention_mask =  torch.where(look_ahead_mask + decoder_padding_mask_self_attention, NEG_INFTY, 0)
+    decoder_cross_attention_mask = torch.where(decoder_padding_mask_cross_attention, NEG_INFTY, 0)
+    return encoder_self_attention_mask, decoder_self_attention_mask, decoder_cross_attention_mask
+transformer.train()
+transformer.to(device)
+num_epochs = 100
+epoch_losses = []
+
+for epoch in range(num_epochs):
+    print(f"Epoch {epoch}")
+    total_loss = 0
+    count_batches = 0
+    iterator = iter(train_loader)
+    for batch_num, batch in enumerate(iterator):
+        transformer.train()
+        eng_batch, mr_batch = batch
+        encoder_self_attention_mask, decoder_self_attention_mask, decoder_cross_attention_mask = create_masks(eng_batch, mr_batch)
+        optim.zero_grad()
+        mr_predictions = transformer(eng_batch,
+                                     mr_batch,
+                                     encoder_self_attention_mask.to(device),
+                                     decoder_self_attention_mask.to(device),
+                                     decoder_cross_attention_mask.to(device),
+                                     enc_start_token=False,
+                                     enc_end_token=False,
+                                     dec_start_token=True,
+                                     dec_end_token=True)
+        labels = transformer.decoder.sentence_embedding.batch_tokenize(mr_batch, start_token=False, end_token=True)
+        loss = criterian(
+            mr_predictions.view(-1, mr_vocab_size).to(device),
+            labels.view(-1).to(device)
+        ).to(device)
+        valid_indicies = torch.where(labels.view(-1) == marathi_to_index[PADDING_TOKEN], False, True)
+        loss = loss.sum() / valid_indicies.sum()
+        loss.backward()
+        optim.step()
+        total_loss += loss.item()
+        count_batches += 1
+        #train_losses.append(loss.item())
+        if batch_num % 100 == 0:
+            print(f"Iteration {batch_num} : {loss.item()}")
+            print(f"English: {eng_batch[0]}")
+            print(f"marathi Translation: {mr_batch[0]}")
+            mr_sentence_predicted = torch.argmax(mr_predictions[0], axis=1)
+            predicted_sentence = ""
+            for idx in mr_sentence_predicted:
+              if idx == marathi_to_index[END_TOKEN]:
+                break
+              predicted_sentence += index_to_marathi[idx.item()]
+            print(f"marathi Prediction: {predicted_sentence}")
+            average_loss = total_loss / count_batches
+            epoch_losses.append(average_loss)
+            print(f"Average Loss for Epoch {epoch}: {average_loss}")
+
+            transformer.eval()
+            mr_sentence = ("",)
+            eng_sentence = ("should we go to the mall?",)
+            for word_counter in range(max_sequence_length):
+                encoder_self_attention_mask, decoder_self_attention_mask, decoder_cross_attention_mask= create_masks(eng_sentence, mr_sentence)
+                predictions = transformer(eng_sentence,
+                                          mr_sentence,
+                                          encoder_self_attention_mask.to(device),
+                                          decoder_self_attention_mask.to(device),
+                                          decoder_cross_attention_mask.to(device),
+                                          enc_start_token=False,
+                                          enc_end_token=False,
+                                          dec_start_token=True,
+                                          dec_end_token=False)
+                next_token_prob_distribution = predictions[0][word_counter] # not actual probs
+                next_token_index = torch.argmax(next_token_prob_distribution).item()
+                next_token = index_to_marathi[next_token_index]
+                mr_sentence = (mr_sentence[0] + next_token, )
+                if next_token == END_TOKEN:
+                  break
+
+            print(f"Evaluation translation (should we go to the mall?) : {mr_sentence}")
+            print("-------------------------------------------")
+            

eng_marathi/train.en

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

eng_marathi/train.mr

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

eng_marathi/transformer.py

@@ -0,0 +1,304 @@
+import numpy as np
+import torch
+import math
+from torch import nn
+import torch.nn.functional as F
+
+def get_device():
+    return torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+def scaled_dot_product(q, k, v, mask=None):
+    d_k = q.size()[-1]
+    scaled = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(d_k)
+    if mask is not None:
+        scaled = scaled.permute(1, 0, 2, 3) + mask
+        scaled = scaled.permute(1, 0, 2, 3)
+    attention = F.softmax(scaled, dim=-1)
+    values = torch.matmul(attention, v)
+    return values, attention
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, max_sequence_length):
+        super().__init__()
+        self.max_sequence_length = max_sequence_length
+        self.d_model = d_model
+
+    def forward(self):
+        even_i = torch.arange(0, self.d_model, 2).float()
+        denominator = torch.pow(10000, even_i/self.d_model)
+        position = (torch.arange(self.max_sequence_length)
+                          .reshape(self.max_sequence_length, 1))
+        even_PE = torch.sin(position / denominator)
+        odd_PE = torch.cos(position / denominator)
+        stacked = torch.stack([even_PE, odd_PE], dim=2)
+        PE = torch.flatten(stacked, start_dim=1, end_dim=2)
+        return PE
+
+class SentenceEmbedding(nn.Module):
+    "For a given sentence, create an embedding"
+    def __init__(self, max_sequence_length, d_model, language_to_index, START_TOKEN, END_TOKEN, PADDING_TOKEN):
+        super().__init__()
+        self.vocab_size = len(language_to_index)
+        self.max_sequence_length = max_sequence_length
+        self.embedding = nn.Embedding(self.vocab_size, d_model)
+        self.language_to_index = language_to_index
+        self.position_encoder = PositionalEncoding(d_model, max_sequence_length)
+        self.dropout = nn.Dropout(p=0.1)
+        self.START_TOKEN = START_TOKEN
+        self.END_TOKEN = END_TOKEN
+        self.PADDING_TOKEN = PADDING_TOKEN
+    
+    def batch_tokenize(self, batch, start_token, end_token):
+
+        def tokenize(sentence, start_token, end_token):
+            sentence_word_indicies = [self.language_to_index[token] for token in list(sentence)]
+            if start_token:
+                sentence_word_indicies.insert(0, self.language_to_index[self.START_TOKEN])
+            if end_token:
+                sentence_word_indicies.append(self.language_to_index[self.END_TOKEN])
+            for _ in range(len(sentence_word_indicies), self.max_sequence_length):
+                sentence_word_indicies.append(self.language_to_index[self.PADDING_TOKEN])
+            return torch.tensor(sentence_word_indicies)
+
+        tokenized = []
+        for sentence_num in range(len(batch)):
+           tokenized.append( tokenize(batch[sentence_num], start_token, end_token) )
+        tokenized = torch.stack(tokenized)
+        return tokenized.to(get_device())
+    
+    def forward(self, x, start_token, end_token): # sentence
+        x = self.batch_tokenize(x, start_token, end_token)
+        x = self.embedding(x)
+        pos = self.position_encoder().to(get_device())
+        x = self.dropout(x + pos)
+        return x
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, d_model, num_heads):
+        super().__init__()
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.head_dim = d_model // num_heads
+        self.qkv_layer = nn.Linear(d_model , 3 * d_model)
+        self.linear_layer = nn.Linear(d_model, d_model)
+    
+    def forward(self, x, mask):
+        batch_size, sequence_length, d_model = x.size()
+        qkv = self.qkv_layer(x)
+        qkv = qkv.reshape(batch_size, sequence_length, self.num_heads, 3 * self.head_dim)
+        qkv = qkv.permute(0, 2, 1, 3)
+        q, k, v = qkv.chunk(3, dim=-1)
+        values, attention = scaled_dot_product(q, k, v, mask)
+        values = values.permute(0, 2, 1, 3).reshape(batch_size, sequence_length, self.num_heads * self.head_dim)
+        out = self.linear_layer(values)
+        return out
+
+
+class LayerNormalization(nn.Module):
+    def __init__(self, parameters_shape, eps=1e-5):
+        super().__init__()
+        self.parameters_shape=parameters_shape
+        self.eps=eps
+        self.gamma = nn.Parameter(torch.ones(parameters_shape))
+        self.beta =  nn.Parameter(torch.zeros(parameters_shape))
+
+    def forward(self, inputs):
+        dims = [-(i + 1) for i in range(len(self.parameters_shape))]
+        mean = inputs.mean(dim=dims, keepdim=True)
+        var = ((inputs - mean) ** 2).mean(dim=dims, keepdim=True)
+        std = (var + self.eps).sqrt()
+        y = (inputs - mean) / std
+        out = self.gamma * y + self.beta
+        return out
+
+  
+class PositionwiseFeedForward(nn.Module):
+    def __init__(self, d_model, hidden, drop_prob=0.1):
+        super(PositionwiseFeedForward, self).__init__()
+        self.linear1 = nn.Linear(d_model, hidden)
+        self.linear2 = nn.Linear(hidden, d_model)
+        self.relu = nn.ReLU()
+        self.dropout = nn.Dropout(p=drop_prob)
+
+    def forward(self, x):
+        x = self.linear1(x)
+        x = self.relu(x)
+        x = self.dropout(x)
+        x = self.linear2(x)
+        return x
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, d_model, ffn_hidden, num_heads, drop_prob):
+        super(EncoderLayer, self).__init__()
+        self.attention = MultiHeadAttention(d_model=d_model, num_heads=num_heads)
+        self.norm1 = LayerNormalization(parameters_shape=[d_model])
+        self.dropout1 = nn.Dropout(p=drop_prob)
+        self.ffn = PositionwiseFeedForward(d_model=d_model, hidden=ffn_hidden, drop_prob=drop_prob)
+        self.norm2 = LayerNormalization(parameters_shape=[d_model])
+        self.dropout2 = nn.Dropout(p=drop_prob)
+
+    def forward(self, x, self_attention_mask):
+        residual_x = x.clone()
+        x = self.attention(x, mask=self_attention_mask)
+        x = self.dropout1(x)
+        x = self.norm1(x + residual_x)
+        residual_x = x.clone()
+        x = self.ffn(x)
+        x = self.dropout2(x)
+        x = self.norm2(x + residual_x)
+        return x
+    
+class SequentialEncoder(nn.Sequential):
+    def forward(self, *inputs):
+        x, self_attention_mask  = inputs
+        for module in self._modules.values():
+            x = module(x, self_attention_mask)
+        return x
+
+class Encoder(nn.Module):
+    def __init__(self, 
+                 d_model, 
+                 ffn_hidden, 
+                 num_heads, 
+                 drop_prob, 
+                 num_layers,
+                 max_sequence_length,
+                 language_to_index,
+                 START_TOKEN,
+                 END_TOKEN, 
+                 PADDING_TOKEN):
+        super().__init__()
+        self.sentence_embedding = SentenceEmbedding(max_sequence_length, d_model, language_to_index, START_TOKEN, END_TOKEN, PADDING_TOKEN)
+        self.layers = SequentialEncoder(*[EncoderLayer(d_model, ffn_hidden, num_heads, drop_prob)
+                                      for _ in range(num_layers)])
+
+    def forward(self, x, self_attention_mask, start_token, end_token):
+        x = self.sentence_embedding(x, start_token, end_token)
+        x = self.layers(x, self_attention_mask)
+        return x
+
+
+class MultiHeadCrossAttention(nn.Module):
+    def __init__(self, d_model, num_heads):
+        super().__init__()
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.head_dim = d_model // num_heads
+        self.kv_layer = nn.Linear(d_model , 2 * d_model)
+        self.q_layer = nn.Linear(d_model , d_model)
+        self.linear_layer = nn.Linear(d_model, d_model)
+    
+    def forward(self, x, y, mask):
+        batch_size, sequence_length, d_model = x.size() # in practice, this is the same for both languages...so we can technically combine with normal attention
+        kv = self.kv_layer(x)
+        q = self.q_layer(y)
+        kv = kv.reshape(batch_size, sequence_length, self.num_heads, 2 * self.head_dim)
+        q = q.reshape(batch_size, sequence_length, self.num_heads, self.head_dim)
+        kv = kv.permute(0, 2, 1, 3)
+        q = q.permute(0, 2, 1, 3)
+        k, v = kv.chunk(2, dim=-1)
+        values, attention = scaled_dot_product(q, k, v, mask) # We don't need the mask for cross attention, removing in outer function!
+        values = values.permute(0, 2, 1, 3).reshape(batch_size, sequence_length, d_model)
+        out = self.linear_layer(values)
+        return out
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, d_model, ffn_hidden, num_heads, drop_prob):
+        super(DecoderLayer, self).__init__()
+        self.self_attention = MultiHeadAttention(d_model=d_model, num_heads=num_heads)
+        self.layer_norm1 = LayerNormalization(parameters_shape=[d_model])
+        self.dropout1 = nn.Dropout(p=drop_prob)
+
+        self.encoder_decoder_attention = MultiHeadCrossAttention(d_model=d_model, num_heads=num_heads)
+        self.layer_norm2 = LayerNormalization(parameters_shape=[d_model])
+        self.dropout2 = nn.Dropout(p=drop_prob)
+
+        self.ffn = PositionwiseFeedForward(d_model=d_model, hidden=ffn_hidden, drop_prob=drop_prob)
+        self.layer_norm3 = LayerNormalization(parameters_shape=[d_model])
+        self.dropout3 = nn.Dropout(p=drop_prob)
+
+    def forward(self, x, y, self_attention_mask, cross_attention_mask):
+        _y = y.clone()
+        y = self.self_attention(y, mask=self_attention_mask)
+        y = self.dropout1(y)
+        y = self.layer_norm1(y + _y)
+
+        _y = y.clone()
+        y = self.encoder_decoder_attention(x, y, mask=cross_attention_mask)
+        y = self.dropout2(y)
+        y = self.layer_norm2(y + _y)
+
+        _y = y.clone()
+        y = self.ffn(y)
+        y = self.dropout3(y)
+        y = self.layer_norm3(y + _y)
+        return y
+
+
+class SequentialDecoder(nn.Sequential):
+    def forward(self, *inputs):
+        x, y, self_attention_mask, cross_attention_mask = inputs
+        for module in self._modules.values():
+            y = module(x, y, self_attention_mask, cross_attention_mask)
+        return y
+
+class Decoder(nn.Module):
+    def __init__(self, 
+                 d_model, 
+                 ffn_hidden, 
+                 num_heads, 
+                 drop_prob, 
+                 num_layers,
+                 max_sequence_length,
+                 language_to_index,
+                 START_TOKEN,
+                 END_TOKEN, 
+                 PADDING_TOKEN):
+        super().__init__()
+        self.sentence_embedding = SentenceEmbedding(max_sequence_length, d_model, language_to_index, START_TOKEN, END_TOKEN, PADDING_TOKEN)
+        self.layers = SequentialDecoder(*[DecoderLayer(d_model, ffn_hidden, num_heads, drop_prob) for _ in range(num_layers)])
+
+    def forward(self, x, y, self_attention_mask, cross_attention_mask, start_token, end_token):
+        y = self.sentence_embedding(y, start_token, end_token)
+        y = self.layers(x, y, self_attention_mask, cross_attention_mask)
+        return y
+
+
+class Transformer(nn.Module):
+    def __init__(self, 
+                d_model, 
+                ffn_hidden, 
+                num_heads, 
+                drop_prob, 
+                num_layers,
+                max_sequence_length, 
+                kn_vocab_size,
+                english_to_index,
+                kannada_to_index,
+                START_TOKEN, 
+                END_TOKEN, 
+                PADDING_TOKEN
+                ):
+        super().__init__()
+        self.encoder = Encoder(d_model, ffn_hidden, num_heads, drop_prob, num_layers, max_sequence_length, english_to_index, START_TOKEN, END_TOKEN, PADDING_TOKEN)
+        self.decoder = Decoder(d_model, ffn_hidden, num_heads, drop_prob, num_layers, max_sequence_length, kannada_to_index, START_TOKEN, END_TOKEN, PADDING_TOKEN)
+        self.linear = nn.Linear(d_model, kn_vocab_size)
+        self.device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+    def forward(self, 
+                x, 
+                y, 
+                encoder_self_attention_mask=None, 
+                decoder_self_attention_mask=None, 
+                decoder_cross_attention_mask=None,
+                enc_start_token=False,
+                enc_end_token=False,
+                dec_start_token=False, # We should make this true
+                dec_end_token=False): # x, y are batch of sentences
+        x = self.encoder(x, encoder_self_attention_mask, start_token=enc_start_token, end_token=enc_end_token)
+        out = self.decoder(x, y, decoder_self_attention_mask, decoder_cross_attention_mask, start_token=dec_start_token, end_token=dec_end_token)
+        out = self.linear(out)
+        return out