fine-tunning-gemma-model-with-kaggle-docs-data.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "02e92d66",
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "papermill": {
     "duration": 0.011778,
     "end_time": "2024-06-30T15:28:29.501435",
     "exception": false,
     "start_time": "2024-06-30T15:28:29.489657",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "<center><h1>Fine-tunning Gemma model with Kaggle Docs data</h1></center>\n",
    "\n",
    "<center><img src=\"https://res.infoq.com/news/2024/02/google-gemma-open-model/en/headerimage/generatedHeaderImage-1708977571481.jpg\" width=\"400\"></center>\n",
    "\n",
    "\n",
    "# Introduction\n",
    "\n",
    "This notebook will demonstrate three things:\n",
    "\n",
    "1. How to fine-tune Gemma model using LoRA\n",
    "2. Creation of a specialised class to query about Kaggle features\n",
    "3. Some results of querying about Kaggle Docs\n",
    "\n",
    "This work is largely based on previous work. Here I list the sources:\n",
    "\n",
    "1. Gemma Model Card, Kaggle Models, https://www.kaggle.com/models/google/gemma\n",
    "2. Kaggle QA with Gemma - KerasNLP Starter, Kaggle Code, https://www.kaggle.com/code/awsaf49/kaggle-qa-with-gemma-kerasnlp-starter (Version 11)  \n",
    "3. Fine-tune Gemma models in Keras using LoRA, Kaggle Code, https://www.kaggle.com/code/nilaychauhan/fine-tune-gemma-models-in-keras-using-lora (Version 1)  \n",
    "4. Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, LoRA: Low-Rank Adaptation of Large Language Models, ArXiv, https://arxiv.org/pdf/2106.09685.pdf\n",
    "5. Abheesht Sharma, Matthew Watson, Parameter-efficient fine-tuning of GPT-2 with LoRA, https://keras.io/examples/nlp/parameter_efficient_finetuning_of_gpt2_with_lora/\n",
    "6. Keras 3 API documentation / KerasNLP / Models / Gemma, https://keras.io/api/keras_nlp/models/gemma/\n",
    "7. Kaggle Docs, Kaggle Dataset, https://www.kaggle.com/datasets/awsaf49/kaggle-docs  \n",
    "\n",
    "\n",
    "**Let's go**!\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0bb8004",
   "metadata": {
    "papermill": {
     "duration": 0.011196,
     "end_time": "2024-06-30T15:28:29.523834",
     "exception": false,
     "start_time": "2024-06-30T15:28:29.512638",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# What is Gemma?\n",
    "\n",
    "\n",
    "Gemma is a collection of lightweight source generative AI models designed to be used mostly by developers and researchers. Created by Google DeepMind research lab that also developed Gemini, Gemma is available in several versions, with 2B and 7B parameters, as following:\n",
    "\n",
    "\n",
    "| Model                  | Parameters      | Tuned versions    | Description                                    | Recomemnded target platforms       |\n",
    "|------------------------|-----------------|-------------------|------------------------------------------------|------------------------------------|\n",
    "| `gemma_2b_en`          | 2.51B           | Pretrained        | 18-layer Gemma model (Gemma with 2B parameters)|Mobile devices and laptops          |\n",
    "| `gemma_instruct_2b_en` | 2.51B           | Instruction tuned | 18-layer Gemma model (Gemma with 2B parameters)| Mobile devices and laptops         | \n",
    "| `gemma_7b_en`          | 8.54B           | Pretrained        | 28-layer Gemma model (Gemma with 7B parameters)| Desktop computers and small servers|\n",
    "| `gemma_instruct_7b_en` | 8.54B           | Instruction tuned | 28-layer Gemma model (Gemma with 7B parameters)| Desktop computers and small servers|\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27466746",
   "metadata": {
    "papermill": {
     "duration": 0.010845,
     "end_time": "2024-06-30T15:28:29.546314",
     "exception": false,
     "start_time": "2024-06-30T15:28:29.535469",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# What is LoRA?  \n",
    "\n",
    "LoRA stands for Low-Rank Adaptation. It is a method used to fine-tune large language models (LLMs) by freezing the weights of the LLM and injecting trainable rank-decomposition matrices. The number of trainable parameters during fine-tunning will decrease therefore considerably. According to LoRA paper, this number decreases 10,000 times, and the computational resources size decreases 3 times. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cee508e",
   "metadata": {
    "papermill": {
     "duration": 0.010837,
     "end_time": "2024-06-30T15:28:29.569112",
     "exception": false,
     "start_time": "2024-06-30T15:28:29.558275",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# How we proceed?\n",
    "\n",
    "For fine-tunning with LoRA, we will follow the steps:\n",
    "\n",
    "1. Install prerequisites\n",
    "2. Load and process the data for fine-tuning\n",
    "3. Initialize the code for Gemma causal language model (Gemma Causal LM)\n",
    "4. Perform fine-tuning\n",
    "5. Test the fine-tunned model with questions from the data used for fine-tuning and with aditional questions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8db28e46",
   "metadata": {
    "papermill": {
     "duration": 0.010804,
     "end_time": "2024-06-30T15:28:29.591020",
     "exception": false,
     "start_time": "2024-06-30T15:28:29.580216",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Prerequisites\n",
    "\n",
    "\n",
    "## Install packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4de551a4",
   "metadata": {
    "_kg_hide-output": true,
    "execution": {
     "iopub.execute_input": "2024-06-30T15:28:29.615034Z",
     "iopub.status.busy": "2024-06-30T15:28:29.614679Z",
     "iopub.status.idle": "2024-06-30T15:28:59.952416Z",
     "shell.execute_reply": "2024-06-30T15:28:59.951479Z"
    },
    "papermill": {
     "duration": 30.35287,
     "end_time": "2024-06-30T15:28:59.954916",
     "exception": false,
     "start_time": "2024-06-30T15:28:29.602046",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\r\n",
      "tensorflow-decision-forests 1.8.1 requires wurlitzer, which is not installed.\u001b[0m\u001b[31m\r\n",
      "\u001b[0m\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\r\n",
      "tensorflow-decision-forests 1.8.1 requires wurlitzer, which is not installed.\r\n",
      "tensorflow 2.15.0 requires keras<2.16,>=2.15.0, but you have keras 3.4.1 which is incompatible.\u001b[0m\u001b[31m\r\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "# Install Keras 3 last. See https://keras.io/getting_started/ for more details.\n",
    "!pip install -q -U keras-nlp\n",
    "!pip install -q -U keras>=3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea669e64",
   "metadata": {
    "papermill": {
     "duration": 0.011309,
     "end_time": "2024-06-30T15:28:59.978034",
     "exception": false,
     "start_time": "2024-06-30T15:28:59.966725",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "affc1eda",
   "metadata": {
    "_kg_hide-output": true,
    "execution": {
     "iopub.execute_input": "2024-06-30T15:29:00.001950Z",
     "iopub.status.busy": "2024-06-30T15:29:00.001622Z",
     "iopub.status.idle": "2024-06-30T15:29:14.544585Z",
     "shell.execute_reply": "2024-06-30T15:29:14.543615Z"
    },
    "papermill": {
     "duration": 14.557812,
     "end_time": "2024-06-30T15:29:14.547094",
     "exception": false,
     "start_time": "2024-06-30T15:28:59.989282",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-06-30 15:29:04.547281: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2024-06-30 15:29:04.547414: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2024-06-30 15:29:04.677591: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ[\"KERAS_BACKEND\"] = \"jax\" # you can also use tensorflow or torch\n",
    "os.environ[\"XLA_PYTHON_CLIENT_MEM_FRACTION\"] = \"1.00\" # avoid memory fragmentation on JAX backend.\n",
    "os.environ[\"JAX_PLATFORMS\"] = \"\"\n",
    "import keras\n",
    "import keras_nlp\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm.notebook import tqdm\n",
    "tqdm.pandas() # progress bar for pandas\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from IPython.display import display, Markdown"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1cfe680",
   "metadata": {
    "papermill": {
     "duration": 0.011362,
     "end_time": "2024-06-30T15:29:14.570250",
     "exception": false,
     "start_time": "2024-06-30T15:29:14.558888",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6b93ee9f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:29:14.594964Z",
     "iopub.status.busy": "2024-06-30T15:29:14.594176Z",
     "iopub.status.idle": "2024-06-30T15:29:14.599643Z",
     "shell.execute_reply": "2024-06-30T15:29:14.598805Z"
    },
    "papermill": {
     "duration": 0.019965,
     "end_time": "2024-06-30T15:29:14.601574",
     "exception": false,
     "start_time": "2024-06-30T15:29:14.581609",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Config:\n",
    "    seed = 42\n",
    "    dataset_path = \"/kaggle/input/kaggle-docs/questions_answers\"\n",
    "    preset = \"gemma_2b_en\" # name of pretrained Gemma\n",
    "    sequence_length = 512 # max size of input sequence for training\n",
    "    batch_size = 1 # size of the input batch in training, x 2 as two GPUs\n",
    "    lora_rank = 4 # rank for LoRA, higher means more trainable parameters\n",
    "    learning_rate=8e-5 # learning rate used in train\n",
    "    epochs = 15 # number of epochs to train"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "775149f5",
   "metadata": {
    "papermill": {
     "duration": 0.01102,
     "end_time": "2024-06-30T15:29:14.624145",
     "exception": false,
     "start_time": "2024-06-30T15:29:14.613125",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Set a random seed for results reproducibility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6436910a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:29:14.648220Z",
     "iopub.status.busy": "2024-06-30T15:29:14.647404Z",
     "iopub.status.idle": "2024-06-30T15:29:14.652217Z",
     "shell.execute_reply": "2024-06-30T15:29:14.651507Z"
    },
    "papermill": {
     "duration": 0.018877,
     "end_time": "2024-06-30T15:29:14.654200",
     "exception": false,
     "start_time": "2024-06-30T15:29:14.635323",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "keras.utils.set_random_seed(Config.seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d30e25f",
   "metadata": {
    "papermill": {
     "duration": 0.011103,
     "end_time": "2024-06-30T15:29:14.676617",
     "exception": false,
     "start_time": "2024-06-30T15:29:14.665514",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Load the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1c3b9781",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:29:14.700920Z",
     "iopub.status.busy": "2024-06-30T15:29:14.700173Z",
     "iopub.status.idle": "2024-06-30T15:29:14.741530Z",
     "shell.execute_reply": "2024-06-30T15:29:14.740714Z"
    },
    "papermill": {
     "duration": 0.055506,
     "end_time": "2024-06-30T15:29:14.743443",
     "exception": false,
     "start_time": "2024-06-30T15:29:14.687937",
     "status": "completed"
    },
    "tags": []
   },
   "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>Question</th>\n",
       "      <th>Answer</th>\n",
       "      <th>Category</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>What are the different types of competitions a...</td>\n",
       "      <td># Types of Competitions\\n\\nKaggle Competitions...</td>\n",
       "      <td>competition</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>What are the different competition formats on ...</td>\n",
       "      <td>There are handful of different formats competi...</td>\n",
       "      <td>competition</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>How to join a competition?</td>\n",
       "      <td>Before you start, navigate to the [Competition...</td>\n",
       "      <td>competition</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>How to form, manage, and disband teams in a co...</td>\n",
       "      <td>Everyone that competes in a Competition does s...</td>\n",
       "      <td>competition</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>How do I make a submission in a competition?</td>\n",
       "      <td>You will need to submit your model predictions...</td>\n",
       "      <td>competition</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            Question  \\\n",
       "0  What are the different types of competitions a...   \n",
       "1  What are the different competition formats on ...   \n",
       "2                         How to join a competition?   \n",
       "3  How to form, manage, and disband teams in a co...   \n",
       "4       How do I make a submission in a competition?   \n",
       "\n",
       "                                              Answer     Category  \n",
       "0  # Types of Competitions\\n\\nKaggle Competitions...  competition  \n",
       "1  There are handful of different formats competi...  competition  \n",
       "2  Before you start, navigate to the [Competition...  competition  \n",
       "3  Everyone that competes in a Competition does s...  competition  \n",
       "4  You will need to submit your model predictions...  competition  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.read_csv(f\"{Config.dataset_path}/data.csv\")\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5471fca9",
   "metadata": {
    "papermill": {
     "duration": 0.011495,
     "end_time": "2024-06-30T15:29:14.767295",
     "exception": false,
     "start_time": "2024-06-30T15:29:14.755800",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's check the total number of rows in this dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c285e2ae",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:29:14.791586Z",
     "iopub.status.busy": "2024-06-30T15:29:14.791322Z",
     "iopub.status.idle": "2024-06-30T15:29:14.796403Z",
     "shell.execute_reply": "2024-06-30T15:29:14.795612Z"
    },
    "papermill": {
     "duration": 0.019489,
     "end_time": "2024-06-30T15:29:14.798364",
     "exception": false,
     "start_time": "2024-06-30T15:29:14.778875",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "60"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.shape[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf9ac01a",
   "metadata": {
    "papermill": {
     "duration": 0.011602,
     "end_time": "2024-06-30T15:29:14.821869",
     "exception": false,
     "start_time": "2024-06-30T15:29:14.810267",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "For easiness, we will create the following template for QA: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6c4f406f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:29:14.846312Z",
     "iopub.status.busy": "2024-06-30T15:29:14.846072Z",
     "iopub.status.idle": "2024-06-30T15:29:14.855980Z",
     "shell.execute_reply": "2024-06-30T15:29:14.855119Z"
    },
    "papermill": {
     "duration": 0.024301,
     "end_time": "2024-06-30T15:29:14.857894",
     "exception": false,
     "start_time": "2024-06-30T15:29:14.833593",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "template = \"\\n\\nCategory:\\nkaggle-{Category}\\n\\nQuestion:\\n{Question}\\n\\nAnswer:\\n{Answer}\"\n",
    "df[\"prompt\"] = df.apply(lambda row: template.format(Category=row.Category,\n",
    "                                                             Question=row.Question,\n",
    "                                                             Answer=row.Answer), axis=1)\n",
    "data = df.prompt.tolist()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6404730f",
   "metadata": {
    "papermill": {
     "duration": 0.011611,
     "end_time": "2024-06-30T15:29:14.881322",
     "exception": false,
     "start_time": "2024-06-30T15:29:14.869711",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Template utility function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f6509ffa",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:29:14.906398Z",
     "iopub.status.busy": "2024-06-30T15:29:14.905659Z",
     "iopub.status.idle": "2024-06-30T15:29:14.910510Z",
     "shell.execute_reply": "2024-06-30T15:29:14.909715Z"
    },
    "papermill": {
     "duration": 0.019328,
     "end_time": "2024-06-30T15:29:14.912389",
     "exception": false,
     "start_time": "2024-06-30T15:29:14.893061",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def colorize_text(text):\n",
    "    for word, color in zip([\"Category\", \"Question\", \"Answer\"], [\"blue\", \"red\", \"green\"]):\n",
    "        text = text.replace(f\"\\n\\n{word}:\", f\"\\n\\n**<font color='{color}'>{word}:</font>**\")\n",
    "    return text"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "707b65ed",
   "metadata": {
    "papermill": {
     "duration": 0.011445,
     "end_time": "2024-06-30T15:29:14.935491",
     "exception": false,
     "start_time": "2024-06-30T15:29:14.924046",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Specialized class to query Gemma\n",
    "\n",
    "\n",
    "We define a specialized class to query Gemma. But first, we need to initialize an object of GemmaCausalLM class."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5f6ea2c",
   "metadata": {
    "papermill": {
     "duration": 0.01147,
     "end_time": "2024-06-30T15:29:14.958672",
     "exception": false,
     "start_time": "2024-06-30T15:29:14.947202",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Initialize the code for Gemma Causal LM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f8db7598",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:29:15.022047Z",
     "iopub.status.busy": "2024-06-30T15:29:15.021697Z",
     "iopub.status.idle": "2024-06-30T15:30:22.027459Z",
     "shell.execute_reply": "2024-06-30T15:30:22.026647Z"
    },
    "papermill": {
     "duration": 67.020562,
     "end_time": "2024-06-30T15:30:22.029410",
     "exception": false,
     "start_time": "2024-06-30T15:29:15.008848",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Attaching 'model.safetensors' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'model.safetensors.index.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'metadata.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'metadata.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'task.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'config.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'model.safetensors' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'model.safetensors.index.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'metadata.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'metadata.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'config.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'config.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'model.weights.h5' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'model.safetensors' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'model.safetensors.index.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'metadata.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'metadata.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'preprocessor.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'tokenizer.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'tokenizer.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "Attaching 'assets/tokenizer/vocabulary.spm' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\n",
      "normalizer.cc(51) LOG(INFO) precompiled_charsmap is empty. use identity normalization.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Preprocessor: \"gemma_causal_lm_preprocessor\"</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1mPreprocessor: \"gemma_causal_lm_preprocessor\"\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\"> Tokenizer (type)                                   </span>┃<span style=\"font-weight: bold\">                                             Vocab # </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│ gemma_tokenizer (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">GemmaTokenizer</span>)                   │                                             <span style=\"color: #00af00; text-decoration-color: #00af00\">256,000</span> │\n",
       "└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1mTokenizer (type)                                  \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m                                            Vocab #\u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│ gemma_tokenizer (\u001b[38;5;33mGemmaTokenizer\u001b[0m)                   │                                             \u001b[38;5;34m256,000\u001b[0m │\n",
       "└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Model: \"gemma_causal_lm\"</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1mModel: \"gemma_causal_lm\"\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\"> Layer (type)                  </span>┃<span style=\"font-weight: bold\"> Output Shape              </span>┃<span style=\"font-weight: bold\">         Param # </span>┃<span style=\"font-weight: bold\"> Connected to               </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│ padding_mask (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)     │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ token_ids (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)        │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ gemma_backbone                │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">2048</span>)        │   <span style=\"color: #00af00; text-decoration-color: #00af00\">2,506,172,416</span> │ padding_mask[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>],        │\n",
       "│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">GemmaBackbone</span>)               │                           │                 │ token_ids[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>]            │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ token_embedding               │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">256000</span>)      │     <span style=\"color: #00af00; text-decoration-color: #00af00\">524,288,000</span> │ gemma_backbone[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>]       │\n",
       "│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">ReversibleEmbedding</span>)         │                           │                 │                            │\n",
       "└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1mLayer (type)                 \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mOutput Shape             \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m        Param #\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mConnected to              \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│ padding_mask (\u001b[38;5;33mInputLayer\u001b[0m)     │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ token_ids (\u001b[38;5;33mInputLayer\u001b[0m)        │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ gemma_backbone                │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m2048\u001b[0m)        │   \u001b[38;5;34m2,506,172,416\u001b[0m │ padding_mask[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m],        │\n",
       "│ (\u001b[38;5;33mGemmaBackbone\u001b[0m)               │                           │                 │ token_ids[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m]            │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ token_embedding               │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m256000\u001b[0m)      │     \u001b[38;5;34m524,288,000\u001b[0m │ gemma_backbone[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m]       │\n",
       "│ (\u001b[38;5;33mReversibleEmbedding\u001b[0m)         │                           │                 │                            │\n",
       "└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Total params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">2,506,172,416</span> (9.34 GB)\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m Total params: \u001b[0m\u001b[38;5;34m2,506,172,416\u001b[0m (9.34 GB)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">2,506,172,416</span> (9.34 GB)\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m Trainable params: \u001b[0m\u001b[38;5;34m2,506,172,416\u001b[0m (9.34 GB)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Non-trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> (0.00 B)\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m Non-trainable params: \u001b[0m\u001b[38;5;34m0\u001b[0m (0.00 B)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "gemma_causal_lm = keras_nlp.models.GemmaCausalLM.from_preset(\"gemma_2b_en\")\n",
    "gemma_causal_lm.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e14d3524",
   "metadata": {
    "papermill": {
     "duration": 0.014582,
     "end_time": "2024-06-30T15:30:22.058777",
     "exception": false,
     "start_time": "2024-06-30T15:30:22.044195",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Define the specialized class\n",
    "\n",
    "Here we define the special class `GemmaQA`. \n",
    "in the `__init__` we pass the `GemmaCausalLM` object created before.\n",
    "The `query` member function uses `GemmaCausalLM` member function `generate` to generate the answer, based on a prompt that includes the category and the question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7932e6dc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:30:22.089851Z",
     "iopub.status.busy": "2024-06-30T15:30:22.089530Z",
     "iopub.status.idle": "2024-06-30T15:30:22.095456Z",
     "shell.execute_reply": "2024-06-30T15:30:22.094622Z"
    },
    "papermill": {
     "duration": 0.023836,
     "end_time": "2024-06-30T15:30:22.097280",
     "exception": false,
     "start_time": "2024-06-30T15:30:22.073444",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class GemmaQA:\n",
    "    def __init__(self, max_length=512):\n",
    "        self.max_length = max_length\n",
    "        self.prompt = template\n",
    "        self.gemma_causal_lm = gemma_causal_lm\n",
    "        \n",
    "    def query(self, category, question):\n",
    "        response = self.gemma_causal_lm.generate(\n",
    "            self.prompt.format(\n",
    "                Category=category,\n",
    "                Question=question,\n",
    "                Answer=\"\"), \n",
    "            max_length=self.max_length)\n",
    "        display(Markdown(colorize_text(response)))\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf1bc092",
   "metadata": {
    "papermill": {
     "duration": 0.014361,
     "end_time": "2024-06-30T15:30:22.126237",
     "exception": false,
     "start_time": "2024-06-30T15:30:22.111876",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Gemma preprocessor\n",
    "\n",
    "\n",
    "This preprocessing layer will take in batches of strings, and return outputs in a ```(x, y, sample_weight)``` format, where the y label is the next token id in the x sequence.\n",
    "\n",
    "From the code below, we can see that, after the preprocessor, the data shape is ```(num_samples, sequence_length)```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "62ff6211",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:30:22.156566Z",
     "iopub.status.busy": "2024-06-30T15:30:22.156301Z",
     "iopub.status.idle": "2024-06-30T15:30:22.572828Z",
     "shell.execute_reply": "2024-06-30T15:30:22.572061Z"
    },
    "papermill": {
     "duration": 0.434389,
     "end_time": "2024-06-30T15:30:22.575119",
     "exception": false,
     "start_time": "2024-06-30T15:30:22.140730",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "x, y, sample_weight = gemma_causal_lm.preprocessor(data[0:2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "08a0f630",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:30:22.606203Z",
     "iopub.status.busy": "2024-06-30T15:30:22.605894Z",
     "iopub.status.idle": "2024-06-30T15:30:22.611988Z",
     "shell.execute_reply": "2024-06-30T15:30:22.610918Z"
    },
    "papermill": {
     "duration": 0.023789,
     "end_time": "2024-06-30T15:30:22.613835",
     "exception": false,
     "start_time": "2024-06-30T15:30:22.590046",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'token_ids': Array([[   2,  109, 8606, ...,    0,    0,    0],\n",
      "       [   2,  109, 8606, ...,    0,    0,    0]], dtype=int32), 'padding_mask': Array([[ True,  True,  True, ..., False, False, False],\n",
      "       [ True,  True,  True, ..., False, False, False]], dtype=bool)} [[   109   8606 235292 ...      0      0      0]\n",
      " [   109   8606 235292 ...      0      0      0]]\n"
     ]
    }
   ],
   "source": [
    "print(x, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e51a4cd1",
   "metadata": {
    "papermill": {
     "duration": 0.014522,
     "end_time": "2024-06-30T15:30:22.642980",
     "exception": false,
     "start_time": "2024-06-30T15:30:22.628458",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Perform fine-tuning with LoRA"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ccaeca5",
   "metadata": {
    "papermill": {
     "duration": 0.014506,
     "end_time": "2024-06-30T15:30:22.672031",
     "exception": false,
     "start_time": "2024-06-30T15:30:22.657525",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Enable LoRA for the model\n",
    "\n",
    "LoRA rank is setting the number of trainable parameters. A larger rank will result in a larger number of parameters to train."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "da9f27e9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:30:22.702301Z",
     "iopub.status.busy": "2024-06-30T15:30:22.701995Z",
     "iopub.status.idle": "2024-06-30T15:30:23.156083Z",
     "shell.execute_reply": "2024-06-30T15:30:23.155201Z"
    },
    "papermill": {
     "duration": 0.471411,
     "end_time": "2024-06-30T15:30:23.157894",
     "exception": false,
     "start_time": "2024-06-30T15:30:22.686483",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Preprocessor: \"gemma_causal_lm_preprocessor\"</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1mPreprocessor: \"gemma_causal_lm_preprocessor\"\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\"> Tokenizer (type)                                   </span>┃<span style=\"font-weight: bold\">                                             Vocab # </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│ gemma_tokenizer (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">GemmaTokenizer</span>)                   │                                             <span style=\"color: #00af00; text-decoration-color: #00af00\">256,000</span> │\n",
       "└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1mTokenizer (type)                                  \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m                                            Vocab #\u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│ gemma_tokenizer (\u001b[38;5;33mGemmaTokenizer\u001b[0m)                   │                                             \u001b[38;5;34m256,000\u001b[0m │\n",
       "└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Model: \"gemma_causal_lm\"</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1mModel: \"gemma_causal_lm\"\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\"> Layer (type)                  </span>┃<span style=\"font-weight: bold\"> Output Shape              </span>┃<span style=\"font-weight: bold\">         Param # </span>┃<span style=\"font-weight: bold\"> Connected to               </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│ padding_mask (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)     │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ token_ids (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)        │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ gemma_backbone                │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">2048</span>)        │   <span style=\"color: #00af00; text-decoration-color: #00af00\">2,507,536,384</span> │ padding_mask[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>],        │\n",
       "│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">GemmaBackbone</span>)               │                           │                 │ token_ids[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>]            │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ token_embedding               │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">256000</span>)      │     <span style=\"color: #00af00; text-decoration-color: #00af00\">524,288,000</span> │ gemma_backbone[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>]       │\n",
       "│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">ReversibleEmbedding</span>)         │                           │                 │                            │\n",
       "└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1mLayer (type)                 \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mOutput Shape             \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m        Param #\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mConnected to              \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│ padding_mask (\u001b[38;5;33mInputLayer\u001b[0m)     │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ token_ids (\u001b[38;5;33mInputLayer\u001b[0m)        │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ gemma_backbone                │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m2048\u001b[0m)        │   \u001b[38;5;34m2,507,536,384\u001b[0m │ padding_mask[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m],        │\n",
       "│ (\u001b[38;5;33mGemmaBackbone\u001b[0m)               │                           │                 │ token_ids[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m]            │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ token_embedding               │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m256000\u001b[0m)      │     \u001b[38;5;34m524,288,000\u001b[0m │ gemma_backbone[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m]       │\n",
       "│ (\u001b[38;5;33mReversibleEmbedding\u001b[0m)         │                           │                 │                            │\n",
       "└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Total params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">2,507,536,384</span> (9.34 GB)\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m Total params: \u001b[0m\u001b[38;5;34m2,507,536,384\u001b[0m (9.34 GB)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">1,363,968</span> (5.20 MB)\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m Trainable params: \u001b[0m\u001b[38;5;34m1,363,968\u001b[0m (5.20 MB)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Non-trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">2,506,172,416</span> (9.34 GB)\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m Non-trainable params: \u001b[0m\u001b[38;5;34m2,506,172,416\u001b[0m (9.34 GB)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Enable LoRA for the model and set the LoRA rank to the lora_rank as set in Config (4).\n",
    "gemma_causal_lm.backbone.enable_lora(rank=Config.lora_rank)\n",
    "gemma_causal_lm.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a784deb5",
   "metadata": {
    "papermill": {
     "duration": 0.015639,
     "end_time": "2024-06-30T15:30:23.189843",
     "exception": false,
     "start_time": "2024-06-30T15:30:23.174204",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We see that only a small part of the parameters are trainable. 2.5 billions parameters total, and only 1.3 Millions parameters trainable."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4aea04c9",
   "metadata": {
    "papermill": {
     "duration": 0.015613,
     "end_time": "2024-06-30T15:30:23.221206",
     "exception": false,
     "start_time": "2024-06-30T15:30:23.205593",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Run the training sequence\n",
    "\n",
    "We set the `sequence_length` for the `GemmaCausalLM` (from configuration, will be 512).\n",
    "We compile the model, with the loss, optimizer and metric.\n",
    "For the metric, it is used `SparseCategoricalAccuracy`. This metric calculates how often predictions match integer labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "01636240",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:30:23.253712Z",
     "iopub.status.busy": "2024-06-30T15:30:23.253415Z",
     "iopub.status.idle": "2024-06-30T15:42:23.009442Z",
     "shell.execute_reply": "2024-06-30T15:42:23.008524Z"
    },
    "papermill": {
     "duration": 719.77451,
     "end_time": "2024-06-30T15:42:23.011378",
     "exception": false,
     "start_time": "2024-06-30T15:30:23.236868",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/15\n",
      "\u001b[1m60/60\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m65s\u001b[0m 735ms/step - loss: 1.7209 - sparse_categorical_accuracy: 0.5241\n",
      "Epoch 2/15\n",
      "\u001b[1m60/60\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m83s\u001b[0m 1s/step - loss: 1.6869 - sparse_categorical_accuracy: 0.5313\n",
      "Epoch 3/15\n",
      "\u001b[1m60/60\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m44s\u001b[0m 728ms/step - loss: 1.6175 - sparse_categorical_accuracy: 0.5417\n",
      "Epoch 4/15\n",
      "\u001b[1m60/60\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m44s\u001b[0m 729ms/step - loss: 1.5770 - sparse_categorical_accuracy: 0.5509\n",
      "Epoch 5/15\n",
      "\u001b[1m60/60\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m44s\u001b[0m 729ms/step - loss: 1.5537 - sparse_categorical_accuracy: 0.5552\n",
      "Epoch 6/15\n",
      "\u001b[1m60/60\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m44s\u001b[0m 729ms/step - loss: 1.5304 - sparse_categorical_accuracy: 0.5568\n",
      "Epoch 7/15\n",
      "\u001b[1m60/60\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m44s\u001b[0m 728ms/step - loss: 1.5028 - sparse_categorical_accuracy: 0.5630\n",
      "Epoch 8/15\n",
      "\u001b[1m60/60\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m44s\u001b[0m 728ms/step - loss: 1.4733 - sparse_categorical_accuracy: 0.5682\n",
      "Epoch 9/15\n",
      "\u001b[1m60/60\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m44s\u001b[0m 728ms/step - loss: 1.4444 - sparse_categorical_accuracy: 0.5747\n",
      "Epoch 10/15\n",
      "\u001b[1m60/60\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m44s\u001b[0m 728ms/step - loss: 1.4025 - sparse_categorical_accuracy: 0.5873\n",
      "Epoch 11/15\n",
      "\u001b[1m60/60\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m44s\u001b[0m 728ms/step - loss: 1.3607 - sparse_categorical_accuracy: 0.5960\n",
      "Epoch 12/15\n",
      "\u001b[1m60/60\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m44s\u001b[0m 728ms/step - loss: 1.3163 - sparse_categorical_accuracy: 0.6079\n",
      "Epoch 13/15\n",
      "\u001b[1m60/60\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m44s\u001b[0m 728ms/step - loss: 1.2684 - sparse_categorical_accuracy: 0.6198\n",
      "Epoch 14/15\n",
      "\u001b[1m60/60\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m44s\u001b[0m 728ms/step - loss: 1.2149 - sparse_categorical_accuracy: 0.6325\n",
      "Epoch 15/15\n",
      "\u001b[1m60/60\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m44s\u001b[0m 728ms/step - loss: 1.1585 - sparse_categorical_accuracy: 0.6451\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.src.callbacks.history.History at 0x7ce99c332800>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#set sequence length cf. config (512)\n",
    "gemma_causal_lm.preprocessor.sequence_length = Config.sequence_length \n",
    "\n",
    "# Compile the model with loss, optimizer, and metric\n",
    "gemma_causal_lm.compile(\n",
    "    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
    "    optimizer=keras.optimizers.Adam(learning_rate=Config.learning_rate),\n",
    "    weighted_metrics=[keras.metrics.SparseCategoricalAccuracy()],\n",
    ")\n",
    "\n",
    "# Train model\n",
    "gemma_causal_lm.fit(data, epochs=Config.epochs, batch_size=Config.batch_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da9acf5c",
   "metadata": {
    "papermill": {
     "duration": 0.087939,
     "end_time": "2024-06-30T15:42:23.187800",
     "exception": false,
     "start_time": "2024-06-30T15:42:23.099861",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Test the fine-tuned model\n",
    "\n",
    "We instantiate an object of class GemmaQA. Because `gemma_causal_lm` was fine-tuned using LoRA, `gemma_qa` defined here will use the fine-tuned model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a2a043b4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:42:23.380738Z",
     "iopub.status.busy": "2024-06-30T15:42:23.380113Z",
     "iopub.status.idle": "2024-06-30T15:42:23.384456Z",
     "shell.execute_reply": "2024-06-30T15:42:23.383615Z"
    },
    "papermill": {
     "duration": 0.108343,
     "end_time": "2024-06-30T15:42:23.386268",
     "exception": false,
     "start_time": "2024-06-30T15:42:23.277925",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "gemma_qa = GemmaQA()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1e25876",
   "metadata": {
    "papermill": {
     "duration": 0.095407,
     "end_time": "2024-06-30T15:42:23.579541",
     "exception": false,
     "start_time": "2024-06-30T15:42:23.484134",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "For start, we are testing the model with some of the data from the training set itself."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e1a933b",
   "metadata": {
    "papermill": {
     "duration": 0.088466,
     "end_time": "2024-06-30T15:42:23.756729",
     "exception": false,
     "start_time": "2024-06-30T15:42:23.668263",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Sample 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b57b2892",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:42:23.934569Z",
     "iopub.status.busy": "2024-06-30T15:42:23.934212Z",
     "iopub.status.idle": "2024-06-30T15:42:43.322742Z",
     "shell.execute_reply": "2024-06-30T15:42:43.321842Z"
    },
    "papermill": {
     "duration": 19.479903,
     "end_time": "2024-06-30T15:42:43.324967",
     "exception": false,
     "start_time": "2024-06-30T15:42:23.845064",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "\n",
       "**<font color='blue'>Category:</font>**\n",
       "kaggle-competition\n",
       "\n",
       "**<font color='red'>Question:</font>**\n",
       "What are the different types of competitions available on Kaggle?\n",
       "\n",
       "**<font color='green'>Answer:</font>**\n",
       "## Types of Competitions\n",
       "\n",
       "There are two main types of competitions on Kaggle:\n",
       "\n",
       "- **Data science challenges**: These are typically data science challenges where the goal is to build a model that outperforms a provided baseline. Examples include building a model to predict the likelihood of a fraudulent transaction or building a model to predict the likelihood of a heart attack.\n",
       "\n",
       "- **Classification challenges**: These are typically classification challenges where the goal is to build a model that outperforms a provided baseline. Examples include building a model to classify emails as spam or ham or building a model to classify images as cats or dogs.\n",
       "\n",
       "## How Are the Kaggle Competitions Ranked?\n",
       "\n",
       "The rankings on Kaggle are based on a combination of two metrics:\n",
       "\n",
       "- **Submissions**: The total number of submissions made by a user.\n",
       "\n",
       "- **Leaderboard Rank**: The rank of a user on the leaderboard of a given competition.\n",
       "\n",
       "The leaderboard rank is calculated by taking the median rank of the top 100 users on the leaderboard. The median rank is calculated by first sorting the leaderboard by leaderboard rank and then taking the middle value of the leaderboard.\n",
       "\n",
       "The combination of these two metrics is designed to encourage users to make high-quality submissions and to reward users who are consistently at the top of the leaderboard."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "row = df.iloc[0]\n",
    "gemma_qa.query(row.Category,row.Question)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b11ce04",
   "metadata": {
    "papermill": {
     "duration": 0.107633,
     "end_time": "2024-06-30T15:42:43.527880",
     "exception": false,
     "start_time": "2024-06-30T15:42:43.420247",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Sample 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "76db8d50",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:42:43.706552Z",
     "iopub.status.busy": "2024-06-30T15:42:43.705862Z",
     "iopub.status.idle": "2024-06-30T15:42:57.220689Z",
     "shell.execute_reply": "2024-06-30T15:42:57.219705Z"
    },
    "papermill": {
     "duration": 13.605999,
     "end_time": "2024-06-30T15:42:57.222774",
     "exception": false,
     "start_time": "2024-06-30T15:42:43.616775",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "\n",
       "**<font color='blue'>Category:</font>**\n",
       "kaggle-tpu\n",
       "\n",
       "**<font color='red'>Question:</font>**\n",
       "How to load and save model on TPU?\n",
       "\n",
       "**<font color='green'>Answer:</font>**\n",
       "## Overview\n",
       "\n",
       "TPU models are saved in SavedModel format. To load a model from SavedModel, you can use the `tpu.Model.from_saved_model` method.\n",
       "\n",
       "To save a model, you can use the `tpu.Model.save_main_session()` method.\n",
       "\n",
       "## Example\n",
       "\n",
       "```python\n",
       "# Load a model from SavedModel\n",
       "model = tpu.Model.from_saved_model(tpu.saved_model_fn)\n",
       "```\n",
       "\n",
       "## TPU Model Format\n",
       "\n",
       "TPU models are saved in SavedModel format. To load a model from SavedModel, you can use the tpu.Model.from_saved_model method.\n",
       "\n",
       "To save a model, you can use the tpu.Model.save_main_session() method.\n",
       "\n",
       "## TPU Model Format Overview\n",
       "\n",
       "The SavedModel format is a model archive that contains all the information needed to run a model on TPU. It is a zip archive with the following structure:\n",
       "\n",
       "- `model.pb` is the TensorFlow graph that describes the model.\n",
       "- `checkpoint.index` is a binary file that contains the metadata for the checkpoints in the model.\n",
       "- `config.json` is a JSON file that contains the configuration for the model.\n",
       "- `variables.add_file_io_ops.json` is a JSON file that contains the metadata for the file I/O ops in the model.\n",
       "- `variables.add_file_io_ops.pb` is a binary file that contains the code for the file I/O ops in the model.\n",
       "- `variables.add_op_for_each_checkpoint.index` is a binary file that contains the metadata for the checkpoints in the model.\n",
       "- `variables.add_op_for_each_checkpoint.json` is a JSON file that contains the metadata for the checkpoints in the model.\n",
       "- `variables.add_op_for_each_checkpoint.pb` is a binary file that contains the code for the checkpoints in the model.\n",
       "- `variables.add_op_for_each_checkpoint.summary.tsv` is a tab-separated-values file that contains the summary of the checkpoints in the model.\n",
       "- `variables.add_op_for_each_checkpoint.summary.tsv` is a tab-separated-values"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "row = df.iloc[15]\n",
    "gemma_qa.query(row.Category,row.Question)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6370b91",
   "metadata": {
    "papermill": {
     "duration": 0.089198,
     "end_time": "2024-06-30T15:42:57.401650",
     "exception": false,
     "start_time": "2024-06-30T15:42:57.312452",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Sample 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "2028ab33",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:42:57.581162Z",
     "iopub.status.busy": "2024-06-30T15:42:57.580797Z",
     "iopub.status.idle": "2024-06-30T15:43:10.988169Z",
     "shell.execute_reply": "2024-06-30T15:43:10.987212Z"
    },
    "papermill": {
     "duration": 13.499385,
     "end_time": "2024-06-30T15:43:10.990134",
     "exception": false,
     "start_time": "2024-06-30T15:42:57.490749",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "\n",
       "**<font color='blue'>Category:</font>**\n",
       "kaggle-noteboook\n",
       "\n",
       "**<font color='red'>Question:</font>**\n",
       "What are the different types of notebooks available on Kaggle?\n",
       "\n",
       "**<font color='green'>Answer:</font>**\n",
       "## Notebooks\n",
       "\n",
       "Kaggle Notebooks are interactive markdown documents that allow users to run code and see the results. Notebooks are a key part of the Kaggle experience, and are used by the community to share data science workflows, discuss techniques, and explore new ideas.\n",
       "\n",
       "Kaggle Notebooks are hosted on nbviewer.com. You can access them from any web browser, or from the Kaggle mobile app.\n",
       "\n",
       "You can create a Notebook from the Kaggle website or from the command line using the kaggle notebook command.\n",
       "\n",
       "## Versions\n",
       "\n",
       "Each Notebook has multiple versions. Each version is a snapshot of the Notebook at a specific point in time. You can view and compare the versions of a Notebook to see how it has changed over time.\n",
       "\n",
       "You can access the version history of a Notebook from the Notebook editor.\n",
       "\n",
       "You can create a new version of a Notebook from the Notebook editor.\n",
       "\n",
       "You can delete a Notebook version from the Notebook editor.\n",
       "\n",
       "## Versions and Sharing\n",
       "\n",
       "Each Notebook version has a URL that looks like this:\n",
       "\n",
       "```\n",
       "https://www.kaggle.com/kylegordon/kaggle-noteboook-versions?version=1666666\n",
       "```\n",
       "\n",
       "You can share a Notebook version URL with others, and they will be prompted to view the Notebook as a guest.\n",
       "\n",
       "You can also create a Notebook version from a URL that someone else has shared with you.\n",
       "\n",
       "## Versions and Permissions\n",
       "\n",
       "Each Notebook version has a URL that looks like this:\n",
       "\n",
       "```\n",
       "https://www.kaggle.com/kylegordon/kaggle-noteboook-versions?version=1666666\n",
       "```\n",
       "\n",
       "You can access the version history of a Notebook from the Notebook editor.\n",
       "\n",
       "You can create a new version of a Notebook from the Notebook editor.\n",
       "\n",
       "You can delete a Notebook version from the Notebook editor.\n",
       "\n",
       "You can view the Notebook from any web browser, or from the Kaggle mobile app.\n",
       "\n",
       "## Notebooks and Privacy\n",
       "\n",
       "By default, Notebooks are public. You can make a Notebook private by toggling the “Make Public” button to “Private”.\n",
       "\n",
       "If you make a Notebook private, anyone will need to be invited to view the Notebook. You can find the private Notebook URL in the URL bar of your browser.\n",
       "\n",
       "If you make a Notebook private, anyone will need to be invited to"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "row = df.iloc[25]\n",
    "gemma_qa.query(row.Category,row.Question)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34510d2c",
   "metadata": {
    "papermill": {
     "duration": 0.089865,
     "end_time": "2024-06-30T15:43:11.170404",
     "exception": false,
     "start_time": "2024-06-30T15:43:11.080539",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Not seen question(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "e8df7f02",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:43:11.351062Z",
     "iopub.status.busy": "2024-06-30T15:43:11.350678Z",
     "iopub.status.idle": "2024-06-30T15:43:24.974412Z",
     "shell.execute_reply": "2024-06-30T15:43:24.973428Z"
    },
    "papermill": {
     "duration": 13.716244,
     "end_time": "2024-06-30T15:43:24.976563",
     "exception": false,
     "start_time": "2024-06-30T15:43:11.260319",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "\n",
       "**<font color='blue'>Category:</font>**\n",
       "kaggle-notebook\n",
       "\n",
       "**<font color='red'>Question:</font>**\n",
       "How to run a notebook?\n",
       "\n",
       "**<font color='green'>Answer:</font>**\n",
       "## Installation\n",
       "\n",
       "You can run notebooks on Kaggle without installing anything, but if you want to make use of the community-built add-ons marketplace or you have a team of users that you want to collaborate with on a notebook, then you should consider creating a notebook on Kaggle.\n",
       "\n",
       "To install Kaggle and Notebook features in your own conda environment, run:\n",
       "\n",
       "```\n",
       "kaggle datasets login-or-create <USERNAME>\n",
       "kaggle kernels --help\n",
       "```\n",
       "\n",
       "For a full list of command line arguments, run `kaggle kernels help <command` from the Kaggle command line.\n",
       "\n",
       "## Run Locally\n",
       "\n",
       "If you are running locally, you can run the notebook locally by running the following command from the notebook directory:\n",
       "\n",
       "```python\n",
       "!kaggle run --kernels /kaggle/notebooks/tutorial-saving-and-loading-models: Cells 1-X\n",
       "```\n",
       "\n",
       "The command will run the notebook from the first cell to the most recently executed cell. If you want to run a specific cell, you can specify the cell number after the `--kernels` flag.\n",
       "\n",
       "If you want to run the notebook from scratch, you can run the notebook locally by running the following command from the notebook directory:\n",
       "\n",
       "```python\n",
       "kaggle run --kernels <USERNAME>/notebooks/tutorial-saving-and-loading-models:latest\n",
       "```\n",
       "\n",
       "The command will run the notebook from the first cell to the most recently executed cell.\n",
       "\n",
       "If you want to run a specific cell, you can specify the cell number after the `--kernels` flag.\n",
       "\n",
       "If you want to run the notebook from scratch, you can run the notebook locally by running the following command from the notebook directory:\n",
       "\n",
       "```python\n",
       "kaggle run --kernels <USERNAME>/notebooks/tutorial-saving-and-loading-models:latest\n",
       "```\n",
       "\n",
       "## Run on Kaggle\n",
       "\n",
       "If you are running on Kaggle, you can run the notebook on Kaggle by clicking the “Run Notebook” button in the top right corner of the notebook editor.\n",
       "\n",
       "If you are running locally, you can upload the notebook to Kaggle and run the notebook locally by running the following command from the notebook directory:\n",
       "\n",
       "```python\n",
       "kaggle runnotebooks <USERNAME>/notebooks/tutorial-saving-and-loading-models:latest\n",
       "```\n",
       "\n",
       "If you are running on Kaggle, you can run the notebook"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "category = \"notebook\"\n",
    "question = \"How to run a notebook?\"\n",
    "gemma_qa.query(category,question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "257f40e8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:43:25.161564Z",
     "iopub.status.busy": "2024-06-30T15:43:25.160738Z",
     "iopub.status.idle": "2024-06-30T15:43:29.396227Z",
     "shell.execute_reply": "2024-06-30T15:43:29.395280Z"
    },
    "papermill": {
     "duration": 4.329437,
     "end_time": "2024-06-30T15:43:29.398303",
     "exception": false,
     "start_time": "2024-06-30T15:43:25.068866",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "\n",
       "**<font color='blue'>Category:</font>**\n",
       "kaggle-discussions\n",
       "\n",
       "**<font color='red'>Question:</font>**\n",
       "How to create a discussion topic?\n",
       "\n",
       "**<font color='green'>Answer:</font>**\n",
       "## Create a Discussion\n",
       "\n",
       "To create a new discussion, click on the \"Discussions\" tab and click on the \"New Discussion\" button.\n",
       "\n",
       "You can create a new discussion by either:\n",
       "\n",
       "- Selecting a template from the \"New Discussion\" dropdown menu and clicking \"Create Discussion\"\n",
       "- Clicking \"Create Discussion\" directly on a blank discussion page\n",
       "\n",
       "You can also create a new discussion from the \"Discussions\" tab by clicking on the \"New Discussion\" button.\n",
       "\n",
       "## Edit a Discussion\n",
       "\n",
       "To edit an existing discussion, click on the title of the discussion you want to edit.\n",
       "\n",
       "You can also edit a discussion from the \"Discussions\" tab by clicking on the title of the discussion you want to edit."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "category = \"discussions\"\n",
    "question = \"How to create a discussion topic?\"\n",
    "gemma_qa.query(category,question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "edf45b55",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:43:29.581917Z",
     "iopub.status.busy": "2024-06-30T15:43:29.581162Z",
     "iopub.status.idle": "2024-06-30T15:43:43.177918Z",
     "shell.execute_reply": "2024-06-30T15:43:43.176950Z"
    },
    "papermill": {
     "duration": 13.690486,
     "end_time": "2024-06-30T15:43:43.179964",
     "exception": false,
     "start_time": "2024-06-30T15:43:29.489478",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "\n",
       "**<font color='blue'>Category:</font>**\n",
       "kaggle-competitions\n",
       "\n",
       "**<font color='red'>Question:</font>**\n",
       "What is a code competition?\n",
       "\n",
       "**<font color='green'>Answer:</font>**\n",
       "Code competitions are a type of competition where participants are given a dataset and a set of tools to build machine learning models. The goal is to build the best model for a given task. Code competitions are a popular format for Kaggle Datasets because they encourage participants to use the tools and techniques available on Kaggle.\n",
       "\n",
       "Code competitions are a great way to learn new tools and techniques, and to see how they can be applied to a real-world problem. They also provide a structured way to evaluate and improve your model, as you'll be competing against other participants who are using the same dataset and tools.\n",
       "\n",
       "There are many benefits to participating in code competitions. First, they're a great way to learn how to use a dataset and a task to drive model development. Second, they're a great way to improve your model's performance. Third, they're a great way to meet other data scientists and machine learning enthusiasts. Fourth, they're a lot of fun!\n",
       "\n",
       "## Types of Code Competitions\n",
       "\n",
       "There are two main types of code competitions on Kaggle: public and private.\n",
       "\n",
       "Public competitions are open to anyone and everyone. They're a great way to see what other people are working on in the community, and to learn from others' experiences. Public competitions are a great way to get exposure and recognition for your work.\n",
       "\n",
       "Private competitions are invitation-only and typically invitees are part of a larger community or organization. They're a great way to collaborate with others who share your interest in the same topic. Private competitions are a great way to get feedback on your work from experts in the field.\n",
       "\n",
       "## How to Participate in a Code Competition\n",
       "\n",
       "To participate in a code competition, you'll need to create a competition notebook. This notebook will contain all of the code you use to build your model. You can then submit your notebook to the competition.\n",
       "\n",
       "The first step is to find a competition that interests you. You can browse all of the current competitions on Kaggle Datasets by visiting the [code competitions page](https://www.kaggle.com/competitions/code/overview).\n",
       "\n",
       "Once you've found a competition you'd like to enter, click on the \"Participate\" button. This will take you to a page where you can sign up for the competition.\n",
       "\n",
       "The next step is to familiarize yourself with the dataset. This is the collection"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "category = \"competitions\"\n",
    "question = \"What is a code competition?\"\n",
    "gemma_qa.query(category,question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "13e50094",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:43:43.367049Z",
     "iopub.status.busy": "2024-06-30T15:43:43.366204Z",
     "iopub.status.idle": "2024-06-30T15:43:56.852572Z",
     "shell.execute_reply": "2024-06-30T15:43:56.851674Z"
    },
    "papermill": {
     "duration": 13.58407,
     "end_time": "2024-06-30T15:43:56.854553",
     "exception": false,
     "start_time": "2024-06-30T15:43:43.270483",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "\n",
       "**<font color='blue'>Category:</font>**\n",
       "kaggle-datasets\n",
       "\n",
       "**<font color='red'>Question:</font>**\n",
       "What are the steps to create a Kaggle dataset?\n",
       "\n",
       "**<font color='green'>Answer:</font>**\n",
       "## Introduction\n",
       "\n",
       "Kaggle Datasets is a repository of public datasets that are created by members of the Kaggle community. Datasets are used by the Kaggle community for a variety of purposes, including training machine learning models, evaluating models, and sharing data.\n",
       "\n",
       "Kaggle Datasets is a great resource for anyone working with data.\n",
       "\n",
       "## Getting Started\n",
       "\n",
       "To get started creating a new dataset, navigate to the \"Create Dataset\" button in the left-hand navigation bar.\n",
       "\n",
       "You can also access the \"Create Dataset\" page from your profile by clicking on your username in the top-right corner of any page on Kaggle and selecting \"Create Dataset\" from the drop-down menu.\n",
       "\n",
       "There are two ways to create a new dataset:\n",
       "\n",
       "- Upload a file: If you have a large dataset that you want to upload, this is the way to go. Simply browse to the location of the file on your computer and select \"Upload\" to begin the upload process.\n",
       "- Generate a new dataset: If you don't have a large dataset to upload, this is the way to go. You can choose from a variety of template datasets to get started, or you can start from scratch by selecting \"Create Blank Dataset\" and editing the dataset settings to your liking.\n",
       "\n",
       "Once you have created your dataset, you will be taken to the \"Manage Dataset\" page. Here, you can edit the metadata (e.g. title, description, tags), permissions (who can view and edit the dataset), and the dataset files themselves.\n",
       "\n",
       "## Analytics\n",
       "\n",
       "Analytics are available to help you understand how your dataset is being used. You can see the number of downloads, clones, and edits your dataset has received over time, as well as the number of unique users who have accessed the dataset.\n",
       "\n",
       "To access analytics, navigate to the \"Manage Dataset\" page and click on the \"Analytics\" tab.\n",
       "\n",
       "## Sharing\n",
       "\n",
       "Kaggle Datasets is a community of data scientists, machine learning engineers, and data enthusiasts. By sharing your datasets with others in the community, you can help advance the state of the art in machine learning and data science.\n",
       "\n",
       "To share your dataset, navigate to the \"Manage Dataset\" page and click on the \"Share\" tab. Here, you can choose from a variety of sharing options, including public, private, and unlisted.\n",
       "\n",
       "## Search\n",
       "\n",
       "Kaggle Datasets is a"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "category = \"datasets\"\n",
    "question = \"What are the steps to create a Kaggle dataset?\"\n",
    "gemma_qa.query(category,question)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc14e3b2",
   "metadata": {
    "papermill": {
     "duration": 0.089043,
     "end_time": "2024-06-30T15:43:57.034166",
     "exception": false,
     "start_time": "2024-06-30T15:43:56.945123",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Save the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "b81425da",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-06-30T15:43:57.258633Z",
     "iopub.status.busy": "2024-06-30T15:43:57.258279Z",
     "iopub.status.idle": "2024-06-30T15:45:31.377134Z",
     "shell.execute_reply": "2024-06-30T15:45:31.372497Z"
    },
    "papermill": {
     "duration": 94.216508,
     "end_time": "2024-06-30T15:45:31.382721",
     "exception": false,
     "start_time": "2024-06-30T15:43:57.166213",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "gemma_causal_lm.save(\"gemma_2b_en_kaggle_docs.keras\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "783bde85",
   "metadata": {
    "papermill": {
     "duration": 0.091805,
     "end_time": "2024-06-30T15:45:31.574549",
     "exception": false,
     "start_time": "2024-06-30T15:45:31.482744",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Conclusions\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79f32be4",
   "metadata": {
    "papermill": {
     "duration": 0.090925,
     "end_time": "2024-06-30T15:45:31.759038",
     "exception": false,
     "start_time": "2024-06-30T15:45:31.668113",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We demonstated how to fine-tune a Gemma model using LoRA.   \n",
    "We also created a class to run queries to the Gemma model and tested it with some examples from the existing training data but also with some new, not seen questions.   \n",
    "We also saved the models as a Keras model."
   ]
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "gpu",
   "dataSources": [
    {
     "databundleVersionId": 7669720,
     "sourceId": 64148,
     "sourceType": "competition"
    },
    {
     "datasetId": 4484051,
     "sourceId": 7711309,
     "sourceType": "datasetVersion"
    },
    {
     "modelInstanceId": 5171,
     "sourceId": 11371,
     "sourceType": "modelInstanceVersion"
    }
   ],
   "dockerImageVersionId": 30674,
   "isGpuEnabled": true,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 1029.683294,
   "end_time": "2024-06-30T15:45:36.334749",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2024-06-30T15:28:26.651455",
   "version": "2.5.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}