training_pairwise.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/beir/datasets/data_loader.py:2: TqdmExperimentalWarning: Using `tqdm.autonotebook.tqdm` in notebook mode. Use `tqdm.tqdm` instead to force console mode (e.g. in jupyter console)\n",
      "  from tqdm.autonotebook import tqdm\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import random\n",
    "import wandb\n",
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from beir.datasets.data_loader import GenericDataLoader\n",
    "\n",
    "from matryoshka import Matryoshka, PairwiseSimilarityLoss, PairwiseSimilarityLossParallel\n",
    "\n",
    "random.seed(42)\n",
    "np.random.seed(42)\n",
    "torch.manual_seed(42)\n",
    "torch.cuda.manual_seed(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bb9972d3c97847fa8b4c0b8669a05666",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3633 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "data_path = \"data/nfcorpus\"\n",
    "corpus, queries, qrels = GenericDataLoader(data_folder=data_path).load(split=\"train\")\n",
    "\n",
    "length = None\n",
    "corpus = {k: v for k, v in list(corpus.items())[:length]}\n",
    "queries = {k: v for k, v in list(queries.items())[:length]}\n",
    "qrels = {k: v for k, v in list(qrels.items())[:length]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
      "  warnings.warn(\n",
      "/opt/conda/lib/python3.10/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "base_model = Matryoshka(matryoshka_dim=384, adaptor=False)\n",
    "model = Matryoshka(matryoshka_dim=384, adaptor=True)\n",
    "tokenizer = model.tokenizer\n",
    "\n",
    "sentences = [\"sentence\"]\n",
    "inputs = tokenizer(sentences, return_tensors=\"pt\", padding=True, truncation=True)\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    model = model.cuda()\n",
    "    base_model = base_model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "cs = [c[\"text\"] for c in corpus.values()]\n",
    "qs = list(queries.values())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(133.9172, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "{64: tensor(35.2179, device='cuda:0', grad_fn=<AddBackward0>), 128: tensor(33.6590, device='cuda:0', grad_fn=<AddBackward0>), 256: tensor(31.9819, device='cuda:0', grad_fn=<AddBackward0>), 384: tensor(33.0584, device='cuda:0', grad_fn=<AddBackward0>)}\n",
      "tensor(128.0457, device='cuda:0', grad_fn=<AddBackward0>)\n",
      "{64: tensor(33.7338, device='cuda:0', grad_fn=<AddBackward0>), 128: tensor(32.2593, device='cuda:0', grad_fn=<AddBackward0>), 256: tensor(30.7107, device='cuda:0', grad_fn=<AddBackward0>), 384: tensor(31.3418, device='cuda:0', grad_fn=<AddBackward0>)}\n"
     ]
    }
   ],
   "source": [
    "inputs = tokenizer(qs[:10], return_tensors=\"pt\", padding=True, truncation=True)\n",
    "if torch.cuda.is_available():\n",
    "    for k, v in inputs.items():\n",
    "        inputs[k] = v.cuda()\n",
    "outputs = model(pooling=True, skip=False, **inputs)\n",
    "target_outputs = base_model(pooling=True, **inputs)\n",
    "\n",
    "loss_fn = PairwiseSimilarityLoss()\n",
    "loss_fn_parallel = PairwiseSimilarityLossParallel()\n",
    "loss, loss_partial = loss_fn(outputs, target_outputs, [64, 128, 256, 384], reduce=False)\n",
    "loss2, loss_partial2 = loss_fn_parallel(outputs, target_outputs, [64, 128, 256, 384], reduce=False)\n",
    "\n",
    "print(loss)\n",
    "print(loss_partial)\n",
    "print(loss2)\n",
    "print(loss_partial2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss\n",
      "tensor(8.4998, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    }
   ],
   "source": [
    "import torch.nn.functional as F\n",
    "\n",
    "loss = 0\n",
    "embeddings = target_outputs\n",
    "embeddings_adapted = outputs\n",
    "for j in range(len(outputs)):\n",
    "    target_similarity = F.cosine_similarity(embeddings[0].squeeze(0), embeddings[j].squeeze(0), dim=0)\n",
    "    for m in [64, 128, 256, 384]:\n",
    "        similarity = F.cosine_similarity(embeddings_adapted[0].squeeze(0)[:m], embeddings[j].squeeze(0)[:m], dim=0)\n",
    "        loss += torch.abs(target_similarity - similarity)\n",
    "print(\"loss\")\n",
    "print(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(8.4998, device='cuda:0', grad_fn=<AddBackward0>)\n"
     ]
    }
   ],
   "source": [
    "loss = 0\n",
    "embeddings = target_outputs\n",
    "embeddings_adapted = outputs\n",
    "cloned_target = embeddings[0].repeat(10, 1)\n",
    "target_similarity = F.cosine_similarity(cloned_target, embeddings, dim=1)\n",
    "cloned_adapted = embeddings_adapted[0].repeat(10, 1)\n",
    "for m in [64, 128, 256, 384]:\n",
    "    similarity = F.cosine_similarity(cloned_adapted[:, :m], embeddings[:, :m], dim=1)\n",
    "    loss += torch.sum(torch.abs(target_similarity - similarity))\n",
    "print(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23.4 ms ± 141 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "loss, loss_partial = loss_fn(outputs, target_outputs, [64, 128, 256, 384], reduce=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5.07 ms ± 12.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "loss2, loss_partial2 = loss_fn_parallel(outputs, target_outputs, [64, 128, 256, 384], reduce=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mlaz4rz\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7e893278a78f4441b12e9cd132983d9b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='Waiting for wandb.init()...\\r'), FloatProgress(value=0.011112051188117929, max=1.0…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.17.7"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/root/matryoshka/wandb/run-20240826_183531-zfzfkghb</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/laz4rz/matryoshka_pairwise/runs/zfzfkghb' target=\"_blank\">pairwise_loss_skip_normalized</a></strong> to <a href='https://wandb.ai/laz4rz/matryoshka_pairwise' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/laz4rz/matryoshka_pairwise' target=\"_blank\">https://wandb.ai/laz4rz/matryoshka_pairwise</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/laz4rz/matryoshka_pairwise/runs/zfzfkghb' target=\"_blank\">https://wandb.ai/laz4rz/matryoshka_pairwise/runs/zfzfkghb</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.38975848257541656\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.3867886245250702\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.3859786123037338\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.3808394134044647\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.37997758388519287\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.3750932842493057\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.3725013703107834\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 0 loss: 0.3812659282188911 eval_loss: 0.37232198683839096\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.36411979496479036\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.3660570025444031\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.3615768730640411\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.3600315451622009\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.3560211807489395\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.35197252929210665\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.34585377275943757\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 1 loss: 0.3567195612114745 eval_loss: 0.34436648927236857\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.33942426443099977\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.33241714239120485\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.33764695525169375\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.32280888855457307\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.3216740429401398\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.311101433634758\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.30887019336223603\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 2 loss: 0.3229202767471214 eval_loss: 0.302460977905675\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.29699728190898894\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.2862587749958038\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.281273353099823\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.27256103307008744\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.2700335308909416\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.2563768431544304\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.23475342392921447\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 3 loss: 0.26715318645749775 eval_loss: 0.23137781572969338\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.21915096789598465\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.21399540156126023\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.18983609229326248\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.1808306396007538\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.17407256215810776\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.15572642385959626\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.14370137751102446\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 4 loss: 0.17796194456614456 eval_loss: 0.13288493023106926\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.1261090323328972\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.1190803050994873\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.11490565091371537\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.11587914302945138\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.11289053708314896\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.11086148098111152\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.1096495345234871\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 5 loss: 0.11497789640705307 eval_loss: 0.09859537137182135\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.10041948556900024\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.09878404140472412\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.09476422816514969\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.09338229224085808\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.09508641734719277\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.09375598728656769\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.09017239436507225\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 6 loss: 0.09426014909109512 eval_loss: 0.08354687612307699\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.08518427833914757\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.08432185351848602\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.08075226992368698\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.0827925406396389\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.079519834369421\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.08240972906351089\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.07842741385102273\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 7 loss: 0.0818034199538169 eval_loss: 0.07373333643925817\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.07600673139095307\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.07434826195240021\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.07366737350821495\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.07496651336550712\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.07218708693981171\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.07064684927463531\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.07196262255311012\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 8 loss: 0.07324535119069087 eval_loss: 0.06797180442433608\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.06919820047914982\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.0693405658006668\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.06908508092164993\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.06682524085044861\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.06858700290322303\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.06746469363570214\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.06761791259050369\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 9 loss: 0.0681458679209282 eval_loss: 0.06366518375120665\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.06587723270058632\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.06616704352200031\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.06536762379109859\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.06489429511129856\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.06409142203629017\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.0643436636775732\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.06438975632190705\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 10 loss: 0.06494999386660465 eval_loss: 0.06171171033852979\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.06445256657898427\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.06270152665674686\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.0630718532949686\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.06256203167140484\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.061821871995925905\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.06307688727974892\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.06252947635948658\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 11 loss: 0.06274598096678784 eval_loss: 0.05932497468433882\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.0627339694648981\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.061760326474905015\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.06218230649828911\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.06074652373790741\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.061000478267669675\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.060573424026370046\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.06148635819554329\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 12 loss: 0.061449884255598115 eval_loss: 0.057912549494128475\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.061158095672726634\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.061031120270490645\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.06092230826616287\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.06049954667687416\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.05943373702466488\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.06023319587111473\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.06051075682044029\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 13 loss: 0.0604112060232596 eval_loss: 0.0569381198208583\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.05987762026488781\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.06038348935544491\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.059453465044498444\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.059129268676042554\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.05962332971394062\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.05853031799197197\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.05872403010725975\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 14 loss: 0.059422185136513275 eval_loss: 0.056020727283076235\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.06053743027150631\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.05810055993497372\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.057860585302114485\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.058801843971014026\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.05866151452064514\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.058558203279972076\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.059622620418667796\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 15 loss: 0.058925517048541605 eval_loss: 0.055218527583699474\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.058775104209780694\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.05843987688422203\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.058484336361289024\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.05775183215737343\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.05749083682894707\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.057723529264330865\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.05790027640759945\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 16 loss: 0.05808762601250178 eval_loss: 0.05433757956090726\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.057487012818455696\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.05771183595061302\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.05756686702370643\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.05786304324865341\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.057784304022789\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.057270390540361406\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.05713200010359287\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 17 loss: 0.05746711553497748 eval_loss: 0.0538826884799882\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.057932083308696744\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.05627211853861809\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.057159971073269845\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.05804916806519032\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.05690103285014629\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.05651534423232078\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.05761745311319828\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 18 loss: 0.05714789865078864 eval_loss: 0.053466729623706716\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.057993557304143906\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.056419705599546434\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.05726307407021523\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.056424865126609804\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.05694531463086605\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.05670414939522743\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.0564150869846344\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 19 loss: 0.056801054820224836 eval_loss: 0.05270253847304143\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.05633576065301895\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.05579679645597935\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.056395990401506425\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.05653482526540756\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.056450647115707395\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.0564983420073986\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.05558919757604599\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 20 loss: 0.05629186288683445 eval_loss: 0.05267483505763506\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.05702760890126228\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.05562193058431149\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.05594151057302952\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.05537833608686924\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.056432831659913064\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.05547347776591778\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.05628661960363388\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 21 loss: 0.056033584252967464 eval_loss: 0.05199772629298662\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.05507693961262703\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.055444838106632234\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.05485650673508644\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.0558669563382864\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.05549939163029194\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.056427692994475365\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.05567287467420101\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 22 loss: 0.05548690525548799 eval_loss: 0.05191877366680848\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.05526545085012913\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.055837782472372054\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.05531158149242401\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.05525295175611973\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.05386478789150715\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.05572314001619816\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.055680985748767856\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 23 loss: 0.05526501985339375 eval_loss: 0.051485018118431695\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 0.05523969978094101\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 0.054640794172883034\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 0.05425774455070496\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 0.055120726302266124\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 0.05484226010739803\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 0.05591348297894001\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 0.05556385479867458\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "Epoch: 24 loss: 0.05500728143500043 eval_loss: 0.051188286393880844\n"
     ]
    }
   ],
   "source": [
    "run_name = \"pairwise_loss_skip_normalized\"\n",
    "epochs = 25\n",
    "batch_size = 64\n",
    "running_loss_step = 10\n",
    "learning_rate = 1e-5\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
    "loss_fn = PairwiseSimilarityLossParallel()\n",
    "# scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)\n",
    "\n",
    "data = qs + cs\n",
    "random.shuffle(data)\n",
    "train_data = data[:int(len(data) * 0.8)]\n",
    "test_data = data[int(len(data) * 0.8):]\n",
    "\n",
    "wandb.init(\n",
    "    project=\"matryoshka_pairwise\",  \n",
    "    name=run_name,        \n",
    "    config={                      \n",
    "        \"learning_rate\": learning_rate,\n",
    "        \"epochs\": epochs,\n",
    "        \"batch_size\": batch_size,\n",
    "        \"loss\": \"PairwiseSimilarityLoss\",\n",
    "        \"model\": model.name,\n",
    "        \"model_card\": model.model_card_data,\n",
    "        \"loss_resolution\": running_loss_step,\n",
    "        \"architecture\": model.__str__(),\n",
    "    }\n",
    ")\n",
    "\n",
    "ls = []\n",
    "for i in range(epochs):\n",
    "    epoch_loss = []\n",
    "    running_loss = []\n",
    "    partial_running_loss = {64: [], 128: [], 256: [], 384: []}\n",
    "\n",
    "    model.train()\n",
    "    random.shuffle(train_data)\n",
    "    for j in range(0, len(train_data), batch_size):\n",
    "        print(j)\n",
    "        if j + batch_size > len(train_data):\n",
    "            break\n",
    "        # c = cs[j : j + 32]\n",
    "        # q = qs[j : j + 32]\n",
    "        q = train_data[j : j + batch_size]\n",
    "\n",
    "        inputs = tokenizer(q, return_tensors=\"pt\", padding=True, truncation=True)\n",
    "        if torch.cuda.is_available():\n",
    "            for k, v in inputs.items():\n",
    "                inputs[k] = v.cuda()\n",
    "        outputs = model(pooling=True, skip=False, **inputs)\n",
    "        target_outputs = base_model(pooling=True, **inputs)\n",
    "\n",
    "        loss, loss_partial = loss_fn(outputs, target_outputs, [64, 128, 256, 384])\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        ls.append(loss.item())\n",
    "        epoch_loss.append(loss.item())\n",
    "        running_loss.append(loss.item())\n",
    "        for k, v in loss_partial.items():\n",
    "            partial_running_loss[k].append(v.item())\n",
    "        if len(running_loss) % 10 == 0:\n",
    "            print(\"Batch:\", j, \"loss:\", np.mean(running_loss))\n",
    "            wandb.log({\"batch\": j, \"loss\": np.mean(running_loss)} | {f\"loss_{k}\": np.mean(v) for k, v in partial_running_loss.items()})\n",
    "            partial_running_loss = {64: [], 128: [], 256: [], 384: []}\n",
    "            running_loss = []\n",
    " \n",
    "    model.eval()\n",
    "    eval_loss = []\n",
    "    partial_eval_loss = {64: [], 128: [], 256: [], 384: []}\n",
    "    for j in range(0, len(test_data), batch_size):\n",
    "        if j + batch_size > len(test_data):\n",
    "            break\n",
    "        q = test_data[j : j + batch_size]\n",
    "\n",
    "        inputs = tokenizer(q, return_tensors=\"pt\", padding=True, truncation=True)\n",
    "        if torch.cuda.is_available():\n",
    "            for k, v in inputs.items():\n",
    "                inputs[k] = v.cuda()\n",
    "        with torch.no_grad():\n",
    "            outputs = model(pooling=True, skip=False, **inputs)\n",
    "            target_outputs = base_model(pooling=True, **inputs)\n",
    "            loss, loss_partial = loss_fn(outputs, target_outputs, [64, 128, 256, 384])\n",
    "            eval_loss.append(loss.item())\n",
    "            for k, v in loss_partial.items():\n",
    "                partial_eval_loss[k].append(v.item())\n",
    "    wandb.log({\"epoch\": i, \"epoch_loss\": np.mean(epoch_loss), \"eval_loss\": np.mean(eval_loss)} | {f\"eval_loss_{k}\": np.mean(v) for k, v in partial_eval_loss.items()})\n",
    "    print(\"Epoch:\", i, \"loss:\", np.mean(epoch_loss), \"eval_loss:\", np.mean(eval_loss))\n",
    "    os.makedirs(f\"models/{run_name}\", exist_ok=True)\n",
    "    torch.save(model, f\"models/{run_name}/{i}.pth\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}