Initial commit: EN-BN Translation Project

TESTAPI.py

@@ -0,0 +1,57 @@
+
+
+import os
+import requests
+
+BASE = os.getenv("ENBN_API_URL", "https://samin7479-en-bn-translator.hf.space")
+HEADERS = {"Content-Type": "application/json"}
+
+def greet():
+    try:
+        r = requests.get(f"{BASE}/greet", headers=HEADERS, timeout=20)
+        r.raise_for_status()
+        return r.json()
+    except Exception as e:
+        return {"error": str(e)}
+
+def translate(text, max_new_tokens=128, num_beams=4):
+    try:
+        payload = {
+            "text": text,
+            "max_new_tokens": max_new_tokens,
+            "num_beams": num_beams,
+            "do_sample": False
+        }
+        r = requests.post(f"{BASE}/translate", json=payload, headers=HEADERS, timeout=60)
+        r.raise_for_status()
+        return r.json().get("translation")
+    except Exception as e:
+        return f"[error] {e}"
+
+def translate_batch(texts, max_new_tokens=128, num_beams=4):
+    try:
+        payload = {
+            "texts": texts,
+            "max_new_tokens": max_new_tokens,
+            "num_beams": num_beams,
+            "do_sample": False
+        }
+        r = requests.post(f"{BASE}/translate_batch", json=payload, headers=HEADERS, timeout=120)
+        r.raise_for_status()
+        return r.json().get("translations", [])
+    except Exception as e:
+        return [f"[error] {e}"]
+
+if __name__ == "__main__":
+    # quick smoke test
+    print("GREET:", greet())
+
+    en = "How are you today?"
+    bn = translate(en)
+    print(f"\nSingle:\nEN: {en}\nBN: {bn}")
+
+    batch = ["Good morning", "Where is the hospital?", "The weather is nice."]
+    outs = translate_batch(batch)
+    print("\nBatch:")
+    for e, b in zip(batch, outs):
+        print(f"EN: {e}\nBN: {b}\n")

app.py

@@ -24,8 +24,7 @@ try:
 except Exception as e:
     raise RuntimeError(f"Failed to load model/tokenizer '{mt_pretrained_model_name}': {e}")
 
-# Optional: be gentle on CPU-only machines
-torch.set_num_threads(max(1, (os.cpu_count() or 1)))
+
 
 # -------------------------
 # FastAPI app + (optional) CORS

project_2_mt_en_bn.ipynb

