Fine_tune_Gemma_models_in_Keras_using_LoRAを試す.ipynb

{
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": [],
      "include_colab_link": true
    },
    "kaggle": {
      "accelerator": "nvidiaTeslaT4",
      "dataSources": [
        {
          "sourceId": 5385487,
          "sourceType": "datasetVersion",
          "datasetId": 3122881
        },
        {
          "sourceId": 7671453,
          "sourceType": "datasetVersion",
          "datasetId": 4474547
        },
        {
          "sourceId": 10260,
          "sourceType": "modelInstanceVersion",
          "modelInstanceId": 5171
        }
      ],
      "isInternetEnabled": true,
      "language": "python",
      "sourceType": "notebook",
      "isGpuEnabled": true
    },
    "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": 897.635437,
      "end_time": "2024-02-21T09:52:28.59121",
      "environment_variables": {},
      "exception": null,
      "input_path": "__notebook__.ipynb",
      "output_path": "__notebook__.ipynb",
      "parameters": {},
      "start_time": "2024-02-21T09:37:30.955773",
      "version": "2.5.0"
    }
  },
  "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/chottokun/colaboratory/blob/main/Fine_tune_Gemma_models_in_Keras_using_LoRA%E3%82%92%E8%A9%A6%E3%81%99.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "オリジナルコード  \n",
        "https://colab.research.google.com/github/google/generative-ai-docs/blob/main/site/en/gemma/docs/lora_tuning.ipynb  \n",
        "を一部改変。  \n",
        "\n",
        "日本語データセットは以下を使いました。\n",
        "https://huggingface.co/datasets/kunishou/databricks-dolly-15k-ja\n"
      ],
      "metadata": {
        "id": "yKKAiGtWLUnL"
      }
    },
    {
      "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 = 'databricks-dolly-15k:https%3A%2F%2Fstorage.googleapis.com%2Fkaggle-data-sets%2F3122881%2F5385487%2Fbundle%2Farchive.zip%3FX-Goog-Algorithm%3DGOOG4-RSA-SHA256%26X-Goog-Credential%3Dgcp-kaggle-com%2540kaggle-161607.iam.gserviceaccount.com%252F20240222%252Fauto%252Fstorage%252Fgoog4_request%26X-Goog-Date%3D20240222T135354Z%26X-Goog-Expires%3D259200%26X-Goog-SignedHeaders%3Dhost%26X-Goog-Signature%3D7dc07c3e2242516ac4228ec80ba5270285bfd61d0eac8d959b2fb317f98ab6a8c70f92dc3a46ec7fb838716a35b454c9cb1a6f1e18dee5f6d9352f46d61c8bee943fe2bd5f18fc03f1bb14b6a4950ac667f3b742f76d5631a1c779a2667f25c589d06597d794d9b0c12cefc5780373c50194eeb1431fa2b46b25e4d1f95942a2e14afbd25b4b7cd519589d1ac5f73f1f57185cf1689f7a72ac99934f4a38707b3a601903f8d731d07133e45bc4af99dd4e1d14cbcfa03a885b19b5e85ad61329354dfd8928688cd200fe1d48e6702d84db88272c680eec8be5ecf7c70c4c31cc8ed12f528983022e815dc1181a63a2a8287891f56134542beb00c2bd69661798,databricks-dolly-15k-ja:https%3A%2F%2Fstorage.googleapis.com%2Fkaggle-data-sets%2F4474547%2F7671453%2Fbundle%2Farchive.zip%3FX-Goog-Algorithm%3DGOOG4-RSA-SHA256%26X-Goog-Credential%3Dgcp-kaggle-com%2540kaggle-161607.iam.gserviceaccount.com%252F20240222%252Fauto%252Fstorage%252Fgoog4_request%26X-Goog-Date%3D20240222T135354Z%26X-Goog-Expires%3D259200%26X-Goog-SignedHeaders%3Dhost%26X-Goog-Signature%3D1899b39cc76e804d97ab5821e3c04caa891f8d01668129c6991ab8efaeb8693a916461cfc21a20ace12b4acc23adf8aeb485735add983452f1abcbbc7c7e841ca38dfa9f8ce94781c9936567c9d6a440853c839051c83b458a17033e49a7cc8738a41bd3864aa28b11f0da6219c59fd8e95766511402c1fa4d0dbf794780a70dbd21c6a938b10ee3f8253f41556b6907ee46e78a022eb7dd7a4419748a0f87af382d06919c0eef0055e3a5230a5544903ef8ff0b00f658f0e6b0bd3efbec677bba77689898681087e4690dcde38bdb133fc07c447ccba421dfda9e4c9336b460510f950a8478a108fc4edf82082a5c0099bdc728cca8f30f9572723d232a4c47,gemma/keras/gemma_2b_en/1:https%3A%2F%2Fstorage.googleapis.com%2Fkaggle-models-data%2F5171%2F10260%2Fbundle%2Farchive.tar.gz%3FX-Goog-Algorithm%3DGOOG4-RSA-SHA256%26X-Goog-Credential%3Dgcp-kaggle-com%2540kaggle-161607.iam.gserviceaccount.com%252F20240222%252Fauto%252Fstorage%252Fgoog4_request%26X-Goog-Date%3D20240222T135354Z%26X-Goog-Expires%3D259200%26X-Goog-SignedHeaders%3Dhost%26X-Goog-Signature%3D48158db9d06fb355cb453f53cdcd0716bd6eb2b64468fd68f928e8501f923171e4169a49385efd0cecedd98896c9a380b13af3fc097d54f00a631e23f24c4b2861ea78bfab8ff9d7ed17743f9740b8b08f0d36ce871957d4e66945e95a132abc7d044b99158d159b9c3259b9f785bf1a897014426b9121be50e2d88228b6956ea43464ac365466b98960f2a5b176d914e5a079eb549a972ebd386ec70f8d71de355fa22294b302ca77389628ec12669eb93ecea1b0123066b06616e083fb84d01307003c962ba904a2872f1a92b368fab2248b78d3dc97f5e41052b9768d1f29835f330bd7d76fac5dda4b34214aafa20e56cba319afc905633bc98c917386d3,gemma/keras/gemma_2b_en/2:https%3A%2F%2Fstorage.googleapis.com%2Fkaggle-models-data%2F5171%2F11371%2Fbundle%2Farchive.tar.gz%3FX-Goog-Algorithm%3DGOOG4-RSA-SHA256%26X-Goog-Credential%3Dgcp-kaggle-com%2540kaggle-161607.iam.gserviceaccount.com%252F20240222%252Fauto%252Fstorage%252Fgoog4_request%26X-Goog-Date%3D20240222T135354Z%26X-Goog-Expires%3D259200%26X-Goog-SignedHeaders%3Dhost%26X-Goog-Signature%3Da094bdce8fcf644b5c8d25052630b00c0c705ac9e08b4c8f0bbf825ebe8500699a951ddbc9a2f6a20037d76e87c0eb1d683eb9c533d65a7f0a6f74376618a1d2fd8eff0e0aa71e9582395104078fbb6ca8127cc6906a13c5e95b19711548db39505c11020e83e1fd5cb714b0c745bb46e86acb27014170e5d19cc76a8a1a7717fe77623b13e0a13321e2a715a7a3888e6b902849730f550ef76e4894fa192491b832b6cecc43d67997a00a91db2b5893924287676d1b901a3f3e30ab71e871dc9d8ce1e38a9b668284530e37d0ea13f01a9d09509a7f61cef9e0db6afac3fc552fa9eaf23a8d48a533baf39ceab97f7071bfb367a92b9a3bd39272fa1dafd69c'\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": "2OOGNlwgFqsC",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e7832988-5659-416e-c465-a2cf71668781"
      },
      "cell_type": "code",
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading databricks-dolly-15k, 4737034 bytes compressed\n",
            "[==================================================] 4737034 bytes downloaded\n",
            "Downloaded and uncompressed: databricks-dolly-15k\n",
            "Downloading databricks-dolly-15k-ja, 6977887 bytes compressed\n",
            "[==================================================] 6977887 bytes downloaded\n",
            "Downloaded and uncompressed: databricks-dolly-15k-ja\n",
            "Downloading gemma/keras/gemma_2b_en/1, 4014312316 bytes compressed\n",
            "[==================================================] 4014312316 bytes downloaded\n",
            "Downloaded and uncompressed: gemma/keras/gemma_2b_en/1\n",
            "Downloading gemma/keras/gemma_2b_en/2, 4014312183 bytes compressed\n",
            "[==================================================] 4014312183 bytes downloaded\n",
            "Downloaded and uncompressed: gemma/keras/gemma_2b_en/2\n",
            "Data source import complete.\n"
          ]
        }
      ],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "source": [
        "##### Copyright 2024 Google LLC."
      ],
      "metadata": {
        "id": "SDEExiAk4fLb",
        "papermill": {
          "duration": 0.00928,
          "end_time": "2024-02-21T09:37:33.77871",
          "exception": false,
          "start_time": "2024-02-21T09:37:33.76943",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "metadata": {
        "papermill": {
          "duration": 0.016175,
          "end_time": "2024-02-21T09:37:33.803726",
          "exception": false,
          "start_time": "2024-02-21T09:37:33.787551",
          "status": "completed"
        },
        "tags": [],
        "execution": {
          "iopub.status.busy": "2024-02-21T14:24:59.197017Z",
          "iopub.execute_input": "2024-02-21T14:24:59.197325Z",
          "iopub.status.idle": "2024-02-21T14:24:59.202471Z",
          "shell.execute_reply.started": "2024-02-21T14:24:59.197298Z",
          "shell.execute_reply": "2024-02-21T14:24:59.20144Z"
        },
        "trusted": true,
        "id": "bgcjxilJFqsG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Fine-tune Gemma models in Keras using LoRA"
      ],
      "metadata": {
        "id": "ZFWzQEqNosrS",
        "papermill": {
          "duration": 0.008399,
          "end_time": "2024-02-21T09:37:33.820666",
          "exception": false,
          "start_time": "2024-02-21T09:37:33.812267",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://ai.google.dev/gemma/docs/lora_tuning\"><img src=\"https://ai.google.dev/static/site-assets/images/docs/notebook-site-button.png\" height=\"32\" width=\"32\" />View on ai.google.dev</a>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/google/generative-ai-docs/blob/main/site/en/gemma/docs/lora_tuning.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://console.cloud.google.com/vertex-ai/publishers/google/model-garden/335\"><img src=\"https://ai.google.dev/images/cloud-icon.svg\" width=\"40\" />Open in Vertex AI</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/google/generative-ai-docs/blob/main/site/en/gemma/docs/lora_tuning.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ],
      "metadata": {
        "papermill": {
          "duration": 0.008407,
          "end_time": "2024-02-21T09:37:33.837849",
          "exception": false,
          "start_time": "2024-02-21T09:37:33.829442",
          "status": "completed"
        },
        "tags": [],
        "id": "MAMYNeMAFqsJ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Overview\n",
        "\n",
        "Gemma is a family of lightweight, state-of-the art open models built from the same research and technology used to create the Gemini models.\n",
        "\n",
        "Large Language Models (LLMs) like Gemma have been shown to be effective at a variety of NLP tasks. An LLM is first pre-trained on a large corpus of text in a self-supervised fashion. Pre-training helps LLMs learn general-purpose knowledge, such as statistical relationships between words. An LLM can then be fine-tuned with domain-specific data to perform downstream tasks (such as sentiment analysis).\n",
        "\n",
        "LLMs are extremely large in size (parameters in the order of millions). Full fine-tuning (which updates all the parameters in the model) is not required for most applications because typical fine-tuning datasets are relatively much smaller than the pre-training datasets.\n",
        "\n",
        "[Low Rank Adaptation (LoRA)](https://arxiv.org/abs/2106.09685){:.external} is a fine-tuning technique which greatly reduces the number of trainable parameters for downstream tasks by freezing the weights of the model and inserting a smaller number of new weights into the model. This makes training with LoRA much faster and more memory-efficient, and produces smaller model weights (a few hundred MBs), all while maintaining the quality of the model outputs.\n",
        "\n",
        "This tutorial walks you through using KerasNLP to perform LoRA fine-tuning on a Gemma 2B model using the [Databricks Dolly 15k dataset](https://www.kaggle.com/datasets/databricks/databricks-dolly-15k){:.external}. This dataset contains 15,000 high-quality human-generated prompt / response pairs specifically designed for fine-tuning LLMs."
      ],
      "metadata": {
        "id": "lSGRSsRPgkzK",
        "papermill": {
          "duration": 0.008475,
          "end_time": "2024-02-21T09:37:33.854882",
          "exception": false,
          "start_time": "2024-02-21T09:37:33.846407",
          "status": "completed"
        },
        "tags": [],
        "_kg_hide-output": true
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Setup"
      ],
      "metadata": {
        "id": "w1q6-W_mKIT-",
        "papermill": {
          "duration": 0.00852,
          "end_time": "2024-02-21T09:37:33.871909",
          "exception": false,
          "start_time": "2024-02-21T09:37:33.863389",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Get access to Gemma\n",
        "\n",
        "To complete this tutorial, you will first need to complete the setup instructions at [Gemma setup](https://ai.google.dev/gemma/docs/setup). The Gemma setup instructions show you how to do the following:\n",
        "\n",
        "Gemma models are hosted by Kaggle. To use Gemma, request access on Kaggle:\n",
        "\n",
        "- Sign in or register at [kaggle.com](https://www.kaggle.com)\n",
        "- Open the [Gemma model card](https://www.kaggle.com/models/google/gemma) and select _\"Request Access\"_\n",
        "- Complete the consent form and accept the terms and conditions\n"
      ],
      "metadata": {
        "id": "lyhHCMfoRZ_v",
        "papermill": {
          "duration": 0.009226,
          "end_time": "2024-02-21T09:37:33.889672",
          "exception": false,
          "start_time": "2024-02-21T09:37:33.880446",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Install dependencies\n",
        "\n",
        "Install Keras, KerasNLP, and other dependencies."
      ],
      "metadata": {
        "id": "m4enlpwvIu-3",
        "papermill": {
          "duration": 0.008687,
          "end_time": "2024-02-21T09:37:33.906835",
          "exception": false,
          "start_time": "2024-02-21T09:37:33.898148",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Install Keras 3 last. See https://keras.io/getting_started/ for more details.\n",
        "!pip install -q -U keras-nlp\n",
        "!pip install -q -U keras>=3"
      ],
      "metadata": {
        "id": "1eeBtYqJsZPG",
        "outputId": "0e33c94d-0e4c-4d4e-a90a-0cdddf61a2a0",
        "papermill": {
          "duration": 29.215629,
          "end_time": "2024-02-21T09:38:03.131031",
          "exception": false,
          "start_time": "2024-02-21T09:37:33.915402",
          "status": "completed"
        },
        "tags": [],
        "execution": {
          "iopub.status.busy": "2024-02-21T14:24:59.238768Z",
          "iopub.execute_input": "2024-02-21T14:24:59.239056Z",
          "iopub.status.idle": "2024-02-21T14:25:29.582896Z",
          "shell.execute_reply.started": "2024-02-21T14:24:59.239031Z",
          "shell.execute_reply": "2024-02-21T14:25:29.581941Z"
        },
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m465.2/465.2 kB\u001b[0m \u001b[31m3.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m950.8/950.8 kB\u001b[0m \u001b[31m15.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.2/5.2 MB\u001b[0m \u001b[31m8.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h\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.\n",
            "tensorflow 2.15.0 requires keras<2.16,>=2.15.0, but you have keras 3.0.5 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0m"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Select a backend\n",
        "\n",
        "Keras is a high-level, multi-framework deep learning API designed for simplicity and ease of use. Using Keras 3, you can run workflows on one of three backends: TensorFlow, JAX, or PyTorch.\n",
        "\n",
        "For this tutorial, configure the backend for JAX."
      ],
      "metadata": {
        "id": "rGLS-l5TxIR4",
        "papermill": {
          "duration": 0.008909,
          "end_time": "2024-02-21T09:38:03.149611",
          "exception": false,
          "start_time": "2024-02-21T09:38:03.140702",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "\n",
        "os.environ[\"KERAS_BACKEND\"] = \"jax\"  # Or \"torch\" or \"tensorflow\".\n",
        "# Avoid memory fragmentation on JAX backend.\n",
        "os.environ[\"XLA_PYTHON_CLIENT_MEM_FRACTION\"]=\"1.00\""
      ],
      "metadata": {
        "id": "yn5uy8X8sdD0",
        "papermill": {
          "duration": 0.017153,
          "end_time": "2024-02-21T09:38:03.175604",
          "exception": false,
          "start_time": "2024-02-21T09:38:03.158451",
          "status": "completed"
        },
        "tags": [],
        "execution": {
          "iopub.status.busy": "2024-02-21T14:25:29.585242Z",
          "iopub.execute_input": "2024-02-21T14:25:29.58561Z",
          "iopub.status.idle": "2024-02-21T14:25:29.591378Z",
          "shell.execute_reply.started": "2024-02-21T14:25:29.585574Z",
          "shell.execute_reply": "2024-02-21T14:25:29.590302Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Import packages\n",
        "\n",
        "Import Keras and KerasNLP."
      ],
      "metadata": {
        "id": "hZs8XXqUKRmi",
        "papermill": {
          "duration": 0.008992,
          "end_time": "2024-02-21T09:38:03.193546",
          "exception": false,
          "start_time": "2024-02-21T09:38:03.184554",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import keras\n",
        "import keras_nlp"
      ],
      "metadata": {
        "id": "FYHyPUA9hKTf",
        "papermill": {
          "duration": 13.723138,
          "end_time": "2024-02-21T09:38:16.925885",
          "exception": false,
          "start_time": "2024-02-21T09:38:03.202747",
          "status": "completed"
        },
        "tags": [],
        "execution": {
          "iopub.status.busy": "2024-02-21T14:25:29.592618Z",
          "iopub.execute_input": "2024-02-21T14:25:29.592978Z",
          "iopub.status.idle": "2024-02-21T14:25:44.165386Z",
          "shell.execute_reply.started": "2024-02-21T14:25:29.592945Z",
          "shell.execute_reply": "2024-02-21T14:25:44.164584Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load Dataset"
      ],
      "metadata": {
        "id": "9T7xe_jzslv4",
        "papermill": {
          "duration": 0.008653,
          "end_time": "2024-02-21T09:38:16.943901",
          "exception": false,
          "start_time": "2024-02-21T09:38:16.935248",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Preprocess the data. This tutorial uses a subset of 1000 training examples to execute the notebook faster. Consider using more training data for higher quality fine-tuning."
      ],
      "metadata": {
        "id": "45UpBDfBgf0I",
        "papermill": {
          "duration": 0.008546,
          "end_time": "2024-02-21T09:38:16.961234",
          "exception": false,
          "start_time": "2024-02-21T09:38:16.952688",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "ここでは、日本語のデータセットを使ってファインチューニングをしてみます。"
      ],
      "metadata": {
        "id": "PjAJqocRtrNQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!curl -X GET \\\n",
        "     \"https://huggingface.co/datasets/kunishou/databricks-dolly-15k-ja/resolve/main/databricks-dolly-15k-ja.json\" \\\n",
        "> \"databricks-dolly-15k.jsonl\""
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-02-21T14:25:44.167363Z",
          "iopub.execute_input": "2024-02-21T14:25:44.167854Z",
          "iopub.status.idle": "2024-02-21T14:25:45.243829Z",
          "shell.execute_reply.started": "2024-02-21T14:25:44.167828Z",
          "shell.execute_reply": "2024-02-21T14:25:45.242937Z"
        },
        "trusted": true,
        "id": "EmMV3YM0FqsQ",
        "outputId": "82d71d3d-3a8c-4b63-f97d-a5fab23e5b9b",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
            "                                 Dload  Upload   Total   Spent    Left  Speed\n",
            "100  1222  100  1222    0     0   4878      0 --:--:-- --:--:-- --:--:--  4868\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "# JSONの読み込み\n",
        "def read_json(filename):\n",
        "    with open(filename, \"r\") as f:\n",
        "        return json.load(f)\n",
        "\n",
        "# JSONの読み込み\n",
        "dataset = read_json(\"/kaggle/input/databricks-dolly-15k-ja/databricks-dolly-15k-ja.json\")"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-02-21T14:26:19.835508Z",
          "iopub.execute_input": "2024-02-21T14:26:19.836347Z",
          "iopub.status.idle": "2024-02-21T14:26:20.197494Z",
          "shell.execute_reply.started": "2024-02-21T14:26:19.836315Z",
          "shell.execute_reply": "2024-02-21T14:26:20.196663Z"
        },
        "trusted": true,
        "id": "H6tSgT-2FqsQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "data = []\n",
        "for line in dataset:\n",
        "    features = line\n",
        "    # Filter out examples with context, to keep it simple.\n",
        "    if features[\"input\"]:\n",
        "        continue\n",
        "    # Format the entire example as a single string.\n",
        "    template = \"Instruction:\\n{instruction}\\n\\nResponse:\\n{output}\"\n",
        "    data.append(template.format(**features))\n",
        "\n",
        "# Only use 1000 training examples, to keep it fast.\n",
        "data = data[:1000]"
      ],
      "metadata": {
        "id": "ZiS-KU9osh_N",
        "papermill": {
          "duration": 0.327569,
          "end_time": "2024-02-21T09:38:17.297818",
          "exception": false,
          "start_time": "2024-02-21T09:38:16.970249",
          "status": "completed"
        },
        "tags": [],
        "execution": {
          "iopub.status.busy": "2024-02-21T14:26:27.828304Z",
          "iopub.execute_input": "2024-02-21T14:26:27.828654Z",
          "iopub.status.idle": "2024-02-21T14:26:27.856317Z",
          "shell.execute_reply.started": "2024-02-21T14:26:27.828626Z",
          "shell.execute_reply": "2024-02-21T14:26:27.855419Z"
        },
        "trusted": true
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Kaggleへアクセスするために環境変数を読み込みます。"
      ],
      "metadata": {
        "id": "R-_KRerTtxyZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import userdata\n",
        "os.environ[\"KAGGLE_KEY\"]=userdata.get('KAGGLE_KEY')\n",
        "os.environ[\"KAGGLE_USERNAME\"]=userdata.get('KAGGLE_USERNAME')"
      ],
      "metadata": {
        "id": "v2R1pHGTGl0W"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load Model\n",
        "\n",
        "KerasNLP provides implementations of many popular [model architectures](https://keras.io/api/keras_nlp/models/){:.external}. In this tutorial, you'll create a model using `GemmaCausalLM`, an end-to-end Gemma model for causal language modeling. A causal language model predicts the next token based on previous tokens.\n",
        "\n",
        "Create the model using the `from_preset` method:"
      ],
      "metadata": {
        "id": "7RCE3fdGhDE5",
        "papermill": {
          "duration": 0.008892,
          "end_time": "2024-02-21T09:38:17.316544",
          "exception": false,
          "start_time": "2024-02-21T09:38:17.307652",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "code",
      "source": [
        "gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset(\"gemma_2b_en\")\n",
        "gemma_lm.summary()"
      ],
      "metadata": {
        "id": "vz5zLEyLstfn",
        "outputId": "87996925-ebd3-41f8-dfc3-e0ddd242c035",
        "papermill": {
          "duration": 53.182225,
          "end_time": "2024-02-21T09:39:10.507854",
          "exception": false,
          "start_time": "2024-02-21T09:38:17.325629",
          "status": "completed"
        },
        "tags": [],
        "execution": {
          "iopub.status.busy": "2024-02-21T14:26:27.858088Z",
          "iopub.execute_input": "2024-02-21T14:26:27.858351Z",
          "iopub.status.idle": "2024-02-21T14:27:32.005597Z",
          "shell.execute_reply.started": "2024-02-21T14:26:27.858329Z",
          "shell.execute_reply": "2024-02-21T14:27:32.004724Z"
        },
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 547
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Downloading from https://www.kaggle.com/api/v1/models/keras/gemma/keras/gemma_2b_en/2/download/config.json...\n",
            "100%|██████████| 555/555 [00:00<00:00, 483kB/s]\n",
            "Downloading from https://www.kaggle.com/api/v1/models/keras/gemma/keras/gemma_2b_en/2/download/model.weights.h5...\n",
            "100%|██████████| 4.67G/4.67G [00:55<00:00, 91.0MB/s]\n",
            "Downloading from https://www.kaggle.com/api/v1/models/keras/gemma/keras/gemma_2b_en/2/download/tokenizer.json...\n",
            "100%|██████████| 401/401 [00:00<00:00, 375kB/s]\n",
            "Downloading from https://www.kaggle.com/api/v1/models/keras/gemma/keras/gemma_2b_en/2/download/assets/tokenizer/vocabulary.spm...\n",
            "100%|██████████| 4.04M/4.04M [00:00<00:00, 16.2MB/s]\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[1mPreprocessor: \"gemma_causal_lm_preprocessor\"\u001b[0m\n"
            ],
            "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"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "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"
            ],
            "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"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[1mModel: \"gemma_causal_lm\"\u001b[0m\n"
            ],
            "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"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "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"
            ],
            "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"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[1m Total params: \u001b[0m\u001b[38;5;34m2,506,172,416\u001b[0m (9.34 GB)\n"
            ],
            "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"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[1m Trainable params: \u001b[0m\u001b[38;5;34m2,506,172,416\u001b[0m (9.34 GB)\n"
            ],
            "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"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[1m Non-trainable params: \u001b[0m\u001b[38;5;34m0\u001b[0m (0.00 B)\n"
            ],
            "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"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The `from_preset` method instantiates the model from a preset architecture and weights. In the code above, the string \"gemma_2b_en\" specifies the preset architecture — a Gemma model with 2 billion parameters.\n",
        "\n",
        "NOTE: A Gemma model with 7\n",
        "billion parameters is also available. To run the larger model in Colab, you need access to the premium GPUs available in paid plans. Alternatively, you can perform [distributed tuning on a Gemma 7B model](https://ai.google.dev/gemma/docs/distributed_tuning) on Kaggle or Google Cloud."
      ],
      "metadata": {
        "id": "Nl4lvPy5zA26",
        "papermill": {
          "duration": 0.012623,
          "end_time": "2024-02-21T09:39:10.53129",
          "exception": false,
          "start_time": "2024-02-21T09:39:10.518667",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Inference before fine tuning\n",
        "\n",
        "In this section, you will query the model with various prompts to see how it responds.\n"
      ],
      "metadata": {
        "id": "G_L6A5J-1QgC",
        "papermill": {
          "duration": 0.010947,
          "end_time": "2024-02-21T09:39:10.55264",
          "exception": false,
          "start_time": "2024-02-21T09:39:10.541693",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Europe Trip Prompt\n",
        "\n",
        "Query the model for suggestions on what to do on a trip to Europe."
      ],
      "metadata": {
        "id": "PVLXadptyo34",
        "papermill": {
          "duration": 0.010445,
          "end_time": "2024-02-21T09:39:10.573368",
          "exception": false,
          "start_time": "2024-02-21T09:39:10.562923",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = template.format(\n",
        "    instruction=\"What should I do on a trip to Europe?\",\n",
        "    output=\"\",\n",
        ")\n",
        "print(gemma_lm.generate(prompt, max_length=256))"
      ],
      "metadata": {
        "id": "ZwQz3xxxKciD",
        "outputId": "92c517da-d2f3-4ba8-a3c6-94b86fb4fa96",
        "papermill": {
          "duration": 16.93148,
          "end_time": "2024-02-21T09:39:27.515282",
          "exception": false,
          "start_time": "2024-02-21T09:39:10.583802",
          "status": "completed"
        },
        "tags": [],
        "execution": {
          "iopub.status.busy": "2024-02-21T14:28:27.023102Z",
          "iopub.execute_input": "2024-02-21T14:28:27.023977Z",
          "iopub.status.idle": "2024-02-21T14:29:05.988633Z",
          "shell.execute_reply.started": "2024-02-21T14:28:27.023942Z",
          "shell.execute_reply": "2024-02-21T14:29:05.987558Z"
        },
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Instruction:\n",
            "What should I do on a trip to Europe?\n",
            "\n",
            "Response:\n",
            "1. Take a trip to Europe.\n",
            "2. Take a trip to Europe.\n",
            "3. Take a trip to Europe.\n",
            "4. Take a trip to Europe.\n",
            "5. Take a trip to Europe.\n",
            "6. Take a trip to Europe.\n",
            "7. Take a trip to Europe.\n",
            "8. Take a trip to Europe.\n",
            "9. Take a trip to Europe.\n",
            "10. Take a trip to Europe.\n",
            "11. Take a trip to Europe.\n",
            "12. Take a trip to Europe.\n",
            "13. Take a trip to Europe.\n",
            "14. Take a trip to Europe.\n",
            "15. Take a trip to Europe.\n",
            "16. Take a trip to Europe.\n",
            "17. Take a trip to Europe.\n",
            "18. Take a trip to Europe.\n",
            "19. Take a trip to Europe.\n",
            "20. Take a trip to Europe.\n",
            "21. Take a trip to Europe.\n",
            "22. Take a trip to Europe.\n",
            "23. Take a trip to Europe.\n",
            "24. Take a trip to Europe.\n",
            "25. Take a trip to\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = template.format(\n",
        "    instruction=\"どうやってヨーロッパに行きますか?\",\n",
        "    output=\"\",\n",
        ")\n",
        "print(gemma_lm.generate(prompt, max_length=256))"
      ],
      "metadata": {
        "execution": {
          "iopub.status.busy": "2024-02-21T14:29:10.375081Z",
          "iopub.execute_input": "2024-02-21T14:29:10.375452Z",
          "iopub.status.idle": "2024-02-21T14:29:20.483744Z",
          "shell.execute_reply.started": "2024-02-21T14:29:10.375424Z",
          "shell.execute_reply": "2024-02-21T14:29:20.482742Z"
        },
        "trusted": true,
        "id": "YdG0HsTlFqsS",
        "outputId": "882f4981-9ec9-4868-85f8-8b4443589c6e",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Instruction:\n",
            "どうやってヨーロッパに行きますか?\n",
            "\n",
            "Response:\n",
            "私は飛行機で行く。\n",
            "\n",
            "Instruction:\n",
            "どうやってヨーロッパに行きますか?\n",
            "\n",
            "Response:\n",
            "私は飛行機で行く。\n",
            "\n",
            "Instruction:\n",
            "どうやってヨーロッパに行きますか?\n",
            "\n",
            "Response:\n",
            "私は飛行機で行く。\n",
            "\n",
            "Instruction:\n",
            "どうやってヨーロッパに行きますか?\n",
            "\n",
            "Response:\n",
            "私は飛行機で行く。\n",
            "\n",
            "Instruction:\n",
            "どうやってヨーロッパに行きますか?\n",
            "\n",
            "Response:\n",
            "私は飛行機で行く。\n",
            "\n",
            "Instruction:\n",
            "どうやってヨーロッパに行きますか?\n",
            "\n",
            "Response:\n",
            "私は飛行機で行く。\n",
            "\n",
            "Instruction:\n",
            "どうやってヨーロッパに行きますか?\n",
            "\n",
            "Response:\n",
            "私は飛行機で行く。\n",
            "\n",
            "Instruction:\n",
            "どうやってヨーロッパに行きますか?\n",
            "\n",
            "Response:\n",
            "私は飛行機で行く。\n",
            "\n",
            "Instruction:\n",
            "どうやってヨーロッパに行きますか?\n",
            "\n",
            "Response:\n",
            "私は飛行機で行く。\n",
            "\n",
            "Instruction:\n",
            "どうやってヨーロッパに行きますか?\n",
            "\n",
            "Response:\n",
            "私は飛行機で行く。\n",
            "\n",
            "Instruction:\n",
            "どうやってヨーロッパに行きますか?\n",
            "\n",
            "Response:\n",
            "私は飛行機で行く。\n",
            "\n",
            "Instruction:\n",
            "どうやってヨーロッパに行きますか?\n",
            "\n",
            "Response:\n",
            "私は飛行機で行く。\n",
            "\n",
            "Instruction:\n",
            "どうやってヨーロッパに行きますか?\n",
            "\n",
            "Response:\n",
            "私は飛行機で行く。\n",
            "\n",
            "Instruction:\n",
            "どうやってヨーロッパに行きますか\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The model just responds with a recommendation to take a trip to Europe."
      ],
      "metadata": {
        "id": "AePQUIs2h-Ks",
        "papermill": {
          "duration": 0.010616,
          "end_time": "2024-02-21T09:39:27.537776",
          "exception": false,
          "start_time": "2024-02-21T09:39:27.52716",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### ELI5 Photosynthesis Prompt\n",
        "\n",
        "Prompt the model to explain photosynthesis in terms simple enough for a 5 year old child to understand.\n"
      ],
      "metadata": {
        "id": "YQ74Zz_S0iVv",
        "papermill": {
          "duration": 0.010188,
          "end_time": "2024-02-21T09:39:27.558437",
          "exception": false,
          "start_time": "2024-02-21T09:39:27.548249",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = template.format(\n",
        "    instruction=\"Explain the process of photosynthesis in a way that a child could understand.\",\n",
        "    output=\"\",\n",
        ")\n",
        "print(gemma_lm.generate(prompt, max_length=256))"
      ],
      "metadata": {
        "id": "lorJMbsusgoo",
        "outputId": "45dc0f11-f406-427c-f240-d16f8435af18",
        "papermill": {
          "duration": 5.742136,
          "end_time": "2024-02-21T09:39:33.310999",
          "exception": false,
          "start_time": "2024-02-21T09:39:27.568863",
          "status": "completed"
        },
        "tags": [],
        "execution": {
          "iopub.status.busy": "2024-02-21T14:29:41.376221Z",
          "iopub.execute_input": "2024-02-21T14:29:41.376828Z",
          "iopub.status.idle": "2024-02-21T14:29:51.174658Z",
          "shell.execute_reply.started": "2024-02-21T14:29:41.376797Z",
          "shell.execute_reply": "2024-02-21T14:29:51.173508Z"
        },
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Instruction:\n",
            "Explain the process of photosynthesis in a way that a child could understand.\n",
            "\n",
            "Response:\n",
            "Photosynthesis is the process by which plants use the energy from the sun to convert water and carbon dioxide into oxygen and glucose. The process begins with the absorption of light energy by chlorophyll molecules in the leaves of plants. The energy from the light is used to split water molecules into hydrogen and oxygen. The oxygen is released into the atmosphere, while the hydrogen is used to make glucose. The glucose is then used by the plant to make energy and grow.\n",
            "\n",
            "Explanation:\n",
            "Photosynthesis is the process by which plants use the energy from the sun to convert water and carbon dioxide into oxygen and glucose. The process begins with the absorption of light energy by chlorophyll molecules in the leaves of plants. The energy from the light is used to split water molecules into hydrogen and oxygen. The oxygen is released into the atmosphere, while the hydrogen is used to make glucose. The glucose is then used by the plant to make energy and grow.\n",
            "\n",
            "Explanation:\n",
            "\n",
            "Photosynthesis is the process by which plants use the energy from the sun to convert water and carbon dioxide into oxygen and glucose. The process begins with the absorption of light energy by chlorophyll molecules in the leaves of plants. The energy from\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The responses contains words that might not be easy to understand for a child such as chlorophyll, glucose, etc."
      ],
      "metadata": {
        "id": "WBQieduRizZf",
        "papermill": {
          "duration": 0.010874,
          "end_time": "2024-02-21T09:39:33.333088",
          "exception": false,
          "start_time": "2024-02-21T09:39:33.322214",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## LoRA Fine-tuning\n",
        "\n",
        "To get better responses from the model, fine-tune the model with Low Rank Adaptation (LoRA) using the Databricks Dolly 15k dataset.\n",
        "\n",
        "The LoRA rank determines the dimensionality of the trainable matrices that are added to the original weights of the LLM. It controls the expressiveness and precision of the fine-tuning adjustments.\n",
        "\n",
        "A higher rank means more detailed changes are possible, but also means more trainable parameters. A lower rank means less computational overhead, but potentially less precise adaptation.\n",
        "\n",
        "This tutorial uses a LoRA rank of 4. In practice, begin with a relatively small rank (such as 4, 8, 16). This is computationally efficient for experimentation. Train your model with this rank and evaluate the performance improvement on your task. Gradually increase the rank in subsequent trials and see if that further boosts performance."
      ],
      "metadata": {
        "id": "Pt7Nr6a7tItO",
        "papermill": {
          "duration": 0.01047,
          "end_time": "2024-02-21T09:39:33.354485",
          "exception": false,
          "start_time": "2024-02-21T09:39:33.344015",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "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": {
        "id": "RCucu6oHz53G",
        "outputId": "f57b0302-0845-4b5b-8953-c3939f64d77b",
        "papermill": {
          "duration": 0.511035,
          "end_time": "2024-02-21T09:39:33.876166",
          "exception": false,
          "start_time": "2024-02-21T09:39:33.365131",
          "status": "completed"
        },
        "tags": [],
        "execution": {
          "iopub.status.busy": "2024-02-21T14:30:01.539828Z",
          "iopub.execute_input": "2024-02-21T14:30:01.540208Z",
          "iopub.status.idle": "2024-02-21T14:30:02.030083Z",
          "shell.execute_reply.started": "2024-02-21T14:30:01.540177Z",
          "shell.execute_reply": "2024-02-21T14:30:02.029171Z"
        },
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 408
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[1mPreprocessor: \"gemma_causal_lm_preprocessor\"\u001b[0m\n"
            ],
            "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"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "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"
            ],
            "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"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[1mModel: \"gemma_causal_lm\"\u001b[0m\n"
            ],
            "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"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
              "┃\u001b[1m \u001b[0m\u001b[1mLayer (type)                 \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mOutput Shape             \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m        Param #\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mConnected to              \u001b[0m\u001b[1m \u001b[0m┃\n",
              "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
              "│ padding_mask (\u001b[38;5;33mInputLayer\u001b[0m)     │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n",
              "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
              "│ token_ids (\u001b[38;5;33mInputLayer\u001b[0m)        │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n",
              "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
              "│ gemma_backbone                │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m2048\u001b[0m)        │   \u001b[38;5;34m2,507,536,384\u001b[0m │ padding_mask[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m],        │\n",
              "│ (\u001b[38;5;33mGemmaBackbone\u001b[0m)               │                           │                 │ token_ids[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m]            │\n",
              "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
              "│ token_embedding               │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m256000\u001b[0m)      │     \u001b[38;5;34m524,288,000\u001b[0m │ gemma_backbone[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m]       │\n",
              "│ (\u001b[38;5;33mReversibleEmbedding\u001b[0m)         │                           │                 │                            │\n",
              "└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
              "┃<span style=\"font-weight: bold\"> Layer (type)                  </span>┃<span style=\"font-weight: bold\"> Output Shape              </span>┃<span style=\"font-weight: bold\">         Param # </span>┃<span style=\"font-weight: bold\"> Connected to               </span>┃\n",
              "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
              "│ padding_mask (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)     │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n",
              "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
              "│ token_ids (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)        │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n",
              "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
              "│ gemma_backbone                │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">2048</span>)        │   <span style=\"color: #00af00; text-decoration-color: #00af00\">2,507,536,384</span> │ padding_mask[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>],        │\n",
              "│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">GemmaBackbone</span>)               │                           │                 │ token_ids[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>]            │\n",
              "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
              "│ token_embedding               │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">256000</span>)      │     <span style=\"color: #00af00; text-decoration-color: #00af00\">524,288,000</span> │ gemma_backbone[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>]       │\n",
              "│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">ReversibleEmbedding</span>)         │                           │                 │                            │\n",
              "└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[1m Total params: \u001b[0m\u001b[38;5;34m2,507,536,384\u001b[0m (9.34 GB)\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Total params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">2,507,536,384</span> (9.34 GB)\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[1m Trainable params: \u001b[0m\u001b[38;5;34m1,363,968\u001b[0m (5.20 MB)\n"
            ],
            "text/html": [
              "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">1,363,968</span> (5.20 MB)\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "\u001b[1m Non-trainable params: \u001b[0m\u001b[38;5;34m2,506,172,416\u001b[0m (9.34 GB)\n"
            ],
            "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"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Note that enabling LoRA reduces the number of trainable parameters significantly (from 2.5 billion to 1.3 million)."
      ],
      "metadata": {
        "id": "hQQ47kcdpbZ9",
        "papermill": {
          "duration": 0.011797,
          "end_time": "2024-02-21T09:39:33.903795",
          "exception": false,
          "start_time": "2024-02-21T09:39:33.891998",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Limit the input sequence length to 512 (to control memory usage).\n",
        "gemma_lm.preprocessor.sequence_length = 512\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=1, batch_size=1)"
      ],
      "metadata": {
        "id": "_Peq7TnLtHse",
        "outputId": "c1d5c107-fafe-4bb7-8a74-c8c2af7ac6df",
        "papermill": {
          "duration": 753.447217,
          "end_time": "2024-02-21T09:52:07.365329",
          "exception": false,
          "start_time": "2024-02-21T09:39:33.918112",
          "status": "completed"
        },
        "tags": [],
        "execution": {
          "iopub.status.busy": "2024-02-21T14:30:02.032033Z",
          "iopub.execute_input": "2024-02-21T14:30:02.032451Z",
          "iopub.status.idle": "2024-02-21T14:41:54.783533Z",
          "shell.execute_reply.started": "2024-02-21T14:30:02.032417Z",
          "shell.execute_reply": "2024-02-21T14:41:54.781988Z"
        },
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[1m1000/1000\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m1404s\u001b[0m 1s/step - loss: 0.6336 - sparse_categorical_accuracy: 0.4726\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<keras.src.callbacks.history.History at 0x7c10e032f8e0>"
            ]
          },
          "metadata": {},
          "execution_count": 16
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Inference after fine-tuning\n",
        "After fine-tuning, responses follow the instruction provided in the prompt."
      ],
      "metadata": {
        "id": "4yd-1cNw1dTn",
        "papermill": {
          "duration": 0.092253,
          "end_time": "2024-02-21T09:52:07.553072",
          "exception": false,
          "start_time": "2024-02-21T09:52:07.460819",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Europe Trip Prompt\n"
      ],
      "metadata": {
        "id": "H55JYJ1a1Kos",
        "papermill": {
          "duration": 0.092097,
          "end_time": "2024-02-21T09:52:07.73732",
          "exception": false,
          "start_time": "2024-02-21T09:52:07.645223",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = template.format(\n",
        "    instruction=\"What should I do on a trip to Europe?\",\n",
        "    output=\"\",\n",
        ")\n",
        "print(gemma_lm.generate(prompt, max_length=256))"
      ],
      "metadata": {
        "id": "Y7cDJHy8WfCB",
        "papermill": {
          "duration": 14.645082,
          "end_time": "2024-02-21T09:52:22.473375",
          "exception": false,
          "start_time": "2024-02-21T09:52:07.828293",
          "status": "completed"
        },
        "tags": [],
        "execution": {
          "iopub.status.busy": "2024-02-21T14:41:54.784384Z",
          "iopub.status.idle": "2024-02-21T14:41:54.784774Z",
          "shell.execute_reply.started": "2024-02-21T14:41:54.784588Z",
          "shell.execute_reply": "2024-02-21T14:41:54.784603Z"
        },
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "9f37f0f5-d5c7-4f76-98fd-ee5650ed2ed2"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Instruction:\n",
            "What should I do on a trip to Europe?\n",
            "\n",
            "Response:\n",
            "I would recommend visiting the major cities in Europe to get the most bang for your buck. Paris, London, Rome, Berlin and Amsterdam all offer unique experiences that will leave you with lasting memories. If you want to experience the outdoors, I would recommend hiking through the Alps in Switzerland or the Scottish Highlands.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The model now recommends places to visit in Europe."
      ],
      "metadata": {
        "id": "OXP6gg2mjs6u",
        "papermill": {
          "duration": 0.104088,
          "end_time": "2024-02-21T09:52:22.681301",
          "exception": false,
          "start_time": "2024-02-21T09:52:22.577213",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = template.format(\n",
        "    instruction=\"どうやってヨーロッパに行きますか?\",\n",
        "    output=\"\",\n",
        ")\n",
        "print(gemma_lm.generate(prompt, max_length=256))"
      ],
      "metadata": {
        "id": "-B3PWFFnFqsW",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "dcfa50ca-e6aa-44df-fd24-78d8c03e8b8c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Instruction:\n",
            "どうやってヨーロッパに行きますか?\n",
            "\n",
            "Response:\n",
            "ヨーロッパは、いくつかの主要な空港と鉄道で easyjet と easyjet が接続されています。\n",
            "ヨーロッパの主要な空港は、ロンドン、ブリュッセル、パリ、アムステルダム、フランクフルト、ベルリン、ブダペスト、ミュンヘン、ウィーン、ストックホルム、ストックホルム、ヘルシンキなどです。\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### ELI5 Photosynthesis Prompt\n"
      ],
      "metadata": {
        "id": "H7nVd8Mi1Yta",
        "papermill": {
          "duration": 0.093425,
          "end_time": "2024-02-21T09:52:22.874292",
          "exception": false,
          "start_time": "2024-02-21T09:52:22.780867",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = template.format(\n",
        "    instruction=\"Explain the process of photosynthesis in a way that a child could understand.\",\n",
        "    output=\"\",\n",
        ")\n",
        "print(gemma_lm.generate(prompt, max_length=256))"
      ],
      "metadata": {
        "id": "X-2sYl2jqwl7",
        "papermill": {
          "duration": 1.343147,
          "end_time": "2024-02-21T09:52:24.310022",
          "exception": false,
          "start_time": "2024-02-21T09:52:22.966875",
          "status": "completed"
        },
        "tags": [],
        "execution": {
          "iopub.status.busy": "2024-02-21T14:41:54.786248Z",
          "iopub.status.idle": "2024-02-21T14:41:54.786642Z",
          "shell.execute_reply.started": "2024-02-21T14:41:54.786451Z",
          "shell.execute_reply": "2024-02-21T14:41:54.786467Z"
        },
        "trusted": true,
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d22d5d87-1cb9-4af1-8687-715571f227c9"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Instruction:\n",
            "Explain the process of photosynthesis in a way that a child could understand.\n",
            "\n",
            "Response:\n",
            "Photosynthesis is the process where green plants, like trees, use the energy of sunlight to convert water into oxygen in order to create food.  The process begins with the use of chlorophyll in the leaves to capture the energy of the sun.  The chlorophyll absorbs the energy and then converts it into chemical energy, called sugar.  This sugar is used as a form of fuel by the plant to grow.  The oxygen produced by this process is then given off into the atmosphere.  This process is what makes the air we breathe.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = template.format(\n",
        "    instruction=\"光合成の仕組みを子どもにもわかるように説明してください。\",\n",
        "    output=\"\",\n",
        ")\n",
        "print(gemma_lm.generate(prompt, max_length=256))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "DKFz63vJfm4X",
        "outputId": "b66a18de-3c93-4586-969a-031825637c1f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Instruction:\n",
            "光合成の仕組みを子どもにもわかるように説明してください。\n",
            "\n",
            "Response:\n",
            "光合成は、植物が光とCO2を反応させて、糖とO2を生成する反応を指します。この反応には光エネルギー、CO2、H2O、O2、糖、ADP、ATP、H+が必要になります。この反応は、光触媒の光触媒の働きで進行します。\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The model now explains photosynthesis in simpler terms."
      ],
      "metadata": {
        "id": "PCmAmqrvkEhc",
        "papermill": {
          "duration": 0.093039,
          "end_time": "2024-02-21T09:52:24.497343",
          "exception": false,
          "start_time": "2024-02-21T09:52:24.404304",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Note that for demonstration purposes, this tutorial fine-tunes the model on a small subset of the dataset for just one epoch and with a low LoRA rank value. To get better responses from the fine-tuned model, you can experiment with:\n",
        "\n",
        "1. Increasing the size of the fine-tuning dataset\n",
        "2. Training for more steps (epochs)\n",
        "3. Setting a higher LoRA rank\n",
        "4. Modifying the hyperparameter values such as `learning_rate` and `weight_decay`.\n"
      ],
      "metadata": {
        "id": "rtj1MWG2HBPv",
        "papermill": {
          "duration": 0.094695,
          "end_time": "2024-02-21T09:52:24.685134",
          "exception": false,
          "start_time": "2024-02-21T09:52:24.590439",
          "status": "completed"
        },
        "tags": []
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Summary and next steps\n",
        "\n",
        "This tutorial covered LoRA fine-tuning on a Gemma model using KerasNLP. Check out the following docs next:\n",
        "\n",
        "* Learn how to [generate text with a Gemma model](https://ai.google.dev/gemma/docs/get_started).\n",
        "* Learn how to perform [distributed fine-tuning and inference on a Gemma model](https://ai.google.dev/gemma/docs/distributed_tuning).\n",
        "* Learn how to [use Gemma open models with Vertex AI](https://cloud.google.com/vertex-ai/docs/generative-ai/open-models/use-gemma).\n",
        "* Learn how to [fine-tune Gemma using KerasNLP and deploy to Vertex AI](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_gemma_kerasnlp_to_vertexai.ipynb).\n"
      ],
      "metadata": {
        "id": "gSsRdeiof_rJ",
        "papermill": {
          "duration": 0.091976,
          "end_time": "2024-02-21T09:52:24.869733",
          "exception": false,
          "start_time": "2024-02-21T09:52:24.777757",
          "status": "completed"
        },
        "tags": []
      }
    }
  ]
}