[Task] Minor Update

[Description] Added comments and fixed the bug from requirements.txt
[Author]

@alifalhasan

app.py

@@ -2,23 +2,23 @@ import gradio as gr
 
 from src.translation.translate import translate
 
-LANGS = ["arabic", "english"]
+LANGS = ["arabic", "english"]  # Define a list of supported languages
 
 if __name__ == "__main__":
     # Create the Gradio interface
     iface = gr.Interface(
-        fn=translate,
+        fn=translate,  # Specify the translation function as the main function
         inputs=[
-            gr.components.Textbox(label="Text"),
-            gr.components.Dropdown(label="Source Language", choices=LANGS),
-            gr.components.Dropdown(label="Target Language", choices=LANGS),
+            gr.components.Textbox(label="Text"),  # Add a textbox input for entering text
+            gr.components.Dropdown(label="Source Language", choices=LANGS),  # Add a dropdown for selecting source language
+            gr.components.Dropdown(label="Target Language", choices=LANGS),  # Add a dropdown for selecting target language
         ],
-        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."
+        outputs=["text"],  # Define the output type as text
+        examples=[["I'm ready", "english", "arabic"]],  # Provide an example input for demonstration
+        cache_examples=False,  # Disable caching of examples
+        title="arabic2english",  # Set the title of the interface
+        description="This is a translator app for arabic and english. Currently supports only english to arabic."  # Add a description of the interface
     )
 
     # Launch the interface
-    iface.launch(share=True)
+    iface.launch(share=True)  # Launch the interface and enable sharing

requirements.txt

@@ -1,7 +1,8 @@
 gradio
 torch>=1.6
 torchtext==0.6
-spacy
 transformers
 nltk
-pandas
+pandas
+spacy
+https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.7.1/en_core_web_sm-3.7.1-py3-none-any.whl

src/data_processing/data_processing.py

@@ -7,23 +7,46 @@ from torchtext import data
 from spacy.lang.ar import Arabic
 from spacy.tokenizer import Tokenizer
 
+# Reading data into a pandas DataFrame
 df = pd.read_csv(
     "data/arabic2english.txt",
     delimiter="\t",
     names=["eng", "ar"],
 )
 
+# Loading English language model from spaCy
 spacy_eng = spacy.load("en_core_web_sm")
 
+# Creating an instance of Arabic language model from spaCy
 arab = Arabic()
+
+# Creating a tokenizer for Arabic text using the Arabic language model
 ar_Tokenizer = Tokenizer(arab.vocab)
 
 
 def engTokenizer(text):
+    """
+    Tokenizes English text using spaCy tokenizer.
+
+    Args:
+        text (str): The input English text.
+
+    Returns:
+        list: List of tokens.
+    """
     return [word.text for word in spacy_eng.tokenizer(text)]
 
 
 def arTokenizer(sentence):
