notebooks/8_latex_gpt2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e044664f",
   "metadata": {},
   "source": [
    "## GPT for code dictation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "65cd4e7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, TextDataset, DataCollatorForLanguageModeling, AutoModelForCausalLM, pipeline, \\\n",
    "                         Trainer, TrainingArguments\n",
    "import pandas as pd\n",
    "from datasets import Dataset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "09b74db5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "MODEL = 'gpt2'\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL)  # load up a standard gpt2 model\n",
    "\n",
    "tokenizer.pad_token = tokenizer.eos_token  # set the pad token to avoid a warning\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cc4a8564",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(50, 2)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>English</th>\n",
       "      <th>LaTeX</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>integral from a to b of x squared</td>\n",
       "      <td>\\int_{a}^{b} x^2 \\,dx</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>integral from negative 1 to 1 of x squared</td>\n",
       "      <td>\\int_{-1}^{1} x^2 \\,dx</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                      English                   LaTeX\n",
       "0           integral from a to b of x squared   \\int_{a}^{b} x^2 \\,dx\n",
       "1  integral from negative 1 to 1 of x squared  \\int_{-1}^{1} x^2 \\,dx"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data = pd.read_csv('../data/english_to_latex.csv')\n",
    "\n",
    "print(data.shape)\n",
    "\n",
    "data.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8a166ea6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>English</th>\n",
       "      <th>LaTeX</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>integral from a to b of x squared</td>\n",
       "      <td>\\int_{a}^{b} x^2 \\,dx</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>integral from negative 1 to 1 of x squared</td>\n",
       "      <td>\\int_{-1}^{1} x^2 \\,dx</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>integral from negative 1 to infinity of x cubed</td>\n",
       "      <td>\\int_{-1}^{\\inf} x^3 \\,dx</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>integral from 0 to infinity of x squared</td>\n",
       "      <td>\\int_{0}^{\\inf} x^2 \\,dx</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>integral from 0 to infinity of y squared</td>\n",
       "      <td>\\int_{0}^{\\inf} y^2 \\,dy</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>integral from 1 to 2 of x over 2</td>\n",
       "      <td>\\int_{1}^{2} \\frac{x}{2} \\,dx</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>f of x equals x squared</td>\n",
       "      <td>f(x) = x^2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>h of x equals x squared</td>\n",
       "      <td>h(x) = x^2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>g of x equals x squared</td>\n",
       "      <td>g(x) = x^2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>g of x equals x to the eighth power</td>\n",
       "      <td>g(x) = x^8</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                           English  \\\n",
       "0                integral from a to b of x squared   \n",
       "1       integral from negative 1 to 1 of x squared   \n",
       "2  integral from negative 1 to infinity of x cubed   \n",
       "3         integral from 0 to infinity of x squared   \n",
       "4         integral from 0 to infinity of y squared   \n",
       "5                 integral from 1 to 2 of x over 2   \n",
       "6                          f of x equals x squared   \n",
       "7                          h of x equals x squared   \n",
       "8                          g of x equals x squared   \n",
       "9              g of x equals x to the eighth power   \n",
       "\n",
       "                           LaTeX  \n",
       "0          \\int_{a}^{b} x^2 \\,dx  \n",
       "1         \\int_{-1}^{1} x^2 \\,dx  \n",
       "2      \\int_{-1}^{\\inf} x^3 \\,dx  \n",
       "3       \\int_{0}^{\\inf} x^2 \\,dx  \n",
       "4       \\int_{0}^{\\inf} y^2 \\,dy  \n",
       "5  \\int_{1}^{2} \\frac{x}{2} \\,dx  \n",
       "6                     f(x) = x^2  \n",
       "7                     h(x) = x^2  \n",
       "8                     g(x) = x^2  \n",
       "9                     g(x) = x^8  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c138752d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "db5ebf1e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Convert English to LaTeX\n",
      "English: integral from a to b of x squared\n",
      "LaTeX: \\int_{a}^{b} x^2 \\,dx\n"
     ]
    }
   ],
   "source": [
    "# Add our singular prompt\n",
    "CONVERSION_PROMPT = 'Convert English to LaTeX\\n'  # LaTeX conversion task\n",
    "\n",
    "CONVERSION_TOKEN = 'LaTeX:'\n",
    "\n",
    "\n",
    "# This is our \"training prompt\" that we want GPT2 to recognize and learn\n",
    "training_examples = f'{CONVERSION_PROMPT}English: ' + data['English'] + '\\n' + CONVERSION_TOKEN + ' ' + data['LaTeX'].astype(str)\n",
    "\n",
    "print(training_examples[0])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "beb1b918",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Convert English to LaTeX\\nEnglish: integral fr...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Convert English to LaTeX\\nEnglish: integral fr...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                text\n",
       "0  Convert English to LaTeX\\nEnglish: integral fr...\n",
       "1  Convert English to LaTeX\\nEnglish: integral fr..."
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "task_df = pd.DataFrame({'text': training_examples})\n",
    "\n",
    "task_df.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c98cd94a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b7617abd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# adding the EOS token at the end so the model knows when to stop predicting\n",
    "\n",
    "task_df['text'] = task_df['text'].map(lambda x: f'{x}{tokenizer.eos_token}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76552d13",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "074754cb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/50 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "latex_data = Dataset.from_pandas(task_df)  # turn a pandas DataFrame into a Dataset\n",
    "\n",
    "def preprocess(examples):  \n",
    "    # tokenize our text but don't pad because our collator will pad for us dynamically\n",
    "    return tokenizer(examples['text'], truncation=True)\n",
    "\n",
    "latex_data = latex_data.map(preprocess, batched=True)\n",
    "\n",
    "latex_data = latex_data.train_test_split(train_size=.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "16bb726f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'text': 'Convert English to LaTeX\\nEnglish: 2 pi r\\nLaTeX: 2 * \\\\pi * r<|endoftext|>',\n",
       " 'input_ids': [3103,\n",
       "  1851,\n",
       "  3594,\n",
       "  284,\n",
       "  4689,\n",
       "  49568,\n",
       "  198,\n",
       "  15823,\n",
       "  25,\n",
       "  362,\n",
       "  31028,\n",
       "  374,\n",
       "  198,\n",
       "  14772,\n",
       "  49568,\n",
       "  25,\n",
       "  362,\n",
       "  1635,\n",
       "  3467,\n",
       "  14415,\n",
       "  1635,\n",
       "  374,\n",
       "  50256],\n",
       " 'attention_mask': [1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1,\n",
       "  1]}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "latex_data['train'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4f1f30f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# standard data collator for auto-regressive language modelling\n",
    "data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e44dd292",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "latex_gpt2 = AutoModelForCausalLM.from_pretrained(MODEL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0712e12d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['text', 'input_ids', 'attention_mask'],\n",
       "        num_rows: 40\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['text', 'input_ids', 'attention_mask'],\n",
       "        num_rows: 10\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "latex_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fd824d6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "80924236",
   "metadata": {},
   "source": [
    "# Attempt 1 at fine-tuning GPT2 at a LaTeX conversion task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "7d7582bd",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The following columns in the evaluation set don't have a corresponding argument in `GPT2LMHeadModel.forward` and have been ignored: text. If text are not expected by `GPT2LMHeadModel.forward`,  you can safely ignore this message.\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 10\n",
      "  Batch size = 20\n",
      "You're using a GPT2TokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='2' max='1' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [1/1 00:19]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Automatic Weights & Biases logging enabled, to disable set os.environ[\"WANDB_DISABLED\"] = \"true\"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mprofoz\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "wandb version 0.15.8 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.14.0"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/Users/sinanozdemir/Teaching/Pearson/quick-start-llms/notebooks/wandb/run-20230803_105136-x9f2la83</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/profoz/huggingface/runs/x9f2la83' target=\"_blank\">confused-smoke-173</a></strong> to <a href='https://wandb.ai/profoz/huggingface' 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/profoz/huggingface' target=\"_blank\">https://wandb.ai/profoz/huggingface</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/profoz/huggingface/runs/x9f2la83' target=\"_blank\">https://wandb.ai/profoz/huggingface/runs/x9f2la83</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "{'eval_loss': 2.7215285301208496,\n",
       " 'eval_runtime': 0.5484,\n",
       " 'eval_samples_per_second': 18.235,\n",
       " 'eval_steps_per_second': 1.823}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./english_to_latex\",\n",
    "    overwrite_output_dir=True, # overwrite the content of the output directory\n",
    "    num_train_epochs=5, # number of training epochs\n",
    "    per_device_train_batch_size=1, # batch size for training\n",
    "    per_device_eval_batch_size=20,  # batch size for evaluation\n",
    "    load_best_model_at_end=True,\n",
    "    logging_steps=5,\n",
    "    log_level='info',\n",
    "    evaluation_strategy='epoch',\n",
    "    save_strategy='epoch',\n",
    "    use_mps_device=True\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=latex_gpt2,\n",
    "    args=training_args,\n",
    "    train_dataset=latex_data[\"train\"],\n",
    "    eval_dataset=latex_data[\"test\"],\n",
    "    data_collator=data_collator,\n",
    ")\n",
    "\n",
    "trainer.evaluate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b78c75d0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The following columns in the training set don't have a corresponding argument in `GPT2LMHeadModel.forward` and have been ignored: text. If text are not expected by `GPT2LMHeadModel.forward`,  you can safely ignore this message.\n",
      "/opt/homebrew/lib/python3.10/site-packages/transformers/optimization.py:411: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n",
      "***** Running training *****\n",
      "  Num examples = 40\n",
      "  Num Epochs = 5\n",
      "  Instantaneous batch size per device = 1\n",
      "  Total train batch size (w. parallel, distributed & accumulation) = 1\n",
      "  Gradient Accumulation steps = 1\n",
      "  Total optimization steps = 200\n",
      "  Number of trainable parameters = 124,439,808\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='200' max='200' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [200/200 01:24, Epoch 5/5]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.886500</td>\n",
       "      <td>0.568916</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>0.795700</td>\n",
       "      <td>0.498212</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>0.757100</td>\n",
       "      <td>0.415177</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>0.493900</td>\n",
       "      <td>0.376287</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5</td>\n",
       "      <td>0.379000</td>\n",
       "      <td>0.381766</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The following columns in the evaluation set don't have a corresponding argument in `GPT2LMHeadModel.forward` and have been ignored: text. If text are not expected by `GPT2LMHeadModel.forward`,  you can safely ignore this message.\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 10\n",
      "  Batch size = 20\n",
      "Saving model checkpoint to ./english_to_latex/checkpoint-40\n",
      "Configuration saved in ./english_to_latex/checkpoint-40/config.json\n",
      "Configuration saved in ./english_to_latex/checkpoint-40/generation_config.json\n",
      "Model weights saved in ./english_to_latex/checkpoint-40/pytorch_model.bin\n",
      "The following columns in the evaluation set don't have a corresponding argument in `GPT2LMHeadModel.forward` and have been ignored: text. If text are not expected by `GPT2LMHeadModel.forward`,  you can safely ignore this message.\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 10\n",
      "  Batch size = 20\n",
      "Saving model checkpoint to ./english_to_latex/checkpoint-80\n",
      "Configuration saved in ./english_to_latex/checkpoint-80/config.json\n",
      "Configuration saved in ./english_to_latex/checkpoint-80/generation_config.json\n",
      "Model weights saved in ./english_to_latex/checkpoint-80/pytorch_model.bin\n",
      "The following columns in the evaluation set don't have a corresponding argument in `GPT2LMHeadModel.forward` and have been ignored: text. If text are not expected by `GPT2LMHeadModel.forward`,  you can safely ignore this message.\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 10\n",
      "  Batch size = 20\n",
      "Saving model checkpoint to ./english_to_latex/checkpoint-120\n",
      "Configuration saved in ./english_to_latex/checkpoint-120/config.json\n",
      "Configuration saved in ./english_to_latex/checkpoint-120/generation_config.json\n",
      "Model weights saved in ./english_to_latex/checkpoint-120/pytorch_model.bin\n",
      "The following columns in the evaluation set don't have a corresponding argument in `GPT2LMHeadModel.forward` and have been ignored: text. If text are not expected by `GPT2LMHeadModel.forward`,  you can safely ignore this message.\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 10\n",
      "  Batch size = 20\n",
      "Saving model checkpoint to ./english_to_latex/checkpoint-160\n",
      "Configuration saved in ./english_to_latex/checkpoint-160/config.json\n",
      "Configuration saved in ./english_to_latex/checkpoint-160/generation_config.json\n",
      "Model weights saved in ./english_to_latex/checkpoint-160/pytorch_model.bin\n",
      "The following columns in the evaluation set don't have a corresponding argument in `GPT2LMHeadModel.forward` and have been ignored: text. If text are not expected by `GPT2LMHeadModel.forward`,  you can safely ignore this message.\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 10\n",
      "  Batch size = 20\n",
      "Saving model checkpoint to ./english_to_latex/checkpoint-200\n",
      "Configuration saved in ./english_to_latex/checkpoint-200/config.json\n",
      "Configuration saved in ./english_to_latex/checkpoint-200/generation_config.json\n",
      "Model weights saved in ./english_to_latex/checkpoint-200/pytorch_model.bin\n",
      "\n",
      "\n",
      "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
      "\n",
      "\n",
      "Loading best model from ./english_to_latex/checkpoint-160 (score: 0.3762868046760559).\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=200, training_loss=0.8274225550889969, metrics={'train_runtime': 85.8452, 'train_samples_per_second': 2.33, 'train_steps_per_second': 2.33, 'total_flos': 3018637440000.0, 'train_loss': 0.8274225550889969, 'epoch': 5.0})"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6010cf42",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "173901dd",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading features from cached file ../data/cached_lm_GPT2TokenizerFast_128_latex-guide-cos423.txt [took 0.011 s]\n",
      "loading configuration file config.json from cache at /Users/sinanozdemir/.cache/huggingface/hub/models--gpt2/snapshots/11c5a3d5811f50298f278a704980280950aedb10/config.json\n",
      "Model config GPT2Config {\n",
      "  \"_name_or_path\": \"gpt2\",\n",
      "  \"activation_function\": \"gelu_new\",\n",
      "  \"architectures\": [\n",
      "    \"GPT2LMHeadModel\"\n",
      "  ],\n",
      "  \"attn_pdrop\": 0.1,\n",
      "  \"bos_token_id\": 50256,\n",
      "  \"embd_pdrop\": 0.1,\n",
      "  \"eos_token_id\": 50256,\n",
      "  \"initializer_range\": 0.02,\n",
      "  \"layer_norm_epsilon\": 1e-05,\n",
      "  \"model_type\": \"gpt2\",\n",
      "  \"n_ctx\": 1024,\n",
      "  \"n_embd\": 768,\n",
      "  \"n_head\": 12,\n",
      "  \"n_inner\": null,\n",
      "  \"n_layer\": 12,\n",
      "  \"n_positions\": 1024,\n",
      "  \"reorder_and_upcast_attn\": false,\n",
      "  \"resid_pdrop\": 0.1,\n",
      "  \"scale_attn_by_inverse_layer_idx\": false,\n",
      "  \"scale_attn_weights\": true,\n",
      "  \"summary_activation\": null,\n",
      "  \"summary_first_dropout\": 0.1,\n",
      "  \"summary_proj_to_labels\": true,\n",
      "  \"summary_type\": \"cls_index\",\n",
      "  \"summary_use_proj\": true,\n",
      "  \"task_specific_params\": {\n",
      "    \"text-generation\": {\n",
      "      \"do_sample\": true,\n",
      "      \"max_length\": 50\n",
      "    }\n",
      "  },\n",
      "  \"transformers_version\": \"4.30.2\",\n",
      "  \"use_cache\": true,\n",
      "  \"vocab_size\": 50257\n",
      "}\n",
      "\n",
      "loading weights file model.safetensors from cache at /Users/sinanozdemir/.cache/huggingface/hub/models--gpt2/snapshots/11c5a3d5811f50298f278a704980280950aedb10/model.safetensors\n",
      "Generate config GenerationConfig {\n",
      "  \"_from_model_config\": true,\n",
      "  \"bos_token_id\": 50256,\n",
      "  \"eos_token_id\": 50256,\n",
      "  \"transformers_version\": \"4.30.2\"\n",
      "}\n",
      "\n",
      "All model checkpoint weights were used when initializing GPT2LMHeadModel.\n",
      "\n",
      "All the weights of GPT2LMHeadModel were initialized from the model checkpoint at gpt2.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use GPT2LMHeadModel for predictions without further training.\n",
      "loading configuration file generation_config.json from cache at /Users/sinanozdemir/.cache/huggingface/hub/models--gpt2/snapshots/11c5a3d5811f50298f278a704980280950aedb10/generation_config.json\n",
      "Generate config GenerationConfig {\n",
      "  \"_from_model_config\": true,\n",
      "  \"bos_token_id\": 50256,\n",
      "  \"eos_token_id\": 50256,\n",
      "  \"transformers_version\": \"4.30.2\"\n",
      "}\n",
      "\n",
      "Found safetensors installation, but --save_safetensors=False. Safetensors should be a preferred weights saving format due to security and performance reasons. If your model cannot be saved by safetensors please feel free to open an issue at https://github.com/huggingface/safetensors!\n",
      "PyTorch: setting up devices\n",
      "The default value for the training argument `--report_to` will change in v5 (from all installed integrations to none). In v5, you will need to use `--report_to all` to get the same behavior as now. You should start updating your code and make this info disappear :-).\n"
     ]
    }
   ],
   "source": [
    "book_data = TextDataset(\n",
    "    tokenizer=tokenizer,\n",
    "    file_path='../data/latex-guide-cos423.txt',  # train on a LaTeX cheat sheet they made\n",
    "    block_size=128\n",
    ")\n",
    "\n",
    "data_collator = DataCollatorForLanguageModeling(\n",
    "    tokenizer=tokenizer, mlm=False,  # MLM is Masked Language Modelling\n",
    ")\n",
    "\n",
    "latex_gpt2 = AutoModelForCausalLM.from_pretrained(MODEL)\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./math_book\",\n",
    "    overwrite_output_dir=True, # overwrite the content of the output directory\n",
    "    num_train_epochs=10, # number of training epochs\n",
    "    per_device_train_batch_size=2, # batch size for training\n",
    "    per_device_eval_batch_size=32,  # batch size for evaluation\n",
    "    load_best_model_at_end=True,\n",
    "    logging_steps=10,\n",
    "    evaluation_strategy='epoch',\n",
    "    save_strategy='epoch',\n",
    "    use_mps_device=True\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=latex_gpt2,\n",
    "    args=training_args,\n",
    "    data_collator=data_collator,\n",
    "    train_dataset=book_data.examples[:int(len(book_data.examples)*.8)],\n",
    "    eval_dataset=book_data.examples[int(len(book_data.examples)*.8):]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "e466aa20",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "***** Running Evaluation *****\n",
      "  Num examples = 12\n",
      "  Batch size = 32\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='2' max='1' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [1/1 00:20]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Automatic Weights & Biases logging enabled, to disable set os.environ[\"WANDB_DISABLED\"] = \"true\"\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'eval_loss': 2.3705315589904785,\n",
       " 'eval_runtime': 8.3843,\n",
       " 'eval_samples_per_second': 1.431,\n",
       " 'eval_steps_per_second': 0.119}"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.evaluate()  # initial loss for the cheat sheet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "388afa38",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/homebrew/lib/python3.10/site-packages/transformers/optimization.py:411: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n",
      "***** Running training *****\n",
      "  Num examples = 47\n",
      "  Num Epochs = 10\n",
      "  Instantaneous batch size per device = 2\n",
      "  Total train batch size (w. parallel, distributed & accumulation) = 2\n",
      "  Gradient Accumulation steps = 1\n",
      "  Total optimization steps = 240\n",
      "  Number of trainable parameters = 124,439,808\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='240' max='240' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [240/240 01:30, Epoch 10/10]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>1.884900</td>\n",
       "      <td>1.815823</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>1.758800</td>\n",
       "      <td>1.791541</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>1.424900</td>\n",
       "      <td>1.782735</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>1.282700</td>\n",
       "      <td>1.818491</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5</td>\n",
       "      <td>1.268700</td>\n",
       "      <td>1.846649</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6</td>\n",
       "      <td>1.091600</td>\n",
       "      <td>1.912194</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7</td>\n",
       "      <td>1.014600</td>\n",
       "      <td>1.939818</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>8</td>\n",
       "      <td>0.756500</td>\n",
       "      <td>1.977578</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>9</td>\n",
       "      <td>0.747800</td>\n",
       "      <td>2.006962</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>0.729000</td>\n",
       "      <td>2.001646</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "***** Running Evaluation *****\n",
      "  Num examples = 12\n",
      "  Batch size = 32\n",
      "Saving model checkpoint to ./math_book/checkpoint-24\n",
      "Configuration saved in ./math_book/checkpoint-24/config.json\n",
      "Configuration saved in ./math_book/checkpoint-24/generation_config.json\n",
      "Model weights saved in ./math_book/checkpoint-24/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 12\n",
      "  Batch size = 32\n",
      "Saving model checkpoint to ./math_book/checkpoint-48\n",
      "Configuration saved in ./math_book/checkpoint-48/config.json\n",
      "Configuration saved in ./math_book/checkpoint-48/generation_config.json\n",
      "Model weights saved in ./math_book/checkpoint-48/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 12\n",
      "  Batch size = 32\n",
      "Saving model checkpoint to ./math_book/checkpoint-72\n",
      "Configuration saved in ./math_book/checkpoint-72/config.json\n",
      "Configuration saved in ./math_book/checkpoint-72/generation_config.json\n",
      "Model weights saved in ./math_book/checkpoint-72/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 12\n",
      "  Batch size = 32\n",
      "Saving model checkpoint to ./math_book/checkpoint-96\n",
      "Configuration saved in ./math_book/checkpoint-96/config.json\n",
      "Configuration saved in ./math_book/checkpoint-96/generation_config.json\n",
      "Model weights saved in ./math_book/checkpoint-96/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 12\n",
      "  Batch size = 32\n",
      "Saving model checkpoint to ./math_book/checkpoint-120\n",
      "Configuration saved in ./math_book/checkpoint-120/config.json\n",
      "Configuration saved in ./math_book/checkpoint-120/generation_config.json\n",
      "Model weights saved in ./math_book/checkpoint-120/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 12\n",
      "  Batch size = 32\n",
      "Saving model checkpoint to ./math_book/checkpoint-144\n",
      "Configuration saved in ./math_book/checkpoint-144/config.json\n",
      "Configuration saved in ./math_book/checkpoint-144/generation_config.json\n",
      "Model weights saved in ./math_book/checkpoint-144/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 12\n",
      "  Batch size = 32\n",
      "Saving model checkpoint to ./math_book/checkpoint-168\n",
      "Configuration saved in ./math_book/checkpoint-168/config.json\n",
      "Configuration saved in ./math_book/checkpoint-168/generation_config.json\n",
      "Model weights saved in ./math_book/checkpoint-168/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 12\n",
      "  Batch size = 32\n",
      "Saving model checkpoint to ./math_book/checkpoint-192\n",
      "Configuration saved in ./math_book/checkpoint-192/config.json\n",
      "Configuration saved in ./math_book/checkpoint-192/generation_config.json\n",
      "Model weights saved in ./math_book/checkpoint-192/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 12\n",
      "  Batch size = 32\n",
      "Saving model checkpoint to ./math_book/checkpoint-216\n",
      "Configuration saved in ./math_book/checkpoint-216/config.json\n",
      "Configuration saved in ./math_book/checkpoint-216/generation_config.json\n",
      "Model weights saved in ./math_book/checkpoint-216/pytorch_model.bin\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 12\n",
      "  Batch size = 32\n",
      "Saving model checkpoint to ./math_book/checkpoint-240\n",
      "Configuration saved in ./math_book/checkpoint-240/config.json\n",
      "Configuration saved in ./math_book/checkpoint-240/generation_config.json\n",
      "Model weights saved in ./math_book/checkpoint-240/pytorch_model.bin\n",
      "\n",
      "\n",
      "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
      "\n",
      "\n",
      "Loading best model from ./math_book/checkpoint-72 (score: 1.7827345132827759).\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=240, training_loss=1.147268941005071, metrics={'train_runtime': 98.3177, 'train_samples_per_second': 4.78, 'train_steps_per_second': 2.441, 'total_flos': 30701813760000.0, 'train_loss': 1.147268941005071, 'epoch': 10.0})"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "270433fa",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to ./math_book\n",
      "Configuration saved in ./math_book/config.json\n",
      "Configuration saved in ./math_book/generation_config.json\n",
      "Model weights saved in ./math_book/pytorch_model.bin\n"
     ]
    }
   ],
   "source": [
    "trainer.save_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "946a90f3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "c38cbe07",
   "metadata": {},
   "source": [
    "# Restart training now with our own pre-trained \"foundation\" model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "0c820b1b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "loading configuration file ./math_book/config.json\n",
      "Model config GPT2Config {\n",
      "  \"_name_or_path\": \"./math_book\",\n",
      "  \"activation_function\": \"gelu_new\",\n",
      "  \"architectures\": [\n",
      "    \"GPT2LMHeadModel\"\n",
      "  ],\n",
      "  \"attn_pdrop\": 0.1,\n",
      "  \"bos_token_id\": 50256,\n",
      "  \"embd_pdrop\": 0.1,\n",
      "  \"eos_token_id\": 50256,\n",
      "  \"initializer_range\": 0.02,\n",
      "  \"layer_norm_epsilon\": 1e-05,\n",
      "  \"model_type\": \"gpt2\",\n",
      "  \"n_ctx\": 1024,\n",
      "  \"n_embd\": 768,\n",
      "  \"n_head\": 12,\n",
      "  \"n_inner\": null,\n",
      "  \"n_layer\": 12,\n",
      "  \"n_positions\": 1024,\n",
      "  \"reorder_and_upcast_attn\": false,\n",
      "  \"resid_pdrop\": 0.1,\n",
      "  \"scale_attn_by_inverse_layer_idx\": false,\n",
      "  \"scale_attn_weights\": true,\n",
      "  \"summary_activation\": null,\n",
      "  \"summary_first_dropout\": 0.1,\n",
      "  \"summary_proj_to_labels\": true,\n",
      "  \"summary_type\": \"cls_index\",\n",
      "  \"summary_use_proj\": true,\n",
      "  \"task_specific_params\": {\n",
      "    \"text-generation\": {\n",
      "      \"do_sample\": true,\n",
      "      \"max_length\": 50\n",
      "    }\n",
      "  },\n",
      "  \"torch_dtype\": \"float32\",\n",
      "  \"transformers_version\": \"4.30.2\",\n",
      "  \"use_cache\": true,\n",
      "  \"vocab_size\": 50257\n",
      "}\n",
      "\n",
      "loading weights file ./math_book/pytorch_model.bin\n",
      "Generate config GenerationConfig {\n",
      "  \"_from_model_config\": true,\n",
      "  \"bos_token_id\": 50256,\n",
      "  \"eos_token_id\": 50256,\n",
      "  \"transformers_version\": \"4.30.2\"\n",
      "}\n",
      "\n",
      "All model checkpoint weights were used when initializing GPT2LMHeadModel.\n",
      "\n",
      "All the weights of GPT2LMHeadModel were initialized from the model checkpoint at ./math_book.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use GPT2LMHeadModel for predictions without further training.\n",
      "loading configuration file ./math_book/generation_config.json\n",
      "Generate config GenerationConfig {\n",
      "  \"_from_model_config\": true,\n",
      "  \"bos_token_id\": 50256,\n",
      "  \"eos_token_id\": 50256,\n",
      "  \"transformers_version\": \"4.30.2\"\n",
      "}\n",
      "\n",
      "Found safetensors installation, but --save_safetensors=False. Safetensors should be a preferred weights saving format due to security and performance reasons. If your model cannot be saved by safetensors please feel free to open an issue at https://github.com/huggingface/safetensors!\n",
      "PyTorch: setting up devices\n",
      "The default value for the training argument `--report_to` will change in v5 (from all installed integrations to none). In v5, you will need to use `--report_to all` to get the same behavior as now. You should start updating your code and make this info disappear :-).\n",
      "The following columns in the evaluation set don't have a corresponding argument in `GPT2LMHeadModel.forward` and have been ignored: text. If text are not expected by `GPT2LMHeadModel.forward`,  you can safely ignore this message.\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 10\n",
      "  Batch size = 20\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='2' max='1' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [1/1 00:13]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Automatic Weights & Biases logging enabled, to disable set os.environ[\"WANDB_DISABLED\"] = \"true\"\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'eval_loss': 2.5888800621032715,\n",
       " 'eval_runtime': 1.1215,\n",
       " 'eval_samples_per_second': 8.917,\n",
       " 'eval_steps_per_second': 0.892}"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# load up our gpt pre-trained on latex cheat sheets\n",
    "math_latex_gpt2 = AutoModelForCausalLM.from_pretrained('./math_book')\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./math_english_to_latex\",\n",
    "    overwrite_output_dir=True, #overwrite the content of the output directory\n",
    "    num_train_epochs=5, # number of training epochs\n",
    "    per_device_train_batch_size=1, # batch size for training\n",
    "    per_device_eval_batch_size=20,  # batch size for evaluation\n",
    "    load_best_model_at_end=True,\n",
    "    logging_steps=5,\n",
    "    log_level='info',\n",
    "    evaluation_strategy='epoch',\n",
    "    save_strategy='epoch',\n",
    "    use_mps_device=True\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=math_latex_gpt2,\n",
    "    args=training_args,\n",
    "    train_dataset=latex_data[\"train\"],\n",
    "    eval_dataset=latex_data[\"test\"],\n",
    "    data_collator=data_collator,\n",
    ")\n",
    "\n",
    "trainer.evaluate()  # loss is starting slightly lower than before"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "67a53cd6",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The following columns in the training set don't have a corresponding argument in `GPT2LMHeadModel.forward` and have been ignored: text. If text are not expected by `GPT2LMHeadModel.forward`,  you can safely ignore this message.\n",
      "/opt/homebrew/lib/python3.10/site-packages/transformers/optimization.py:411: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n",
      "***** Running training *****\n",
      "  Num examples = 40\n",
      "  Num Epochs = 5\n",
      "  Instantaneous batch size per device = 1\n",
      "  Total train batch size (w. parallel, distributed & accumulation) = 1\n",
      "  Gradient Accumulation steps = 1\n",
      "  Total optimization steps = 200\n",
      "  Number of trainable parameters = 124,439,808\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='200' max='200' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [200/200 00:54, Epoch 5/5]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.872800</td>\n",
       "      <td>0.583309</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>0.752600</td>\n",
       "      <td>0.454410</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>0.701400</td>\n",
       "      <td>0.416536</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>0.494500</td>\n",
       "      <td>0.373505</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5</td>\n",
       "      <td>0.368600</td>\n",
       "      <td>0.383265</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The following columns in the evaluation set don't have a corresponding argument in `GPT2LMHeadModel.forward` and have been ignored: text. If text are not expected by `GPT2LMHeadModel.forward`,  you can safely ignore this message.\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 10\n",
      "  Batch size = 20\n",
      "Saving model checkpoint to ./math_english_to_latex/checkpoint-40\n",
      "Configuration saved in ./math_english_to_latex/checkpoint-40/config.json\n",
      "Configuration saved in ./math_english_to_latex/checkpoint-40/generation_config.json\n",
      "Model weights saved in ./math_english_to_latex/checkpoint-40/pytorch_model.bin\n",
      "The following columns in the evaluation set don't have a corresponding argument in `GPT2LMHeadModel.forward` and have been ignored: text. If text are not expected by `GPT2LMHeadModel.forward`,  you can safely ignore this message.\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 10\n",
      "  Batch size = 20\n",
      "Saving model checkpoint to ./math_english_to_latex/checkpoint-80\n",
      "Configuration saved in ./math_english_to_latex/checkpoint-80/config.json\n",
      "Configuration saved in ./math_english_to_latex/checkpoint-80/generation_config.json\n",
      "Model weights saved in ./math_english_to_latex/checkpoint-80/pytorch_model.bin\n",
      "The following columns in the evaluation set don't have a corresponding argument in `GPT2LMHeadModel.forward` and have been ignored: text. If text are not expected by `GPT2LMHeadModel.forward`,  you can safely ignore this message.\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 10\n",
      "  Batch size = 20\n",
      "Saving model checkpoint to ./math_english_to_latex/checkpoint-120\n",
      "Configuration saved in ./math_english_to_latex/checkpoint-120/config.json\n",
      "Configuration saved in ./math_english_to_latex/checkpoint-120/generation_config.json\n",
      "Model weights saved in ./math_english_to_latex/checkpoint-120/pytorch_model.bin\n",
      "The following columns in the evaluation set don't have a corresponding argument in `GPT2LMHeadModel.forward` and have been ignored: text. If text are not expected by `GPT2LMHeadModel.forward`,  you can safely ignore this message.\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 10\n",
      "  Batch size = 20\n",
      "Saving model checkpoint to ./math_english_to_latex/checkpoint-160\n",
      "Configuration saved in ./math_english_to_latex/checkpoint-160/config.json\n",
      "Configuration saved in ./math_english_to_latex/checkpoint-160/generation_config.json\n",
      "Model weights saved in ./math_english_to_latex/checkpoint-160/pytorch_model.bin\n",
      "The following columns in the evaluation set don't have a corresponding argument in `GPT2LMHeadModel.forward` and have been ignored: text. If text are not expected by `GPT2LMHeadModel.forward`,  you can safely ignore this message.\n",
      "***** Running Evaluation *****\n",
      "  Num examples = 10\n",
      "  Batch size = 20\n",
      "Saving model checkpoint to ./math_english_to_latex/checkpoint-200\n",
      "Configuration saved in ./math_english_to_latex/checkpoint-200/config.json\n",
      "Configuration saved in ./math_english_to_latex/checkpoint-200/generation_config.json\n",
      "Model weights saved in ./math_english_to_latex/checkpoint-200/pytorch_model.bin\n",
      "\n",
      "\n",
      "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
      "\n",
      "\n",
      "Loading best model from ./math_english_to_latex/checkpoint-160 (score: 0.3735045790672302).\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=200, training_loss=0.8061065357923508, metrics={'train_runtime': 55.2758, 'train_samples_per_second': 3.618, 'train_steps_per_second': 3.618, 'total_flos': 3018637440000.0, 'train_loss': 0.8061065357923508, 'epoch': 5.0})"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "530c78c5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to ./math_english_to_latex\n",
      "Configuration saved in ./math_english_to_latex/config.json\n",
      "Configuration saved in ./math_english_to_latex/generation_config.json\n",
      "Model weights saved in ./math_english_to_latex/pytorch_model.bin\n"
     ]
    }
   ],
   "source": [
    "trainer.save_model()  # save this model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13153c32",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c34b7af3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "2b2cb777",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "loading configuration file ./math_english_to_latex/config.json\n",
      "Model config GPT2Config {\n",
      "  \"_name_or_path\": \"./math_english_to_latex\",\n",
      "  \"activation_function\": \"gelu_new\",\n",
      "  \"architectures\": [\n",
      "    \"GPT2LMHeadModel\"\n",
      "  ],\n",
      "  \"attn_pdrop\": 0.1,\n",
      "  \"bos_token_id\": 50256,\n",
      "  \"embd_pdrop\": 0.1,\n",
      "  \"eos_token_id\": 50256,\n",
      "  \"initializer_range\": 0.02,\n",
      "  \"layer_norm_epsilon\": 1e-05,\n",
      "  \"model_type\": \"gpt2\",\n",
      "  \"n_ctx\": 1024,\n",
      "  \"n_embd\": 768,\n",
      "  \"n_head\": 12,\n",
      "  \"n_inner\": null,\n",
      "  \"n_layer\": 12,\n",
      "  \"n_positions\": 1024,\n",
      "  \"reorder_and_upcast_attn\": false,\n",
      "  \"resid_pdrop\": 0.1,\n",
      "  \"scale_attn_by_inverse_layer_idx\": false,\n",
      "  \"scale_attn_weights\": true,\n",
      "  \"summary_activation\": null,\n",
      "  \"summary_first_dropout\": 0.1,\n",
      "  \"summary_proj_to_labels\": true,\n",
      "  \"summary_type\": \"cls_index\",\n",
      "  \"summary_use_proj\": true,\n",
      "  \"task_specific_params\": {\n",
      "    \"text-generation\": {\n",
      "      \"do_sample\": true,\n",
      "      \"max_length\": 50\n",
      "    }\n",
      "  },\n",
      "  \"torch_dtype\": \"float32\",\n",
      "  \"transformers_version\": \"4.30.2\",\n",
      "  \"use_cache\": true,\n",
      "  \"vocab_size\": 50257\n",
      "}\n",
      "\n",
      "loading weights file ./math_english_to_latex/pytorch_model.bin\n",
      "Generate config GenerationConfig {\n",
      "  \"_from_model_config\": true,\n",
      "  \"bos_token_id\": 50256,\n",
      "  \"eos_token_id\": 50256,\n",
      "  \"transformers_version\": \"4.30.2\"\n",
      "}\n",
      "\n",
      "All model checkpoint weights were used when initializing GPT2LMHeadModel.\n",
      "\n",
      "All the weights of GPT2LMHeadModel were initialized from the model checkpoint at ./math_english_to_latex.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use GPT2LMHeadModel for predictions without further training.\n",
      "loading configuration file ./math_english_to_latex/generation_config.json\n",
      "Generate config GenerationConfig {\n",
      "  \"_from_model_config\": true,\n",
      "  \"bos_token_id\": 50256,\n",
      "  \"eos_token_id\": 50256,\n",
      "  \"transformers_version\": \"4.30.2\"\n",
      "}\n",
      "\n",
      "Generate config GenerationConfig {\n",
      "  \"_from_model_config\": true,\n",
      "  \"bos_token_id\": 50256,\n",
      "  \"do_sample\": true,\n",
      "  \"eos_token_id\": 50256,\n",
      "  \"max_length\": 50,\n",
      "  \"transformers_version\": \"4.30.2\"\n",
      "}\n",
      "\n",
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Convert English to LaTeX\n",
      "English: g of x equals integral from 0 to 1 of x squared\n",
      "LaTeX: g(x) = \\int_{0}^{1} x^2 \\,dx^2 \\,dx\n"
     ]
    }
   ],
   "source": [
    "loaded_model = AutoModelForCausalLM.from_pretrained('./math_english_to_latex')\n",
    "latex_generator = pipeline('text-generation', model=loaded_model, tokenizer=tokenizer)\n",
    "\n",
    "text_sample = 'g of x equals integral from 0 to 1 of x squared'\n",
    "conversion_text_sample = f'{CONVERSION_PROMPT}English: {text_sample}\\n{CONVERSION_TOKEN}'\n",
    "\n",
    "print(latex_generator(\n",
    "    conversion_text_sample, num_beams=2, early_stopping=True, temperature=0.7,\n",
    "    max_new_tokens=24\n",
    ")[0]['generated_text'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfd233ed",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "46d13e88",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generate config GenerationConfig {\n",
      "  \"_from_model_config\": true,\n",
      "  \"bos_token_id\": 50256,\n",
      "  \"do_sample\": true,\n",
      "  \"eos_token_id\": 50256,\n",
      "  \"max_length\": 50,\n",
      "  \"transformers_version\": \"4.30.2\"\n",
      "}\n",
      "\n",
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Convert English to LaTeX\n",
      "English: r of x is sum from 0 to x of x squared\n",
      "LaTeX: r(x) = \\sum_{0}^{x} x^2 \\,dx^\n"
     ]
    }
   ],
   "source": [
    "# Another example\n",
    "text_sample = 'r of x is sum from 0 to x of x squared'\n",
    "conversion_text_sample = f'{CONVERSION_PROMPT}English: {text_sample}\\n{CONVERSION_TOKEN}'\n",
    "\n",
    "print(latex_generator(\n",
    "    conversion_text_sample, num_beams=5, early_stopping=True, temperature=0.7,\n",
    "    max_length=len(tokenizer.encode(conversion_text_sample)) + 20\n",
    ")[0]['generated_text'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "9ac6c305",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generate config GenerationConfig {\n",
      "  \"_from_model_config\": true,\n",
      "  \"bos_token_id\": 50256,\n",
      "  \"do_sample\": true,\n",
      "  \"eos_token_id\": 50256,\n",
      "  \"max_length\": 50,\n",
      "  \"transformers_version\": \"4.30.2\"\n",
      "}\n",
      "\n",
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "r of x is sum from 0 to x of x squared\n",
      "x^2 \\,dx^2 \\,dy^2 \\,dy^2 \\,dx^2 \\,dy^2 \\,dx^\n"
     ]
    }
   ],
   "source": [
    "print(latex_generator(\n",
    "    text_sample, num_beams=5, early_stopping=True, temperature=0.7,\n",
    "    max_length=len(tokenizer.encode(conversion_text_sample)) + 20\n",
    ")[0]['generated_text'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06c44d01",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "ab316df9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generate config GenerationConfig {\n",
      "  \"_from_model_config\": true,\n",
      "  \"bos_token_id\": 50256,\n",
      "  \"do_sample\": true,\n",
      "  \"eos_token_id\": 50256,\n",
      "  \"max_length\": 50,\n",
      "  \"transformers_version\": \"4.30.2\"\n",
      "}\n",
      "\n",
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Convert English to LaTeX\n",
      "English: f of x is sum from 0 to x of x squared\n",
      "LaTeX: f(x) = \\sum_{0}^{x} x^2 \\,dx ###\n",
      "LCT\n",
      "English: f of x equals integral from 0 to pi of x to the fourth power\n",
      "LaTeX: f(x) = \\int_{0}^{\\pi} x^4 \\,dx ###\n",
      "LCT\n",
      "English: pi to the 8th power\n",
      "LaTeX: pi to the 8th power\n",
      "LCT\n",
      "English: f(x) = \\sum_{\n"
     ]
    }
   ],
   "source": [
    "# try a few shot with standard gpt2\n",
    "few_shot_prompt = CONVERSION_PROMPT+\"\"\"English: f of x is sum from 0 to x of x squared\n",
    "LaTeX: f(x) = \\sum_{0}^{x} x^2 \\,dx \\\n",
    "###\n",
    "LCT\n",
    "English: f of x equals integral from 0 to pi of x to the fourth power\n",
    "LaTeX: f(x) = \\int_{0}^{\\pi} x^4 \\,dx \\\n",
    "###\n",
    "LCT\n",
    "English: pi to the 8th power\n",
    "LaTeX:\"\"\"\n",
    "\n",
    "print(non_finetuned_latex_generator(\n",
    "    few_shot_prompt, num_beams=1, early_stopping=True, temperature=0.1,\n",
    "    max_length=len(tokenizer.encode(few_shot_prompt)) + 20\n",
    ")[0]['generated_text'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9243a0d9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "c87ac747",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generate config GenerationConfig {\n",
      "  \"_from_model_config\": true,\n",
      "  \"bos_token_id\": 50256,\n",
      "  \"do_sample\": true,\n",
      "  \"eos_token_id\": 50256,\n",
      "  \"max_length\": 50,\n",
      "  \"transformers_version\": \"4.30.2\"\n",
      "}\n",
      "\n",
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Convert English to LaTeX\n",
      "English: pi to the 16th power\n",
      "LaTeX: pi to the 16th power\n",
      "LaTeX: pi to the 16th power\n",
      "LaTeX:\n"
     ]
    }
   ],
   "source": [
    "# Just ask with standard gpt2\n",
    "print(non_finetuned_latex_generator(\n",
    "    conversion_text_sample, num_beams=1, early_stopping=True, temperature=0.1,\n",
    "    max_length=len(tokenizer.encode(conversion_text_sample)) + 20\n",
    ")[0]['generated_text'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87615fb5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f280e57b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}