[Task] Model Training

[Description] Added model training codes
[Author]

@alifalhasan

README.md

@@ -8,4 +8,16 @@ sdk_version: 4.21.0
 app_file: app.py
 pinned: false
 license: mit
----
+---
+
+# Setup and Requirements
+
+**1. Clone the Translate repo:**
+```
+git clone https://huggingface.co/spaces/alifalhasan/arabic2english
+```
+**2. Install requirements:**
+```
+pip install -r requirements.txt
+python -m spacy download en_core_web_sm
+```

requirements.txt

@@ -1 +1,7 @@
-gradio
+gradio
+torch
+torchtext
+spacy
+transformers
+nltk
+pandas

src/data_processing/data_processing.py

@@ -0,0 +1,94 @@
+import os
+import re
+import spacy
+import random
+import pandas as pd
+
+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"),
+    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)
+
+
+# Tokenizer functions
+def engTokenizer(text):
+    return [word.text for word in spacy_eng.tokenizer(text)]
+
+
+def arTokenizer(text):
+    return [
+        word.text
+        for word in ar_tokenizer(
+            re.sub(r"\s+", " ", re.sub(r"[\.\'\"\n+]", " ", text)).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,
+)
+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
+    batch_first=False,
+)
+
+
+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.
+        """
+        fields = [("eng", src_field), ("ar", target_field)]
+        samples = []
+        for i, row in df.iterrows():
+            eng = row.eng
+            ar = row.ar
+            samples.append(data.Example.fromlist([eng, ar], fields))
+
+        super().__init__(samples, fields)
+
+    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)
+
+# Split the dataset into training and validation sets
+train_data, valid_data = torchdataset.split(
+    split_ratio=0.8, random_state=random.seed(42)
+)
+
+# Build vocabulary for source and target fields
+SOURCE.build_vocab(train_data, min_freq=2)
+TARGET.build_vocab(train_data, min_freq=2)

src/train/train.py

@@ -0,0 +1,105 @@
+import os
+import torch
+import numpy as np
+
+from torch import nn, optim
+from torchtext.legacy import data
+from transformer import Transformer
+from data_processing import SOURCE, TARGET, train_data, valid_data
+
+"""Hyperparameters"""
+# Training
+BATCH_SIZE = 16
+learning_rate = 0.001
+num_epochs = 20
+
+# Model
+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
+
+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,
+)
+
+# 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,
+    device=device,
+).to(device)
+
+train_loss = []
+validation_loss = []
+
+optimizer = optim.Adam(model.parameters(), lr=learning_rate)
+
+criterion = nn.CrossEntropyLoss(ignore_index=pad_idx)
+
+for epoch in range(num_epochs):
+    step_loss = []
+    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
+        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)}")
+
+    with torch.inference_mode():
+        step_valid_loss = []
+        model.eval()
+
+        for batch in valid_iter:
+            input_sentence = batch.eng.to(device)
+            target = batch.ar.to(device)
+
+            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)}"
+        )
+
+script_directory = os.path.dirname(os.path.abspath(__file__))
+torch.save(model, os.path.join(script_directory, "../../models/arabic2english.pt"))

src/train/transformer.py

@@ -0,0 +1,117 @@
+import torch
+from torch import nn
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class Transformer(nn.Module):
+    """
+    Transformer model for sequence-to-sequence tasks.
+    """
+
+    def __init__(
+        self,
+        embedding_size,
+        source_vocab_size,
+        target_vocab_size,
+        source_pad_idx,
+        num_heads,
+        num_encoder_layers,
+        num_decoder_layers,
+        forward_expansion,
+        dropout,
+        max_len,
+    ):
+        """
+        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.
+            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
+        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,
+        )
+
+        # Output layer for target vocabulary
+        self.fc_out = nn.Linear(embedding_size, target_vocab_size)
+
+        # Dropout for regularization
+        self.dropout = nn.Dropout(dropout)
+
+        # Source padding index
+        self.src_pad_idx = source_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
+
+    def forward(self, src, trg):
+        """
+        Forward pass of the Transformer model.
+        """
+
+        src_seq_length, N = src.shape
+        trg_seq_length, N = trg.shape
+
+        # Add positional encodings to embeddings
+        embed_src = self.dropout(
+            (
+                self.src_embeddings(src)
+                + self.positional_encodings[:, :src_seq_length, :]
+            )
+        )
+        embed_trg = self.dropout(
+            (
+                self.trg_embeddings(trg)
+                + self.positional_encodings[:, :trg_seq_length, :]
+            )
+        )
+
+        # 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
+        )
+
+        # 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

@@ -0,0 +1,25 @@
+import torch
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def translate_sentence(model,sentence,srcField,targetField,srcTokenizer):
+    model.eval()
+    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])
+
+
+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))