update generato

.gitignore

@@ -0,0 +1 @@
+__pycache__/

best_model.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c87d759052debb4e4adb62ef51c9d65671d04bfc6e1f9fd4b2130c66e69b9257
-size 162038291
+oid sha256:f960ed296d73469239b693c1d37c1cef55e92cf3d7a914c75b8c7b9efb9ff701
+size 162076771

generator.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "id": "bae751d8",
    "metadata": {},
    "outputs": [],
@@ -26,7 +26,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 7,
    "id": "c0e30f61",
    "metadata": {},
    "outputs": [
@@ -42,10 +42,10 @@
     {
      "data": {
       "text/plain": [
-       "<models.open_mp_gen.tokenizer.Tokenizer at 0x7a60237a60c0>"
+       "<tokenizer.Tokenizer at 0x7d2bbbafcb90>"
       ]
      },
-     "execution_count": 3,
+     "execution_count": 7,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -57,7 +57,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 8,
    "id": "db130c45",
    "metadata": {},
    "outputs": [
@@ -134,7 +134,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 9,
    "id": "d5747915",
    "metadata": {},
    "outputs": [
@@ -173,7 +173,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 10,
    "id": "5252d457",
    "metadata": {},
    "outputs": [
@@ -190,13 +190,7 @@
       "  input shape: torch.Size([32, 500])\n",
       "  output shape: torch.Size([32, 100])\n",
       "  input_len shape: torch.Size([32])\n",
-      "  First sample input_len: 16\n",
-      "\n",
-      "Sample batch structure:\n",
-      "  input shape: torch.Size([32, 500])\n",
-      "  output shape: torch.Size([32, 100])\n",
-      "  input_len shape: torch.Size([32])\n",
-      "  First sample input_len: 16\n"
+      "  First sample input_len: 12\n"
      ]
     }
    ],
@@ -229,7 +223,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 11,
    "id": "11631bed",
    "metadata": {},
    "outputs": [
@@ -243,6 +237,7 @@
       "    (embedding): Embedding(8002, 128, padding_idx=0)\n",
       "    (lstm): LSTM(128, 256, num_layers=2, batch_first=True, dropout=0.3, bidirectional=True)\n",
       "    (dropout): Dropout(p=0.3, inplace=False)\n",
+      "    (layer_norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (decoder): Decoder(\n",
       "    (attention): BahdanauAttention(\n",
@@ -254,12 +249,13 @@
       "    (lstm): LSTM(640, 256, num_layers=2, batch_first=True, dropout=0.3)\n",
       "    (fc_out): Linear(in_features=896, out_features=8002, bias=True)\n",
       "    (dropout): Dropout(p=0.3, inplace=False)\n",
+      "    (layer_norm): LayerNorm((896,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (hidden_projection): Linear(in_features=512, out_features=256, bias=True)\n",
       "  (cell_projection): Linear(in_features=512, out_features=256, bias=True)\n",
       ")\n",
       "\n",
-      "Total parameters: 13,499,715\n"
+      "Total parameters: 13,502,531\n"
      ]
     }
    ],
@@ -286,7 +282,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 12,
    "id": "2d3125a6",
    "metadata": {},
    "outputs": [],