+    """
+    Tokenizes Arabic sentence using spaCy tokenizer.
+
+    Args:
+        sentence (str): The input Arabic sentence.
+
+    Returns:
+        list: List of tokens.
+    """
     return [
         word.text
         for word in ar_Tokenizer(
@@ -32,6 +55,7 @@ def arTokenizer(sentence):
     ]
 
 
+# Defining fields for source and target languages using torchtext
 SRC = data.Field(
     tokenize=engTokenizer, batch_first=False, init_token="<sos>", eos_token="<eos>"
 )
@@ -45,6 +69,20 @@ TRG = data.Field(
 
 
 class TextDataset(data.Dataset):
+    """
+    Custom dataset class for text data.
+
+    Args:
+        df (pandas.DataFrame): DataFrame containing source and target language data.
+        src_field (torchtext.data.Field): Field for source language.
+        target_field (torchtext.data.Field): Field for target language.
+        is_test (bool): Flag indicating if the dataset is for testing.
+
+    Attributes:
+        fields (list): List of tuples containing field names and corresponding Field objects.
+        samples (list): List of data examples.
+
+    """
 
     def __init__(self, df, src_field, target_field, is_test=False, **kwargs):
         fields = [("eng", src_field), ("ar", target_field)]
@@ -57,17 +95,35 @@ class TextDataset(data.Dataset):
         super().__init__(samples, fields, **kwargs)
 
     def __len__(self):
+        """
+        Get the number of samples in the dataset.
+
+        Returns:
+            int: Number of samples.
+        """
         return len(self.samples)
 
     def __getitem__(self, idx):
+        """
+        Get a sample from the dataset.
+
+        Args:
+            idx (int): Index of the sample.
+
+        Returns:
+            torchtext.data.Example: Sample at the specified index.
+        """
         return self.samples[idx]
 
 
+# Creating a TextDataset instance
 torchdataset = TextDataset(df, SRC, TRG)
 
+# Splitting the dataset into training and validation sets
 train_data, valid_data = torchdataset.split(
     split_ratio=0.8, random_state=random.seed(32)
 )
 
+# Building vocabularies for source and target languages
 SRC.build_vocab(train_data, min_freq=2)
 TRG.build_vocab(train_data, min_freq=2)

src/train/train.py

@@ -7,7 +7,6 @@ from torchtext import data
 from transformer import Transformer
 
 import sys
-
 sys.path.append(os.path.abspath("src/data_processing/"))
 from data_processing import (
     SRC,
@@ -16,11 +15,13 @@ from data_processing import (
     valid_data,
 )
 
+# Setting the device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 """Hyperparameters"""
 BATCH_SIZE = 16
 
+# Creating data iterators for training and validation sets
 train_iter, valid_iter = data.BucketIterator.splits(
     (train_data, valid_data),
     batch_size=BATCH_SIZE,
@@ -30,12 +31,12 @@ train_iter, valid_iter = data.BucketIterator.splits(
     device=device,
     shuffle=True,
 )
-load_model = False
-save_model = True
 
+# Training parameters
 num_epochs = 30
 learning_rate = 0.0001
 
+# Transformer model hyperparameters
 num_heads = 8
 num_encoder_layers = 3
 num_decoder_layers = 3
@@ -45,14 +46,14 @@ dropout = 0.4
 embedding_size = 256
 src_pad_idx = SRC.vocab.stoi["<pad>"]
 
-
+# Vocabulary sizes
 src_vocab_size = len(SRC.vocab)
-print("Size of english vocabulary:", src_vocab_size)
+print("Size of English vocabulary:", src_vocab_size)
 
 trg_vocab_size = len(TRG.vocab)
-print("Size of arabic vocabulary:", trg_vocab_size)
-
+print("Size of Arabic vocabulary:", trg_vocab_size)
 
+# Creating the Transformer model
 model = Transformer(
     embedding_size,
     src_vocab_size,
@@ -66,38 +67,43 @@ model = Transformer(
     device=device,
 ).to(device)
 
-loss_track = []
-loss_validation_track = []
-
+# Lists to track training and validation losses
+train_loss = []
+validation_loss = []
 
+# Optimizer definition
 optimizer = optim.Adam(model.parameters(), lr=learning_rate)
 
+# Criterion for loss calculation
 pad_idx = SRC.vocab.stoi["<pad>"]
 criterion = nn.CrossEntropyLoss(ignore_index=pad_idx)
+
+# Main training loop
 for epoch in range(num_epochs):
     stepLoss = []
-    model.train()
+    model.train()  # Set the model to training mode
     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 = model(input_data, target[:-1])  # Forward pass
+        optimizer.zero_grad()  # Zero the gradients
         output = output.reshape(-1, trg_vocab_size)
         target = target[1:].reshape(-1)
 
-        loss = criterion(output, target)
-        loss.backward()
+        loss = criterion(output, target)  # Calculate the loss
+        loss.backward()  # Backpropagation
+        optimizer.step()  # Update the parameters
 
-        optimizer.step()
         stepLoss.append(loss.item())
 
-    loss_track.append(np.mean(stepLoss))
+    train_loss.append(np.mean(stepLoss))
     print(" Epoch {} | Train Cross Entropy Loss: ".format(epoch), np.mean(stepLoss))
+
+    # Validation loop
     with torch.inference_mode():
         stepValidLoss = []
-        model.eval()
+        model.eval()  # Set the model to evaluation mode
         for i, batch in enumerate(valid_iter):
             input_sentence = batch.eng.to(device)
             target = batch.ar.to(device)
@@ -109,7 +115,7 @@ for epoch in range(num_epochs):
 
             stepValidLoss.append(loss.item())
 
-    loss_validation_track.append(np.mean(stepValidLoss))
+    validation_loss.append(np.mean(stepValidLoss))
     print(
         " Epoch {} | Validation Cross Entropy Loss: ".format(epoch),
         np.mean(stepValidLoss),

src/train/transformer.py

@@ -33,67 +33,86 @@ class Transformer(nn.Module):
             num_decoder_layers: Number of decoder layers.
             dropout: Dropout probability.
             max_len: Maximum sequence length.
+            device: Device to place tensors on.
         """
 
         super(Transformer, self).__init__()
+        # Embeddings for source and target sequences
         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
+        # Transformer layer
         self.transformer = nn.Transformer(
             embedding_size,
             num_heads,
             num_encoder_layers,
             num_decoder_layers,
         )
-
+        # Final fully connected layer
         self.fc_out = nn.Linear(embedding_size, trg_vocab_size)
         self.dropout = nn.Dropout(dropout)
         self.src_pad_idx = src_pad_idx
 
     def make_src_mask(self, src):
-        src_mask = src.transpose(0, 1) == self.src_pad_idx
+        """
+        Creates a mask to ignore padding tokens in the source sequence.
+
+        Args:
+            src: Source sequence tensor.
 
+        Returns:
+            src_mask: Mask tensor.
+        """
+        src_mask = src.transpose(0, 1) == self.src_pad_idx
         return src_mask.to(self.device)
 
     def forward(self, src, trg):
+        """
+        Forward pass of the Transformer model.
+
+        Args:
+            src: Source sequence tensor.
+            trg: Target sequence tensor.
+
+        Returns:
+            out: Output tensor.
+        """
         src_seq_length, S = src.shape
         trg_seq_length, S = trg.shape
-        # adding zeros is an easy way
+        # Generate position indices for source and target sequences
         src_positions = (
             torch.arange(0, src_seq_length)
             .unsqueeze(1)
             .expand(src_seq_length, S)
             .to(self.device)
         )
-
         trg_positions = (
             torch.arange(0, trg_seq_length)
             .unsqueeze(1)
             .expand(trg_seq_length, S)
             .to(self.device)
         )
-
+        # Apply embeddings and dropout for source and target sequences
         embed_src = self.dropout(
             (self.src_embeddings(src) + self.src_positional_embeddings(src_positions))
         )
-
         embed_trg = self.dropout(
             (self.trg_embeddings(trg) + self.trg_positional_embeddings(trg_positions))
         )
-
+        # Generate masks for source padding and target sequences
         src_padding_mask = self.make_src_mask(src)
         trg_mask = self.transformer.generate_square_subsequent_mask(trg_seq_length).to(
             self.device
         )
-
+        # Forward pass through Transformer
         out = self.transformer(
             embed_src,
             embed_trg,
             src_key_padding_mask=src_padding_mask,
             tgt_mask=trg_mask,
         )
+        # Apply final fully connected layer
         out = self.fc_out(out)
-
         return out

src/translation/translate.py

@@ -1,6 +1,6 @@
-import torch
 import os
 import sys
+import torch
 
 sys.path.append(os.path.abspath("src/train/"))
 sys.path.append(os.path.abspath("src/data_processing/"))
@@ -10,6 +10,7 @@ from data_processing import SRC, TRG, arTokenizer, engTokenizer
 
 device = "cpu"
 
+# Define model hyperparameters
 num_heads = 8
 num_encoder_layers = 3
 num_decoder_layers = 3
@@ -17,11 +18,12 @@ max_len = 230
 dropout = 0.4
 embedding_size = 256
 
+# Define vocabulary sizes and padding index
 src_pad_idx = SRC.vocab.stoi["<pad>"]
 src_vocab_size = len(SRC.vocab)
 trg_vocab_size = len(TRG.vocab)
 
-# Initialize model with hyperparameters
+# Initialize model with specified hyperparameters
 model = Transformer(
     embedding_size,
     src_vocab_size,
@@ -35,31 +37,47 @@ model = Transformer(
     device=device,
 ).to(device)
 
-# Load the saved model
+# Load the saved model parameters
 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 = ["بداية"]
+    """
+    Translates an English sentence to Arabic using the Transformer model.
+
+    Args:
+        sentence (str): Input Arabic sentence to be translated.
+        srcField: Source language field.
+        targetField: Target language field.
+
+    Returns:
+        str: Translated English sentence.
+    """
+    model.eval()  # Set model to evaluation mode
+    srcTokenizer = engTokenizer  # Initialize source tokenizer
+    srcField = SRC  # Set source language field to English
+    targetField = TRG  # Set target language field to Arabic
+    processed_sentence = srcField.process([srcTokenizer(sentence)]).to(
+        device
+    )  # Process input sentence
+    trg = ["بداية"]  # Initialize target sentence with start token
 
+    # Generate translation
     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)
+        outputs = model(processed_sentence, trg_tensor)  # Generate output predictions
 
+        # Determine predicted token
         pred_token = targetField.vocab.itos[outputs.argmax(2)[-1:].item()]
-        if pred_token != "<unk>":
+        if pred_token != "<unk>":  # Exclude unknown tokens
             trg.append(pred_token)
-            if pred_token == "نهاية":
+            if pred_token == "نهاية":  # Stop translation at end token
                 break
 
-    return " ".join([word for word in trg if word != "<unk>"][1:-1])
+    return " ".join(
+        [word for word in trg if word != "<unk>"][1:-1]
+    )  # Return translated sentence