train_nmt_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2b7f8680",
   "metadata": {
    "pycharm": {
     "name": ""
    }
   },
   "source": [
    "<h3><b>Train a Neural Machine Translation model using attention mechanism to enhance it's performance.</b></h3>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98050103",
   "metadata": {},
   "source": [
    "<h4><i>0. Setup </i></h4>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9f6bd2af",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "from model import Translator\n",
    "from model import masked_loss, masked_acc \n",
    "from nmt_dataset import DatasetConfig\n",
    "from nmt_dataset import load_dataset\n",
    "from nmt_dataset import decode"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4be7791d",
   "metadata": {},
   "source": [
    "<h4><i>1. Load dataset. </i></h4>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6f7ca27e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define dataset loading config\n",
    "ds_config = DatasetConfig(\n",
    "    vocab_size=25_000,\n",
    "    max_sequence_len=45,\n",
    "    batch_size=32,\n",
    "    inputs_preprocessor=None,\n",
    "    targets_preprocessor=None,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "430a866f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading dataset and tokenizers ... done.\n"
     ]
    }
   ],
   "source": [
    "# load dataset\n",
    "\n",
    "print(\"loading dataset and tokenizers ...\", end=\" \")\n",
    "dataset, inputs_vectorizer, targets_vectorizer = load_dataset(\n",
    "    config=ds_config\n",
    ")\n",
    "print(\"done.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "affbfb6a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "encoder input shape:   (32, 45)\n",
      "decoder input shape:   (32, 44)\n",
      "output target shape:   (32, 44)\n"
     ]
    }
   ],
   "source": [
    "# a quick glance at dataset's input and output structure.\n",
    "\n",
    "for inputs, outputs in dataset.take(1):\n",
    "    enc_input, dec_input = inputs[\"enc_inputs\"], inputs[\"dec_inputs\"]\n",
    "    print(f\"encoder input shape:{'':2} {enc_input.shape}\")\n",
    "    print(f\"decoder input shape:{'':2} {dec_input.shape}\")\n",
    "    print(f\"output target shape:{'':2} {outputs.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5688879",
   "metadata": {},
   "source": [
    "<h4><i>2. Make a new model. </i></h4>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "691e8540",
   "metadata": {},
   "outputs": [],
   "source": [
    "nmt_model = Translator(\n",
    "    inputs_vectorizer,\n",
    "    targets_vectorizer,\n",
    "    units=512,\n",
    "    attn_num_heads=6,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b03d204",
   "metadata": {},
   "source": [
    "<h4><i>3. Compile and train the model. </i></h4>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b3616e7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compile the Model\n",
    "nmt_model.compile(optimizer=\"adam\", loss=masked_loss, metrics=[masked_loss, masked_acc])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5d6e490b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/20\n",
      "236/236 [==============================] - 514s 2s/step - loss: 6.1987 - masked_loss: 6.1977 - masked_acc: 0.1588\n",
      "Epoch 2/20\n",
      "236/236 [==============================] - 448s 2s/step - loss: 4.8464 - masked_loss: 4.8460 - masked_acc: 0.2754\n",
      "Epoch 3/20\n",
      "236/236 [==============================] - 448s 2s/step - loss: 3.6146 - masked_loss: 3.6151 - masked_acc: 0.3939\n",
      "Epoch 4/20\n",
      "236/236 [==============================] - 447s 2s/step - loss: 2.5346 - masked_loss: 2.5354 - masked_acc: 0.5028\n",
      "Epoch 5/20\n",
      "236/236 [==============================] - 460s 2s/step - loss: 1.5997 - masked_loss: 1.6002 - masked_acc: 0.6405\n",
      "Epoch 6/20\n",
      "236/236 [==============================] - 463s 2s/step - loss: 0.9485 - masked_loss: 0.9489 - masked_acc: 0.7691\n",
      "Epoch 7/20\n",
      "236/236 [==============================] - 458s 2s/step - loss: 0.5823 - masked_loss: 0.5821 - masked_acc: 0.8494\n",
      "Epoch 8/20\n",
      "236/236 [==============================] - 451s 2s/step - loss: 0.3750 - masked_loss: 0.3750 - masked_acc: 0.9020\n",
      "Epoch 9/20\n",
      "236/236 [==============================] - 447s 2s/step - loss: 0.2436 - masked_loss: 0.2437 - masked_acc: 0.9368\n",
      "Epoch 10/20\n",
      "236/236 [==============================] - 445s 2s/step - loss: 0.1820 - masked_loss: 0.1820 - masked_acc: 0.9523\n",
      "Epoch 11/20\n",
      "236/236 [==============================] - 448s 2s/step - loss: 0.1392 - masked_loss: 0.1392 - masked_acc: 0.9647\n",
      "Epoch 12/20\n",
      "236/236 [==============================] - 448s 2s/step - loss: 0.1149 - masked_loss: 0.1149 - masked_acc: 0.9712\n",
      "Epoch 13/20\n",
      "236/236 [==============================] - 450s 2s/step - loss: 0.1145 - masked_loss: 0.1144 - masked_acc: 0.9712\n",
      "Epoch 14/20\n",
      "236/236 [==============================] - 448s 2s/step - loss: 0.1150 - masked_loss: 0.1150 - masked_acc: 0.9704\n",
      "Epoch 15/20\n",
      "236/236 [==============================] - 457s 2s/step - loss: 0.1472 - masked_loss: 0.1474 - masked_acc: 0.9607\n",
      "Epoch 16/20\n",
      "236/236 [==============================] - 446s 2s/step - loss: 0.1659 - masked_loss: 0.1661 - masked_acc: 0.9558\n",
      "Epoch 17/20\n",
      "236/236 [==============================] - 445s 2s/step - loss: 0.1538 - masked_loss: 0.1540 - masked_acc: 0.9583\n",
      "Epoch 18/20\n",
      "236/236 [==============================] - 444s 2s/step - loss: 0.1305 - masked_loss: 0.1306 - masked_acc: 0.9654\n",
      "Epoch 19/20\n",
      "236/236 [==============================] - 449s 2s/step - loss: 0.1056 - masked_loss: 0.1056 - masked_acc: 0.9727\n",
      "Epoch 20/20\n",
      "236/236 [==============================] - 447s 2s/step - loss: 0.0912 - masked_loss: 0.0912 - masked_acc: 0.9767\n"
     ]
    }
   ],
   "source": [
    "# Fit the model.\n",
    "history = nmt_model.fit(dataset, epochs=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3b5bc9e",
   "metadata": {},
   "source": [
    "<h4><i>4. Save the model. </i></h4>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98cd14fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "nmt_model.save(\"persian_to_english_nmt_model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb56858e",
   "metadata": {},
   "source": [
    "<h4><i> 5. Evaluate the model. </i></h4>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "bf34f086",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "بگو کجا   : say one tell \n",
      "\n",
      "با تو موافقم: you can be as you for a bit \n",
      "\n",
      "من داشتم قدم می زدم: i love and withered \n",
      "\n",
      " من یک دکتر هستم: i enjoy a doctor \n",
      "\n"
     ]
    }
   ],
   "source": [
    "texts = [\"بگو کجا \", \"با تو موافقم\", \"من داشتم قدم می زدم\", \" من یک دکتر هستم\"]\n",
    "\n",
    "for t in texts:\n",
    "    result = nmt_model.translate([t])\n",
    "    translated = result[0].numpy().decode()\n",
    "    print(f\"{t:10}: {translated}\")\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcf69fa4",
   "metadata": {},
   "source": [
    "It seems that the translation of quality is not satisfying. Some potential reason :\n",
    "    1. Low quality dataset .\n",
    "    2. Small number of samples.\n",
    "    3. Training for just a few epochs."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}