Prompt_Recovery_with_Gemma_KerasNLP_Starter.ipynb

{
  "metadata": {
    "kernelspec": {
      "language": "python",
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.10.13",
      "mimetype": "text/x-python",
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "pygments_lexer": "ipython3",
      "nbconvert_exporter": "python",
      "file_extension": ".py"
    },
    "kaggle": {
      "accelerator": "gpu",
      "dataSources": [
        {
          "sourceId": 67121,
          "databundleVersionId": 7806901,
          "sourceType": "competition"
        },
        {
          "sourceId": 7731345,
          "sourceType": "datasetVersion",
          "datasetId": 4517764
        },
        {
          "sourceId": 7733314,
          "sourceType": "datasetVersion",
          "datasetId": 4518936
        },
        {
          "sourceId": 11372,
          "sourceType": "modelInstanceVersion",
          "modelInstanceId": 5388
        }
      ],
      "isInternetEnabled": false,
      "language": "python",
      "sourceType": "notebook",
      "isGpuEnabled": true
    },
    "colab": {
      "name": "Prompt Recovery with Gemma - KerasNLP Starter",
      "provenance": [],
      "include_colab_link": true
    }
  },
  "nbformat_minor": 0,
  "nbformat": 4,
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/zhaizeyu/test/blob/master/Prompt_Recovery_with_Gemma_KerasNLP_Starter.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "source": [
        "\n",
        "# IMPORTANT: RUN THIS CELL IN ORDER TO IMPORT YOUR KAGGLE DATA SOURCES\n",
        "# TO THE CORRECT LOCATION (/kaggle/input) IN YOUR NOTEBOOK,\n",
        "# THEN FEEL FREE TO DELETE THIS CELL.\n",
        "# NOTE: THIS NOTEBOOK ENVIRONMENT DIFFERS FROM KAGGLE'S PYTHON\n",
        "# ENVIRONMENT SO THERE MAY BE MISSING LIBRARIES USED BY YOUR\n",
        "# NOTEBOOK.\n",
        "\n",
        "import os\n",
        "import sys\n",
        "from tempfile import NamedTemporaryFile\n",
        "from urllib.request import urlopen\n",
        "from urllib.parse import unquote, urlparse\n",
        "from urllib.error import HTTPError\n",
        "from zipfile import ZipFile\n",
        "import tarfile\n",
        "import shutil\n",
        "\n",
        "CHUNK_SIZE = 40960\n",
        "DATA_SOURCE_MAPPING = 'llm-prompt-recovery-synthetic-datastore:https%3A%2F%2Fstorage.googleapis.com%2Fkaggle-data-sets%2F4517764%2F7731345%2Fbundle%2Farchive.zip%3FX-Goog-Algorithm%3DGOOG4-RSA-SHA256%26X-Goog-Credential%3Dgcp-kaggle-com%2540kaggle-161607.iam.gserviceaccount.com%252F20240319%252Fauto%252Fstorage%252Fgoog4_request%26X-Goog-Date%3D20240319T213000Z%26X-Goog-Expires%3D259200%26X-Goog-SignedHeaders%3Dhost%26X-Goog-Signature%3D8c833fad09f34f3adf7aa440e0a25281d86fdfc0ee1de41ea7ce0f2ef6fd0aa5fc85e2801970762eb164a31a34116f513e873b7dcc310904fc1d41dec509d5cd31e956c429fda9d3368c71a259eeaf703adc94db782e00019e37843b41c44b8e0aff2f163866deb5c6013c729ecaca4978e1dd59978df62c4cb10756fd39d57bdbfbb698fc7e7638fc026b0617b72a45b6be99211d881727f6e17faef61572db9c283f59c2d1b37931d025ae8267cba38c380fad9f1640765f40cbd188611832eb1010a6d84aa99d98d7083f2612ba595660917b8d09e04ed232321bda4d26d612335be10becdea47a40394241f1fadd385f8305d3105eaf9fb24588c6a586c8,3000-rewritten-texts-prompt-recovery-challenge:https%3A%2F%2Fstorage.googleapis.com%2Fkaggle-data-sets%2F4518936%2F7733314%2Fbundle%2Farchive.zip%3FX-Goog-Algorithm%3DGOOG4-RSA-SHA256%26X-Goog-Credential%3Dgcp-kaggle-com%2540kaggle-161607.iam.gserviceaccount.com%252F20240319%252Fauto%252Fstorage%252Fgoog4_request%26X-Goog-Date%3D20240319T213000Z%26X-Goog-Expires%3D259200%26X-Goog-SignedHeaders%3Dhost%26X-Goog-Signature%3D19115b36dcab68b8743b96cd1e7da7a277fd3a3371bc96e0053b0149df118c71da8c0ccd91060cc465d1d21dfb60d694d49d19122bf3d37d711d4f5cdee62d94fc26df24c5e1408aca8cb1c38e411f812b6cabf48c3fcdf7fbf31510a85691b7c1691769fa56711059de37c74e3a9df6c1569792acb99ef6d28c63ed4751529721886c7cb5aa2e3c93a51502c9a96f7fe44c4a21f531d52f68d0c1e6ef1ffa8eb6ffc711b26a3a75807c0461a4376b67c6669de3f849d5900ee3f752cf12bdcb9f112cac787c6987ab7f3c0087bd34bb5324812067d72b1cfd68f2759a0e438d34f3fec1ea074756cbc48518e03d1e401318f0b4fb590a16e29eb06a901c3faf'\n",
        "\n",
        "KAGGLE_INPUT_PATH='/kaggle/input'\n",
        "KAGGLE_WORKING_PATH='/kaggle/working'\n",
        "KAGGLE_SYMLINK='kaggle'\n",
        "\n",
        "!umount /kaggle/input/ 2> /dev/null\n",
        "shutil.rmtree('/kaggle/input', ignore_errors=True)\n",
        "os.makedirs(KAGGLE_INPUT_PATH, 0o777, exist_ok=True)\n",
        "os.makedirs(KAGGLE_WORKING_PATH, 0o777, exist_ok=True)\n",
        "\n",
        "try:\n",
        "  os.symlink(KAGGLE_INPUT_PATH, os.path.join(\"..\", 'input'), target_is_directory=True)\n",
        "except FileExistsError:\n",
        "  pass\n",
        "try:\n",
        "  os.symlink(KAGGLE_WORKING_PATH, os.path.join(\"..\", 'working'), target_is_directory=True)\n",
        "except FileExistsError:\n",
        "  pass\n",
        "\n",
        "for data_source_mapping in DATA_SOURCE_MAPPING.split(','):\n",
        "    directory, download_url_encoded = data_source_mapping.split(':')\n",
        "    download_url = unquote(download_url_encoded)\n",
        "    filename = urlparse(download_url).path\n",
        "    destination_path = os.path.join(KAGGLE_INPUT_PATH, directory)\n",
        "    try:\n",
        "        with urlopen(download_url) as fileres, NamedTemporaryFile() as tfile:\n",
        "            total_length = fileres.headers['content-length']\n",
        "            print(f'Downloading {directory}, {total_length} bytes compressed')\n",
        "            dl = 0\n",
        "            data = fileres.read(CHUNK_SIZE)\n",
        "            while len(data) > 0:\n",
        "                dl += len(data)\n",
        "                tfile.write(data)\n",
        "                done = int(50 * dl / int(total_length))\n",
        "                sys.stdout.write(f\"\\r[{'=' * done}{' ' * (50-done)}] {dl} bytes downloaded\")\n",
        "                sys.stdout.flush()\n",
        "                data = fileres.read(CHUNK_SIZE)\n",
        "            if filename.endswith('.zip'):\n",
        "              with ZipFile(tfile) as zfile:\n",
        "                zfile.extractall(destination_path)\n",
        "            else:\n",
        "              with tarfile.open(tfile.name) as tarfile:\n",
        "                tarfile.extractall(destination_path)\n",
        "            print(f'\\nDownloaded and uncompressed: {directory}')\n",
        "    except HTTPError as e:\n",
        "        print(f'Failed to load (likely expired) {download_url} to path {destination_path}')\n",
        "        continue\n",
        "    except OSError as e:\n",
        "        print(f'Failed to load {download_url} to path {destination_path}')\n",
        "        continue\n",
        "\n",
        "print('Data source import complete.')\n"
      ],
      "metadata": {
        "id": "zD4dqErwnxmp"
      },
      "cell_type": "code",
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "source": [
        "<center><img src=\"https://keras.io/img/logo-small.png\" alt=\"Keras logo\" width=\"100\"><br/>\n",
        "This starter notebook is provided by the Keras team.</center>"
      ],
      "metadata": {
        "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
        "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
        "id": "k0lImHoWnxms"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# LLM Prompt Recovery with [KerasNLP](https://github.com/keras-team/keras-nlp) and [Keras](https://github.com/keras-team/keras)\n",
        "\n",
        "<div align=\"center\">\n",
        "    <img src=\"https://i.ibb.co/8xZNc32/Gemma.png\">\n",
        "</div>\n",
        "\n",
        "In this competition, the goal is to find the prompt used to transform a given text. Specifically, we're seeking the prompt or instruction used in the [Gemma 7B-it](https://www.kaggle.com/models/google/gemma/frameworks/pyTorch/variations/7b-it-quant) model to convert one text to another. Typically, large language models are instructed to transform one text to another style, but here we're tasked with the inverse: finding the instruction/prompt used for the transformation. This notebook walks you through fine-tuning the **Gemma 2b-it** model with LoRA for this prompt recovery task using KerasNLP. Witih KerasNLP, we can fine-tune with LoRA using just a few lines of code.\n",
        "\n",
        "**Fun fact**: This notebook is backend-agnostic, supporting TensorFlow, PyTorch, and JAX. However, the best performance can be achieved with `JAX`. KerasNLP and Keras enable the choice of preferred backend. Explore further details on [Keras](https://keras.io/keras_3/).\n",
        "\n",
        "**Note**: For a deeper understanding of KerasNLP, refer to the [KerasNLP guides](https://keras.io/keras_nlp/)."
      ],
      "metadata": {
        "id": "NKmKQLOSnxmt"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Import Libraries"
      ],
      "metadata": {
        "id": "3Vjp153onxmt"
      }
    },
    {
      "cell_type": "code",
      "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",
        "\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 plotly.graph_objs as go\n",
        "import plotly.express as px\n",
        "from IPython.display import display, Markdown"
      ],
      "metadata": {
        "_kg_hide-output": true,
        "execution": {
          "iopub.status.busy": "2024-03-11T04:37:48.365627Z",
          "iopub.execute_input": "2024-03-11T04:37:48.366001Z",
          "iopub.status.idle": "2024-03-11T04:38:02.230539Z",
          "shell.execute_reply.started": "2024-03-11T04:37:48.365972Z",
          "shell.execute_reply": "2024-03-11T04:38:02.229726Z"
        },
        "trusted": true,
        "id": "JtweXbN0nxmt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Configuration"
      ],
      "metadata": {
        "id": "6RFvnUqHnxmu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class CFG:\n",
        "    seed = 42\n",
        "    dataset_path = \"/kaggle/input/llm-prompt-recovery\"\n",
        "    preset = \"gemma_instruct_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\n",
        "    epochs = 1 # number of epochs to train"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-11T04:38:02.232216Z",
          "iopub.execute_input": "2024-03-11T04:38:02.233275Z",
          "iopub.status.idle": "2024-03-11T04:38:02.237947Z",
          "shell.execute_reply.started": "2024-03-11T04:38:02.233239Z",
          "shell.execute_reply": "2024-03-11T04:38:02.237083Z"
        },
        "trusted": true,
        "id": "Mh6Ryrdenxmu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Reproducibility\n",
        "Sets value for random seed to produce similar result in each run."
      ],
      "metadata": {
        "id": "aOfzj74Jnxmu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "keras.utils.set_random_seed(CFG.seed)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-11T04:38:02.239023Z",
          "iopub.execute_input": "2024-03-11T04:38:02.239328Z",
          "iopub.status.idle": "2024-03-11T04:38:02.26896Z",
          "shell.execute_reply.started": "2024-03-11T04:38:02.239279Z",
          "shell.execute_reply": "2024-03-11T04:38:02.268203Z"
        },
        "trusted": true,
        "id": "CxUaVTolnxmu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Data\n",
        "\n",
        "No training data is provided in this competition; in other words, we can use any openly available datasets for this competition. In this notebook, we will use two external datasets that utilize the **Gemma 7B** model to transform texts using prompts.\n",
        "\n",
        "**Data Format:**\n",
        "\n",
        "These datasets includes:\n",
        "- `original_text`: Input text/essay that needs to be transformed.\n",
        "- `rewrite_prompt`: Prompt/Instruction that was used in the Gemma LM to transform `original_text`. This is also our **target** for this competition.\n",
        "- `rewritten_text`: Output text that was generated by the Gemma model."
      ],
      "metadata": {
        "id": "UKknRUYcnxmv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# `LLM Prompt Recovery - Synthetic Datastore dataset` by @dschettler8845\n",
        "df1 = pd.read_csv(\"/kaggle/input/llm-prompt-recovery-synthetic-datastore/gemma1000_w7b.csv\")\n",
        "df1 = df1[[\"original_text\", \"rewrite_prompt\", \"gemma_7b_rewritten_text_temp0\"]]\n",
        "df1 = df1.rename(columns={\"gemma_7b_rewritten_text_temp0\":\"rewritten_text\"})\n",
        "df1.head(2)\n",
        "\n",
        "# `3000 Rewritten texts - Prompt recovery Challenge` by @dipamc77\n",
        "df2 = pd.read_csv(\"/kaggle/input/3000-rewritten-texts-prompt-recovery-challenge/prompts_0_500_wiki_first_para_3000.csv\")\n",
        "df2.head(2)\n",
        "\n",
        "# Merge all datasets\n",
        "df = pd.concat([df1, df2], axis=0)\n",
        "df = df.sample(2000).reset_index(drop=True) # to reduce training time we are only using 2k samples\n",
        "df.head(5)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-11T04:38:02.271577Z",
          "iopub.execute_input": "2024-03-11T04:38:02.271992Z",
          "iopub.status.idle": "2024-03-11T04:38:02.512551Z",
          "shell.execute_reply.started": "2024-03-11T04:38:02.27196Z",
          "shell.execute_reply": "2024-03-11T04:38:02.51163Z"
        },
        "trusted": true,
        "id": "GlV0VRFVnxmv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Prompt Engineering\n",
        "\n",
        "Here's a simple prompt template we'll use to create instruction-response pairs from the `original_text`, `rewritten_text`, and `rewritten_prompt`:\n",
        "\n",
        "```\n",
        "Instruction:\n",
        "Below, the `Original Text` passage has been rewritten/transformed/improved into `Rewritten Text` by the `Gemma 7b-it` LLM with a certain prompt/instruction. Your task is to carefully analyze the differences between the \"Original Text\" and \"Rewritten Text\", and try to infer the specific prompt or instruction that was likely given to the LLM to rewrite/transform/improve the text in this way.\n",
        "\n",
        "Original Text:\n",
        "...\n",
        "\n",
        "Rewritten Text:\n",
        "...\n",
        "\n",
        "Response:\n",
        "...\n",
        "```\n",
        "\n",
        "This template will help the model to follow instruction and respond accurately. You can explore more advanced prompt templates for better results."
      ],
      "metadata": {
        "id": "xFBzTiW_nxmv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "template = \"\"\"Instruction:\\nBelow, the `Original Text` passage has been rewritten/transformed/improved into `Rewritten Text` by the `Gemma 7b-it` LLM with a certain prompt/instruction. Your task is to carefully analyze the differences between the `Original Text` and `Rewritten Text`, and try to infer the specific prompt or instruction that was likely given to the LLM to rewrite/transform/improve the text in this way.\\n\\nOriginal Text:\\n{original_text}\\n\\nRewriten Text:\\n{rewritten_text}\\n\\nResponse:\\n{rewrite_prompt}\"\"\""
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-11T04:39:38.532219Z",
          "iopub.execute_input": "2024-03-11T04:39:38.532624Z",
          "iopub.status.idle": "2024-03-11T04:39:38.537933Z",
          "shell.execute_reply.started": "2024-03-11T04:39:38.532589Z",
          "shell.execute_reply": "2024-03-11T04:39:38.536727Z"
        },
        "trusted": true,
        "id": "X6KoL1Gxnxmv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "df[\"prompt\"] = df.progress_apply(lambda row: template.format(original_text=row.original_text,\n",
        "                                                             rewritten_text=row.rewritten_text,\n",
        "                                                             rewrite_prompt=row.rewrite_prompt), axis=1)\n",
        "data = df.prompt.tolist()"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-11T04:39:38.819832Z",
          "iopub.execute_input": "2024-03-11T04:39:38.820168Z",
          "iopub.status.idle": "2024-03-11T04:39:38.927952Z",
          "shell.execute_reply.started": "2024-03-11T04:39:38.820142Z",
          "shell.execute_reply": "2024-03-11T04:39:38.927032Z"
        },
        "trusted": true,
        "id": "hLbwzSfYnxmv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's examine a sample prompt. As the answers in our dataset are curated with **markdown** format, we will render the sample using `Markdown()` to properly visualize the formatting."
      ],
      "metadata": {
        "id": "kiEXlctcnxmv"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Sample"
      ],
      "metadata": {
        "id": "oQMKaDiwnxmv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def colorize_text(text):\n",
        "    for word, color in zip([\"Instruction\", \"Original Text\", \"Rewriten Text\", \"Response\"],\n",
        "                           [\"red\", \"yellow\", \"blue\", \"green\"]):\n",
        "        text = text.replace(f\"{word}:\", f\"\\n\\n**<font color='{color}'>{word}:</font>**\")\n",
        "    return text"
      ],
      "metadata": {
        "_kg_hide-input": true,
        "execution": {
          "iopub.status.busy": "2024-03-11T04:39:40.990029Z",
          "iopub.execute_input": "2024-03-11T04:39:40.990701Z",
          "iopub.status.idle": "2024-03-11T04:39:40.995568Z",
          "shell.execute_reply.started": "2024-03-11T04:39:40.990667Z",
          "shell.execute_reply": "2024-03-11T04:39:40.994691Z"
        },
        "trusted": true,
        "id": "a4UuuthYnxmw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Take a random sample\n",
        "sample = data[10]\n",
        "\n",
        "# Give colors to Instruction, Response and Category\n",
        "sample = colorize_text(sample)\n",
        "\n",
        "# Show sample in markdown\n",
        "display(Markdown(sample))"
      ],
      "metadata": {
        "_kg_hide-input": true,
        "execution": {
          "iopub.status.busy": "2024-03-11T04:39:41.205072Z",
          "iopub.execute_input": "2024-03-11T04:39:41.205738Z",
          "iopub.status.idle": "2024-03-11T04:39:41.211687Z",
          "shell.execute_reply.started": "2024-03-11T04:39:41.205708Z",
          "shell.execute_reply": "2024-03-11T04:39:41.210768Z"
        },
        "trusted": true,
        "id": "GVTHUFa4nxmw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Modeling\n",
        "\n",
        "<div align=\"center\"><img src=\"https://i.ibb.co/Bqg9w3g/Gemma-Logo-no-background.png\" width=\"300\"></div>\n",
        "\n",
        "**Gemma** is a collection of advanced open LLMs developed by **Google DeepMind** and other **Google teams**, derived from the same research and technology behind the **Gemini** models. They can be integrated into applications and run on various platforms including mobile devices and hosted services. Developers can customize Gemma models using tuning techniques to enhance their performance for specific tasks, offering more targeted and efficient generative AI solutions beyond text generation.\n",
        "\n",
        "Gemma models are available in several sizes so we can build generative AI solutions based on your available computing resources, the capabilities you need, and where you want to run them.\n",
        "\n",
        "| Parameters size | Tuned versions    | Intended platforms                 | Preset                 |\n",
        "|-----------------|-------------------|------------------------------------|------------------------|\n",
        "| 2B              | Pretrained        | Mobile devices and laptops         | `gemma_2b_en`          |\n",
        "| 2B              | Instruction tuned | Mobile devices and laptops         | `gemma_instruct_2b_en` |\n",
        "| 7B              | Pretrained        | Desktop computers and small servers| `gemma_7b_en`          |\n",
        "| 7B              | Instruction tuned | Desktop computers and small servers| `gemma_instruct_7b_en` |\n",
        "\n",
        "In this notebook, we will utilize the `Gemma 2b-it` model from KerasNLP's pretrained models to recover the prompt. We are using the \"Instruction tuned\" model instead of the \"Pretrained\" one because the test data was generated from an instruction-tuned Gemma model. Additionally, we will fine-tune our model using instruction-response pairs thus fine-tuning an instruction-tuned model will likely yield better results.\n",
        "\n",
        "To explore other available models, you can simply adjust the `preset` value in the `CFG` (config). You can find a list of other pretrained models on the [KerasNLP website](https://keras.io/api/keras_nlp/models/)."
      ],
      "metadata": {
        "id": "1VhhPZRDnxmw"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Gemma Causal LM\n",
        "\n",
        "The code below will build an end-to-end Gemma model for causal language modeling (hence the name `GemmaCausalLM`). A causal language model (LM) predicts the next token based on previous tokens. This task setup can be used to train the model unsupervised on plain text input or to autoregressively generate plain text similar to the data used for training. This task can be used for pre-training or fine-tuning a Gemma model simply by calling `fit()`.\n",
        "\n",
        "This model has a `generate()` method, which generates text based on a prompt. The generation strategy used is controlled by an additional sampler argument on `compile()`. You can recompile the model with different `keras_nlp.samplers` objects to control the generation. By default, `\"greedy\"` sampling will be used.\n",
        "\n",
        "> The `from_preset` method instantiates the model from a preset architecture and weights."
      ],
      "metadata": {
        "id": "zkwSdaJgnxmw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset(CFG.preset)\n",
        "gemma_lm.summary()"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-11T04:39:51.30769Z",
          "iopub.execute_input": "2024-03-11T04:39:51.308137Z",
          "iopub.status.idle": "2024-03-11T04:40:50.868354Z",
          "shell.execute_reply.started": "2024-03-11T04:39:51.308097Z",
          "shell.execute_reply": "2024-03-11T04:40:50.867375Z"
        },
        "trusted": true,
        "id": "Jv-aohJOnxmw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Gemma LM Preprocessor\n",
        "\n",
        "An important part of the Gemma model is the **Preprocessor** layer, which under the hood uses **Tokenizer**.\n",
        "\n",
        "**What it does:** The preprocessor takes input strings and transforms them into a dictionary (`token_ids`, `padding_mask`) containing preprocessed tensors. This process starts with tokenization, where input strings are converted into sequences of token IDs.\n",
        "\n",
        "**Why it's important:** Initially, raw text data is complex and challenging for modeling due to its high dimensionality. By converting text into a compact set of tokens, such as transforming `\"The quick brown fox\"` into `[\"the\", \"qu\", \"##ick\", \"br\", \"##own\", \"fox\"]`, we simplify the data. Many models rely on special tokens and additional tensors to understand input. These tokens help divide input and identify padding, among other tasks. Making all sequences the same length through padding boosts computational efficiency, making subsequent steps smoother.\n",
        "\n",
        "Explore the following pages to access the available preprocessing and tokenizer layers in **KerasNLP**:\n",
        "- [Preprocessing](https://keras.io/api/keras_nlp/preprocessing_layers/)\n",
        "- [Tokenizers](https://keras.io/api/keras_nlp/tokenizers/)"
      ],
      "metadata": {
        "id": "xmIw38iwnxmw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "x, y, sample_weight = gemma_lm.preprocessor(data[0:2])"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-11T04:41:04.359091Z",
          "iopub.execute_input": "2024-03-11T04:41:04.359811Z",
          "iopub.status.idle": "2024-03-11T04:41:04.716064Z",
          "shell.execute_reply.started": "2024-03-11T04:41:04.359775Z",
          "shell.execute_reply": "2024-03-11T04:41:04.714948Z"
        },
        "trusted": true,
        "id": "Yr0h-Ogrnxmw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "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)`."
      ],
      "metadata": {
        "id": "aeYdq5y7nxmw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Display the shape of each processed output\n",
        "for k, v in x.items():\n",
        "    print(k, \":\", v.shape)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-11T04:41:07.078851Z",
          "iopub.execute_input": "2024-03-11T04:41:07.079534Z",
          "iopub.status.idle": "2024-03-11T04:41:07.084376Z",
          "shell.execute_reply.started": "2024-03-11T04:41:07.079501Z",
          "shell.execute_reply": "2024-03-11T04:41:07.083482Z"
        },
        "trusted": true,
        "id": "Ifk2IE1anxmw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Inference before Fine-Tuning\n",
        "\n",
        "Before we do fine-tuning, let's try to recover the prompt using the Gemma model with some prepared prompts and see how it responds.\n",
        "\n",
        "> As this model is not yet fine-tuned for instruction, you will notice that the model's responses are inaccurate."
      ],
      "metadata": {
        "id": "aWrHPDwTnxmw"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Sample 1"
      ],
      "metadata": {
        "id": "GxPHQZ22nxmw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Take one sample\n",
        "row = df.iloc[10]\n",
        "\n",
        "# Generate Prompt using template\n",
        "prompt = template.format(\n",
        "    original_text=row.original_text,\n",
        "    rewritten_text=row.rewritten_text,\n",
        "    rewrite_prompt=\"\",\n",
        ")\n",
        "\n",
        "# Infer\n",
        "output = gemma_lm.generate(prompt, max_length=512)\n",
        "\n",
        "# Colorize\n",
        "output = colorize_text(output)\n",
        "\n",
        "# Display in markdown\n",
        "display(Markdown(output))\n"
      ],
      "metadata": {
        "_kg_hide-input": true,
        "execution": {
          "iopub.status.busy": "2024-03-11T04:41:11.165803Z",
          "iopub.execute_input": "2024-03-11T04:41:11.166534Z",
          "iopub.status.idle": "2024-03-11T04:41:24.532716Z",
          "shell.execute_reply.started": "2024-03-11T04:41:11.166502Z",
          "shell.execute_reply": "2024-03-11T04:41:24.53178Z"
        },
        "trusted": true,
        "id": "aKhcvbVhnxmw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Sample 2"
      ],
      "metadata": {
        "id": "x_z1bLSinxmx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Take one sample\n",
        "row = df.iloc[20]\n",
        "\n",
        "# Generate Prompt using template\n",
        "prompt = template.format(\n",
        "    original_text=row.original_text,\n",
        "    rewritten_text=row.rewritten_text,\n",
        "    rewrite_prompt=\"\",\n",
        ")\n",
        "\n",
        "# Infer\n",
        "output = gemma_lm.generate(prompt, max_length=512)\n",
        "\n",
        "# Colorize\n",
        "output = colorize_text(output)\n",
        "\n",
        "# Display in markdown\n",
        "display(Markdown(output))\n"
      ],
      "metadata": {
        "_kg_hide-input": true,
        "execution": {
          "iopub.status.busy": "2024-03-11T04:42:24.871902Z",
          "iopub.execute_input": "2024-03-11T04:42:24.872608Z",
          "iopub.status.idle": "2024-03-11T04:42:29.324289Z",
          "shell.execute_reply.started": "2024-03-11T04:42:24.872577Z",
          "shell.execute_reply": "2024-03-11T04:42:29.323386Z"
        },
        "trusted": true,
        "id": "MroQOLFXnxmx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Fine-tuning with LoRA\n",
        "\n",
        "To get better responses from the model, we will fine-tune the model with Low Rank Adaptation (LoRA).\n",
        "\n",
        "**What exactly is LoRA?**\n",
        "\n",
        "LoRA is a method used to fine-tune large language models (LLMs) in an efficient way. It involves freezing the weights of the LLM and injecting trainable rank-decomposition matrices.\n",
        "\n",
        "Imagine in an LLM, we have a pre-trained dense layer, represented by a $d \\times d$ weight matrix, denoted as $W_0$. We then initialize two additional dense layers, labeled as $A$ and $B$, with shapes $d \\times r$ and $r \\times d$, respectively. Here, $r$ denotes the rank, which is typically **much smaller than** $d$. Prior to LoRA, the model's output was computed using the equation $output = W_0 \\cdot x + b_0$, where $x$ represents the input and $b_0$ denotes the bias term associated with the original dense layer, which remains frozen. After applying LoRA, the equation becomes $output = (W_0 \\cdot x + b_0) + (B \\cdot A \\cdot x)$, where $A$ and $B$ denote the trainable rank-decomposition matrices that have been introduced.\n",
        "\n",
        "<center><img src=\"https://i.ibb.co/DWsbhLg/LoRA.png\" width=\"300\"><br/>\n",
        "Credit: <a href=\"https://arxiv.org/abs/2106.09685\">LoRA: Low-Rank Adaptation of Large Language Models</a> Paper</center>\n",
        "\n",
        "\n",
        "In the LoRA paper, $A$ is initialized with $\\mathcal{N} (0, \\sigma^2)$ and $B$ with $0$, where $\\mathcal{N}$ denotes the normal distribution, and $\\sigma^2$ is the variance.\n",
        "\n",
        "**Why does LoRA save memory?**\n",
        "\n",
        "Even though we're adding more layers to the model with LoRA, it actually helps save memory. This is because the smaller layers (A and B) have fewer parameters to learn compared to the big model and fewer trainable parameters mean fewer optimizer variables to store. So, even though the overall model might seem bigger, it's actually more efficient in terms of memory usage.\n",
        "\n",
        "> This notebook uses a LoRA rank of `4`. A higher rank means more detailed changes are possible, but also means more trainable parameters."
      ],
      "metadata": {
        "id": "NLHiRvsenxmx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Enable LoRA for the model and set the LoRA rank to 4.\n",
        "gemma_lm.backbone.enable_lora(rank=4)\n",
        "gemma_lm.summary()"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-11T04:42:52.13053Z",
          "iopub.execute_input": "2024-03-11T04:42:52.130875Z",
          "iopub.status.idle": "2024-03-11T04:42:52.595752Z",
          "shell.execute_reply.started": "2024-03-11T04:42:52.130849Z",
          "shell.execute_reply": "2024-03-11T04:42:52.594706Z"
        },
        "trusted": true,
        "id": "45U_4t-wnxmx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Notice** that, the number of trainable parameters is reduced from ~$2.5$ billions to ~$1.3$ millions after enabling LoRA."
      ],
      "metadata": {
        "id": "j28JDNYbnxmx"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Training"
      ],
      "metadata": {
        "id": "vJolkUKinxmx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Limit the input sequence length to 512 (to control memory usage).\n",
        "gemma_lm.preprocessor.sequence_length = CFG.sequence_length\n",
        "\n",
        "# Compile the model with loss, optimizer, and metric\n",
        "gemma_lm.compile(\n",
        "    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "    optimizer=keras.optimizers.Adam(learning_rate=3e-5),\n",
        "    weighted_metrics=[keras.metrics.SparseCategoricalAccuracy()],\n",
        ")\n",
        "\n",
        "# Train model\n",
        "gemma_lm.fit(data, epochs=CFG.epochs, batch_size=CFG.batch_size)"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-11T04:42:56.600925Z",
          "iopub.execute_input": "2024-03-11T04:42:56.601601Z"
        },
        "trusted": true,
        "id": "fCvFqfCXnxmx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Inference after fine-tuning\n",
        "\n",
        "Let's see how our fine-tuned model responds to the same questions we asked before fine-tuning the model."
      ],
      "metadata": {
        "id": "MnZZiaMwnxmx"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Sample 1"
      ],
      "metadata": {
        "id": "ktPyq8uonxmx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Take one sample\n",
        "row = df.iloc[10]\n",
        "\n",
        "# Generate Prompt using template\n",
        "prompt = template.format(\n",
        "    original_text=row.original_text,\n",
        "    rewritten_text=row.rewritten_text,\n",
        "    rewrite_prompt=\"\",\n",
        ")\n",
        "\n",
        "# Infer\n",
        "output = gemma_lm.generate(prompt, max_length=512)\n",
        "\n",
        "# Colorize\n",
        "output = colorize_text(output)\n",
        "\n",
        "# Display in markdown\n",
        "display(Markdown(output))"
      ],
      "metadata": {
        "_kg_hide-input": true,
        "execution": {
          "iopub.status.busy": "2024-03-08T19:25:50.542437Z",
          "iopub.execute_input": "2024-03-08T19:25:50.54274Z",
          "iopub.status.idle": "2024-03-08T19:26:06.485306Z",
          "shell.execute_reply.started": "2024-03-08T19:25:50.542714Z",
          "shell.execute_reply": "2024-03-08T19:26:06.484216Z"
        },
        "trusted": true,
        "id": "tWdRQBI6nxmx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Sample 2"
      ],
      "metadata": {
        "id": "m_qML9GVnxmx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Take one sample\n",
        "row = df.iloc[20]\n",
        "\n",
        "# Generate Prompt using template\n",
        "prompt = template.format(\n",
        "    original_text=row.original_text,\n",
        "    rewritten_text=row.rewritten_text,\n",
        "    rewrite_prompt=\"\",\n",
        ")\n",
        "\n",
        "# Infer\n",
        "output = gemma_lm.generate(prompt, max_length=512)\n",
        "\n",
        "# Colorize\n",
        "output = colorize_text(output)\n",
        "\n",
        "# Display in markdown\n",
        "display(Markdown(output))\n"
      ],
      "metadata": {
        "_kg_hide-input": true,
        "execution": {
          "iopub.status.busy": "2024-03-08T19:26:06.486876Z",
          "iopub.execute_input": "2024-03-08T19:26:06.487352Z",
          "iopub.status.idle": "2024-03-08T19:26:11.834985Z",
          "shell.execute_reply.started": "2024-03-08T19:26:06.487319Z",
          "shell.execute_reply": "2024-03-08T19:26:11.833948Z"
        },
        "trusted": true,
        "id": "2Bt3aEM8nxm1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Test Data"
      ],
      "metadata": {
        "id": "A6ZRGvONnxm1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "test_df = pd.read_csv(\"/kaggle/input/llm-prompt-recovery/test.csv\")\n",
        "test_df['original_text'] = test_df['original_text'].fillna(\"\")\n",
        "test_df['rewritten_text'] = test_df['rewritten_text'].fillna(\"\")\n",
        "test_df.head()"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-08T19:26:54.434613Z",
          "iopub.execute_input": "2024-03-08T19:26:54.43503Z",
          "iopub.status.idle": "2024-03-08T19:26:54.448795Z",
          "shell.execute_reply.started": "2024-03-08T19:26:54.434999Z",
          "shell.execute_reply": "2024-03-08T19:26:54.447604Z"
        },
        "trusted": true,
        "id": "T9iwojwHnxm1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Test Sample\n",
        "\n",
        "Now, let's try out a sample from test data that model hasn't seen during training."
      ],
      "metadata": {
        "id": "T2NpgleUnxm1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "row = test_df.iloc[0]\n",
        "\n",
        "# Generate Prompt using template\n",
        "prompt = template.format(\n",
        "    original_text=row.original_text,\n",
        "    rewritten_text=row.rewritten_text,\n",
        "    rewrite_prompt=\"\",\n",
        ")\n",
        "\n",
        "# Infer\n",
        "output = gemma_lm.generate(prompt, max_length=512)\n",
        "\n",
        "# Colorize\n",
        "output = colorize_text(output)\n",
        "\n",
        "# Display in markdown\n",
        "display(Markdown(output))"
      ],
      "metadata": {
        "_kg_hide-input": true,
        "execution": {
          "iopub.status.busy": "2024-03-08T19:27:29.022534Z",
          "iopub.execute_input": "2024-03-08T19:27:29.023234Z",
          "iopub.status.idle": "2024-03-08T19:27:30.390404Z",
          "shell.execute_reply.started": "2024-03-08T19:27:29.023199Z",
          "shell.execute_reply": "2024-03-08T19:27:30.389393Z"
        },
        "trusted": true,
        "id": "fYhWSBiunxm1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Submission"
      ],
      "metadata": {
        "id": "O6Pb6Rr8nxm1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "preds = []\n",
        "for i in tqdm(range(len(test_df))):\n",
        "    row = test_df.iloc[i]\n",
        "\n",
        "    # Generate Prompt using template\n",
        "    prompt = template.format(\n",
        "        original_text=row.original_text,\n",
        "        rewritten_text=row.rewritten_text,\n",
        "        rewrite_prompt=\"\"\n",
        "    )\n",
        "\n",
        "    # Infer\n",
        "    output = gemma_lm.generate(prompt, max_length=512)\n",
        "    pred = output.replace(prompt, \"\") # remove the prompt from output\n",
        "\n",
        "    # Store predictions\n",
        "    preds.append([row.id, pred])"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-08T19:28:43.847896Z",
          "iopub.execute_input": "2024-03-08T19:28:43.848333Z",
          "iopub.status.idle": "2024-03-08T19:28:45.843319Z",
          "shell.execute_reply.started": "2024-03-08T19:28:43.848301Z",
          "shell.execute_reply": "2024-03-08T19:28:45.842148Z"
        },
        "trusted": true,
        "id": "EuAB71w1nxm1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "While preparing the submission file, we must keep in mind that, leaving any `rewrite_prompt` blank as null answers will throw an error."
      ],
      "metadata": {
        "id": "aUx8fK8Fnxm1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "sub_df = pd.DataFrame(preds, columns=[\"id\", \"rewrite_prompt\"])\n",
        "sub_df['rewrite_prompt'] = sub_df['rewrite_prompt'].fillna(\"\")\n",
        "sub_df['rewrite_prompt'] = sub_df['rewrite_prompt'].map(lambda x: \"Improve the essay\" if len(x) == 0 else x)\n",
        "sub_df.to_csv(\"submission.csv\",index=False)\n",
        "sub_df.head()"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-03-08T19:29:44.702671Z",
          "iopub.execute_input": "2024-03-08T19:29:44.703073Z",
          "iopub.status.idle": "2024-03-08T19:29:44.721731Z",
          "shell.execute_reply.started": "2024-03-08T19:29:44.703039Z",
          "shell.execute_reply": "2024-03-08T19:29:44.720264Z"
        },
        "trusted": true,
        "id": "VLVh5fpFnxm1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Conclusion\n",
        "\n",
        "The result is pretty good. Still there is ample room for improvement. Here are some tips to improve performance:\n",
        "\n",
        "- Try using the larger version of **Gemma** (7B).\n",
        "- Increase `sequence_length`.\n",
        "- Experiment with advanced prompt engineering techniques.\n",
        "- Implement augmentation to increase the number of samples.\n",
        "- Utilize a learning rate scheduler."
      ],
      "metadata": {
        "id": "Lpo5wUm6nxm2"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Reference\n",
        "* [Fine-tune Gemma models in Keras using LoRA](https://www.kaggle.com/code/nilaychauhan/fine-tune-gemma-models-in-keras-using-lora)\n",
        "* [Parameter-efficient fine-tuning of GPT-2 with LoRA](https://keras.io/examples/nlp/parameter_efficient_finetuning_of_gpt2_with_lora/)\n",
        "* [Gemma - KerasNLP](https://keras.io/api/keras_nlp/models/gemma/)"
      ],
      "metadata": {
        "id": "zSQkPka3nxm2"
      }
    }
  ]
}