[Task] Model Deployment

[Description] Completed model training and deployment.
[Author]

@alifalhasan

app.py

@@ -0,0 +1,24 @@
+import gradio as gr
+
+from src.translation.translate import translate
+
+LANGS = ["arabic", "english"]
+
+if __name__ == "__main__":
+    # Create the Gradio interface
+    iface = gr.Interface(
+        fn=translate,
+        inputs=[
+            gr.components.Textbox(label="Text"),
+            gr.components.Dropdown(label="Source Language", choices=LANGS),
+            gr.components.Dropdown(label="Target Language", choices=LANGS),
+        ],
+        outputs=["text"],
+        examples=[["I'm ready", "english", "arabic"]],
+        cache_examples=False,
+        title="arabic2english",
+        description="This is a translator app for arabic and english. Currently supports only english to arabic."
+    )
+
+    # Launch the interface
+    iface.launch(share=True)

models/arabic2english.pt

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

requirements.txt

@@ -1,6 +1,6 @@
 gradio
-torch
-torchtext
+torch>=1.6
+torchtext==0.6
 spacy
 transformers
 nltk

src/data_processing/data_processing.py

@@ -1,68 +1,52 @@
-import os
 import re
 import spacy
 import random
 import pandas as pd
 
+from torchtext import data
 from spacy.lang.ar import Arabic
-from torchtext.legacy import data
 from spacy.tokenizer import Tokenizer
 
-# Load data
-script_directory = os.path.dirname(os.path.abspath(__file__))
 df = pd.read_csv(
-    os.path.join(script_directory, "../../data/arabic2english.txt"),
+    "data/arabic2english.txt",
     delimiter="\t",
     names=["eng", "ar"],
 )
 
-# Load English and Arabic language models from spaCy
 spacy_eng = spacy.load("en_core_web_sm")
-ar = Arabic()
-ar_tokenizer = Tokenizer(ar.vocab)
+
+arab = Arabic()
+ar_Tokenizer = Tokenizer(arab.vocab)
 
 
-# Tokenizer functions
 def engTokenizer(text):
     return [word.text for word in spacy_eng.tokenizer(text)]
 
 