@@ -0,0 +1,1207 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "id": "initial_id",
+   "metadata": {
+    "collapsed": true,
+    "id": "initial_id",
+    "executionInfo": {
+     "status": "error",
+     "timestamp": 1757400199832,
+     "user_tz": -360,
+     "elapsed": 136,
+     "user": {
+      "displayName": "KARABI KUMARI MEDHA 1604062",
+      "userId": "02676772162340716864"
+     }
+    },
+    "outputId": "4ca67866-ac7a-4f35-9d5f-11d09460a5ef",
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 383
+    },
+    "ExecuteTime": {
+     "end_time": "2025-09-14T07:00:59.517452Z",
+     "start_time": "2025-09-14T07:00:50.959173Z"
+    }
+   },
+   "source": [
+    "from typing import Any\n",
+    "\n",
+    "from pytorch_lightning.utilities.types import STEP_OUTPUT\n",
+    "\n",
+    "\"\"\" Class 25 | Project 2 | Machine Translation using Pretrained Model\n",
+    "\n",
+    "Objectives:\n",
+    "1. End-to-end machine translation training pipeline\n",
+    "2. Fine-tune a pre-trained model for the custom dataset\n",
+    "\"\"\"\n",
+    "\n",
+    "import pytorch_lightning as pl\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "import pandas as pd\n",
+    "from torchmetrics.text import BLEUScore\n",
+    "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM"
+   ],
+   "outputs": [],
+   "execution_count": 2
+  },
+  {
+   "metadata": {
+    "id": "cd6712aad1b548d7",
+    "outputId": "824c43b6-2d5e-480e-da6d-35338820f1fe",
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 159
+    },
+    "executionInfo": {
+     "status": "error",
+     "timestamp": 1757187729883,
+     "user_tz": -360,
+     "elapsed": 187,
+     "user": {
+      "displayName": "Chironjit Banerjee",
+      "userId": "04428016465669976257"
+     }
+    },
+    "ExecuteTime": {
+     "end_time": "2025-09-14T07:00:59.576987Z",
+     "start_time": "2025-09-14T07:00:59.523970Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "device"
+   ],
+   "id": "cd6712aad1b548d7",
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "device(type='cuda')"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "execution_count": 3
+  },
+  {
+   "metadata": {
+    "id": "eaa67c3f07ec30e2",
+    "ExecuteTime": {
+     "end_time": "2025-09-14T07:00:59.585691Z",
+     "start_time": "2025-09-14T07:00:59.581345Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "\"\"\"Task: English to Bangla \"\"\"\n",
+    "\n",
+    "mt_pretrained_model_name = \"shhossain/opus-mt-en-to-bn\""
+   ],
+   "id": "eaa67c3f07ec30e2",
+   "outputs": [],
+   "execution_count": 4
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "outputs": [],
+   "execution_count": null,
+   "source": [
+    "ROOT_DIR = \"E:\\Projects\\DS & ML\"\n",
+    "DATA_DIR = os.path.join(ROOT_DIR, \"DS\")\n",
+    "dataset_file = os.path.join(DATA_DIR, \"digit_train.csv\")\n",
+    "\n",
+    "ARTIFACT_FOLDER_NAME = \"model\" # Directory to save models\n",
+    "SOURCE_CODE_PATH = os.path.join(\n",
+    "        os.getcwd(),\n",
+    "        \"project_2_mt_en_bn.ipynb\",\n",
+    ") # Our current notebook file path\n",
+    "\n",
+    "SOURCE_CODE_ARTIFACT = \"trainer.ipynb\"\n"
+   ],
+   "id": "eddf18ffeb5bab6c"
+  },
+  {
+   "metadata": {
+    "id": "a0d805fe4a8ab875",
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 193
+    },
+    "executionInfo": {
+     "status": "error",
+     "timestamp": 1757187729919,
+     "user_tz": -360,
+     "elapsed": 18,
+     "user": {
+      "displayName": "Chironjit Banerjee",
+      "userId": "04428016465669976257"
+     }
+    },
+    "outputId": "cae2f309-54cf-49f1-dee7-786aedc7622d",
+    "ExecuteTime": {
+     "end_time": "2025-09-14T07:01:02.538952Z",
+     "start_time": "2025-09-14T07:00:59.593206Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "\"\"\" For NLP tasks, we basically need two entities:\n",
+    "1. Tokenizer\n",
+    "2. Model\n",
+    "\"\"\"\n",
+    "\n",
+    "tokenizer = AutoTokenizer.from_pretrained(mt_pretrained_model_name)\n",
+    "mt_pretrained_model = AutoModelForSeq2SeqLM.from_pretrained(mt_pretrained_model_name)"
+   ],
+   "id": "a0d805fe4a8ab875",
+   "outputs": [],
+   "execution_count": 5
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "outputs": [],
+   "execution_count": null,
+   "source": [
+    "\"\"\" Hyperparameters: Parameters that are not for neural networks but use to train\n",
+    "models. \"\"\"\n",
+    "EPOCHS = 3\n",
+    "BATCH_SIZE = 32\n",
+    "LEARNING_RATE = 2e-5"
+   ],
+   "id": "adece7449c00450c"
+  },
+  {
+   "metadata": {
+    "id": "e93068cfc700f5f8"
+   },
+   "cell_type": "markdown",
+   "source": [
+    "# Data"
+   ],
+   "id": "e93068cfc700f5f8"
+  },
+  {
+   "metadata": {
+    "id": "89449c4bacc42140",
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 211
+    },
+    "executionInfo": {
+     "status": "error",
+     "timestamp": 1757187730045,
+     "user_tz": -360,
+     "elapsed": 29,
+     "user": {
+      "displayName": "Chironjit Banerjee",
+      "userId": "04428016465669976257"
+     }
+    },
+    "outputId": "a4dc0573-6e6d-4908-ac0f-8813ed1dc901",
+    "ExecuteTime": {
+     "end_time": "2025-09-14T07:01:02.555974Z",
+     "start_time": "2025-09-14T07:01:02.547968Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "\"\"\"\n",
+    "Sentence: How are you, dude?\n",
+    "Tokens: 'How', 'are', 'you', 'dude?'\n",
+    "ids: 125, 14, 145, 78\n",
+    "max_length = 3\n",
+    "ids: [125, 14, 145]\n",
+    "\"\"\"\n",
+    "\n",
+    "class MTDataset(Dataset):\n",
+    "    def __init__(self, csv_file):\n",
+    "        self.data = pd.read_csv(csv_file)\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.data)\n",
+    "\n",
+    "    def __getitem__(self, idx):\n",
+    "        src_text = str(self.data.iloc[idx]['en'])\n",
+    "        tgt_text = str(self.data.iloc[idx]['bn'])\n",
+    "\n",
+    "        src_encoding = tokenizer(\n",
+    "            src_text,\n",
+    "            max_length=128,\n",
+    "            padding='max_length',\n",
+    "            truncation=True,\n",
+    "            return_tensors='pt',\n",
+    "        )\n",
+    "\n",
+    "        tgt_encoding = tokenizer(\n",
+    "            tgt_text,\n",
+    "            max_length=128,\n",
+    "            padding='max_length',\n",
+    "            truncation=True,\n",
+    "            return_tensors='pt'\n",
+    "        )\n",
+    "\n",
+    "        return {\n",
+    "            'src_input_ids': src_encoding['input_ids'].squeeze(),\n",
+    "            'src_attention_mask': src_encoding['attention_mask'].squeeze(),\n",
+    "            'tgt_input_ids': tgt_encoding['input_ids'].squeeze(),\n",
+    "            'tgt_attention_mask': tgt_encoding['attention_mask'].squeeze()\n",
+    "        }\n",
+    "\n",
+    "\"\"\"\n",
+    "example: How are you, dude?\n",
+    "input_ids: 125, 14, 145, 78\n",
+    "max_length = 7\n",
+    "input_ids: [125, 14, 145, 147, 0, 0, 0]\n",
+    "attention_mask: [1, 1, 1, 1, 0, 0, 0]\n",
+    "\"\"\""
+   ],
+   "id": "89449c4bacc42140",
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'\\nexample: How are you, dude?\\ninput_ids: 125, 14, 145, 78\\nmax_length = 7\\ninput_ids: [125, 14, 145, 147, 0, 0, 0]\\nattention_mask: [1, 1, 1, 1, 0, 0, 0]\\n'"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "execution_count": 6
+  },
+  {
+   "metadata": {
+    "id": "7dec7cfe5693f5f1",
+    "ExecuteTime": {
+     "end_time": "2025-09-14T07:01:02.571975Z",
+     "start_time": "2025-09-14T07:01:02.567999Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "class MTDataModule(pl.LightningDataModule):\n",
+    "    def __init__(self, train_csv, val_csv, test_csv, batch_size=BATCH_SIZE):\n",
+    "        super().__init__()\n",
+    "        self.train_csv = train_csv\n",
+    "        self.val_csv = val_csv\n",
+    "        self.test_csv = test_csv\n",
+    "        self.batch_size = BATCH_SIZE\n",
+    "\n",
+    "    def setup(self, stage=None):\n",
+    "        self.train_dataset = MTDataset(self.train_csv)\n",
+    "        self.val_dataset = MTDataset(self.val_csv)\n",
+    "        self.test_dataset = MTDataset(self.test_csv)\n",
+    "\n",
+    "    def train_dataloader(self):\n",
+    "        return DataLoader(\n",
+    "            self.train_dataset,\n",
+    "            batch_size=self.BATCH_SIZE,\n",
+    "            shuffle=True\n",
+    "        )\n",
+    "\n",
+    "    def val_dataloader(self):\n",
+    "        return DataLoader(\n",
+    "            self.val_dataset,\n",
+    "            batch_size=self.BATCH_SIZE,\n",
+    "            shuffle=False\n",
+    "        )\n",
+    "\n",
+    "    def test_dataloader(self):\n",
+    "        return DataLoader(\n",
+    "            self.test_dataset,\n",
+    "            batch_size=self.BATCH_SIZE,\n",
+    "            shuffle=False\n",
+    "        )"
+   ],
+   "id": "7dec7cfe5693f5f1",
+   "outputs": [],
+   "execution_count": 7
+  },
+  {
+   "metadata": {
+    "id": "ef2deed7494ec4b4",
+    "ExecuteTime": {
+     "end_time": "2025-09-14T07:01:02.583942Z",
+     "start_time": "2025-09-14T07:01:02.580979Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "data_module = MTDataModule(\n",
+    "    train_csv=r'E:\\Projects\\DS & ML\\EN to BN ML Project\\train.csv',\n",
+    "    val_csv=r'E:\\Projects\\DS & ML\\EN to BN ML Project\\val.csv',\n",
+    "    test_csv=r'E:\\Projects\\DS & ML\\EN to BN ML Project\\test.csv',\n",
+    "    batch_size= BATCH_SIZE\n",
+    ")"
+   ],
+   "id": "ef2deed7494ec4b4",
+   "outputs": [],
+   "execution_count": 8
+  },
+  {
+   "metadata": {
+    "id": "86e90bfb5b63dafe"
+   },
+   "cell_type": "markdown",
+   "source": [
+    "# Model"
+   ],
+   "id": "86e90bfb5b63dafe"
+  },
+  {
+   "metadata": {
+    "id": "70ac9ff9786267a5",
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 211
+    },
+    "executionInfo": {
+     "status": "error",
+     "timestamp": 1757187730281,
+     "user_tz": -360,
+     "elapsed": 204,
+     "user": {
+      "displayName": "Chironjit Banerjee",
+      "userId": "04428016465669976257"
+     }
+    },
+    "outputId": "7a8e4dc0-2b1e-46ec-8c32-cef34f48ee96",
+    "ExecuteTime": {
+     "end_time": "2025-09-14T07:01:02.599947Z",
+     "start_time": "2025-09-14T07:01:02.591950Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "class MTModel(pl.LightningModule):\n",
+    "    def __init__(self):\n",
+    "        super().__init__()\n",
+    "        # load pretrained model\n",
+    "        self.model = AutoModelForSeq2SeqLM.from_pretrained(mt_pretrained_model_name)\n",
+    "        # load pretrained tokenizer\n",
+    "        self.tokenizer = AutoTokenizer.from_pretrained(mt_pretrained_model_name)\n",
+    "        # learning rate\n",
+    "        self.learning_rate = 2e-5\n",
+    "        # loss function\n",
+    "        self.loss_fn = nn.CrossEntropyLoss(\n",
+    "            ignore_index=self.tokenizer.pad_token_id\n",
+    "        )\n",
+    "        # evaluation metric\n",
+    "        self.bleu = BLEUScore()\n",
+    "\n",
+    "    def forward(self,\n",
+    "                src_input_ids,\n",
+    "                src_attention_mask,\n",
+    "                tgt_input_ids,\n",
+    "                tgt_attention_mask\n",
+    "        ):\n",
+    "        outputs = self.model(\n",
+    "            input_ids=src_input_ids,\n",
+    "            attention_mask=src_attention_mask,\n",
+    "            decoder_input_ids=tgt_input_ids[:, :-1],\n",
+    "            decoder_attention_mask=tgt_attention_mask[:, :-1]\n",
+    "        )\n",
+    "        return outputs\n",
+    "\n",
+    "    def training_step(self, batch, batch_idx):\n",
+    "        loss = self.compute_loss(batch, batch_idx, 'train')\n",
+    "        self.log('train_loss', loss, prog_bar=True)\n",
+    "        return loss\n",
+    "\n",
+    "    def validation_step(self, batch, batch_idx):\n",
+    "        loss = self.compute_loss(batch, batch_idx, 'val')\n",
+    "        self.log('val_loss', loss, prog_bar=True)\n",
+    "        return loss\n",
+    "\n",
+    "    def test_step(self, batch, batch_idx):\n",
+    "        loss = self.compute_loss(batch, batch_idx, 'test')\n",
+    "        self.log('test_loss', loss, prog_bar=True)\n",
+    "        return loss\n",
+    "\n",
+    "    def configure_optimizers(self):\n",
+    "        optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate)\n",
+    "        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n",
+    "            optimizer,\n",
+    "            T_max=10\n",
+    "        )\n",
+    "        return {'optimizer': optimizer, 'lr_scheduler': scheduler}\n",
+    "\n",
+    "    def compute_loss(self, batch, batch_idx, stage):\n",
+    "        src_input_ids = batch['src_input_ids']\n",
+    "        src_attention_mask = batch['src_attention_mask']\n",
+    "        tgt_input_ids = batch['tgt_input_ids']\n",
+    "        tgt_attention_mask = batch['tgt_attention_mask']\n",
+    "\n",
+    "        outputs = self(\n",
+    "            src_input_ids,\n",
+    "            src_attention_mask,\n",
+    "            tgt_input_ids,\n",
+    "            tgt_attention_mask\n",
+    "        )\n",
+    "        logits = outputs.logits\n",
+    "        loss = self.loss_fn(\n",
+    "            logits.view(-1, logits.size(-1)),\n",
+    "            tgt_input_ids[:, 1:].contiguous().view(-1)\n",
+    "        )\n",
+    "\n",
+    "        if stage == 'val' or stage == 'test':\n",
+    "            preds = torch.argmax(logits, dim=-1)\n",
+    "            pred_texts = self.tokenizer.batch_decode(preds, skip_special_tokens=True)\n",
+    "            tgt_texts = self.tokenizer.batch_decode(tgt_input_ids[:, 1:], skip_special_tokens=True)\n",
+    "            bleu_score = self.bleu(pred_texts, [[tgt] for tgt in tgt_texts])\n",
+    "            self.log(f'{stage}_bleu', bleu_score, prog_bar=True)\n",
+    "\n",
+    "        return loss\n"
+   ],
+   "id": "70ac9ff9786267a5",
+   "outputs": [],
+   "execution_count": 9
+  },
+  {
+   "metadata": {
+    "id": "76dccd8fe08376a4",
+    "ExecuteTime": {
+     "end_time": "2025-09-14T07:01:05.089548Z",
+     "start_time": "2025-09-14T07:01:02.605071Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "model = MTModel()"
+   ],
+   "id": "76dccd8fe08376a4",
+   "outputs": [],
+   "execution_count": 10
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "outputs": [],
+   "execution_count": null,
+   "source": [
+    "early_stopping = EarlyStopping(\n",
+    "    monitor='val_loss', # Should match with the validation step log key\n",
+    "    patience=2,\n",
+    "    verbose=True,\n",
+    ")\n",
+    "\n",
+    "checkpoint_callback = ModelCheckpoint(\n",
+    "    monitor='val_accuracy', # Should match with the validation step log key\n",
+    "    save_top_k=1, # Saves top one model\n",
+    "    mode='max', # top means max validation accuracy\n",
+    ")\n",
+    "\n",
+    "checkpoint_path = os.path.join(\n",
+    "    os.getcwd(), \"checkpoints\", \"best_model.pth\"\n",
+    ")\n"
+   ],
+   "id": "b280d211a42ceeee"
+  },
+  {
+   "metadata": {
+    "id": "c037b19d321b93ff"
+   },
+   "cell_type": "markdown",
+   "source": [
+    "# Train"
+   ],
+   "id": "c037b19d321b93ff"
+  },
+  {
+   "metadata": {
+    "id": "1bd38416398d770a",
+    "ExecuteTime": {
+     "end_time": "2025-09-14T07:01:05.135876Z",
+     "start_time": "2025-09-14T07:01:05.096305Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "trainer = pl.Trainer(\n",
+    "    max_epochs=5,\n",
+    "    accelerator='gpu' if torch.cuda.is_available() else 'cpu',\n",
+    "    devices=1,\n",
+    "    precision=\"16-mixed\",\n",
+    "    log_every_n_steps=10,\n",
+    "    val_check_interval=0.25\n",
+    ")"
+   ],
+   "id": "1bd38416398d770a",
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Using 16bit Automatic Mixed Precision (AMP)\n",
+      "💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.\n",
+      "GPU available: True (cuda), used: True\n",
+      "TPU available: False, using: 0 TPU cores\n",
+      "HPU available: False, using: 0 HPUs\n"
+     ]
+    }
+   ],
+   "execution_count": 11
+  },
+  {
+   "metadata": {
+    "id": "add377254e158c86",
+    "jupyter": {
+     "is_executing": true
+    },
+    "ExecuteTime": {
+     "start_time": "2025-09-14T07:01:05.145883Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "trainer.fit(model, data_module)"
+   ],
+   "id": "add377254e158c86",
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "You are using a CUDA device ('NVIDIA GeForce RTX 3060') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision\n",
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
+      "C:\\Users\\User\\PyCharmMiscProject\\.venv\\Lib\\site-packages\\pytorch_lightning\\utilities\\model_summary\\model_summary.py:231: Precision 16-mixed is not supported by the model summary.  Estimated model size in MB will not be accurate. Using 32 bits instead.\n",
+      "\n",
+      "  | Name    | Type             | Params | Mode \n",
+      "-----------------------------------------------------\n",
+      "0 | model   | MarianMTModel    | 76.3 M | eval \n",
+      "1 | loss_fn | CrossEntropyLoss | 0      | train\n",
+      "2 | bleu    | BLEUScore        | 0      | train\n",
+      "-----------------------------------------------------\n",
+      "75.8 M    Trainable params\n",
+      "524 K     Non-trainable params\n",
+      "76.3 M    Total params\n",
+      "305.136   Total estimated model params size (MB)\n",
+      "2         Modules in train mode\n",
+      "178       Modules in eval mode\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "Sanity Checking: |          | 0/? [00:00<?, ?it/s]"
+      ],
+      "application/vnd.jupyter.widget-view+json": {
+       "version_major": 2,
+       "version_minor": 0,
+       "model_id": "7653d36a0abd4e27bea488b14b89d42b"
+      }
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "C:\\Users\\User\\PyCharmMiscProject\\.venv\\Lib\\site-packages\\pytorch_lightning\\trainer\\connectors\\data_connector.py:433: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.\n",
+      "C:\\Users\\User\\PyCharmMiscProject\\.venv\\Lib\\site-packages\\pytorch_lightning\\trainer\\connectors\\data_connector.py:433: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "Training: |          | 0/? [00:00<?, ?it/s]"
+      ],
+      "application/vnd.jupyter.widget-view+json": {
+       "version_major": 2,
+       "version_minor": 0,
+       "model_id": "b3323116f1b44cd9aacdfe59272a1310"
+      }
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ],
+      "application/vnd.jupyter.widget-view+json": {
+       "version_major": 2,
+       "version_minor": 0,
+       "model_id": "e1f76950fb134014a65f59bd58f85541"
+      }
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ],
+      "application/vnd.jupyter.widget-view+json": {
+       "version_major": 2,
+       "version_minor": 0,
+       "model_id": "52f93b059f014695ac2984b1ab8ab6e0"
+      }
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ],
+      "application/vnd.jupyter.widget-view+json": {
+       "version_major": 2,
+       "version_minor": 0,
+       "model_id": "b738ac5554f9427ea2fe44f82b159593"
+      }
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ],
+      "application/vnd.jupyter.widget-view+json": {
+       "version_major": 2,
+       "version_minor": 0,
+       "model_id": "85d57ade665a4808a138b0364dc40d0d"
+      }
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "execution_count": null
+  },
+  {
+   "metadata": {
+    "id": "652b3f73247ae77c",
+    "ExecuteTime": {
+     "end_time": "2025-09-12T11:07:00.608309Z",
+     "start_time": "2025-09-12T11:06:38.554751Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "trainer.test(model, data_module)"
+   ],
+   "id": "652b3f73247ae77c",
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
+      "C:\\Users\\User\\PyCharmMiscProject\\.venv\\Lib\\site-packages\\pytorch_lightning\\trainer\\connectors\\data_connector.py:433: The 'test_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "Testing: |          | 0/? [00:00<?, ?it/s]"
+      ],
+      "application/vnd.jupyter.widget-view+json": {
+       "version_major": 2,
+       "version_minor": 0,
+       "model_id": "37db78bcf4f646bb995c4a725d9126eb"
+      }
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
+       "┃\u001B[1m \u001B[0m\u001B[1m       Test metric       \u001B[0m\u001B[1m \u001B[0m┃\u001B[1m \u001B[0m\u001B[1m      DataLoader 0       \u001B[0m\u001B[1m \u001B[0m┃\n",
+       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
+       "│\u001B[36m \u001B[0m\u001B[36m        test_bleu        \u001B[0m\u001B[36m \u001B[0m│\u001B[35m \u001B[0m\u001B[35m   0.13064175844192505   \u001B[0m\u001B[35m \u001B[0m│\n",
+       "│\u001B[36m \u001B[0m\u001B[36m        test_loss        \u001B[0m\u001B[36m \u001B[0m│\u001B[35m \u001B[0m\u001B[35m   0.5454719662666321    \u001B[0m\u001B[35m \u001B[0m│\n",
+       "└───────────────────────────┴───────────────────────────┘\n"
+      ],
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
+       "┃<span style=\"font-weight: bold\">        Test metric        </span>┃<span style=\"font-weight: bold\">       DataLoader 0        </span>┃\n",
+       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">         test_bleu         </span>│<span style=\"color: #800080; text-decoration-color: #800080\">    0.13064175844192505    </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">         test_loss         </span>│<span style=\"color: #800080; text-decoration-color: #800080\">    0.5454719662666321     </span>│\n",
+       "└───────────────────────────┴───────────────────────────┘\n",
+       "</pre>\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "[{'test_bleu': 0.13064175844192505, 'test_loss': 0.5454719662666321}]"
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "execution_count": 12
+  },
+  {
+   "metadata": {
+    "id": "9f115c718388b7f9",
+    "outputId": "2e3186ce-b52e-4344-cfd5-8fea0ee7b28a",
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 141
+    },
+    "executionInfo": {
+     "status": "error",
+     "timestamp": 1757187730430,
+     "user_tz": -360,
+     "elapsed": 104,
+     "user": {
+      "displayName": "Chironjit Banerjee",
+      "userId": "04428016465669976257"
+     }
+    },
+    "ExecuteTime": {
+     "end_time": "2025-09-12T11:07:00.671344Z",
+     "start_time": "2025-09-12T11:07:00.665343Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "model.model.config"
+   ],
+   "id": "9f115c718388b7f9",
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "MarianConfig {\n",
+       "  \"activation_dropout\": 0.0,\n",
+       "  \"activation_function\": \"swish\",\n",
+       "  \"add_bias_logits\": false,\n",
+       "  \"add_final_layer_norm\": false,\n",
+       "  \"architectures\": [\n",
+       "    \"MarianMTModel\"\n",
+       "  ],\n",
+       "  \"attention_dropout\": 0.0,\n",
+       "  \"bad_words_ids\": [\n",
+       "    [\n",
+       "      61759\n",
+       "    ]\n",
+       "  ],\n",
+       "  \"bos_token_id\": 0,\n",
+       "  \"classif_dropout\": 0.0,\n",
+       "  \"classifier_dropout\": 0.0,\n",
+       "  \"d_model\": 512,\n",
+       "  \"decoder_attention_heads\": 8,\n",
+       "  \"decoder_ffn_dim\": 2048,\n",
+       "  \"decoder_layerdrop\": 0.0,\n",
+       "  \"decoder_layers\": 6,\n",
+       "  \"decoder_start_token_id\": 61759,\n",
+       "  \"decoder_vocab_size\": 61760,\n",
+       "  \"dropout\": 0.1,\n",
+       "  \"dtype\": \"float32\",\n",
+       "  \"encoder_attention_heads\": 8,\n",
+       "  \"encoder_ffn_dim\": 2048,\n",
+       "  \"encoder_layerdrop\": 0.0,\n",
+       "  \"encoder_layers\": 6,\n",
+       "  \"eos_token_id\": 0,\n",
+       "  \"extra_pos_embeddings\": 61760,\n",
+       "  \"forced_eos_token_id\": 0,\n",
+       "  \"id2label\": {\n",
+       "    \"0\": \"LABEL_0\",\n",
+       "    \"1\": \"LABEL_1\",\n",
+       "    \"2\": \"LABEL_2\"\n",
+       "  },\n",
+       "  \"init_std\": 0.02,\n",
+       "  \"is_encoder_decoder\": true,\n",
+       "  \"label2id\": {\n",
+       "    \"LABEL_0\": 0,\n",
+       "    \"LABEL_1\": 1,\n",
+       "    \"LABEL_2\": 2\n",
+       "  },\n",
+       "  \"max_length\": 512,\n",
+       "  \"max_position_embeddings\": 512,\n",
+       "  \"model_type\": \"marian\",\n",
+       "  \"normalize_before\": false,\n",
+       "  \"normalize_embedding\": false,\n",
+       "  \"num_beams\": 4,\n",
+       "  \"num_hidden_layers\": 6,\n",
+       "  \"pad_token_id\": 61759,\n",
+       "  \"scale_embedding\": true,\n",
+       "  \"share_encoder_decoder_embeddings\": true,\n",
+       "  \"static_position_embeddings\": true,\n",
+       "  \"transformers_version\": \"4.56.1\",\n",
+       "  \"use_cache\": true,\n",
+       "  \"vocab_size\": 61760\n",
+       "}"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "execution_count": 13
+  },
+  {
+   "metadata": {
+    "id": "4db952c7f44ec3b2",
+    "outputId": "1c8af67a-9c64-4017-c39b-5bf773a0e8ad",
+    "colab": {
+     "base_uri": "https://localhost:8080/",
+     "height": 159
+    },
+    "executionInfo": {
+     "status": "error",
+     "timestamp": 1757187730464,
+     "user_tz": -360,
+     "elapsed": 20,
+     "user": {
+      "displayName": "Chironjit Banerjee",
+      "userId": "04428016465669976257"
+     }
+    },
+    "ExecuteTime": {
+     "end_time": "2025-09-12T11:07:00.694215Z",
+     "start_time": "2025-09-12T11:07:00.687372Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "for name, module in model.model.named_modules():\n",
+    "    print(name)"
+   ],
+   "id": "4db952c7f44ec3b2",
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "model\n",
+      "model.shared\n",
+      "model.encoder\n",
+      "model.encoder.embed_positions\n",
+      "model.encoder.layers\n",
+      "model.encoder.layers.0\n",
+      "model.encoder.layers.0.self_attn\n",
+      "model.encoder.layers.0.self_attn.k_proj\n",
+      "model.encoder.layers.0.self_attn.v_proj\n",
+      "model.encoder.layers.0.self_attn.q_proj\n",
+      "model.encoder.layers.0.self_attn.out_proj\n",
+      "model.encoder.layers.0.self_attn_layer_norm\n",
+      "model.encoder.layers.0.activation_fn\n",
+      "model.encoder.layers.0.fc1\n",
+      "model.encoder.layers.0.fc2\n",
+      "model.encoder.layers.0.final_layer_norm\n",
+      "model.encoder.layers.1\n",
+      "model.encoder.layers.1.self_attn\n",
+      "model.encoder.layers.1.self_attn.k_proj\n",
+      "model.encoder.layers.1.self_attn.v_proj\n",
+      "model.encoder.layers.1.self_attn.q_proj\n",
+      "model.encoder.layers.1.self_attn.out_proj\n",
+      "model.encoder.layers.1.self_attn_layer_norm\n",
+      "model.encoder.layers.1.activation_fn\n",
+      "model.encoder.layers.1.fc1\n",
+      "model.encoder.layers.1.fc2\n",
+      "model.encoder.layers.1.final_layer_norm\n",
+      "model.encoder.layers.2\n",
+      "model.encoder.layers.2.self_attn\n",
+      "model.encoder.layers.2.self_attn.k_proj\n",
+      "model.encoder.layers.2.self_attn.v_proj\n",
+      "model.encoder.layers.2.self_attn.q_proj\n",
+      "model.encoder.layers.2.self_attn.out_proj\n",
+      "model.encoder.layers.2.self_attn_layer_norm\n",
+      "model.encoder.layers.2.activation_fn\n",
+      "model.encoder.layers.2.fc1\n",
+      "model.encoder.layers.2.fc2\n",
+      "model.encoder.layers.2.final_layer_norm\n",
+      "model.encoder.layers.3\n",
+      "model.encoder.layers.3.self_attn\n",
+      "model.encoder.layers.3.self_attn.k_proj\n",
+      "model.encoder.layers.3.self_attn.v_proj\n",
+      "model.encoder.layers.3.self_attn.q_proj\n",
+      "model.encoder.layers.3.self_attn.out_proj\n",
+      "model.encoder.layers.3.self_attn_layer_norm\n",
+      "model.encoder.layers.3.activation_fn\n",
+      "model.encoder.layers.3.fc1\n",
+      "model.encoder.layers.3.fc2\n",
+      "model.encoder.layers.3.final_layer_norm\n",
+      "model.encoder.layers.4\n",
+      "model.encoder.layers.4.self_attn\n",
+      "model.encoder.layers.4.self_attn.k_proj\n",
+      "model.encoder.layers.4.self_attn.v_proj\n",
+      "model.encoder.layers.4.self_attn.q_proj\n",
+      "model.encoder.layers.4.self_attn.out_proj\n",
+      "model.encoder.layers.4.self_attn_layer_norm\n",
+      "model.encoder.layers.4.activation_fn\n",
+      "model.encoder.layers.4.fc1\n",
+      "model.encoder.layers.4.fc2\n",
+      "model.encoder.layers.4.final_layer_norm\n",
+      "model.encoder.layers.5\n",
+      "model.encoder.layers.5.self_attn\n",
+      "model.encoder.layers.5.self_attn.k_proj\n",
+      "model.encoder.layers.5.self_attn.v_proj\n",
+      "model.encoder.layers.5.self_attn.q_proj\n",
+      "model.encoder.layers.5.self_attn.out_proj\n",
+      "model.encoder.layers.5.self_attn_layer_norm\n",
+      "model.encoder.layers.5.activation_fn\n",
+      "model.encoder.layers.5.fc1\n",
+      "model.encoder.layers.5.fc2\n",
+      "model.encoder.layers.5.final_layer_norm\n",
+      "model.decoder\n",
+      "model.decoder.embed_positions\n",
+      "model.decoder.layers\n",
+      "model.decoder.layers.0\n",
+      "model.decoder.layers.0.self_attn\n",
+      "model.decoder.layers.0.self_attn.k_proj\n",
+      "model.decoder.layers.0.self_attn.v_proj\n",
+      "model.decoder.layers.0.self_attn.q_proj\n",
+      "model.decoder.layers.0.self_attn.out_proj\n",
+      "model.decoder.layers.0.activation_fn\n",
+      "model.decoder.layers.0.self_attn_layer_norm\n",
+      "model.decoder.layers.0.encoder_attn\n",
+      "model.decoder.layers.0.encoder_attn.k_proj\n",
+      "model.decoder.layers.0.encoder_attn.v_proj\n",
+      "model.decoder.layers.0.encoder_attn.q_proj\n",
+      "model.decoder.layers.0.encoder_attn.out_proj\n",
+      "model.decoder.layers.0.encoder_attn_layer_norm\n",
+      "model.decoder.layers.0.fc1\n",
+      "model.decoder.layers.0.fc2\n",
+      "model.decoder.layers.0.final_layer_norm\n",
+      "model.decoder.layers.1\n",
+      "model.decoder.layers.1.self_attn\n",
+      "model.decoder.layers.1.self_attn.k_proj\n",
+      "model.decoder.layers.1.self_attn.v_proj\n",
+      "model.decoder.layers.1.self_attn.q_proj\n",
+      "model.decoder.layers.1.self_attn.out_proj\n",
+      "model.decoder.layers.1.activation_fn\n",
+      "model.decoder.layers.1.self_attn_layer_norm\n",
+      "model.decoder.layers.1.encoder_attn\n",
+      "model.decoder.layers.1.encoder_attn.k_proj\n",
+      "model.decoder.layers.1.encoder_attn.v_proj\n",
+      "model.decoder.layers.1.encoder_attn.q_proj\n",
+      "model.decoder.layers.1.encoder_attn.out_proj\n",
+      "model.decoder.layers.1.encoder_attn_layer_norm\n",
+      "model.decoder.layers.1.fc1\n",
+      "model.decoder.layers.1.fc2\n",
+      "model.decoder.layers.1.final_layer_norm\n",
+      "model.decoder.layers.2\n",
+      "model.decoder.layers.2.self_attn\n",
+      "model.decoder.layers.2.self_attn.k_proj\n",
+      "model.decoder.layers.2.self_attn.v_proj\n",
+      "model.decoder.layers.2.self_attn.q_proj\n",
+      "model.decoder.layers.2.self_attn.out_proj\n",
+      "model.decoder.layers.2.activation_fn\n",
+      "model.decoder.layers.2.self_attn_layer_norm\n",
+      "model.decoder.layers.2.encoder_attn\n",
+      "model.decoder.layers.2.encoder_attn.k_proj\n",
+      "model.decoder.layers.2.encoder_attn.v_proj\n",
+      "model.decoder.layers.2.encoder_attn.q_proj\n",
+      "model.decoder.layers.2.encoder_attn.out_proj\n",
+      "model.decoder.layers.2.encoder_attn_layer_norm\n",
+      "model.decoder.layers.2.fc1\n",
+      "model.decoder.layers.2.fc2\n",
+      "model.decoder.layers.2.final_layer_norm\n",
+      "model.decoder.layers.3\n",
+      "model.decoder.layers.3.self_attn\n",
+      "model.decoder.layers.3.self_attn.k_proj\n",
+      "model.decoder.layers.3.self_attn.v_proj\n",
+      "model.decoder.layers.3.self_attn.q_proj\n",
+      "model.decoder.layers.3.self_attn.out_proj\n",
+      "model.decoder.layers.3.activation_fn\n",
+      "model.decoder.layers.3.self_attn_layer_norm\n",
+      "model.decoder.layers.3.encoder_attn\n",
+      "model.decoder.layers.3.encoder_attn.k_proj\n",
+      "model.decoder.layers.3.encoder_attn.v_proj\n",
+      "model.decoder.layers.3.encoder_attn.q_proj\n",
+      "model.decoder.layers.3.encoder_attn.out_proj\n",
+      "model.decoder.layers.3.encoder_attn_layer_norm\n",
+      "model.decoder.layers.3.fc1\n",
+      "model.decoder.layers.3.fc2\n",
+      "model.decoder.layers.3.final_layer_norm\n",
+      "model.decoder.layers.4\n",
+      "model.decoder.layers.4.self_attn\n",
+      "model.decoder.layers.4.self_attn.k_proj\n",
+      "model.decoder.layers.4.self_attn.v_proj\n",
+      "model.decoder.layers.4.self_attn.q_proj\n",
+      "model.decoder.layers.4.self_attn.out_proj\n",
+      "model.decoder.layers.4.activation_fn\n",
+      "model.decoder.layers.4.self_attn_layer_norm\n",
+      "model.decoder.layers.4.encoder_attn\n",
+      "model.decoder.layers.4.encoder_attn.k_proj\n",
+      "model.decoder.layers.4.encoder_attn.v_proj\n",
+      "model.decoder.layers.4.encoder_attn.q_proj\n",
+      "model.decoder.layers.4.encoder_attn.out_proj\n",
+      "model.decoder.layers.4.encoder_attn_layer_norm\n",
+      "model.decoder.layers.4.fc1\n",
+      "model.decoder.layers.4.fc2\n",
+      "model.decoder.layers.4.final_layer_norm\n",
+      "model.decoder.layers.5\n",
+      "model.decoder.layers.5.self_attn\n",
+      "model.decoder.layers.5.self_attn.k_proj\n",
+      "model.decoder.layers.5.self_attn.v_proj\n",
+      "model.decoder.layers.5.self_attn.q_proj\n",
+      "model.decoder.layers.5.self_attn.out_proj\n",
+      "model.decoder.layers.5.activation_fn\n",
+      "model.decoder.layers.5.self_attn_layer_norm\n",
+      "model.decoder.layers.5.encoder_attn\n",
+      "model.decoder.layers.5.encoder_attn.k_proj\n",
+      "model.decoder.layers.5.encoder_attn.v_proj\n",
+      "model.decoder.layers.5.encoder_attn.q_proj\n",
+      "model.decoder.layers.5.encoder_attn.out_proj\n",
+      "model.decoder.layers.5.encoder_attn_layer_norm\n",
+      "model.decoder.layers.5.fc1\n",
+      "model.decoder.layers.5.fc2\n",
+      "model.decoder.layers.5.final_layer_norm\n",
+      "lm_head\n"
+     ]
+    }
+   ],
+   "execution_count": 14
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2025-09-14T07:00:16.535340Z",
+     "start_time": "2025-09-14T07:00:14.981559Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "import mlflow\n",
+    "mlflow.set_experiment(experiment_name= \"BongoDev Projects\")\n"
+   ],
+   "id": "1c8e5f9092cfe5c6",
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2025/09/14 13:00:16 INFO mlflow.tracking.fluent: Experiment with name 'BongoDev Projects' does not exist. Creating a new experiment.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "<Experiment: artifact_location='file:///C:/Users/User/PyCharmMiscProject/mlruns/793621701339965882', creation_time=1757833216529, experiment_id='793621701339965882', last_update_time=1757833216529, lifecycle_stage='active', name='BongoDev Projects', tags={}>"
+      ]
+     },
+     "execution_count": 1,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "execution_count": 1
+  },
+  {
+   "metadata": {},
+   "cell_type": "markdown",
+   "source": "### Experiment Tracking using MLFlow",
+   "id": "1f356c18c4d0a1f9"
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "outputs": [],
+   "execution_count": null,
+   "source": [
+    "with mlflow.start_run():\n",
+    "    # Log Hyperparameters\n",
+    "    mlflow.log_param(\"learning_rate\", LEARNING_RATE)\n",
+    "    mlflow.log_param(\"batch_size\", BATCH_SIZE)\n",
+    "    mlflow.log_param(\"epochs\", EPOCHS)\n",
+    "\n",
+    "\n",
+    "\n",
+    "    trainer.fit(\n",
+    "        model=model,\n",
+    "        datamodule=data_module\n",
+    "    )\n",
+    "\n",
+    "    # Get the best model\n",
+    "    best_model_path = checkpoint_callback.best_model_path\n",
+    "    best_model = DigitClassifier.load_from_checkpoint(best_model_path)\n",
+    "\n",
+    "    # Evaluate the model on the test set\n",
+    "    evaluation_score = trainer.test(\n",
+    "        best_model,\n",
+    "        datamodule= data_module,\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "    mlflow.log_metric(\"test_accuracy\", evaluation_score[0][\"test_accuracy\"])\n",
+    "    mlflow.log_metric(\"test_loss\", evaluation_score[0][\"test_loss\"])\n",
+    "\n",
+    "\n",
+    "    # Save the model\n",
+    "    # Prepare a small input_example from the test loader\n",
+    "    test_loader = data_module.test_dataloader()\n",
+    "    first_batch = next(iter(test_loader))\n",
+    "    src_input_ids_example = first_batch[\"src_input_ids\"].cpu().numpy()\n",
+    "\n",
+    "    signature = infer_signature(src_input_ids_example, src_input_ids_example)\n",
+    "\n",
+    "    # Log the underlying HF Seq2Seq model (nn.Module) to keep it simple\n",
+    "    import mlflow.pytorch\n",
+    "    mlflow.pytorch.log_model(\n",
+    "        pytorch_model=best_model.model,\n",
+    "        artifact_path=ARTIFACT_FOLDER_NAME,\n",
+    "        input_example=src_input_ids_example,\n",
+    "        signature=signature\n",
+    "    )\n",
+    "\n",
+    "    import shutil\n",
+    "    shutil.copyfile(SOURCE_CODE_PATH, SOURCE_CODE_ARTIFACT)\n",
+    "    mlflow.log_artifact(SOURCE_CODE_ARTIFACT)\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n",
+    "\n"
+   ],
+   "id": "23776af915ea20ae"
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  },
+  "colab": {
+   "provenance": []
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}