training_pairwise_regloss_topk.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "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, TopKSimilarityLoss\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": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5f0d7db2208c4cd987e8e8853bb64d08",
       "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": 8,
   "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"
     ]
    }
   ],
   "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": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "cs = [c[\"text\"] for c in corpus.values()]\n",
    "qs = list(queries.values())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "Finishing last run (ID:iy606665) before initializing another..."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "34c673fa4f784832b958adc77bb3347d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='0.013 MB of 0.013 MB uploaded\\r'), FloatProgress(value=1.0, max=1.0)))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run <strong style=\"color:#cdcd00\">pairwise_reg_skip_layernorm</strong> at: <a href='https://wandb.ai/laz4rz/matryoshka_pairwise/runs/iy606665' target=\"_blank\">https://wandb.ai/laz4rz/matryoshka_pairwise/runs/iy606665</a><br/> View project at: <a href='https://wandb.ai/laz4rz/matryoshka_pairwise' target=\"_blank\">https://wandb.ai/laz4rz/matryoshka_pairwise</a><br/>Synced 6 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>./wandb/run-20240826_203508-iy606665/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "The new W&B backend becomes opt-out in version 0.18.0; try it out with `wandb.require(\"core\")`! See https://wandb.me/wandb-core for more information."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Successfully finished last run (ID:iy606665). Initializing new run:<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d6236e992fff4165a1eb1e4cc804500a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='Waiting for wandb.init()...\\r'), FloatProgress(value=0.011111914697620604, 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_203549-8xz6e0fi</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/8xz6e0fi' target=\"_blank\">pairwise_reg_topk_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/8xz6e0fi' target=\"_blank\">https://wandb.ai/laz4rz/matryoshka_pairwise/runs/8xz6e0fi</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.0714909553527834\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.0424808740615843\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.017916536331177\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.9935804724693298\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.973906433582306\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.9541754961013793\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.932689893245697\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.9012663960456848\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.8807628393173217\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 0 loss: 1.9717461091020834 eval_loss: 0.24278169125318527\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.8423680782318115\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.8209438800811768\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.8033029675483703\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.7826645255088807\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.7369699716567992\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.730507755279541\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.706375288963318\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.6769646286964417\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.655928099155426\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 1 loss: 1.7481176555156708 eval_loss: 0.19160200282931328\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.6190183639526368\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.6058061838150024\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.5732953190803527\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.5440703868865966\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.5253185033798218\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.4994437098503113\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.4622482538223267\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.449941599369049\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.415003228187561\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 2 loss: 1.5192180288874584 eval_loss: 0.1370648294687271\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.3878769755363465\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.3679767727851868\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.341284441947937\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.3138965010643004\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.3024012327194214\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.2716529965400696\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.2423068284988403\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.2342555522918701\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.2098871231079102\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 3 loss: 1.2952672424523726 eval_loss: 0.08826634287834167\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.1984980940818786\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.177922296524048\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.1567253828048707\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.1374552369117736\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.1240820169448853\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.1075989365577699\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.0838119268417359\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.0762299060821534\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.0557409405708313\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 4 loss: 1.1226705040620721 eval_loss: 0.061398536898195744\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.0607634902000427\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.0363744497299194\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.032974773645401\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.027301549911499\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.0197879731655122\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.004603236913681\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.9994607269763947\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.0056794166564942\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.9954966068267822\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 5 loss: 1.0198362385449202 eval_loss: 0.05260169506072998\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.9699135243892669\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.9867300093173981\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.9764590799808502\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.965187257528305\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.9454527854919433\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.9764470100402832\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.9474615633487702\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.9504186153411865\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.9509088873863221\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 6 loss: 0.9625698366890783 eval_loss: 0.049728672951459885\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.9345678210258483\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.9252578318119049\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.9361412107944489\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.9381246626377105\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.9175156831741333\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.9322361528873444\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.9349242746829987\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.9097214937210083\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.9167334794998169\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 7 loss: 0.9265106877555018 eval_loss: 0.0483147157356143\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.9058258771896363\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.9047136425971984\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.8932783126831054\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.9093328237533569\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.9186095774173737\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.9051150679588318\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.9032449066638947\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.8984106004238128\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.890012800693512\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 8 loss: 0.9030188088831694 eval_loss: 0.047691903077065945\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.8893606424331665\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.892059862613678\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.892192953824997\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.8998451113700867\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.8950607299804687\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.8932998061180115\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.8755385398864746\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.876514983177185\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.8836945474147797\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 9 loss: 0.887922175552534 eval_loss: 0.047188843600451946\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.8867657899856567\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.8888525903224945\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.8681677043437958\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.8798153817653656\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.869137567281723\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.8797962069511414\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.865215641260147\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.8769360959529877\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.8723880231380463\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 10 loss: 0.875957833684009 eval_loss: 0.04701343551278114\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.8859590351581573\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.8739035844802856\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.8647003173828125\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.8614374101161957\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.8667252004146576\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.8625456333160401\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.8643939018249511\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.853375369310379\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.8619961857795715\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 11 loss: 0.8657998028008834 eval_loss: 0.04682436678558588\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.8718172132968902\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.8537865519523621\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.8485569834709168\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.8585374236106873\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.8596360087394714\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.8635716617107392\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.8400083124637604\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.8508914113044739\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.8488603711128235\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 12 loss: 0.8555191735858503 eval_loss: 0.04656581953167915\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.861348170042038\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.8503109157085419\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.8424749910831452\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.8505492210388184\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.8327948451042175\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.8520928263664246\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.8364994168281555\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.8522222936153412\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.8509903192520142\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 13 loss: 0.8481355912011602 eval_loss: 0.04641688335686922\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.8527559220790863\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.8505757808685303\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.8473664581775665\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.8362796843051911\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.8357851386070252\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.829872316122055\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.8467422366142273\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.8371345698833466\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.8230834066867828\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 14 loss: 0.8400040994519773 eval_loss: 0.04631275311112404\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.829780638217926\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.8326732575893402\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.825967139005661\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.8562900841236114\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.8363583147525787\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.8274723827838898\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.8255309343338013\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.8446797132492065\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.8256119430065155\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 15 loss: 0.8336659456076829 eval_loss: 0.04611305147409439\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.8248075664043426\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.8230002224445343\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.8422209620475769\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.8193593502044678\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.8266858518123626\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.8294976770877838\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.8423159122467041\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.8193506479263306\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.8181253492832183\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 16 loss: 0.8276208425345628 eval_loss: 0.04603450279682875\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.8128633737564087\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.8173852980136871\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.8266371130943299\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.8280156135559082\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.8327736020088196\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.8248191714286804\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.8190480053424836\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.8160410404205323\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.8309522092342376\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 17 loss: 0.8232500468907149 eval_loss: 0.04597785975784063\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.8156634449958802\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.825304526090622\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.8221731781959534\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.8143977046012878\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.8211024940013886\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.8121215105056763\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.8156988441944122\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.8233668565750122\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.8147494077682496\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 18 loss: 0.8186585935561553 eval_loss: 0.04592515807598829\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.824483722448349\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.8201368391513825\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.8126678943634034\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.8062148869037629\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.7985086679458618\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.8213733613491059\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.8057322859764099\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.8245919585227967\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.808605694770813\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 19 loss: 0.8142048373170521 eval_loss: 0.04579660575836897\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.8031754851341247\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.807389211654663\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.8119488775730133\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.8060495436191559\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.8061563014984131\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.8091408491134644\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.8149797916412354\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.8141852498054505\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.8019773781299591\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 20 loss: 0.8081537880327391 eval_loss: 0.045685685239732265\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.8114294886589051\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.7957885921001434\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.8138441801071167\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.8086468696594238\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.8007187962532043\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.7974467992782592\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.8082451760768891\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.8065548539161682\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.8043025195598602\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 21 loss: 0.8048521215501039 eval_loss: 0.045633633621037006\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.7895053386688232\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.8189716815948487\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.7847348272800445\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.8053452014923096\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.7959512829780578\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.8089095473289489\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.8115442276000977\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.8023633897304535\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.8006280422210693\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 22 loss: 0.8017927129631457 eval_loss: 0.04556760657578707\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.7949838638305664\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.8059546291828156\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.8030726730823516\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.8120013415813446\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.7873087048530578\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.7962858557701111\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.7904390454292297\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.7843383252620697\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.7916596055030822\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 23 loss: 0.7958612020896829 eval_loss: 0.04554499126970768\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.7929193258285523\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.7989086985588074\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.7935534596443177\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.8001280903816224\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.7889819443225861\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.8015326678752899\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.7958470046520233\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.7819853007793427\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.7971375405788421\n",
      "5760\n",
      "5824\n",
      "5888\n",
      "Epoch: 24 loss: 0.7940995829260867 eval_loss: 0.04555346630513668\n"
     ]
    }
   ],
   "source": [
    "run_name = \"pairwise_reg_topk_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",
    "loss_fn_topk = TopKSimilarityLoss(k=10)\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",
    "        loss_topk = loss_fn_topk(target_outputs, outputs, [64, 128, 256, 384])\n",
    "\n",
    "        loss = loss + loss_reg + loss_topk\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
}