fine-tune-gemma-using-lora-and-keras.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "45bc477c",
   "metadata": {
    "papermill": {
     "duration": 0.00884,
     "end_time": "2024-11-12T21:55:13.962049",
     "exception": false,
     "start_time": "2024-11-12T21:55:13.953209",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Introduction\n",
    "\n",
    "\n",
    "This notebook is based on a tutorial from goo.gle/ai-kaggle-keras-gemma about how to fine-tune **Gemma** on a **Kaggle** dataset and share the model with the community.  \n",
    "\n",
    "## Model used\n",
    "\n",
    "We are using **Gemma** model with 2B parameters, Keras, English version, v2.\n",
    "\n",
    "## Dataset\n",
    "\n",
    "We will fine-tune **Gemma** using a [Medical Q & A](https://www.kaggle.com/datasets/gpreda/medquad/) dataset. This is a subset of the full public dataset [Healthcare NLP: LLMs, Transformers, Datasets](https://www.kaggle.com/datasets/jpmiller/layoutlm).\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a071760",
   "metadata": {
    "papermill": {
     "duration": 0.008254,
     "end_time": "2024-11-12T21:55:13.979908",
     "exception": false,
     "start_time": "2024-11-12T21:55:13.971654",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Prepare packages\n",
    "\n",
    "\n",
    "We will install updated version of Keras, KerasNLP, which we need for fine-tuning, and other dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "93a70eef",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-12T21:55:13.997825Z",
     "iopub.status.busy": "2024-11-12T21:55:13.997488Z",
     "iopub.status.idle": "2024-11-12T21:57:27.459019Z",
     "shell.execute_reply": "2024-11-12T21:57:27.457826Z"
    },
    "papermill": {
     "duration": 133.473451,
     "end_time": "2024-11-12T21:57:27.461551",
     "exception": false,
     "start_time": "2024-11-12T21:55:13.988100",
     "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.\r\n",
      "tensorflow-decision-forests 1.8.1 requires tensorflow~=2.15.0, but you have tensorflow 2.18.0 which is incompatible.\r\n",
      "tensorflow-text 2.15.0 requires tensorflow<2.16,>=2.15.0; platform_machine != \"arm64\" or platform_system != \"Darwin\", but you have tensorflow 2.18.0 which is incompatible.\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.\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.1 requires keras<2.16,>=2.15.0, but you have keras 3.6.0 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 tf-keras\n",
    "!pip install -q -U keras-nlp==0.10.0\n",
    "!pip install -q -U keras>=3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fc0dbdb",
   "metadata": {
    "papermill": {
     "duration": 0.008454,
     "end_time": "2024-11-12T21:57:27.479202",
     "exception": false,
     "start_time": "2024-11-12T21:57:27.470748",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "With Keras 3, we can run workflows on one of three backends: **TensorFlow**, **JAX**, and **PyTorch**.\n",
    "\n",
    "For this Notebook, we will configure the backend for **JAX**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "63e56b8c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-12T21:57:27.498975Z",
     "iopub.status.busy": "2024-11-12T21:57:27.498658Z",
     "iopub.status.idle": "2024-11-12T21:57:27.503774Z",
     "shell.execute_reply": "2024-11-12T21:57:27.502920Z"
    },
    "papermill": {
     "duration": 0.017899,
     "end_time": "2024-11-12T21:57:27.505616",
     "exception": false,
     "start_time": "2024-11-12T21:57:27.487717",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "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\"] = \"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66f2f061",
   "metadata": {
    "papermill": {
     "duration": 0.008441,
     "end_time": "2024-11-12T21:57:27.522620",
     "exception": false,
     "start_time": "2024-11-12T21:57:27.514179",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Once installed the additional packages, we can now include them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "89ab87ba",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-12T21:57:27.541861Z",
     "iopub.status.busy": "2024-11-12T21:57:27.541412Z",
     "iopub.status.idle": "2024-11-12T21:57:34.524008Z",
     "shell.execute_reply": "2024-11-12T21:57:34.523097Z"
    },
    "papermill": {
     "duration": 6.994541,
     "end_time": "2024-11-12T21:57:34.526132",
     "exception": false,
     "start_time": "2024-11-12T21:57:27.531591",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-11-12 21:57:28.019458: 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-11-12 21:57:28.019511: 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-11-12 21:57:28.021030: 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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "KerasNLP version:  0.10.0\n",
      "Keras version:  3.6.0\n"
     ]
    }
   ],
   "source": [
    "import keras_nlp\n",
    "import keras\n",
    "import csv\n",
    "\n",
    "print(\"KerasNLP version: \", keras_nlp.__version__)\n",
    "print(\"Keras version: \", keras.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8942845",
   "metadata": {
    "papermill": {
     "duration": 0.008753,
     "end_time": "2024-11-12T21:57:34.544043",
     "exception": false,
     "start_time": "2024-11-12T21:57:34.535290",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Load the model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbe86bcc",
   "metadata": {
    "papermill": {
     "duration": 0.008661,
     "end_time": "2024-11-12T21:57:34.561531",
     "exception": false,
     "start_time": "2024-11-12T21:57:34.552870",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We load the model `gemma_2b_en` using `keras_nlp`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "afe9c29a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-12T21:57:34.580739Z",
     "iopub.status.busy": "2024-11-12T21:57:34.580191Z",
     "iopub.status.idle": "2024-11-12T21:58:16.194209Z",
     "shell.execute_reply": "2024-11-12T21:58:16.193453Z"
    },
    "papermill": {
     "duration": 41.626169,
     "end_time": "2024-11-12T21:58:16.196550",
     "exception": false,
     "start_time": "2024-11-12T21:57:34.570381",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "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 '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 '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"
     ]
    }
   ],
   "source": [
    "gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset(\"gemma_2b_en\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bf0cc6b",
   "metadata": {
    "papermill": {
     "duration": 0.009964,
     "end_time": "2024-11-12T21:58:16.216923",
     "exception": false,
     "start_time": "2024-11-12T21:58:16.206959",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's print the summary of the model loaded."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3d4633a7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-12T21:58:16.238766Z",
     "iopub.status.busy": "2024-11-12T21:58:16.238463Z",
     "iopub.status.idle": "2024-11-12T21:58:16.266559Z",
     "shell.execute_reply": "2024-11-12T21:58:16.265729Z"
    },
    "papermill": {
     "duration": 0.041732,
     "end_time": "2024-11-12T21:58:16.268646",
     "exception": false,
     "start_time": "2024-11-12T21:58:16.226914",
     "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,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_lm.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17683caf",
   "metadata": {
    "papermill": {
     "duration": 0.010891,
     "end_time": "2024-11-12T21:58:16.290914",
     "exception": false,
     "start_time": "2024-11-12T21:58:16.280023",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Prepare the training data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9024f5cb",
   "metadata": {
    "papermill": {
     "duration": 0.01106,
     "end_time": "2024-11-12T21:58:16.313391",
     "exception": false,
     "start_time": "2024-11-12T21:58:16.302331",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We prepare the **Medical Q & A** data for training. We will load the data using the template where, for each data that will be included in the training set, we provide pairs of questions and answers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7cf84199",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-12T21:58:16.337383Z",
     "iopub.status.busy": "2024-11-12T21:58:16.336691Z",
     "iopub.status.idle": "2024-11-12T21:58:16.808806Z",
     "shell.execute_reply": "2024-11-12T21:58:16.808013Z"
    },
    "papermill": {
     "duration": 0.486542,
     "end_time": "2024-11-12T21:58:16.811089",
     "exception": false,
     "start_time": "2024-11-12T21:58:16.324547",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "data = []\n",
    "\n",
    "# The CSV file contains two columns 'question' and 'answer'\n",
    "with open(\"//kaggle/input/medquad/medquad.csv\", mode='r', encoding='utf-8') as file:\n",
    "    reader = csv.DictReader(file)\n",
    "    for row in reader:\n",
    "        # Use a template to format the questions and answers in the CSV into\n",
    "        # questions and answers in the data.\n",
    "        template = \"Question:\\n{question}\\n\\nAnswer:\\n{answer}\"\n",
    "        data.append(template.format(**row))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e68baab8",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-12T21:58:16.835367Z",
     "iopub.status.busy": "2024-11-12T21:58:16.835085Z",
     "iopub.status.idle": "2024-11-12T21:58:16.840689Z",
     "shell.execute_reply": "2024-11-12T21:58:16.839845Z"
    },
    "papermill": {
     "duration": 0.019808,
     "end_time": "2024-11-12T21:58:16.842550",
     "exception": false,
     "start_time": "2024-11-12T21:58:16.822742",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "16412"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13dfb336",
   "metadata": {
    "papermill": {
     "duration": 0.011198,
     "end_time": "2024-11-12T21:58:16.865064",
     "exception": false,
     "start_time": "2024-11-12T21:58:16.853866",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's limit the data to **200 rows** of questions and answers. We do this to limit the training time for this demonstration Notebook. You can get back and fine-tune using entire dataset or use your own data size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9cb561c0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-12T21:58:16.889141Z",
     "iopub.status.busy": "2024-11-12T21:58:16.888374Z",
     "iopub.status.idle": "2024-11-12T21:58:16.894618Z",
     "shell.execute_reply": "2024-11-12T21:58:16.893769Z"
    },
    "papermill": {
     "duration": 0.020324,
     "end_time": "2024-11-12T21:58:16.896587",
     "exception": false,
     "start_time": "2024-11-12T21:58:16.876263",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "data = data[:300]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37ca9a99",
   "metadata": {
    "papermill": {
     "duration": 0.011113,
     "end_time": "2024-11-12T21:58:16.918960",
     "exception": false,
     "start_time": "2024-11-12T21:58:16.907847",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Check model inference before fine tuning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea333369",
   "metadata": {
    "papermill": {
     "duration": 0.011165,
     "end_time": "2024-11-12T21:58:16.941739",
     "exception": false,
     "start_time": "2024-11-12T21:58:16.930574",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "We wil first check the model before proceeding to fine-tuning. We will test it with some questions about medical matters.  \n",
    "\n",
    "First, we will define an utility function to display the query and answer from LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9706006d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-12T21:58:16.965346Z",
     "iopub.status.busy": "2024-11-12T21:58:16.965095Z",
     "iopub.status.idle": "2024-11-12T21:58:16.970062Z",
     "shell.execute_reply": "2024-11-12T21:58:16.969179Z"
    },
    "papermill": {
     "duration": 0.019107,
     "end_time": "2024-11-12T21:58:16.972041",
     "exception": false,
     "start_time": "2024-11-12T21:58:16.952934",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from IPython.display import display, Markdown\n",
    "def colorize_text(text):\n",
    "    for word, color in zip([\"Category\", \"Question\", \"Answer\"], [\"blue\", \"red\", \"green\"]):\n",
    "        text = text.replace(f\"{word}:\", f\"\\n\\n**<font color='{color}'>{word}:</font>**\")\n",
    "    return text"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92f6fee5",
   "metadata": {
    "papermill": {
     "duration": 0.011087,
     "end_time": "2024-11-12T21:58:16.994429",
     "exception": false,
     "start_time": "2024-11-12T21:58:16.983342",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Let's check how we can display the content of one data input using the `colorize_text` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "62643412",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-12T21:58:17.018732Z",
     "iopub.status.busy": "2024-11-12T21:58:17.017915Z",
     "iopub.status.idle": "2024-11-12T21:58:17.022776Z",
     "shell.execute_reply": "2024-11-12T21:58:17.021899Z"
    },
    "papermill": {
     "duration": 0.019227,
     "end_time": "2024-11-12T21:58:17.024878",
     "exception": false,
     "start_time": "2024-11-12T21:58:17.005651",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question:\n",
      "What are the treatments for Glaucoma ?\n",
      "\n",
      "Answer:\n",
      "Although open-angle glaucoma cannot be cured, it can usually be controlled. While treatments may save remaining vision, they do not improve sight already lost from glaucoma. The most common treatments for glaucoma are medication and surgery. Medications  Medications for glaucoma may be either in the form of eye drops or pills. Some drugs reduce pressure by slowing the flow of fluid into the eye. Others help to improve fluid drainage. (Watch the video to learn more about coping with glaucoma. To enlarge the video, click the brackets in the lower right-hand corner. To reduce the video, press the Escape (Esc) button on your keyboard.) For most people with glaucoma, regular use of medications will control the increased fluid pressure. But, these drugs may stop working over time. Or, they may cause side effects. If a problem occurs, the eye care professional may select other drugs, change the dose, or suggest other ways to deal with the problem.  Read or listen to ways some patients are coping with glaucoma. Surgery Laser surgery is another treatment for glaucoma. During laser surgery, a strong beam of light is focused on the part of the anterior chamber where the fluid leaves the eye. This results in a series of small changes that makes it easier for fluid to exit the eye. Over time, the effect of laser surgery may wear off. Patients who have this form of surgery may need to keep taking glaucoma drugs. Researching Causes and Treatments Through studies in the laboratory and with patients, NEI is seeking better ways to detect, treat, and prevent vision loss in people with glaucoma. For example, researchers have discovered genes that could help explain how glaucoma damages the eye. NEI also is supporting studies to learn more about who is likely to get glaucoma, when to treat people who have increased eye pressure, and which treatment to use first.\n"
     ]
    }
   ],
   "source": [
    "print(data[3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "bd43ebe5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-12T21:58:17.048662Z",
     "iopub.status.busy": "2024-11-12T21:58:17.048416Z",
     "iopub.status.idle": "2024-11-12T21:58:17.054240Z",
     "shell.execute_reply": "2024-11-12T21:58:17.053394Z"
    },
    "papermill": {
     "duration": 0.019929,
     "end_time": "2024-11-12T21:58:17.056171",
     "exception": false,
     "start_time": "2024-11-12T21:58:17.036242",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "\n",
       "**<font color='red'>Question:</font>**\n",
       "What are the treatments for Glaucoma ?\n",
       "\n",
       "\n",
       "\n",
       "**<font color='green'>Answer:</font>**\n",
       "Although open-angle glaucoma cannot be cured, it can usually be controlled. While treatments may save remaining vision, they do not improve sight already lost from glaucoma. The most common treatments for glaucoma are medication and surgery. Medications  Medications for glaucoma may be either in the form of eye drops or pills. Some drugs reduce pressure by slowing the flow of fluid into the eye. Others help to improve fluid drainage. (Watch the video to learn more about coping with glaucoma. To enlarge the video, click the brackets in the lower right-hand corner. To reduce the video, press the Escape (Esc) button on your keyboard.) For most people with glaucoma, regular use of medications will control the increased fluid pressure. But, these drugs may stop working over time. Or, they may cause side effects. If a problem occurs, the eye care professional may select other drugs, change the dose, or suggest other ways to deal with the problem.  Read or listen to ways some patients are coping with glaucoma. Surgery Laser surgery is another treatment for glaucoma. During laser surgery, a strong beam of light is focused on the part of the anterior chamber where the fluid leaves the eye. This results in a series of small changes that makes it easier for fluid to exit the eye. Over time, the effect of laser surgery may wear off. Patients who have this form of surgery may need to keep taking glaucoma drugs. Researching Causes and Treatments Through studies in the laboratory and with patients, NEI is seeking better ways to detect, treat, and prevent vision loss in people with glaucoma. For example, researchers have discovered genes that could help explain how glaucoma damages the eye. NEI also is supporting studies to learn more about who is likely to get glaucoma, when to treat people who have increased eye pressure, and which treatment to use first."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(Markdown(colorize_text(data[3])))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aec03901",
   "metadata": {
    "papermill": {
     "duration": 0.01146,
     "end_time": "2024-11-12T21:58:17.079247",
     "exception": false,
     "start_time": "2024-11-12T21:58:17.067787",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Now we will ask the model to answer to a question for which we know the expected answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a68bec20",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-12T21:58:17.103490Z",
     "iopub.status.busy": "2024-11-12T21:58:17.103218Z",
     "iopub.status.idle": "2024-11-12T21:58:27.468701Z",
     "shell.execute_reply": "2024-11-12T21:58:27.467743Z"
    },
    "papermill": {
     "duration": 10.380058,
     "end_time": "2024-11-12T21:58:27.470838",
     "exception": false,
     "start_time": "2024-11-12T21:58:17.090780",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "\n",
       "**<font color='red'>Question:</font>**\n",
       "What are the complications of Paget's Disease of Bone ?\n",
       "\n",
       "\n",
       "\n",
       "**<font color='green'>Answer:</font>**\n",
       "Paget's disease of bone is a disorder of bone metabolism characterized by excessive bone resorption and osteoclast activity. The disease is characterized by the formation of new bone in the form of osteoid, which is subsequently replaced by mature bone. The disease is characterized by the formation of new bone in the form of osteoid, which is subsequently replaced by mature bone. The disease is characterized by the formation of new bone in the form of osteoid, which is subsequently replaced by mature bone. The disease is characterized by the formation"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "prompt = template.format(\n",
    "    question=\"What are the complications of Paget's Disease of Bone ?\",\n",
    "    answer=\"\",\n",
    ")\n",
    "response = gemma_lm.generate(prompt, max_length=128)\n",
    "display(Markdown(colorize_text(response)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7ce6c7d",
   "metadata": {
    "papermill": {
     "duration": 0.011417,
     "end_time": "2024-11-12T21:58:27.494588",
     "exception": false,
     "start_time": "2024-11-12T21:58:27.483171",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Another example below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "127333c6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-12T21:58:27.519274Z",
     "iopub.status.busy": "2024-11-12T21:58:27.518921Z",
     "iopub.status.idle": "2024-11-12T21:58:30.264302Z",
     "shell.execute_reply": "2024-11-12T21:58:30.263413Z"
    },
    "papermill": {
     "duration": 2.759866,
     "end_time": "2024-11-12T21:58:30.266242",
     "exception": false,
     "start_time": "2024-11-12T21:58:27.506376",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "\n",
       "**<font color='red'>Question:</font>**\n",
       "What are the treatments for Diabetes ?\n",
       "\n",
       "\n",
       "\n",
       "**<font color='green'>Answer:</font>**\n",
       "Diabetes is a chronic disease that affects the way your body uses blood sugar (glucose). It is caused by a lack of insulin, a hormone that helps glucose enter your cells to be used for energy.\n",
       "\n",
       "There are two types of diabetes:\n",
       "\n",
       "Type 1 diabetes: This is an autoimmune disease in which the body’s immune system attacks the cells that produce insulin. It is often called “insulin-dependent diabetes” or “juvenile diabetes.”\n",
       "\n",
       "Type 2 diabetes: This is the most common type of diabetes. It is caused by a combination of factors"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "prompt = template.format(\n",
    "    question=\"What are the treatments for Diabetes ?\",\n",
    "    answer=\"\",\n",
    ")\n",
    "response = gemma_lm.generate(prompt, max_length=128)\n",
    "display(Markdown(colorize_text(response)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c9b55aa",
   "metadata": {
    "papermill": {
     "duration": 0.012046,
     "end_time": "2024-11-12T21:58:30.290544",
     "exception": false,
     "start_time": "2024-11-12T21:58:30.278498",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Fine-tunning with LoRA   \n",
    "\n",
    "\n",
    "We are using now **LoRA** for fine-tunning. **LoRA** stands for **Low Rank Adaptation** and is a method for modifying a pretrained model (for example, an LLM or vision transformer) to better suit a specific, often smaller, dataset by **adjusting only a small, low-rank subset of the model's parameters**.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a611fcf",
   "metadata": {
    "papermill": {
     "duration": 0.011656,
     "end_time": "2024-11-12T21:58:30.314010",
     "exception": false,
     "start_time": "2024-11-12T21:58:30.302354",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "The rank used here for LoRA controls the number of parameters that will be recalculated during fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0b9016c0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-12T21:58:30.339218Z",
     "iopub.status.busy": "2024-11-12T21:58:30.338878Z",
     "iopub.status.idle": "2024-11-12T21:58:30.689855Z",
     "shell.execute_reply": "2024-11-12T21:58:30.688319Z"
    },
    "papermill": {
     "duration": 0.366316,
     "end_time": "2024-11-12T21:58:30.692174",
     "exception": false,
     "start_time": "2024-11-12T21:58:30.325858",
     "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,195,392</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,195,392\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,195,392</span> (9.34 GB)\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m Total params: \u001b[0m\u001b[38;5;34m2,507,195,392\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,022,976</span> (3.90 MB)\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m Trainable params: \u001b[0m\u001b[38;5;34m1,022,976\u001b[0m (3.90 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 3.\n",
    "gemma_lm.backbone.enable_lora(rank=3)\n",
    "gemma_lm.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "30cbe12a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-12T21:58:30.720490Z",
     "iopub.status.busy": "2024-11-12T21:58:30.720166Z",
     "iopub.status.idle": "2024-11-12T22:09:07.120693Z",
     "shell.execute_reply": "2024-11-12T22:09:07.119774Z"
    },
    "papermill": {
     "duration": 636.417162,
     "end_time": "2024-11-12T22:09:07.123134",
     "exception": false,
     "start_time": "2024-11-12T21:58:30.705972",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/10\n",
      "\u001b[1m300/300\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m83s\u001b[0m 228ms/step - loss: 1.8024 - sparse_categorical_accuracy: 0.5780\n",
      "Epoch 2/10\n",
      "\u001b[1m300/300\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m69s\u001b[0m 201ms/step - loss: 1.5297 - sparse_categorical_accuracy: 0.6032\n",
      "Epoch 3/10\n",
      "\u001b[1m300/300\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m60s\u001b[0m 201ms/step - loss: 1.4629 - sparse_categorical_accuracy: 0.6131\n",
      "Epoch 4/10\n",
      "\u001b[1m300/300\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m60s\u001b[0m 201ms/step - loss: 1.4265 - sparse_categorical_accuracy: 0.6188\n",
      "Epoch 5/10\n",
      "\u001b[1m300/300\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m60s\u001b[0m 201ms/step - loss: 1.3921 - sparse_categorical_accuracy: 0.6255\n",
      "Epoch 6/10\n",
      "\u001b[1m300/300\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m60s\u001b[0m 201ms/step - loss: 1.3586 - sparse_categorical_accuracy: 0.6314\n",
      "Epoch 7/10\n",
      "\u001b[1m300/300\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m60s\u001b[0m 201ms/step - loss: 1.3259 - sparse_categorical_accuracy: 0.6384\n",
      "Epoch 8/10\n",
      "\u001b[1m300/300\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m60s\u001b[0m 201ms/step - loss: 1.2897 - sparse_categorical_accuracy: 0.6434\n",
      "Epoch 9/10\n",
      "\u001b[1m300/300\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m60s\u001b[0m 201ms/step - loss: 1.2519 - sparse_categorical_accuracy: 0.6522\n",
      "Epoch 10/10\n",
      "\u001b[1m300/300\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m60s\u001b[0m 201ms/step - loss: 1.2149 - sparse_categorical_accuracy: 0.6614\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.src.callbacks.history.History at 0x7ce3203707c0>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Fine-tune on the Medical QA dataset.\n",
    "\n",
    "# Limit the input sequence length to 128 (to control memory usage).\n",
    "gemma_lm.preprocessor.sequence_length = 128\n",
    "# Use AdamW (a common optimizer for transformer models).\n",
    "optimizer = keras.optimizers.AdamW(\n",
    "    learning_rate=5e-5,\n",
    "    weight_decay=0.01,\n",
    ")\n",
    "# Exclude layernorm and bias terms from decay.\n",
    "optimizer.exclude_from_weight_decay(var_names=[\"bias\", \"scale\"])\n",
    "\n",
    "gemma_lm.compile(\n",
    "    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
    "    optimizer=optimizer,\n",
    "    weighted_metrics=[keras.metrics.SparseCategoricalAccuracy()],\n",
    ")\n",
    "gemma_lm.fit(data, epochs=10, batch_size=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39a5f3a7",
   "metadata": {
    "papermill": {
     "duration": 0.258641,
     "end_time": "2024-11-12T22:09:07.694545",
     "exception": false,
     "start_time": "2024-11-12T22:09:07.435904",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Inference after fine tuning\n",
    "\n",
    "\n",
    "Let's try again now the two examples from above. We will run now the queries through the fine-tuned model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "74c60b0b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-12T22:09:08.211171Z",
     "iopub.status.busy": "2024-11-12T22:09:08.210775Z",
     "iopub.status.idle": "2024-11-12T22:09:18.850322Z",
     "shell.execute_reply": "2024-11-12T22:09:18.849381Z"
    },
    "papermill": {
     "duration": 10.898319,
     "end_time": "2024-11-12T22:09:18.852560",
     "exception": false,
     "start_time": "2024-11-12T22:09:07.954241",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "\n",
       "**<font color='red'>Question:</font>**\n",
       "What are the complications of Paget's Disease of Bone ?\n",
       "\n",
       "\n",
       "\n",
       "**<font color='green'>Answer:</font>**\n",
       "Paget's disease of bone is a disorder that causes the bones to become enlarged and deformed. It is a bone disease that affects people of all ages, but it is most common in people over 60. The disease causes the bones to become enlarged and deformed. It is a bone disease that affects people of all ages, but it is most common in people over 60. The disease causes the bones to become enlarged and deformed. It is a bone disease that affects people of all ages, but it is most common"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "prompt = template.format(\n",
    "    question=\"What are the complications of Paget's Disease of Bone ?\",\n",
    "    answer=\"\",\n",
    ")\n",
    "response = gemma_lm.generate(prompt, max_length=128)\n",
    "display(Markdown(colorize_text(response)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "ccdda820",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-12T22:09:19.378900Z",
     "iopub.status.busy": "2024-11-12T22:09:19.378549Z",
     "iopub.status.idle": "2024-11-12T22:09:22.456520Z",
     "shell.execute_reply": "2024-11-12T22:09:22.455560Z"
    },
    "papermill": {
     "duration": 3.335225,
     "end_time": "2024-11-12T22:09:22.458618",
     "exception": false,
     "start_time": "2024-11-12T22:09:19.123393",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "\n",
       "\n",
       "**<font color='red'>Question:</font>**\n",
       "What are the treatments for Diabetes ?\n",
       "\n",
       "\n",
       "\n",
       "**<font color='green'>Answer:</font>**\n",
       "Treatments for Diabetes There is no cure for diabetes, but there are many ways to manage it. The goal of treatment is to keep blood glucose levels as close to normal as possible. This can be done with diet, exercise, and medications. The type of treatment a person needs depends on the type of diabetes he or she has. The two main types of diabetes are type 1 and type 2. Type 1 diabetes is an autoimmune disease. The body makes no insulin, a hormone that allows glucose to enter cells and be used for energy. Type 2"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "prompt = template.format(\n",
    "    question=\"What are the treatments for Diabetes ?\",\n",
    "    answer=\"\",\n",
    ")\n",
    "response = gemma_lm.generate(prompt, max_length=128)\n",
    "display(Markdown(colorize_text(response)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "089fd090",
   "metadata": {
    "papermill": {
     "duration": 0.260405,
     "end_time": "2024-11-12T22:09:22.981206",
     "exception": false,
     "start_time": "2024-11-12T22:09:22.720801",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Save the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "0091fb75",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-12T22:09:23.493882Z",
     "iopub.status.busy": "2024-11-12T22:09:23.493416Z",
     "iopub.status.idle": "2024-11-12T22:10:13.473626Z",
     "shell.execute_reply": "2024-11-12T22:10:13.472425Z"
    },
    "papermill": {
     "duration": 50.242585,
     "end_time": "2024-11-12T22:10:13.477654",
     "exception": false,
     "start_time": "2024-11-12T22:09:23.235069",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "preset = \"./medical_gemma\"\n",
    "# Save the model to the preset directory.\n",
    "gemma_lm.save_to_preset(preset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b14c8cf8",
   "metadata": {
    "papermill": {
     "duration": 0.251912,
     "end_time": "2024-11-12T22:10:14.014228",
     "exception": false,
     "start_time": "2024-11-12T22:10:13.762316",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Conclusions\n",
    "\n",
    "We saw how to fine tune Gemma model using **LoRA** and **KerasNLP**.\n",
    "\n",
    "We only used a subset of the **Medical Q & A** data to perform the fine tuning. If you have access to more computational resources, please feel free to improve the fine-tuned model by:\n",
    "* Adding more data from the available dataset.  \n",
    "* Increasing the rank parameter used in fine-tuning.\n",
    "\n",
    "We will use the saved model to publish a Kaggle Model."
   ]
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "gpu",
   "dataSources": [
    {
     "datasetId": 5314260,
     "sourceId": 8831896,
     "sourceType": "datasetVersion"
    },
    {
     "modelId": 3533,
     "modelInstanceId": 5171,
     "sourceId": 11371,
     "sourceType": "modelInstanceVersion"
    }
   ],
   "dockerImageVersionId": 30733,
   "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": 905.616141,
   "end_time": "2024-11-12T22:10:16.795528",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2024-11-12T21:55:11.179387",
   "version": "2.5.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}