template_nb_question1_relu_final_latest_nodropout.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "F2UvNJlIZXa2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-27T13:00:39.936796Z",
     "start_time": "2024-10-27T13:00:39.934202Z"
    },
    "execution": {
     "iopub.execute_input": "2024-10-27T18:35:44.559164Z",
     "iopub.status.busy": "2024-10-27T18:35:44.558777Z",
     "iopub.status.idle": "2024-10-27T18:35:44.563310Z",
     "shell.execute_reply": "2024-10-27T18:35:44.562426Z",
     "shell.execute_reply.started": "2024-10-27T18:35:44.559125Z"
    },
    "id": "F2UvNJlIZXa2",
    "trusted": true
   },
   "outputs": [],
   "source": [
    "import re\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6xsIy6cSZXa3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-27T12:57:46.625439Z",
     "start_time": "2024-10-27T12:57:46.621767Z"
    },
    "execution": {
     "iopub.execute_input": "2024-10-27T18:35:44.829446Z",
     "iopub.status.busy": "2024-10-27T18:35:44.828775Z",
     "iopub.status.idle": "2024-10-27T18:35:44.833305Z",
     "shell.execute_reply": "2024-10-27T18:35:44.832411Z",
     "shell.execute_reply.started": "2024-10-27T18:35:44.829406Z"
    },
    "id": "6xsIy6cSZXa3",
    "trusted": true
   },
   "outputs": [],
   "source": [
    "text = \"\"\"Taehyung is trying to divide 21/11 liters (L) of water into 7/11 liters (L) per cup, and Hoseok is trying to divide 8/17 liters (L) of water into 2/17 liters (L) per cup. How many cups will Taehyung and Hoseok need in order to divide all the water they each have into the cups?\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "FBvq4XjhZXa3",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-27T13:05:58.576215Z",
     "start_time": "2024-10-27T13:05:58.572440Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2024-10-27T18:35:45.222632Z",
     "iopub.status.busy": "2024-10-27T18:35:45.221803Z",
     "iopub.status.idle": "2024-10-27T18:35:45.227854Z",
     "shell.execute_reply": "2024-10-27T18:35:45.226957Z",
     "shell.execute_reply.started": "2024-10-27T18:35:45.222582Z"
    },
    "id": "FBvq4XjhZXa3",
    "outputId": "ecd1e5cd-9b57-4fda-923b-88020bb03634",
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Taehyung', 'is', 'trying', 'to', 'divide', '2', '1', '/', '1', '1', 'liters', '(', 'L', ')', 'of', 'water', 'into', '7', '/', '1', '1', 'liters', '(', 'L', ')', 'per', 'cup', ',', 'and', 'Hoseok', 'is', 'trying', 'to', 'divide', '8', '/', '1', '7', 'liters', '(', 'L', ')', 'of', 'water', 'into', '2', '/', '1', '7', 'liters', '(', 'L', ')', 'per', 'cup', '.', 'How', 'many', 'cups', 'will', 'Taehyung', 'and', 'Hoseok', 'need', 'in', 'order', 'to', 'divide', 'all', 'the', 'water', 'they', 'each', 'have', 'into', 'the', 'cups', '?']\n"
     ]
    }
   ],
   "source": [
    "tokens = [i for i in re.split('(\\s|[^a-zA-Z])', text) if i and not i.isspace()]\n",
    "print(tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "Lv_h3VnXZ7iG",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Lv_h3VnXZ7iG",
    "outputId": "2ee59a2b-1fdc-4f7c-e619-39263fc0c675"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: datasets in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (3.0.2)\n",
      "Requirement already satisfied: filelock in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from datasets) (3.15.4)\n",
      "Requirement already satisfied: numpy>=1.17 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from datasets) (1.26.4)\n",
      "Requirement already satisfied: pyarrow>=15.0.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from datasets) (17.0.0)\n",
      "Requirement already satisfied: dill<0.3.9,>=0.3.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from datasets) (0.3.8)\n",
      "Requirement already satisfied: pandas in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from datasets) (2.1.4)\n",
      "Requirement already satisfied: requests>=2.32.2 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from datasets) (2.32.3)\n",
      "Requirement already satisfied: tqdm>=4.66.3 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from datasets) (4.66.5)\n",
      "Requirement already satisfied: xxhash in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from datasets) (3.5.0)\n",
      "Requirement already satisfied: multiprocess<0.70.17 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from datasets) (0.70.16)\n",
      "Requirement already satisfied: fsspec<=2024.9.0,>=2023.1.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from fsspec[http]<=2024.9.0,>=2023.1.0->datasets) (2024.9.0)\n",
      "Requirement already satisfied: aiohttp in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from datasets) (3.10.5)\n",
      "Requirement already satisfied: huggingface-hub>=0.23.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from datasets) (0.26.1)\n",
      "Requirement already satisfied: packaging in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from datasets) (24.1)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from datasets) (6.0.2)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from aiohttp->datasets) (2.4.0)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from aiohttp->datasets) (1.3.1)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from aiohttp->datasets) (24.2.0)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from aiohttp->datasets) (1.4.1)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from aiohttp->datasets) (6.0.5)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from aiohttp->datasets) (1.9.11)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from aiohttp->datasets) (4.0.3)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from huggingface-hub>=0.23.0->datasets) (4.12.2)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from requests>=2.32.2->datasets) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from requests>=2.32.2->datasets) (3.8)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from requests>=2.32.2->datasets) (2.2.2)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from requests>=2.32.2->datasets) (2024.8.30)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from pandas->datasets) (2.9.0.post0)\n",
      "Requirement already satisfied: pytz>=2020.1 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from pandas->datasets) (2024.1)\n",
      "Requirement already satisfied: tzdata>=2022.1 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from pandas->datasets) (2024.1)\n",
      "Requirement already satisfied: six>=1.5 in /system/conda/miniconda3/envs/cloudspace/lib/python3.10/site-packages (from python-dateutil>=2.8.2->pandas->datasets) (1.16.0)\n"
     ]
    }
   ],
   "source": [
    "!pip install datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "aHw0YOurZXa4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-27T13:08:31.787862Z",
     "start_time": "2024-10-27T13:08:11.280944Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2024-10-27T18:35:45.385882Z",
     "iopub.status.busy": "2024-10-27T18:35:45.385577Z",
     "iopub.status.idle": "2024-10-27T18:35:46.679193Z",
     "shell.execute_reply": "2024-10-27T18:35:46.678444Z",
     "shell.execute_reply.started": "2024-10-27T18:35:45.385849Z"
    },
    "id": "aHw0YOurZXa4",
    "outputId": "0dc0f632-b447-433f-baa1-0d41c6302a86",
    "trusted": true
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "ds = load_dataset(\"microsoft/orca-math-word-problems-200k\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "m2DJP8uUZXa4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-27T13:14:28.885715Z",
     "start_time": "2024-10-27T13:14:28.875906Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 66,
     "referenced_widgets": [
      "92afd065b44c4103ae1ab107810c0be3",
      "011102cd7bbb4f7b8dce44e8b4029431",
      "88b25582c551459a804b210b842ce294",
      "34fb85e4fc2944f18605d8d8c2bdf2c7",
      "ff54f7ca4bd14b49897ec982f5e45c00",
      "1e5cd1d27a374d6fb199e3189dafd731",
      "22ca91d86e3b4915be281faca2fe6a2e",
      "1a6b4d686aaa40d4b57339a7ef0730cb",
      "1578c9b9e6db41a08d596194a87dd2bf",
      "01548ad02d0a492baff8e9b85ca4d356",
      "4d642ae9a8f64261befd29cf45db20b1"
     ]
    },
    "execution": {
     "iopub.execute_input": "2024-10-27T18:35:46.680897Z",
     "iopub.status.busy": "2024-10-27T18:35:46.680597Z",
     "iopub.status.idle": "2024-10-27T18:35:49.107590Z",
     "shell.execute_reply": "2024-10-27T18:35:49.106809Z",
     "shell.execute_reply.started": "2024-10-27T18:35:46.680864Z"
    },
    "id": "m2DJP8uUZXa4",
    "outputId": "99ae2f68-d44b-41d9-de4d-77a86feaf4d6",
    "trusted": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d6ed520f1adf4ef8ac4fbf798461b859",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating json from Arrow format:   0%|          | 0/201 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "236118589"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds['train'].to_json(\"dataset_question.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "taqOcB-9ZXa4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-27T13:09:08.985348Z",
     "start_time": "2024-10-27T13:09:08.982278Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2024-10-27T18:35:49.108972Z",
     "iopub.status.busy": "2024-10-27T18:35:49.108673Z",
     "iopub.status.idle": "2024-10-27T18:35:49.113618Z",
     "shell.execute_reply": "2024-10-27T18:35:49.112740Z",
     "shell.execute_reply.started": "2024-10-27T18:35:49.108940Z"
    },
    "id": "taqOcB-9ZXa4",
    "outputId": "158ef52f-3ab8-42a9-e8d5-372f09ad25d3",
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DatasetDict({\n",
      "    train: Dataset({\n",
      "        features: ['question', 'answer'],\n",
      "        num_rows: 200035\n",
      "    })\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "print(ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "pwj9NhzhZXa5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-27T13:13:07.649275Z",
     "start_time": "2024-10-27T13:13:06.582446Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 49,
     "referenced_widgets": [
      "46e1900d9f8842c88ee0c04ed66d624a",
      "b5bb01b4575d438e9c22ec36a8aa7a6d",
      "8e3c704dcfc540f98df5a3e11b98cf4a",
      "25a993e98f6c41e8b154052648898245",
      "4e2f2f8654a9444ea23022333df6edd0",
      "c81057bec3a241c5a615a4ceae93a10c",
      "2fd3ecbcfc44432580c13b1327c80ccd",
      "1cb047915d2b415b94d4f705bbcd40ec",
      "3596b1318099434a9cdd8dd2a13247cc",
      "ccc51efd3907409c987a71735bc87ce7",
      "7a184a73fe454a17a69db6c937d4f1df"
     ]
    },
    "execution": {
     "iopub.execute_input": "2024-10-27T18:35:49.115825Z",
     "iopub.status.busy": "2024-10-27T18:35:49.115540Z",
     "iopub.status.idle": "2024-10-27T18:35:50.149493Z",
     "shell.execute_reply": "2024-10-27T18:35:50.148726Z",
     "shell.execute_reply.started": "2024-10-27T18:35:49.115793Z"
    },
    "id": "pwj9NhzhZXa5",
    "outputId": "556c6597-af7a-42f6-e06a-f26a8c3aa248",
    "trusted": true
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f886c636e44f4a8b9aada55c7f1a18ed",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "ds = load_dataset('json', data_files={'train': 'dataset_question.json'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "98wcCoYpZXa6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-27T13:23:50.253863Z",
     "start_time": "2024-10-27T13:23:47.466257Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2024-10-27T18:35:50.150962Z",
     "iopub.status.busy": "2024-10-27T18:35:50.150658Z",
     "iopub.status.idle": "2024-10-27T18:36:01.326264Z",
     "shell.execute_reply": "2024-10-27T18:36:01.325315Z",
     "shell.execute_reply.started": "2024-10-27T18:35:50.150929Z"
    },
    "id": "98wcCoYpZXa6",
    "outputId": "215acf4f-1749-45b0-9d92-baefa2a1fa87",
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vocabulary size: 21644\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import re\n",
    "\n",
    "def tokenizer(text):\n",
    "    tokens = [i for i in re.split(r'(\\s|[^a-zA-Z])', text) if i and not i.isspace()]\n",
    "    return tokens\n",
    "\n",
    "def build_vocab(dataset):\n",
    "    vocab = {'EOQ'}\n",
    "    count = 0\n",
    "    for string in dataset['train']['question']:\n",
    "        if not string:\n",
    "            continue\n",
    "\n",
    "        tokens = tokenizer(string.lower())\n",
    "        vocab.update(tokens)\n",
    "\n",
    "    vocab = sorted(vocab)\n",
    "    token_to_idx = {token: idx for idx, token in enumerate(vocab)}\n",
    "    idx_to_token = {idx: token for token, idx in token_to_idx.items()}\n",
    "    return token_to_idx, idx_to_token\n",
    "\n",
    "\n",
    "token_to_idx, idx_to_token = build_vocab(ds)\n",
    "vocab_size = len(token_to_idx)\n",
    "print(\"Vocabulary size:\", vocab_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "PYafCA33ZXa6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-27T13:26:44.085956Z",
     "start_time": "2024-10-27T13:26:44.083314Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2024-10-27T18:36:01.327797Z",
     "iopub.status.busy": "2024-10-27T18:36:01.327478Z",
     "iopub.status.idle": "2024-10-27T18:36:01.332137Z",
     "shell.execute_reply": "2024-10-27T18:36:01.331198Z",
     "shell.execute_reply.started": "2024-10-27T18:36:01.327763Z"
    },
    "id": "PYafCA33ZXa6",
    "outputId": "e6a9195d-19d2-44ea-fa8f-0f02f3f8d09f",
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "13659\n"
     ]
    }
   ],
   "source": [
    "print(token_to_idx['jungkook'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "F2gUcamlZXa7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-27T13:31:34.597530Z",
     "start_time": "2024-10-27T13:31:34.591966Z"
    },
    "execution": {
     "iopub.execute_input": "2024-10-27T18:36:01.333722Z",
     "iopub.status.busy": "2024-10-27T18:36:01.333364Z",
     "iopub.status.idle": "2024-10-27T18:36:01.344738Z",
     "shell.execute_reply": "2024-10-27T18:36:01.343948Z",
     "shell.execute_reply.started": "2024-10-27T18:36:01.333678Z"
    },
    "id": "F2gUcamlZXa7",
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class TokenizedDataset(Dataset):\n",
    "    def __init__(self, text_data, token_to_idx, context_size=5):\n",
    "        self.samples = []\n",
    "        self.context_size = context_size\n",
    "        self.token_to_idx = token_to_idx\n",
    "\n",
    "        counter = 0\n",
    "\n",
    "        for text in text_data:\n",
    "            if not text:\n",
    "                continue\n",
    "\n",
    "            counter += 1\n",
    "            if counter % 10000 == 0:\n",
    "                print(f'{counter*100/len(text_data)}% processed')\n",
    "\n",
    "            tokens = tokenizer(text.strip().lower())\n",
    "            indexed_tokens = [self.token_to_idx.get(token, self.token_to_idx['EOQ']) for token in tokens]\n",
    "            indexed_tokens = [self.token_to_idx['EOQ']] * self.context_size + indexed_tokens + [self.token_to_idx['EOQ']]\n",
    "\n",
    "            for i in range(context_size - 1, len(indexed_tokens) - 1):\n",
    "                context = indexed_tokens[i - context_size + 1: i + 1]\n",
    "                target = indexed_tokens[i + 1]\n",
    "                self.samples.append((context, target))\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.samples)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        context, target = self.samples[idx]\n",
    "        return torch.tensor(context), torch.tensor(target)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6k0OxD6MZXa7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-27T18:36:01.346034Z",
     "iopub.status.busy": "2024-10-27T18:36:01.345755Z",
     "iopub.status.idle": "2024-10-27T18:36:01.360656Z",
     "shell.execute_reply": "2024-10-27T18:36:01.359928Z",
     "shell.execute_reply.started": "2024-10-27T18:36:01.346003Z"
    },
    "id": "6k0OxD6MZXa7",
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# HYPER PARAMS\n",
    "context_size = 5\n",
    "emb_dim = 64\n",
    "activation = 'relu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6mUGNPsw5WYK",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "6mUGNPsw5WYK",
    "outputId": "a070839d-655e-4a38-ff79-49d2baafb22e"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "context_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "VuzVmUJBZXa7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-27T13:31:45.902918Z",
     "start_time": "2024-10-27T13:31:34.734110Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2024-10-27T18:36:01.361982Z",
     "iopub.status.busy": "2024-10-27T18:36:01.361684Z",
     "iopub.status.idle": "2024-10-27T18:36:36.385155Z",
     "shell.execute_reply": "2024-10-27T18:36:36.384154Z",
     "shell.execute_reply.started": "2024-10-27T18:36:01.361950Z"
    },
    "id": "VuzVmUJBZXa7",
    "outputId": "d0b6030e-ca2d-4c85-91e1-26c171c794cc",
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4.999125153098208% processed\n",
      "9.998250306196416% processed\n",
      "14.997375459294624% processed\n",
      "19.996500612392833% processed\n",
      "24.995625765491038% processed\n",
      "29.994750918589247% processed\n",
      "34.99387607168745% processed\n",
      "39.993001224785665% processed\n",
      "44.99212637788387% processed\n",
      "49.991251530982076% processed\n",
      "54.99037668408029% processed\n",
      "59.989501837178494% processed\n",
      "64.9886269902767% processed\n",
      "69.9877521433749% processed\n",
      "74.98687729647311% processed\n",
      "79.98600244957133% processed\n",
      "84.98512760266954% processed\n",
      "89.98425275576774% processed\n",
      "94.98337790886595% processed\n",
      "99.98250306196415% processed\n"
     ]
    }
   ],
   "source": [
    "train_texts = ds['train']['question']\n",
    "tokenized_dataset = TokenizedDataset(train_texts, token_to_idx, context_size=context_size)\n",
    "train_loader = DataLoader(tokenized_dataset, batch_size=4096, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "h9ZzkPge50Tr",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "h9ZzkPge50Tr",
    "outputId": "15ae9d73-5df2-4f37-b187-13c7c9cbb7f5"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "11892783"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(tokenized_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "h3MakPbjZXa7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-27T13:32:30.518188Z",
     "start_time": "2024-10-27T13:32:29.791638Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2024-10-27T18:36:36.388507Z",
     "iopub.status.busy": "2024-10-27T18:36:36.388176Z",
     "iopub.status.idle": "2024-10-27T18:36:37.922958Z",
     "shell.execute_reply": "2024-10-27T18:36:37.922030Z",
     "shell.execute_reply.started": "2024-10-27T18:36:36.388473Z"
    },
    "id": "h3MakPbjZXa7",
    "outputId": "da3a5fd7-38a9-4eb8-bf46-962efcd4ab42",
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Context: tensor([[20342, 11552,  3788, 18250,  8917],\n",
      "        [18364,    25, 11105, 16171,  7070],\n",
      "        [ 2175, 11726,  8917,  7070,    25],\n",
      "        ...,\n",
      "        [20342, 15964,  8833, 17158,  8072],\n",
      "        [20342,  9047,  9176,  8183,  3550],\n",
      "        [ 6912,  8318,   513,   427, 16171]])\n",
      "Target: tensor([ 3718,    25, 11105,  ..., 16920, 20337, 16920])\n"
     ]
    }
   ],
   "source": [
    "for context, target in train_loader:\n",
    "    print(\"Context:\", context)\n",
    "    print(\"Target:\", target)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "c2Dzmria6UEG",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "c2Dzmria6UEG",
    "outputId": "4876f1bf-1c67-40c1-cfff-1cc31d7ac0c9"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<__main__.TokenizedDataset at 0x7f5e77028bb0>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_loader.dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "blxJcisI6NoX",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "blxJcisI6NoX",
    "outputId": "4f983401-91f7-4ca1-fe67-ac4cc800c5f0"
   },
   "outputs": [],
   "source": [
    "# import torch\n",
    "# from torch.utils.data import DataLoader, Subset\n",
    "\n",
    "# torch.manual_seed(31)\n",
    "# subset_size = int(0.1 * len(train_loader.dataset))\n",
    "\n",
    "# subset_indices = torch.randperm(len(train_loader.dataset))[:subset_size]\n",
    "\n",
    "# subset = Subset(train_loader.dataset, subset_indices)\n",
    "\n",
    "# subset_loader = DataLoader(subset, batch_size=4096, shuffle=True, num_workers=2, pin_memory=True)\n",
    "\n",
    "# for context, target in subset_loader:\n",
    "#     print(\"Context:\", context)\n",
    "#     print(\"Target:\", target)\n",
    "#     print(target.shape)\n",
    "#     print(context.shape)\n",
    "#     break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "fHehA2ASZXa7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-27T13:32:45.707864Z",
     "start_time": "2024-10-27T13:32:45.703461Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2024-10-27T18:36:37.924856Z",
     "iopub.status.busy": "2024-10-27T18:36:37.924537Z",
     "iopub.status.idle": "2024-10-27T18:36:37.932403Z",
     "shell.execute_reply": "2024-10-27T18:36:37.931477Z",
     "shell.execute_reply.started": "2024-10-27T18:36:37.924821Z"
    },
    "id": "fHehA2ASZXa7",
    "outputId": "d6ec77c6-bc0d-416d-cee8-8b2403f98fc7",
    "trusted": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([4096]), torch.int64, torch.Size([4096, 5]), torch.int64)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "target.shape, target.dtype, context.shape, context.dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "Q_Hejn0kZXa7",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2024-10-27T18:36:37.933992Z",
     "iopub.status.busy": "2024-10-27T18:36:37.933624Z",
     "iopub.status.idle": "2024-10-27T18:36:37.943119Z",
     "shell.execute_reply": "2024-10-27T18:36:37.942341Z",
     "shell.execute_reply.started": "2024-10-27T18:36:37.933949Z"
    },
    "id": "Q_Hejn0kZXa7",
    "outputId": "a1ba3f98-028c-412a-b6ec-a4ffb323a72b",
    "trusted": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "06aIP6jAZXa7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-27T18:36:37.944697Z",
     "iopub.status.busy": "2024-10-27T18:36:37.944340Z",
     "iopub.status.idle": "2024-10-27T18:36:37.953566Z",
     "shell.execute_reply": "2024-10-27T18:36:37.952778Z",
     "shell.execute_reply.started": "2024-10-27T18:36:37.944655Z"
    },
    "id": "06aIP6jAZXa7",
    "trusted": true
   },
   "outputs": [],
   "source": [
    "def get_default_device():\n",
    "    \"\"\"Pick GPU if available, else CPU\"\"\"\n",
    "    if torch.cuda.is_available():\n",
    "        return torch.device('cuda')\n",
    "    else:\n",
    "        return torch.device('cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "L4JBMx61ZXa7",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2024-10-27T18:36:37.954868Z",
     "iopub.status.busy": "2024-10-27T18:36:37.954583Z",
     "iopub.status.idle": "2024-10-27T18:36:37.966470Z",
     "shell.execute_reply": "2024-10-27T18:36:37.965648Z",
     "shell.execute_reply.started": "2024-10-27T18:36:37.954838Z"
    },
    "id": "L4JBMx61ZXa7",
    "outputId": "98d5d60d-a9e9-45d0-8b51-502e429cdb9e",
    "trusted": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda')"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = get_default_device()\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "k3LX3ltCZXa8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-27T18:36:37.968024Z",
     "iopub.status.busy": "2024-10-27T18:36:37.967764Z",
     "iopub.status.idle": "2024-10-27T18:36:37.976413Z",
     "shell.execute_reply": "2024-10-27T18:36:37.975569Z",
     "shell.execute_reply.started": "2024-10-27T18:36:37.967995Z"
    },
    "id": "k3LX3ltCZXa8",
    "trusted": true
   },
   "outputs": [],
   "source": [
    "def to_device(data, device):\n",
    "    \"\"\"Move tensor(s) to chosen device\"\"\"\n",
    "    if isinstance(data, (list,tuple)):\n",
    "        return [to_device(x, device) for x in data]\n",
    "    return data.to(device, non_blocking=True)\n",
    "\n",
    "class DeviceDataLoader():\n",
    "    \"\"\"Wrap a dataloader to move data to a device\"\"\"\n",
    "    def __init__(self, dl, device):\n",
    "        self.dl = dl\n",
    "        self.device = device\n",
    "\n",
    "    def __iter__(self):\n",
    "        \"\"\"Yield a batch of data after moving it to device\"\"\"\n",
    "        for b in self.dl:\n",
    "            yield to_device(b, self.device)\n",
    "\n",
    "    def __len__(self):\n",
    "        \"\"\"Number of batches\"\"\"\n",
    "        return len(self.dl)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "K9Ld7dp0ZXa8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-27T18:36:37.977643Z",
     "iopub.status.busy": "2024-10-27T18:36:37.977358Z",
     "iopub.status.idle": "2024-10-27T18:36:37.989500Z",
     "shell.execute_reply": "2024-10-27T18:36:37.988638Z",
     "shell.execute_reply.started": "2024-10-27T18:36:37.977612Z"
    },
    "id": "K9Ld7dp0ZXa8",
    "trusted": true
   },
   "outputs": [],
   "source": [
    "train_loader = DeviceDataLoader(train_loader, device)\n",
    "# subset_loader = DeviceDataLoader(subset_loader, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "sItBcXbeZXa8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-27T13:33:17.420977Z",
     "start_time": "2024-10-27T13:33:17.416310Z"
    },
    "execution": {
     "iopub.execute_input": "2024-10-27T18:36:37.990899Z",
     "iopub.status.busy": "2024-10-27T18:36:37.990591Z",
     "iopub.status.idle": "2024-10-27T18:36:38.001248Z",
     "shell.execute_reply": "2024-10-27T18:36:38.000489Z",
     "shell.execute_reply.started": "2024-10-27T18:36:37.990868Z"
    },
    "id": "sItBcXbeZXa8",
    "trusted": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "vocab_size = len(token_to_idx)\n",
    "num_epochs = 15\n",
    "learning_rate = 0.001\n",
    "\n",
    "class NextWord(nn.Module):\n",
    "    def __init__(self, context_size, vocab_size=len(token_to_idx), emb_dim=64, hidden_size=256, activation='sine'):\n",
    "        super().__init__()\n",
    "        self.emb = nn.Embedding(vocab_size, emb_dim)\n",
    "        self.lin1 = nn.Linear(context_size * emb_dim, hidden_size)\n",
    "        self.lin2 = nn.Linear(hidden_size, hidden_size)\n",
    "        self.lin3 = nn.Linear(hidden_size, hidden_size)\n",
    "        self.lin4 = nn.Linear(hidden_size, vocab_size)\n",
    "        self.activation = activation\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.emb(x)\n",
    "        x = x.view(x.shape[0], -1)\n",
    "        if self.activation == 'sine':\n",
    "            x = torch.sin(self.lin1(x))\n",
    "            x = torch.sin(self.lin2(x))\n",
    "            x = torch.sin(self.lin3(x))\n",
    "        elif self.activation == 'relu':\n",
    "            x = nn.functional.relu(self.lin1(x))\n",
    "            x = nn.functional.relu(self.lin2(x))\n",
    "            x = nn.functional.relu(self.lin3(x))\n",
    "        x = self.lin4(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "jwgWleT4AtkB",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "jwgWleT4AtkB",
    "outputId": "30ca5d53-6c32-4bd1-e197-9fa0c83f94cf"
   },
   "outputs": [],
   "source": [
    "# !pip install triton --upgrade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "-jUniRPZZXa8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-27T13:33:21.309502Z",
     "start_time": "2024-10-27T13:33:20.258568Z"
    },
    "execution": {
     "iopub.execute_input": "2024-10-27T18:36:38.002603Z",
     "iopub.status.busy": "2024-10-27T18:36:38.002277Z",
     "iopub.status.idle": "2024-10-27T18:36:38.133657Z",
     "shell.execute_reply": "2024-10-27T18:36:38.132838Z",
     "shell.execute_reply.started": "2024-10-27T18:36:38.002565Z"
    },
    "id": "-jUniRPZZXa8",
    "trusted": true
   },
   "outputs": [],
   "source": [
    "model = NextWord(context_size, vocab_size, emb_dim, hidden_size=512, activation=activation)\n",
    "to_device(model, device)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "k-UG0tev-fAX",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "k-UG0tev-fAX",
    "outputId": "5f745644-2a4b-4552-cc90-1b72053e4273"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "emb.weight torch.Size([21644, 64])\n",
      "lin1.weight torch.Size([512, 320])\n",
      "lin1.bias torch.Size([512])\n",
      "lin2.weight torch.Size([512, 512])\n",
      "lin2.bias torch.Size([512])\n",
      "lin3.weight torch.Size([512, 512])\n",
      "lin3.bias torch.Size([512])\n",
      "lin4.weight torch.Size([21644, 512])\n",
      "lin4.bias torch.Size([21644])\n"
     ]
    }
   ],
   "source": [
    "for param_name, param in model.named_parameters():\n",
    "    print(param_name, param.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "-eGh5JSZZXa8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-27T18:36:38.135136Z",
     "iopub.status.busy": "2024-10-27T18:36:38.134828Z",
     "iopub.status.idle": "2024-10-27T18:36:38.143298Z",
     "shell.execute_reply": "2024-10-27T18:36:38.142359Z",
     "shell.execute_reply.started": "2024-10-27T18:36:38.135103Z"
    },
    "id": "-eGh5JSZZXa8",
    "trusted": true
   },
   "outputs": [],
   "source": [
    "def reconstruct_text(tokens):\n",
    "    result = []\n",
    "    print(tokens)\n",
    "\n",
    "    for i, token in enumerate(tokens):\n",
    "        result.append(token)\n",
    "\n",
    "        if i < len(tokens) - 1:\n",
    "            if token.isalpha() and tokens[i + 1].isalpha():\n",
    "                result.append(\" \")\n",
    "            elif token.isalpha() and tokens[i + 1].isdigit():\n",
    "                result.append(\" \")\n",
    "            elif token.isdigit() and tokens[i + 1].isalpha():\n",
    "                result.append(\" \")\n",
    "            elif token in [',','.','!','?',':','%','&'] and tokens[i + 1].isalnum():\n",
    "                result.append(\" \")\n",
    "\n",
    "    return ''.join(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "G33BD7UxZXa8",
   "metadata": {
    "id": "G33BD7UxZXa8"
   },
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "4JEkEgCJKycL",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 445
    },
    "id": "4JEkEgCJKycL",
    "outputId": "b69a1409-cc15-42c9-dece-720134bce18e"
   },
   "outputs": [],
   "source": [
    "# checkpoint = torch.load('/content/next_word_model_emb64_cont10_relu_layernorm_batchnorm_hidden512.pth', weights_only=True)\n",
    "# model.load_state_dict(checkpoint['model_state_dict'], strict=False)\n",
    "# optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "# epoch = checkpoint['epoch']\n",
    "# loss = checkpoint['loss']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf0b8c1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.manifold import TSNE\n",
    "\n",
    "\n",
    "indices = torch.arange(0, 100)\n",
    "embeddings = model.emb.weight.data.cpu().numpy()\n",
    "idx_to_token = {idx: token for token, idx in token_to_idx.items()}\n",
    "\n",
    "tsne = TSNE(n_components=2, random_state=42)\n",
    "embeddings_2d = tsne.fit_transform(embeddings)\n",
    "\n",
    "for i in range(embeddings_2d.shape[0]):\n",
    "    plt.annotate(idx_to_token[i], (embeddings_2d[i, 0], embeddings_2d[i, 1]), fontsize=8, alpha=0.7)\n",
    "\n",
    "\n",
    "plt.figure(figsize=(10, 10))\n",
    "plt.scatter(embeddings_2d[:, 0], embeddings_2d[:, 1], alpha=0.7)\n",
    "plt.title(\"t-SNE Visualization of Embeddings\")\n",
    "plt.xlabel(\"t-SNE Component 1\")\n",
    "plt.ylabel(\"t-SNE Component 2\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "TWfKZKSKLLGg",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "TWfKZKSKLLGg",
    "outputId": "235a38dc-dfc8-4537-b248-84ee016a2f1a"
   },
   "outputs": [],
   "source": [
    "# def compare_weights(model1, model2):\n",
    "#     for (name1, param1), (name2, param2) in zip(model1.items(), model2.items()):\n",
    "#         if name1 != name2:\n",
    "#             print(f\"Mismatch in parameter names: {name1} vs {name2}\")\n",
    "#             return False\n",
    "\n",
    "#         # Use torch.allclose to check if tensors are nearly identical\n",
    "#         if not torch.allclose(param1, param2, atol=1e-6):\n",
    "#             print(f\"Mismatch found in layer: {name1}\")\n",
    "#             return False\n",
    "\n",
    "#     print(\"All weights match!\")\n",
    "#     return True\n",
    "\n",
    "# compare_weights(model.state_dict(), old_model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "FHet8bztZXa9",
   "metadata": {
    "ExecuteTime": {
     "start_time": "2024-10-27T13:33:39.430419Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2024-10-27T18:36:38.145226Z",
     "iopub.status.busy": "2024-10-27T18:36:38.144591Z"
    },
    "id": "FHet8bztZXa9",
    "jupyter": {
     "is_executing": true
    },
    "outputId": "1dac0a51-030d-4a8c-8cf1-0d0e83385e92",
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/15\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Batch 0, Loss: 9.9844\n",
      "Batch 100, Loss: 5.1179\n",
      "Batch 200, Loss: 4.6700\n",
      "Batch 300, Loss: 4.3862\n",
      "Batch 400, Loss: 4.2689\n",
      "Batch 500, Loss: 4.1547\n",
      "Batch 600, Loss: 4.0234\n",
      "Batch 700, Loss: 3.9107\n",
      "Batch 800, Loss: 3.8488\n",
      "Batch 900, Loss: 3.7860\n",
      "Batch 1000, Loss: 3.7830\n",
      "Batch 1100, Loss: 3.6536\n",
      "Batch 1200, Loss: 3.6501\n",
      "Batch 1300, Loss: 3.6166\n",
      "Batch 1400, Loss: 3.5084\n",
      "Batch 1500, Loss: 3.5253\n",
      "Batch 1600, Loss: 3.5141\n",
      "Batch 1700, Loss: 3.4825\n",
      "Batch 1800, Loss: 3.3861\n",
      "Batch 1900, Loss: 3.3909\n",
      "Batch 2000, Loss: 3.3072\n",
      "Batch 2100, Loss: 3.2982\n",
      "Batch 2200, Loss: 3.3028\n",
      "Batch 2300, Loss: 3.3392\n",
      "Batch 2400, Loss: 3.2424\n",
      "Batch 2500, Loss: 3.2639\n",
      "Batch 2600, Loss: 3.2054\n",
      "Batch 2700, Loss: 3.2320\n",
      "Batch 2800, Loss: 3.2138\n",
      "Batch 2900, Loss: 3.1308\n",
      "Epoch 1 Average Loss: 3.7108\n",
      "Epoch 2/15\n",
      "Batch 0, Loss: 3.0589\n",
      "Batch 100, Loss: 3.0550\n",
      "Batch 200, Loss: 3.1506\n",
      "Batch 300, Loss: 2.9994\n",
      "Batch 400, Loss: 3.0450\n",
      "Batch 500, Loss: 3.0101\n",
      "Batch 600, Loss: 2.9976\n",
      "Batch 700, Loss: 3.0111\n",
      "Batch 800, Loss: 2.9537\n",
      "Batch 900, Loss: 2.9648\n",
      "Batch 1000, Loss: 3.0131\n",
      "Batch 1100, Loss: 2.9586\n",
      "Batch 1200, Loss: 3.0187\n",
      "Batch 1300, Loss: 2.8537\n",
      "Batch 1400, Loss: 2.8709\n",
      "Batch 1500, Loss: 2.9149\n",
      "Batch 1600, Loss: 2.9291\n",
      "Batch 1700, Loss: 2.8597\n",
      "Batch 1800, Loss: 2.9303\n",
      "Batch 1900, Loss: 2.9248\n",
      "Batch 2000, Loss: 2.8778\n",
      "Batch 2100, Loss: 2.8863\n",
      "Batch 2200, Loss: 2.9079\n",
      "Batch 2300, Loss: 2.8645\n",
      "Batch 2400, Loss: 2.8469\n",
      "Batch 2500, Loss: 2.7751\n",
      "Batch 2600, Loss: 2.8386\n",
      "Batch 2700, Loss: 2.7560\n",
      "Batch 2800, Loss: 2.8403\n",
      "Batch 2900, Loss: 2.7724\n",
      "Epoch 2 Average Loss: 2.9264\n",
      "Epoch 3/15\n",
      "Batch 0, Loss: 2.7340\n",
      "Batch 100, Loss: 2.7509\n",
      "Batch 200, Loss: 2.6976\n",
      "Batch 300, Loss: 2.7160\n",
      "Batch 400, Loss: 2.7210\n",
      "Batch 500, Loss: 2.7275\n",
      "Batch 600, Loss: 2.7229\n",
      "Batch 700, Loss: 2.7243\n",
      "Batch 800, Loss: 2.6643\n",
      "Batch 900, Loss: 2.7407\n",
      "Batch 1000, Loss: 2.6895\n",
      "Batch 1100, Loss: 2.6639\n",
      "Batch 1200, Loss: 2.7172\n",
      "Batch 1300, Loss: 2.6480\n",
      "Batch 1400, Loss: 2.6545\n",
      "Batch 1500, Loss: 2.7047\n",
      "Batch 1600, Loss: 2.7874\n",
      "Batch 1700, Loss: 2.6602\n",
      "Batch 1800, Loss: 2.6290\n",
      "Batch 1900, Loss: 2.7108\n",
      "Batch 2000, Loss: 2.7541\n",
      "Batch 2100, Loss: 2.6890\n",
      "Batch 2200, Loss: 2.6927\n",
      "Batch 2300, Loss: 2.6984\n",
      "Batch 2400, Loss: 2.6258\n",
      "Batch 2500, Loss: 2.6995\n",
      "Batch 2600, Loss: 2.6531\n",
      "Batch 2700, Loss: 2.6486\n",
      "Batch 2800, Loss: 2.6573\n",
      "Batch 2900, Loss: 2.6618\n",
      "Epoch 3 Average Loss: 2.6885\n",
      "Epoch 4/15\n",
      "Batch 0, Loss: 2.5181\n",
      "Batch 100, Loss: 2.5705\n",
      "Batch 200, Loss: 2.5833\n",
      "Batch 300, Loss: 2.6347\n",
      "Batch 400, Loss: 2.5948\n",
      "Batch 500, Loss: 2.5326\n",
      "Batch 600, Loss: 2.4972\n",
      "Batch 700, Loss: 2.5977\n",
      "Batch 800, Loss: 2.5384\n",
      "Batch 900, Loss: 2.5572\n",
      "Batch 1000, Loss: 2.4322\n",
      "Batch 1100, Loss: 2.5631\n",
      "Batch 1200, Loss: 2.5342\n",
      "Batch 1300, Loss: 2.5583\n",
      "Batch 1400, Loss: 2.6479\n",
      "Batch 1500, Loss: 2.5607\n",
      "Batch 1600, Loss: 2.5101\n",
      "Batch 1700, Loss: 2.4903\n",
      "Batch 1800, Loss: 2.5086\n",
      "Batch 1900, Loss: 2.5389\n",
      "Batch 2000, Loss: 2.5236\n",
      "Batch 2100, Loss: 2.5126\n",
      "Batch 2200, Loss: 2.6072\n",
      "Batch 2300, Loss: 2.4998\n",
      "Batch 2400, Loss: 2.5984\n",
      "Batch 2500, Loss: 2.6158\n",
      "Batch 2600, Loss: 2.6236\n",
      "Batch 2700, Loss: 2.6012\n",
      "Batch 2800, Loss: 2.5416\n",
      "Batch 2900, Loss: 2.5453\n",
      "Epoch 4 Average Loss: 2.5594\n",
      "Epoch 5/15\n",
      "Batch 0, Loss: 2.3855\n",
      "Batch 100, Loss: 2.4728\n",
      "Batch 200, Loss: 2.4344\n",
      "Batch 300, Loss: 2.4591\n",
      "Batch 400, Loss: 2.3789\n",
      "Batch 500, Loss: 2.4789\n",
      "Batch 600, Loss: 2.4581\n",
      "Batch 700, Loss: 2.4949\n",
      "Batch 800, Loss: 2.4731\n",
      "Batch 900, Loss: 2.4717\n",
      "Batch 1000, Loss: 2.4428\n",
      "Batch 1100, Loss: 2.4726\n",
      "Batch 1200, Loss: 2.3999\n",
      "Batch 1300, Loss: 2.4776\n",
      "Batch 1400, Loss: 2.4264\n",
      "Batch 1500, Loss: 2.4941\n",
      "Batch 1600, Loss: 2.4331\n",
      "Batch 1700, Loss: 2.5092\n",
      "Batch 1800, Loss: 2.5005\n",
      "Batch 1900, Loss: 2.4831\n",
      "Batch 2000, Loss: 2.4958\n",
      "Batch 2100, Loss: 2.4527\n"
     ]
    }
   ],
   "source": [
    "model.train()\n",
    "for epoch in range(num_epochs):\n",
    "    total_loss = 0\n",
    "    print(f\"Epoch {epoch + 1}/{num_epochs}\")\n",
    "    for i, (context, target) in enumerate(train_loader):\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(context)\n",
    "        loss = criterion(outputs, target)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss += loss.item()\n",
    "\n",
    "        if i % 100 == 0:\n",
    "            print(f\"Batch {i}, Loss: {loss.item():.4f}\")\n",
    "\n",
    "    avg_loss = total_loss / len(train_loader)\n",
    "    print(f\"Epoch {epoch + 1} Average Loss: {avg_loss:.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "AwEHsd2TJYSl",
   "metadata": {
    "id": "AwEHsd2TJYSl"
   },
   "outputs": [],
   "source": [
    "# old_model = model.state_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "Vnpp8bwhKZ1d",
   "metadata": {
    "id": "Vnpp8bwhKZ1d"
   },
   "outputs": [],
   "source": [
    "torch.save({\n",
    "            'epoch': epoch,\n",
    "            'model_state_dict': model.state_dict(),\n",
    "            'optimizer_state_dict': optimizer.state_dict(),\n",
    "            'loss': loss,\n",
    "            }, 'next_word_model_emb64_cont5_relu_hidden512.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "mKX9HCsqZXa-",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "mKX9HCsqZXa-",
    "outputId": "e09774fe-c98c-45c7-83b0-2f6ccb5b0a47",
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['ramesh', 'has', '5', 'pencils', 'and', '1', '0', 'pens', 'with', 'him', '.', 'the', 'ratio', 'of', 'marbles', 'to', '.', 'mark', 'now', 'has', 'a', '4', '0', '%', 'chance', 'that', 'she', 'gets', 'jamie', 'and', 'ends', 'up', 'her', 'three', 'thirds', 'of', 'the', 'raspberries', 'as', 'many', 'blueberries', 'as', 'oranges', '.', 'she', 'has', '4', 'pounds', 'of', 'oranges', '.', 'how', 'many', 'peanuts', 'does', 'kimberly', 'have', 'initially', 'in', 'the']\n",
      "ramesh has 5 pencils and 10 pens with him. the ratio of marbles to. mark now has a 40% chance that she gets jamie and ends up her three thirds of the raspberries as many blueberries as oranges. she has 4 pounds of oranges. how many peanuts does kimberly have initially in the\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "def generate_text(model, start_text, token_to_idx, idx_to_token, context_size=5, max_length=50):\n",
    "    model.eval()\n",
    "    tokens = tokenizer(start_text.lower())\n",
    "    input_indices = [token_to_idx.get(token, token_to_idx['EOQ']) for token in tokens]\n",
    "\n",
    "    if len(input_indices) < context_size:\n",
    "        input_indices = [token_to_idx['EOQ']] * (context_size - len(input_indices)) + input_indices\n",
    "\n",
    "    context = input_indices[-context_size:]\n",
    "    generated_tokens = tokens\n",
    "\n",
    "\n",
    "    with torch.no_grad():\n",
    "\n",
    "        for _ in range(max_length):\n",
    "            context_tensor = torch.tensor(context, device=device).unsqueeze(0)\n",
    "\n",
    "            output_logits = model(context_tensor)\n",
    "            predicted_token_idx = torch.distributions.categorical.Categorical(logits=output_logits).sample().item()\n",
    "\n",
    "            predicted_token = idx_to_token.get(predicted_token_idx, '<unk>')\n",
    "            generated_tokens.append(predicted_token)\n",
    "\n",
    "            context = context[1:] + [predicted_token_idx]\n",
    "\n",
    "            if predicted_token == 'EOQ':\n",
    "                break\n",
    "    return reconstruct_text(generated_tokens)\n",
    "#     return ' '.join(generated_tokens)\n",
    "\n",
    "start_text = \"Ramesh has 5 pencils and 10 pens with him\"\n",
    "generated_text = generate_text(model, start_text, token_to_idx, idx_to_token, context_size=5, max_length=50)\n",
    "print(generated_text)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "oEPWAOGYb5aY",
   "metadata": {
    "id": "oEPWAOGYb5aY"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kaggle": {
   "accelerator": "gpu",
   "dataSources": [
    {
     "datasetId": 5960275,
     "sourceId": 9738015,
     "sourceType": "datasetVersion"
    },
    {
     "isSourceIdPinned": true,
     "modelId": 148868,
     "modelInstanceId": 125886,
     "sourceId": 148352,
     "sourceType": "modelInstanceVersion"
    }
   ],
   "dockerImageVersionId": 30787,
   "isGpuEnabled": true,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "011102cd7bbb4f7b8dce44e8b4029431": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_1e5cd1d27a374d6fb199e3189dafd731",
      "placeholder": "​",
      "style": "IPY_MODEL_22ca91d86e3b4915be281faca2fe6a2e",
      "value": "Creating json from Arrow format: 100%"
     }
    },
    "01548ad02d0a492baff8e9b85ca4d356": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "1578c9b9e6db41a08d596194a87dd2bf": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "1a6b4d686aaa40d4b57339a7ef0730cb": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "1cb047915d2b415b94d4f705bbcd40ec": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": "20px"
     }
    },
    "1e5cd1d27a374d6fb199e3189dafd731": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "22ca91d86e3b4915be281faca2fe6a2e": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "25a993e98f6c41e8b154052648898245": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_ccc51efd3907409c987a71735bc87ce7",
      "placeholder": "​",
      "style": "IPY_MODEL_7a184a73fe454a17a69db6c937d4f1df",
      "value": " 200035/0 [00:03&lt;00:00, 64064.45 examples/s]"
     }
    },
    "2fd3ecbcfc44432580c13b1327c80ccd": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "34fb85e4fc2944f18605d8d8c2bdf2c7": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_01548ad02d0a492baff8e9b85ca4d356",
      "placeholder": "​",
      "style": "IPY_MODEL_4d642ae9a8f64261befd29cf45db20b1",
      "value": " 201/201 [00:04&lt;00:00, 23.96ba/s]"
     }
    },
    "3596b1318099434a9cdd8dd2a13247cc": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "46e1900d9f8842c88ee0c04ed66d624a": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_b5bb01b4575d438e9c22ec36a8aa7a6d",
       "IPY_MODEL_8e3c704dcfc540f98df5a3e11b98cf4a",
       "IPY_MODEL_25a993e98f6c41e8b154052648898245"
      ],
      "layout": "IPY_MODEL_4e2f2f8654a9444ea23022333df6edd0"
     }
    },
    "4d642ae9a8f64261befd29cf45db20b1": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "4e2f2f8654a9444ea23022333df6edd0": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "7a184a73fe454a17a69db6c937d4f1df": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "88b25582c551459a804b210b842ce294": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_1a6b4d686aaa40d4b57339a7ef0730cb",
      "max": 201,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_1578c9b9e6db41a08d596194a87dd2bf",
      "value": 201
     }
    },
    "8e3c704dcfc540f98df5a3e11b98cf4a": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_1cb047915d2b415b94d4f705bbcd40ec",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_3596b1318099434a9cdd8dd2a13247cc",
      "value": 1
     }
    },
    "92afd065b44c4103ae1ab107810c0be3": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_011102cd7bbb4f7b8dce44e8b4029431",
       "IPY_MODEL_88b25582c551459a804b210b842ce294",
       "IPY_MODEL_34fb85e4fc2944f18605d8d8c2bdf2c7"
      ],
      "layout": "IPY_MODEL_ff54f7ca4bd14b49897ec982f5e45c00"
     }
    },
    "b5bb01b4575d438e9c22ec36a8aa7a6d": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_c81057bec3a241c5a615a4ceae93a10c",
      "placeholder": "​",
      "style": "IPY_MODEL_2fd3ecbcfc44432580c13b1327c80ccd",
      "value": "Generating train split: "
     }
    },
    "c81057bec3a241c5a615a4ceae93a10c": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "ccc51efd3907409c987a71735bc87ce7": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "ff54f7ca4bd14b49897ec982f5e45c00": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}