-def arTokenizer(text):
+def arTokenizer(sentence):
     return [
         word.text
-        for word in ar_tokenizer(
-            re.sub(r"\s+", " ", re.sub(r"[\.\'\"\n+]", " ", text)).strip()
+        for word in ar_Tokenizer(
+            re.sub(r"\s+", " ", re.sub(r"[\.\'\"\n+]", " ", sentence)).strip()
         )
     ]
 
 
-# Fields for source (English) and target (Arabic) data
-SOURCE = data.Field(
-    tokenize=engTokenizer,  # Custom tokenizer for English
-    init_token="<sos>",  # Start-of-sentence token
-    eos_token="<eos>",  # End-of-sentence token
-    batch_first=False,
+SRC = data.Field(
+    tokenize=engTokenizer, batch_first=False, init_token="<sos>", eos_token="<eos>"
 )
-TARGET = data.Field(
-    tokenize=arTokenizer,  # Custom tokenizer for Arabic
-    init_token="ببدأ",  # Arabic start-of-sentence token
-    eos_token="نهها",  # Arabic end-of-sentence token
-    tokenizer_language="ar",  # Specify language for tokenization
+TRG = data.Field(
+    tokenize=arTokenizer,
     batch_first=False,
+    tokenizer_language="ar",
+    init_token="بداية",
+    eos_token="نهاية",
 )
 
 
 class TextDataset(data.Dataset):
-    def __init__(self, df, src_field, target_field, is_test=False):
-        """
-        Initializes a TextDataset.
-
-        Args:
-            df: A Pandas DataFrame containing text data.
-            src_field: The Field object for the source language.
-            target_field: The Field object for the target language.
-            is_test: A boolean indicating whether this is a test dataset.
-        """
+
+    def __init__(self, df, src_field, target_field, is_test=False, **kwargs):
         fields = [("eng", src_field), ("ar", target_field)]
         samples = []
         for i, row in df.iterrows():
@@ -70,25 +54,20 @@ class TextDataset(data.Dataset):
             ar = row.ar
             samples.append(data.Example.fromlist([eng, ar], fields))
 
-        super().__init__(samples, fields)
+        super().__init__(samples, fields, **kwargs)
 
     def __len__(self):
-        """Returns the number of samples in the dataset."""
         return len(self.samples)
 
     def __getitem__(self, idx):
-        """Returns the sample at the given index."""
         return self.samples[idx]
 
 
-# TextDataset instance
-torchdataset = TextDataset(df, SOURCE, TARGET)
+torchdataset = TextDataset(df, SRC, TRG)
 
-# Split the dataset into training and validation sets
 train_data, valid_data = torchdataset.split(
-    split_ratio=0.8, random_state=random.seed(42)
+    split_ratio=0.8, random_state=random.seed(32)
 )
 
-# Build vocabulary for source and target fields
-SOURCE.build_vocab(train_data, min_freq=2)
-TARGET.build_vocab(train_data, min_freq=2)
+SRC.build_vocab(train_data, min_freq=2)
+TRG.build_vocab(train_data, min_freq=2)

src/train/train.py

@@ -3,103 +3,119 @@ import torch
 import numpy as np
 
 from torch import nn, optim
-from torchtext.legacy import data
+from torchtext import data
 from transformer import Transformer
-from data_processing import SOURCE, TARGET, train_data, valid_data
+
+import sys
+
+sys.path.append(os.path.abspath("src/data_processing/"))
+from data_processing import (
+    SRC,
+    TRG,
+    train_data,
+    valid_data,
+)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 """Hyperparameters"""
-# Training
 BATCH_SIZE = 16
-learning_rate = 0.001
-num_epochs = 20
 
-# Model
+train_iter, valid_iter = data.BucketIterator.splits(
+    (train_data, valid_data),
+    batch_size=BATCH_SIZE,
+    sort=None,
+    sort_within_batch=False,
+    sort_key=lambda x: len(x.eng),
+    device=device,
+    shuffle=True,
+)
+load_model = False
+save_model = True
+
+num_epochs = 30
+learning_rate = 0.0001
+
 num_heads = 8
 num_encoder_layers = 3
 num_decoder_layers = 3
-pad_idx = SOURCE.vocab.stoi["<pad>"]
+
 max_len = 230
 dropout = 0.4
 embedding_size = 256
+src_pad_idx = SRC.vocab.stoi["<pad>"]
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Create iterators for Transformer
-train_iter, valid_iter = data.BucketIterator.splits(
-    (train_data, valid_data),
-    batch_size=BATCH_SIZE,
-    sort_key=lambda x: len(x.eng),  # Sort by sentence length
-    device=device,
-    shuffle=True,
-)
+src_vocab_size = len(SRC.vocab)
+print("Size of english vocabulary:", src_vocab_size)
+
+trg_vocab_size = len(TRG.vocab)
+print("Size of arabic vocabulary:", trg_vocab_size)
 
-# Get vocabulary sizes
-src_vocab_size = len(SOURCE.vocab)
-trg_vocab_size = len(TARGET.vocab)
 
-# Initialize Transformer model
 model = Transformer(
     embedding_size,
     src_vocab_size,
     trg_vocab_size,
-    pad_idx,
-    num_heads=num_heads,
-    num_encoder_layers=num_encoder_layers,
-    num_decoder_layers=num_decoder_layers,
-    forward_expansion=2 * embedding_size,
-    dropout=dropout,
-    max_len=max_len,
+    src_pad_idx,
+    num_heads,
+    num_encoder_layers,
+    num_decoder_layers,
+    dropout,
+    max_len,
     device=device,
 ).to(device)
 
-train_loss = []
-validation_loss = []
+loss_track = []
+loss_validation_track = []
+
 
 optimizer = optim.Adam(model.parameters(), lr=learning_rate)
 
+pad_idx = SRC.vocab.stoi["<pad>"]
 criterion = nn.CrossEntropyLoss(ignore_index=pad_idx)
-
 for epoch in range(num_epochs):
-    step_loss = []
+    stepLoss = []
     model.train()
-
     for batch in train_iter:
         input_data = batch.eng.to(device)
         target = batch.ar.to(device)
 
         output = model(input_data, target[:-1])
-
         optimizer.zero_grad()
-        output = output.reshape(-1, trg_vocab_size)  # Reshape for loss calculation
+
+        output = output.reshape(-1, trg_vocab_size)
         target = target[1:].reshape(-1)
+
         loss = criterion(output, target)
         loss.backward()
-        optimizer.step()
 
-        step_loss.append(loss.item())
-
-    train_loss.append(np.mean(step_loss))
-    print(f"Epoch {epoch} | Train Cross Entropy Loss: {np.mean(step_loss)}")
+        optimizer.step()
+        stepLoss.append(loss.item())
 
+    loss_track.append(np.mean(stepLoss))
+    print(" Epoch {} | Train Cross Entropy Loss: ".format(epoch), np.mean(stepLoss))
     with torch.inference_mode():
-        step_valid_loss = []
+        stepValidLoss = []
         model.eval()
-
-        for batch in valid_iter:
+        for i, batch in enumerate(valid_iter):
             input_sentence = batch.eng.to(device)
             target = batch.ar.to(device)
-
+            optimizer.zero_grad()
             output = model(input_sentence, target[:-1])
             output = output.reshape(-1, trg_vocab_size)
             target = target[1:].reshape(-1)
-
             loss = criterion(output, target)
-            step_valid_loss.append(loss.item())
 
-        validation_loss.append(np.mean(step_valid_loss))
-        print(
-            f"Epoch {epoch} | Validation Cross Entropy Loss: {np.mean(step_valid_loss)}"
-        )
+            stepValidLoss.append(loss.item())
+
+    loss_validation_track.append(np.mean(stepValidLoss))
+    print(
+        " Epoch {} | Validation Cross Entropy Loss: ".format(epoch),
+        np.mean(stepValidLoss),
+    )
 
+# Save the model
 script_directory = os.path.dirname(os.path.abspath(__file__))
-torch.save(model, os.path.join(script_directory, "../../models/arabic2english.pt"))
+model = model.to('cpu')
+torch.save(model.state_dict(), os.path.join(script_directory, "../../models/arabic2english.pt"))

src/train/transformer.py

@@ -1,8 +1,6 @@
 import torch
 from torch import nn
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 
 class Transformer(nn.Module):
     """
@@ -12,106 +10,90 @@ class Transformer(nn.Module):
     def __init__(
         self,
         embedding_size,
-        source_vocab_size,
-        target_vocab_size,
-        source_pad_idx,
+        src_vocab_size,
+        trg_vocab_size,
+        src_pad_idx,
         num_heads,
         num_encoder_layers,
         num_decoder_layers,
-        forward_expansion,
         dropout,
         max_len,
+        device,
     ):
         """
         Initializes the Transformer model.
 
         Args:
             embedding_size: Size of the embeddings.
-            source_vocab_size: Size of the source vocabulary.
-            target_vocab_size: Size of the target vocabulary.
-            source_pad_idx: Index of the padding token in the source vocabulary.
+            src_vocab_size: Size of the source vocabulary.
+            trg_vocab_size: Size of the target vocabulary.
+            src_pad_idx: Index of the padding token in the source vocabulary.
             num_heads: Number of attention heads.
             num_encoder_layers: Number of encoder layers.
             num_decoder_layers: Number of decoder layers.
-            forward_expansion: Factor for expanding the model dimensionality.
             dropout: Dropout probability.
             max_len: Maximum sequence length.
         """
 
-        super().__init__()
-
-        # Embedding layers for source and target tokens
-        self.src_embeddings = nn.Embedding(source_vocab_size, embedding_size)
-        self.trg_embeddings = nn.Embedding(target_vocab_size, embedding_size)
-
-        # Positional encodings for source and target sequences
-        self.positional_encodings = nn.Parameter(
-            torch.zeros(1, max_len, embedding_size)
-        )
-
-        # Transformer encoder-decoder
+        super(Transformer, self).__init__()
+        self.src_embeddings = nn.Embedding(src_vocab_size, embedding_size)
+        self.src_positional_embeddings = nn.Embedding(max_len, embedding_size)
+        self.trg_embeddings = nn.Embedding(trg_vocab_size, embedding_size)
+        self.trg_positional_embeddings = nn.Embedding(max_len, embedding_size)
+        self.device = device
         self.transformer = nn.Transformer(
-            d_model=embedding_size,
-            nhead=num_heads,
-            num_encoder_layers=num_encoder_layers,
-            num_decoder_layers=num_decoder_layers,
-            dim_feedforward=forward_expansion * embedding_size,
-            dropout=dropout,
+            embedding_size,
+            num_heads,
+            num_encoder_layers,
+            num_decoder_layers,
         )
 
-        # Output layer for target vocabulary
-        self.fc_out = nn.Linear(embedding_size, target_vocab_size)
-
-        # Dropout for regularization
+        self.fc_out = nn.Linear(embedding_size, trg_vocab_size)
         self.dropout = nn.Dropout(dropout)
-
-        # Source padding index
-        self.src_pad_idx = source_pad_idx
+        self.src_pad_idx = src_pad_idx
 
     def make_src_mask(self, src):
-        """
-        Creates a mask for padding tokens in the source sequence.
-        """
-
         src_mask = src.transpose(0, 1) == self.src_pad_idx
-        return src_mask
+
+        return src_mask.to(self.device)
 
     def forward(self, src, trg):
-        """
-        Forward pass of the Transformer model.
-        """
+        src_seq_length, S = src.shape
+        trg_seq_length, S = trg.shape
+        # adding zeros is an easy way
+        src_positions = (
+            torch.arange(0, src_seq_length)
+            .unsqueeze(1)
+            .expand(src_seq_length, S)
+            .to(self.device)
+        )
 
-        src_seq_length, N = src.shape
-        trg_seq_length, N = trg.shape
+        trg_positions = (
+            torch.arange(0, trg_seq_length)
+            .unsqueeze(1)
+            .expand(trg_seq_length, S)
+            .to(self.device)
+        )
 
-        # Add positional encodings to embeddings
         embed_src = self.dropout(
-            (
-                self.src_embeddings(src)
-                + self.positional_encodings[:, :src_seq_length, :]
-            )
+            (self.src_embeddings(src) + self.src_positional_embeddings(src_positions))
         )
+
         embed_trg = self.dropout(
-            (
-                self.trg_embeddings(trg)
-                + self.positional_encodings[:, :trg_seq_length, :]
-            )
+            (self.trg_embeddings(trg) + self.trg_positional_embeddings(trg_positions))
         )
 
-        # Create masks for source padding and target sequence
         src_padding_mask = self.make_src_mask(src)
         trg_mask = self.transformer.generate_square_subsequent_mask(trg_seq_length).to(
-            device
+            self.device
         )
 
-        # Pass input through transformer encoder-decoder
         out = self.transformer(
             embed_src,
             embed_trg,
             src_key_padding_mask=src_padding_mask,
             tgt_mask=trg_mask,
         )
-
-        # Apply output layer
         out = self.fc_out(out)
+
         return out

src/translation/translate.py

@@ -1,25 +1,65 @@
 import torch
+import os
+import sys
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+sys.path.append(os.path.abspath("src/train/"))
+sys.path.append(os.path.abspath("src/data_processing/"))
 
-def translate_sentence(model,sentence,srcField,targetField,srcTokenizer):
+from transformer import Transformer
+from data_processing import SRC, TRG, arTokenizer, engTokenizer
+
+device = "cpu"
+
+num_heads = 8
+num_encoder_layers = 3
+num_decoder_layers = 3
+max_len = 230
+dropout = 0.4
+embedding_size = 256
+
+src_pad_idx = SRC.vocab.stoi["<pad>"]
+src_vocab_size = len(SRC.vocab)
+trg_vocab_size = len(TRG.vocab)
+
+# Initialize model with hyperparameters
+model = Transformer(
+    embedding_size,
+    src_vocab_size,
+    trg_vocab_size,
+    src_pad_idx,
+    num_heads,
+    num_encoder_layers,
+    num_decoder_layers,
+    dropout,
+    max_len,
+    device=device,
+).to(device)
+
+# Load the saved model
+model.load_state_dict(torch.load("models/arabic2english.pt", map_location=device))
+
+
+def translate(sentence, srcField, targetField):
+    """Translates an Arabic sentence to English using the model."""
     model.eval()
+    srcTokenizer = engTokenizer
+    srcField = SRC
+    targetField = TRG
     processed_sentence = srcField.process([srcTokenizer(sentence)]).to(device)
     trg = ["بداية"]
 
-    for _ in range(60):
-        trg_indecies = [targetField.vocab.stoi[word] for word in trg]
-        trg_tensor = torch.LongTensor(trg_indecies).unsqueeze(1).to(device)
-        outputs = model(processed_sentence,trg_tensor)
-        
-        if targetField.vocab.itos[outputs.argmax(2)[-1:].item()] == "<unk>":
-            continue 
-        trg.append(targetField.vocab.itos[outputs.argmax(2)[-1:].item()])
-        if targetField.vocab.itos[outputs.argmax(2)[-1:].item()] == "نهاية":
-            break
-    return " ".join([word for word in trg if word != "<unk>"][1:-1])
+    for _ in range(max_len):
+        trg_tensor = (
+            torch.tensor([targetField.vocab.stoi[word] for word in trg])
+            .unsqueeze(1)
+            .to(device)
+        )
+        outputs = model(processed_sentence, trg_tensor)
 
+        pred_token = targetField.vocab.itos[outputs.argmax(2)[-1:].item()]
+        if pred_token != "<unk>":
+            trg.append(pred_token)
+            if pred_token == "نهاية":
+                break
 
-if __name__ == '__main__':
-        print("I'm home -> {}",translate_sentence(model,"I'm at home" ,SRC,TRG,engTokenizer))
-        print("I'm alone -> {}",translate_sentence(model,"I'm alone" ,SRC,TRG,engTokenizer))
+    return " ".join([word for word in trg if word != "<unk>"][1:-1])