training_pairwise_regloss.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, RegularizingLoss\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": "5f20b653935948f38c1982e6413a3f10",
       "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": "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": {
      "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_201548-d02j5n8k</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/d02j5n8k' target=\"_blank\">pairwise_reg_skip_layernorm</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/d02j5n8k' target=\"_blank\">https://wandb.ai/laz4rz/matryoshka_pairwise/runs/d02j5n8k</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: 2.07143087387085\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 2.0424219608306884\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 2.017858695983887\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 1.993524181842804\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 1.973851215839386\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 1.954120934009552\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 1.9326368689537048\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 1.9012141227722168\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 1.8807119965553283\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 0 loss: 1.9716907752596813 eval_loss: 0.24278249591588974\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 1.8423197507858275\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 1.8208955883979798\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 1.803256595134735\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 1.7826189041137694\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 1.7369270324707031\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 1.7304641366004945\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 1.7063336849212647\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 1.6769247174263\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 1.6558892726898193\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 1 loss: 1.7480738512847735 eval_loss: 0.1916041597723961\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 1.6189814805984497\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 1.6057694554328918\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 1.5732603073120117\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 1.5440366744995118\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 1.5252868294715882\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 1.4994136095046997\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 1.4622195959091187\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 1.449914014339447\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 1.4149765968322754\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 2 loss: 1.5191862622032994 eval_loss: 0.1370687186717987\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 1.387851893901825\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 1.3679529666900634\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 1.3412614464759827\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 1.313875424861908\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 1.3023807525634765\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 1.2716333389282226\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 1.242289423942566\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 1.234238040447235\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 1.209871804714203\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 3 loss: 1.2952469807604086 eval_loss: 0.08827047049999237\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 1.198484218120575\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 1.1779089212417602\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 1.1567129969596863\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 1.1374436497688294\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 1.124072015285492\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 1.107588255405426\n",
      "3840\n",
      "3904\n",
      "3968\n",
      "4032\n",
      "4096\n",
      "4160\n",
      "4224\n",
      "4288\n",
      "4352\n",
      "4416\n",
      "Batch: 4416 loss: 1.083800494670868\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 1.0762203454971313\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 1.0557317852973938\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 4 loss: 1.1226592050946278 eval_loss: 0.06139976345002651\n",
      "0\n",
      "64\n",
      "128\n",
      "192\n",
      "256\n",
      "320\n",
      "384\n",
      "448\n",
      "512\n",
      "576\n",
      "Batch: 576 loss: 1.060754132270813\n",
      "640\n",
      "704\n",
      "768\n",
      "832\n",
      "896\n",
      "960\n",
      "1024\n",
      "1088\n",
      "1152\n",
      "1216\n",
      "Batch: 1216 loss: 1.0363656401634216\n",
      "1280\n",
      "1344\n",
      "1408\n",
      "1472\n",
      "1536\n",
      "1600\n",
      "1664\n",
      "1728\n",
      "1792\n",
      "1856\n",
      "Batch: 1856 loss: 1.0329653024673462\n",
      "1920\n",
      "1984\n",
      "2048\n",
      "2112\n",
      "2176\n",
      "2240\n",
      "2304\n",
      "2368\n",
      "2432\n",
      "2496\n",
      "Batch: 2496 loss: 1.0272923946380614\n",
      "2560\n",
      "2624\n",
      "2688\n",
      "2752\n",
      "2816\n",
      "2880\n",
      "2944\n",
      "3008\n",
      "3072\n",
      "3136\n",
      "Batch: 3136 loss: 1.019779086112976\n",
      "3200\n",
      "3264\n",
      "3328\n",
      "3392\n",
      "3456\n",
      "3520\n",
      "3584\n",
      "3648\n",
      "3712\n",
      "3776\n",
      "Batch: 3776 loss: 1.004595297574997\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.9994521796703338\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 1.0056699573993684\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.9954874813556671\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 5 loss: 1.0198272661022518 eval_loss: 0.05260183196514845\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.9699057519435883\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.9867214977741241\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.9764498829841614\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.9651783168315887\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.9454442441463471\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.976436972618103\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.9474526643753052\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 0.9504084408283233\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.9509001016616822\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 6 loss: 0.9625608661900396 eval_loss: 0.049727657809853554\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.9345594704151153\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.9252503991127015\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.9361324310302734\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.9381145536899567\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.9175065279006958\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.9322273313999176\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.934915429353714\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 0.909714549779892\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.9167232096195221\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 7 loss: 0.9265019278163495 eval_loss: 0.04831430595368147\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.9058156669139862\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.904705411195755\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.8932705521583557\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.9093243837356567\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.918600070476532\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.9051078259944916\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.903236347436905\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 0.8984029471874238\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.8900034070014954\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 8 loss: 0.9030102037865183 eval_loss: 0.04769071005284786\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.8893527746200561\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.8920509874820709\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.8921830534934998\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.8998368263244629\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.8950519323348999\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.8932918846607208\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.875530618429184\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 0.8765076994895935\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.8836865425109863\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 9 loss: 0.887913843859797 eval_loss: 0.04719036631286144\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.8867589294910431\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.8888465702533722\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.8681585729122162\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.8798071861267089\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.8691290557384491\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.8797869741916656\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.86520716547966\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 0.8769276082515717\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.8723781049251557\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 10 loss: 0.8759494812592216 eval_loss: 0.04701513051986694\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.8859518527984619\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.8738971650600433\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.8646934509277344\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.8614292085170746\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.8667176425457\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.8625382900238037\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.864384937286377\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 0.8533701837062836\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.8619885385036469\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 11 loss: 0.8657926055400268 eval_loss: 0.04682258144021034\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.8718105137348175\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.8537777304649353\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.8485504269599915\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.8585318505764008\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.859628027677536\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.863562536239624\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.8400034308433533\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 0.8508822321891785\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.8488543093204498\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 12 loss: 0.8555119102415831 eval_loss: 0.04656588099896908\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.8613408803939819\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.8503037691116333\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.8424672365188599\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.8505404472351075\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.8327888965606689\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.852086991071701\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.8364926218986511\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 0.8522137045860291\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.8509844660758972\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 13 loss: 0.8481285688669785 eval_loss: 0.04641709104180336\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.8527490556240082\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.8505683839321136\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.8473581433296203\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.836273443698883\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.8357780039310455\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.8298637390136718\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.8467368245124817\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 0.8371263742446899\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.8230774521827697\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 14 loss: 0.8399970388930776 eval_loss: 0.0463138734921813\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.8297750353813171\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.8326666593551636\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.8259604215621948\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.8562809526920319\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.8363502800464631\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.8274668693542481\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.8255228698253632\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 0.8446728467941285\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.8256034255027771\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 15 loss: 0.8336586978124536 eval_loss: 0.04611244332045317\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.824799758195877\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.8229940414428711\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.842213386297226\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.8193539381027222\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.8266808688640594\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.8294912815093994\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.8423081994056701\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 0.8193435430526733\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.8181199431419373\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 16 loss: 0.82761436117732 eval_loss: 0.046035634353756905\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.8128562450408936\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.8173796474933624\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.8266320586204529\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.8280090987682343\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.8327681064605713\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.8248124837875366\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.819040733575821\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 0.8160342752933503\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.8309447228908539\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 17 loss: 0.8232435882091522 eval_loss: 0.045976645313203335\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.8156577825546265\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.8253005802631378\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.822168493270874\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.8143926739692688\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.8210972785949707\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.8121142446994781\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.8156929194927216\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 0.8233616530895234\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.8147450268268586\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 18 loss: 0.81865336065707 eval_loss: 0.04592460300773382\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.8244767963886261\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.8201315820217132\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.8126620948314667\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.8062079787254334\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.7985038757324219\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.8213655948638916\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.8057250261306763\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 0.8245842754840851\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.8085979580879211\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 19 loss: 0.8141982024130614 eval_loss: 0.04579675476998091\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.8031678915023803\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.8073818743228912\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.8119424700736999\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.8060428678989411\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.8061486661434174\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.8091342329978943\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.8149737358093262\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 0.8141780078411103\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.8019692718982696\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 20 loss: 0.8081467119248017 eval_loss: 0.04568643681704998\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.8114240169525146\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.7957820236682892\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.8138363838195801\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.8086400687694549\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.800710940361023\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.7974393308162689\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.8082393050193787\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 0.8065479099750519\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.8042962312698364\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 21 loss: 0.8048453803943552 eval_loss: 0.04563468415290117\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.7894973933696747\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.8189643740653991\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.7847303450107574\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.8053371787071228\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.7959429800510407\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.8089022040367126\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.8115385472774506\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 0.8023563385009765\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.8006205916404724\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 22 loss: 0.8017854904351027 eval_loss: 0.04556785523891449\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.7949769496917725\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.8059463679790497\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.8030638217926025\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.8119937598705291\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.7873017847537994\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.7962831676006317\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.7904304146766663\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 0.7843301296234131\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.7916539132595062\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 23 loss: 0.7958541195029798 eval_loss: 0.04554487392306328\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.7929144024848938\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.7989008069038391\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.7935472965240479\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.8001224458217621\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.7889769613742829\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.8015250563621521\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.7958392262458801\n",
      "4480\n",
      "4544\n",
      "4608\n",
      "4672\n",
      "4736\n",
      "4800\n",
      "4864\n",
      "4928\n",
      "4992\n",
      "5056\n",
      "Batch: 5056 loss: 0.7819815039634704\n",
      "5120\n",
      "5184\n",
      "5248\n",
      "5312\n",
      "5376\n",
      "5440\n",
      "5504\n",
      "5568\n",
      "5632\n",
      "5696\n",
      "Batch: 5696 loss: 0.7971316635608673\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 24 loss: 0.7940934300422668 eval_loss: 0.04555210657417774\n"
     ]
    }
   ],
   "source": [
    "run_name = \"pairwise_reg_skip_layernorm\"\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",
    "loss_fn_reg = RegularizingLoss()\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.95)]\n",
    "test_data = data[int(len(data) * 0.95):]\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+RegLoss\",\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, **inputs)\n",
    "        target_outputs = base_model(pooling=True, **inputs)\n",
    "\n",
    "        loss, loss_partial = loss_fn(target_outputs, outputs, [64, 128, 256, 384])\n",
    "        loss_reg = loss_fn_reg(target_outputs, outputs)\n",
    "\n",
    "        loss = loss + loss_reg\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, **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
}