@@ -304,12 +300,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 13,
    "id": "794c40e7",
    "metadata": {},
    "outputs": [],
    "source": [
-    "def train(model, iterator, optimizer, criterion, clip=1.0, teacher_forcing_ratio=0.5):\n",
+    "def train(model, iterator, optimizer, criterion, clip=1.0, teacher_forcing_ratio=0.8):\n",
     "    model.train()\n",
     "    epoch_loss = 0\n",
     "    \n",
@@ -361,7 +357,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 14,
    "id": "d4bb0e92",
    "metadata": {},
    "outputs": [
@@ -369,236 +365,192 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "                                                           "
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Epoch: 01/15 | Time: 7m 39s | TF Ratio: 0.50\n",
-      "\tTrain Loss: 4.5316 | Val Loss: 4.2697 | Best Val: 4.2697 ✓ SAVED\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "                                                           "
+      "                                                           \r"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch: 02/15 | Time: 7m 33s | TF Ratio: 0.45\n",
-      "\tTrain Loss: 3.6810 | Val Loss: 4.0286 | Best Val: 4.0286 ✓ SAVED\n"
+      "Epoch: 01/25 | Time: 8m 1s | TF Ratio: 0.50\n",
+      "\tTrain Loss: 4.1408 | Val Loss: 3.8033 | Best Val: 3.8033 ✓ SAVED\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "                                                           "
+      "                                                           \r"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch: 03/15 | Time: 7m 40s | TF Ratio: 0.41\n",
-      "\tTrain Loss: 3.4275 | Val Loss: 3.8817 | Best Val: 3.8817 ✓ SAVED\n"
+      "Epoch: 02/25 | Time: 7m 41s | TF Ratio: 0.45\n",
+      "\tTrain Loss: 3.0543 | Val Loss: 3.5220 | Best Val: 3.5220 ✓ SAVED\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "                                                           "
+      "                                                           \r"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch: 04/15 | Time: 7m 40s | TF Ratio: 0.36\n",
-      "\tTrain Loss: 3.2257 | Val Loss: 3.7254 | Best Val: 3.7254 ✓ SAVED\n"
+      "Epoch: 03/25 | Time: 7m 40s | TF Ratio: 0.41\n",
+      "\tTrain Loss: 2.6443 | Val Loss: 3.2353 | Best Val: 3.2353 ✓ SAVED\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "                                                           "
+      "                                                           \r"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch: 05/15 | Time: 7m 38s | TF Ratio: 0.33\n",
-      "\tTrain Loss: 3.0585 | Val Loss: 3.6210 | Best Val: 3.6210 ✓ SAVED\n"
+      "Epoch: 04/25 | Time: 7m 44s | TF Ratio: 0.36\n",
+      "\tTrain Loss: 2.3818 | Val Loss: 3.1132 | Best Val: 3.1132 ✓ SAVED\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "                                                           "
+      "                                                           \r"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch: 06/15 | Time: 7m 37s | TF Ratio: 0.30\n",
-      "\tTrain Loss: 2.9102 | Val Loss: 3.4103 | Best Val: 3.4103 ✓ SAVED\n"
+      "Epoch: 05/25 | Time: 7m 42s | TF Ratio: 0.33\n",
+      "\tTrain Loss: 2.2041 | Val Loss: 2.9274 | Best Val: 2.9274 ✓ SAVED\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "                                                           "
+      "                                                           \r"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch: 07/15 | Time: 7m 39s | TF Ratio: 0.27\n",
-      "\tTrain Loss: 2.7814 | Val Loss: 3.3304 | Best Val: 3.3304 ✓ SAVED\n"
+      "Epoch: 06/25 | Time: 7m 36s | TF Ratio: 0.30\n",
+      "\tTrain Loss: 2.0576 | Val Loss: 2.8356 | Best Val: 2.8356 ✓ SAVED\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "                                                           "
+      "                                                           \r"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch: 08/15 | Time: 7m 38s | TF Ratio: 0.24\n",
-      "\tTrain Loss: 2.6669 | Val Loss: 3.2644 | Best Val: 3.2644 ✓ SAVED\n"
+      "Epoch: 07/25 | Time: 7m 41s | TF Ratio: 0.27\n",
+      "\tTrain Loss: 1.9377 | Val Loss: 2.8092 | Best Val: 2.8092 ✓ SAVED\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "                                                           "
+      "                                                           \r"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch: 09/15 | Time: 7m 38s | TF Ratio: 0.22\n",
-      "\tTrain Loss: 2.5686 | Val Loss: 3.2038 | Best Val: 3.2038 ✓ SAVED\n"
+      "Epoch: 08/25 | Time: 7m 39s | TF Ratio: 0.24\n",
+      "\tTrain Loss: 1.8034 | Val Loss: 2.8102 | Best Val: 2.8092  \n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "                                                           "
+      "                                                           \r"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch: 10/15 | Time: 7m 38s | TF Ratio: 0.19\n",
-      "\tTrain Loss: 2.4794 | Val Loss: 3.0976 | Best Val: 3.0976 ✓ SAVED\n"
+      "Epoch: 09/25 | Time: 7m 39s | TF Ratio: 0.22\n",
+      "\tTrain Loss: 1.7125 | Val Loss: 2.7772 | Best Val: 2.7772 ✓ SAVED\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "                                                           "
+      "                                                           \r"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch: 11/15 | Time: 7m 37s | TF Ratio: 0.17\n",
-      "\tTrain Loss: 2.4153 | Val Loss: 3.0713 | Best Val: 3.0713 ✓ SAVED\n"
+      "Epoch: 10/25 | Time: 7m 38s | TF Ratio: 0.19\n",
+      "\tTrain Loss: 1.6454 | Val Loss: 2.8247 | Best Val: 2.7772  \n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "                                                           "
+      "                                                           \r"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch: 12/15 | Time: 7m 35s | TF Ratio: 0.16\n",
-      "\tTrain Loss: 2.3247 | Val Loss: 2.9971 | Best Val: 2.9971 ✓ SAVED\n"
+      "Epoch: 11/25 | Time: 7m 42s | TF Ratio: 0.17\n",
+      "\tTrain Loss: 1.5686 | Val Loss: 2.8969 | Best Val: 2.7772  \n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "                                                           "
+      "                                                          \r"
      ]
     },
     {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Epoch: 13/15 | Time: 7m 38s | TF Ratio: 0.14\n",
-      "\tTrain Loss: 2.2682 | Val Loss: 2.9529 | Best Val: 2.9529 ✓ SAVED\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "                                                           "
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Epoch: 14/15 | Time: 7m 38s | TF Ratio: 0.13\n",
-      "\tTrain Loss: 2.2045 | Val Loss: 2.9489 | Best Val: 2.9489 ✓ SAVED\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "                                                           "
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Epoch: 15/15 | Time: 7m 39s | TF Ratio: 0.11\n",
-      "\tTrain Loss: 2.1487 | Val Loss: 2.9050 | Best Val: 2.9050 ✓ SAVED\n",
-      "\n",
-      "======================================================================\n",
-      "✓ TRAINING COMPLETE!\n",
-      "Best validation loss: 2.9050\n",
-      "Model saved to 'best_model.pth'\n",
-      "======================================================================\n"
+     "ename": "KeyboardInterrupt",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
+      "\u001b[31mKeyboardInterrupt\u001b[39m                         Traceback (most recent call last)",
+      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[14]\u001b[39m\u001b[32m, line 10\u001b[39m\n\u001b[32m      7\u001b[39m start_time = time.time()\n\u001b[32m      9\u001b[39m tf_ratio = \u001b[38;5;28mmax\u001b[39m(\u001b[32m0.1\u001b[39m, \u001b[32m0.5\u001b[39m * (\u001b[32m0.9\u001b[39m ** epoch))\n\u001b[32m---> \u001b[39m\u001b[32m10\u001b[39m train_loss = \u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrain_loader\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptimizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcriterion\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mCLIP\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtf_ratio\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m     11\u001b[39m valid_loss = evaluate(model, val_loader, criterion)\n\u001b[32m     12\u001b[39m scheduler.step(valid_loss)\n",
+      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[13]\u001b[39m\u001b[32m, line 17\u001b[39m, in \u001b[36mtrain\u001b[39m\u001b[34m(model, iterator, optimizer, criterion, clip, teacher_forcing_ratio)\u001b[39m\n\u001b[32m     14\u001b[39m trg = trg[\u001b[32m1\u001b[39m:].reshape(-\u001b[32m1\u001b[39m)\n\u001b[32m     16\u001b[39m loss = criterion(output, trg)\n\u001b[32m---> \u001b[39m\u001b[32m17\u001b[39m \u001b[43mloss\u001b[49m\u001b[43m.\u001b[49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m     19\u001b[39m torch.nn.utils.clip_grad_norm_(model.parameters(), clip)\n\u001b[32m     21\u001b[39m optimizer.step()\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/Desktop/projects/env/lib/python3.12/site-packages/torch/_tensor.py:630\u001b[39m, in \u001b[36mTensor.backward\u001b[39m\u001b[34m(self, gradient, retain_graph, create_graph, inputs)\u001b[39m\n\u001b[32m    620\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_unary(\u001b[38;5;28mself\u001b[39m):\n\u001b[32m    621\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[32m    622\u001b[39m         Tensor.backward,\n\u001b[32m    623\u001b[39m         (\u001b[38;5;28mself\u001b[39m,),\n\u001b[32m   (...)\u001b[39m\u001b[32m    628\u001b[39m         inputs=inputs,\n\u001b[32m    629\u001b[39m     )\n\u001b[32m--> \u001b[39m\u001b[32m630\u001b[39m \u001b[43mtorch\u001b[49m\u001b[43m.\u001b[49m\u001b[43mautograd\u001b[49m\u001b[43m.\u001b[49m\u001b[43mbackward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m    631\u001b[39m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mgradient\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m=\u001b[49m\u001b[43minputs\u001b[49m\n\u001b[32m    632\u001b[39m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/Desktop/projects/env/lib/python3.12/site-packages/torch/autograd/__init__.py:364\u001b[39m, in \u001b[36mbackward\u001b[39m\u001b[34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[39m\n\u001b[32m    359\u001b[39m     retain_graph = create_graph\n\u001b[32m    361\u001b[39m \u001b[38;5;66;03m# The reason we repeat the same comment below is that\u001b[39;00m\n\u001b[32m    362\u001b[39m \u001b[38;5;66;03m# some Python versions print out the first line of a multi-line function\u001b[39;00m\n\u001b[32m    363\u001b[39m \u001b[38;5;66;03m# calls in the traceback and some print out the last line\u001b[39;00m\n\u001b[32m--> \u001b[39m\u001b[32m364\u001b[39m \u001b[43m_engine_run_backward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m    365\u001b[39m \u001b[43m    \u001b[49m\u001b[43mtensors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    366\u001b[39m \u001b[43m    \u001b[49m\u001b[43mgrad_tensors_\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    367\u001b[39m \u001b[43m    \u001b[49m\u001b[43mretain_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    368\u001b[39m \u001b[43m    \u001b[49m\u001b[43mcreate_graph\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    369\u001b[39m \u001b[43m    \u001b[49m\u001b[43minputs_tuple\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m    370\u001b[39m \u001b[43m    \u001b[49m\u001b[43mallow_unreachable\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[32m    371\u001b[39m \u001b[43m    \u001b[49m\u001b[43maccumulate_grad\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[32m    372\u001b[39m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
+      "\u001b[36mFile \u001b[39m\u001b[32m~/Desktop/projects/env/lib/python3.12/site-packages/torch/autograd/graph.py:865\u001b[39m, in \u001b[36m_engine_run_backward\u001b[39m\u001b[34m(t_outputs, *args, **kwargs)\u001b[39m\n\u001b[32m    863\u001b[39m     unregister_hooks = _register_logging_hooks_on_whole_graph(t_outputs)\n\u001b[32m    864\u001b[39m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[32m--> \u001b[39m\u001b[32m865\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mVariable\u001b[49m\u001b[43m.\u001b[49m\u001b[43m_execution_engine\u001b[49m\u001b[43m.\u001b[49m\u001b[43mrun_backward\u001b[49m\u001b[43m(\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# Calls into the C++ engine to run the backward pass\u001b[39;49;00m\n\u001b[32m    866\u001b[39m \u001b[43m        \u001b[49m\u001b[43mt_outputs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\n\u001b[32m    867\u001b[39m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# Calls into the C++ engine to run the backward pass\u001b[39;00m\n\u001b[32m    868\u001b[39m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[32m    869\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m attach_logging_hooks:\n",
+      "\u001b[31mKeyboardInterrupt\u001b[39m: "
      ]
     }
    ],
    "source": [
-    "EPOCHS = 15\n",
+    "EPOCHS = 25\n",
     "CLIP = 1.0\n",
     "best_valid_loss = float('inf')\n",
     "training_history = {'train_loss': [], 'valid_loss': []}\n",
@@ -645,7 +597,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 32,
    "id": "6d9a8e25",
    "metadata": {},
    "outputs": [
@@ -653,12 +605,14 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Loaded checkpoint from best_model.pth (epoch 14)\n",
-      "Sample input (truncated): [CLS:parallel_for] for (i = 0; i < 16; ++i)\n",
-      "  ;\n",
+      "Loaded checkpoint from best_model.pth (epoch 8)\n",
+      "Sample input (truncated): [CLS:reduction] for (i = 0; i < 1000; ++i)\n",
+      "{\n",
+      "  logic_and = logic_and && logics[i];\n",
+      "}\n",
       "\n",
-      "Reference pragma: omp target parallel for simd simdlen(4 4)\n",
-      "Model prediction: omp parallel for simd lastprivate(\n"
+      "Reference pragma: omp parallel for schedule(dynamic,1) private(i) reduction(&&:logic_and)\n",
+      "Model prediction: omp parallel for schedule(dynamic,1) private(i) reduction(&&:logic_and)\n"
      ]
     }
    ],
@@ -712,8 +666,8 @@
     "    return tokenizer.decode(generated)\n",
     "\n",
     "# Quick sanity check on a validation example\n",
-    "sample_input = val_inputs[0]\n",
-    "reference = val_outputs[0]\n",
+    "sample_input = val_inputs[18]\n",
+    "reference = val_outputs[18]\n",
     "prediction = greedy_generate(sample_input)\n",
     "print(\"Sample input (truncated):\", sample_input[:140] + \"...\" if len(sample_input) > 140 else sample_input)\n",
     "print(\"Reference pragma:\", reference)\n",

model/decoder.py

@@ -39,6 +39,7 @@ class Decoder(nn.Module):
         )
         
         self.dropout = nn.Dropout(dropout)
+        self.layer_norm = nn.LayerNorm(hidden_size + hidden_size * 2 + embed_size)
     
     def forward(
         self,
@@ -53,6 +54,7 @@ class Decoder(nn.Module):
         context, attention_weights = self.attention(
             top_hidden, encoder_outputs, mask
         )
+        context = self.dropout(context)
         
         lstm_input = torch.cat((embedded, context.unsqueeze(1)), dim=2)
         
@@ -65,7 +67,8 @@ class Decoder(nn.Module):
         embedded = embedded.squeeze(1)
         
         output_context = torch.cat((output, context, embedded), dim=1)
+        output_context = self.layer_norm(output_context)
         
         prediction = self.fc_out(output_context)
         
-        return prediction, decoder_hidden, decoder_cell, attention_weights
+        return prediction, decoder_hidden, decoder_cell, attention_weights

model/encoder.py

@@ -32,8 +32,8 @@ class Encoder(nn.Module):
             dropout=dropout if num_layers > 1 else 0,
             bidirectional=True
         )
-        
         self.dropout = nn.Dropout(dropout)
+        self.layer_norm = nn.LayerNorm(hidden_size * 2)
     
     def forward(
         self,
@@ -52,5 +52,5 @@ class Encoder(nn.Module):
             packed_output,
             batch_first=True
         )
-        
-        return outputs, hidden, cell
+        outputs = self.layer_norm(outputs)
+        return outputs, hidden, cell