notebooks/community/model_garden/model_garden_gemma_kerasnlp_to_vertexai.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 Google LLC\n",
        "#\n",
        "# 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."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Finetune Gemma using KerasNLP and deploy to Vertex AI\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_gemma_kerasnlp_to_vertexai.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Run in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_gemma_kerasnlp_to_vertexai.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://upload.wikimedia.org/wikipedia/commons/9/91/Octicons-mark-github.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iYQE9Mza89tj"
      },
      "source": [
        "> This notebook was tested in the following environment:\n",
        ">\n",
        "> - Python 3.10\n",
        "> - Colab Enterprise with a `g2-standard-8` runtime:\n",
        ">   - 32 GB of system RAM\n",
        ">   - 24 GB of GPU RAM (NVIDIA L4)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eIVltG2O2qgF"
      },
      "source": [
        "## Overview\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "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",
        "This notebook demonstrates loading, finetuning, converting, and deploying Gemma to Vertex AI.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d975e698c9a4"
      },
      "source": [
        "### Objective\n",
        "\n",
        "- Load Gemma using KerasNLP\n",
        "- Finetune Gemma using KerasNLP\n",
        "- Convert Gemma to Hugging Face Transformers\n",
        "- Deploy Gemma to Vertex AI\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aed92deeb4a0"
      },
      "source": [
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "- Vertex AI\n",
        "- Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage](https://cloud.google.com/storage/pricing) pricings,\n",
        "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OrDeaJ_3vL0-"
      },
      "source": [
        "## Installation\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i7EUnXsZhAGF"
      },
      "source": [
        "Install the following packages required to execute this notebook:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SpeKcQcw5E5T"
      },
      "outputs": [],
      "source": [
        "# Keras & KerasNLP\n",
        "# Install Keras 3 last, see https://keras.io/getting_started\n",
        "%pip install --upgrade --quiet keras-nlp\n",
        "%pip install --upgrade --quiet keras\n",
        "\n",
        "# Hugging Face Transformers\n",
        "%pip install --upgrade --quiet accelerate sentencepiece transformers\n",
        "\n",
        "# Vertex AI SDK\n",
        "%pip install --upgrade --quiet google-cloud-aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JaKPnjQsWiLn"
      },
      "source": [
        "## Before you begin\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9xWLhlsd09zY"
      },
      "source": [
        "### Kaggle credentials\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2y1Zf6OcvDxU"
      },
      "source": [
        "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",
        "\n",
        "Then, to use the Kaggle API, create an API token:\n",
        "\n",
        "- Open the [Kaggle settings](https://www.kaggle.com/settings)\n",
        "- Select _\"Create New Token\"_\n",
        "- A `kaggle.json` file is downloaded. It contains your Kaggle credentials\n",
        "\n",
        "Run the following cell and enter your Kaggle credentials.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TxmBScR4vjkf"
      },
      "outputs": [],
      "source": [
        "import kagglehub\n",
        "\n",
        "kagglehub.login()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y8hNHhFkjMaf"
      },
      "source": [
        "> Note: If `kagglehub.login()` doesn't work for you, an alternative way is to set `KAGGLE_USERNAME` and `KAGGLE_KEY` environment variables.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_WqZXDDQqrgW"
      },
      "source": [
        "### Google Cloud setup\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa",
        "tags": []
      },
      "source": [
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
        "\n",
        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PRKmPKBnqjaY"
      },
      "source": [
        "### Google Cloud authentication\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sBCra4QMA2wR"
      },
      "source": [
        "If you run this notebook from Colab Enterprise, the Cloud SDK, code, and other libraries already run using your Google Cloud account.\n",
        "\n",
        "Check your active account:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "40e71J5Go2b7"
      },
      "outputs": [],
      "source": [
        "!gcloud config get core/account"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "espvMVsIph-p"
      },
      "source": [
        "If your account is not defined, you need to authenticate:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tqNbW2pmphbG"
      },
      "outputs": [],
      "source": [
        "# Authenticate the Cloud SDK with your credentials\n",
        "# !gcloud auth login\n",
        "\n",
        "# Authenticate code and libraries with your credentials\n",
        "# !gcloud auth application-default login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6TNLOa5JqmVI"
      },
      "source": [
        "### Google Cloud project\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "If you run this notebook in Colab Enterprise, the default project is automatically defined:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oM1iC_MfAts1"
      },
      "outputs": [],
      "source": [
        "res = !gcloud config get core/project\n",
        "PROJECT_ID = res[0]\n",
        "\n",
        "print(f\"{PROJECT_ID=}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1onIbGdaclHa"
      },
      "source": [
        "Otherwise, list your projects and define the default project manually:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O6ymkb2dZjFN"
      },
      "outputs": [],
      "source": [
        "# List your projects\n",
        "# !gcloud projects list\n",
        "\n",
        "# Define the default project\n",
        "# PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
        "# !gcloud config set core/project $PROJECT_ID"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6CGQ1lxEr_WQ"
      },
      "source": [
        "### Vertex AI region\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "region"
      },
      "source": [
        "Define your default Vertex AI region. See available [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kfe_T6dv5E5Y"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type: \"string\"}\n",
        "\n",
        "!gcloud config set ai/region $REGION"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8dEmh_icsNQI"
      },
      "source": [
        "> Note: This notebook deploys a Gemma model to a single region. In production, you can deploy to multiple regions, to serve your worldwide users with optimal latencies.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ETnL6RT1sGKM"
      },
      "source": [
        "### Cloud Storage bucket\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "Create a storage bucket (or use an existing one) to store artifacts such as model weights or datasets.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "# Define a bucket related to your project\n",
        "BUCKET_URI = f\"gs://gemma-{PROJECT_ID}-unique\"\n",
        "# Or use an existing one\n",
        "# BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
        "\n",
        "res = !gcloud storage buckets describe $BUCKET_URI --format \"value(name)\"\n",
        "if len(res) == 1 and \"ERROR\" not in res[0]:\n",
        "    print(\"✔️ The bucket exists\")\n",
        "else:\n",
        "    print(\"⚙️ Creating the bucket…\")\n",
        "    !gcloud storage buckets create $BUCKET_URI --project $PROJECT_ID --location $REGION"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4OqNY6-Xwj2W"
      },
      "source": [
        "### Service account\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vRAvXQlxubf7"
      },
      "source": [
        "When deploying Gemma to a Vertex AI endpoint, the model service will require a service account with \"Storage Object Admin\" and \"Vertex AI User\" roles.\n",
        "\n",
        "Create the service account (or use an existing one):\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FuFyhcsSri-O"
      },
      "outputs": [],
      "source": [
        "# Create the service account for the Vertex AI endpoint\n",
        "SERVICE_ACCOUNT_NAME = \"gemma-vertexai\"\n",
        "SERVICE_ACCOUNT_DISPLAY_NAME = \"Gemma Vertex AI endpoint\"\n",
        "SERVICE_ACCOUNT = f\"{SERVICE_ACCOUNT_NAME}@{PROJECT_ID}.iam.gserviceaccount.com\"\n",
        "# Or use an existing one\n",
        "# SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}\n",
        "assert SERVICE_ACCOUNT.endswith(f\"@{PROJECT_ID}.iam.gserviceaccount.com\")\n",
        "\n",
        "res = !gcloud iam service-accounts describe $SERVICE_ACCOUNT --format \"value(email)\"\n",
        "if len(res) == 1 and \"ERROR\" not in res[0]:\n",
        "    print(\"✔️ The service account exists\")\n",
        "else:\n",
        "    print(\"⚙️ Creating the service account…\")\n",
        "    !gcloud iam service-accounts create $SERVICE_ACCOUNT_NAME --display-name \"$SERVICE_ACCOUNT_DISPLAY_NAME\"\n",
        "    # Grant \"Storage Object Admin\" role\n",
        "    !gcloud projects add-iam-policy-binding $PROJECT_ID --member \"serviceAccount:$SERVICE_ACCOUNT\" --role \"roles/storage.objectAdmin\"\n",
        "    # Grant \"Vertex AI User\" role\n",
        "    !gcloud projects add-iam-policy-binding $PROJECT_ID --member \"serviceAccount:$SERVICE_ACCOUNT\" --role \"roles/aiplatform.user\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uQDEhMB_3kbE"
      },
      "source": [
        "### Dependencies\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DVQhjWte3s7M"
      },
      "outputs": [],
      "source": [
        "import datetime\n",
        "import json\n",
        "import locale\n",
        "\n",
        "import keras\n",
        "import keras_nlp\n",
        "import torch\n",
        "import transformers\n",
        "from google.cloud import aiplatform\n",
        "from numba import cuda"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zd9-P50rzNNW"
      },
      "source": [
        "### Model constants\n",
        "\n",
        "Gemma models are available in several sizes and variants. This notebook uses the `gemma_2b_en` version, which has lower resource requirements. To learn more about Gemma, see the [Gemma Model Garden card](https://console.cloud.google.com/vertex-ai/publishers/google/model-garden/gemma).\n",
        "\n",
        "Define the model and related constants:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M6xCNXc8ROnU"
      },
      "outputs": [],
      "source": [
        "MODEL_NAME = \"gemma_2b_en\"\n",
        "# MODEL_NAME = \"gemma_instruct_2b_en\"\n",
        "# MODEL_NAME = \"gemma_7b_en\"\n",
        "# MODEL_NAME = \"gemma_instruct_7b_en\"\n",
        "\n",
        "# Deduce model size from name format: \"gemma[_instruct]_{2b,7b}_en\"\n",
        "MODEL_SIZE = MODEL_NAME.split(\"_\")[-2]\n",
        "assert MODEL_SIZE in (\"2b\", \"7b\")\n",
        "\n",
        "# Dataset\n",
        "DATASET_NAME = \"databricks-dolly-15k\"\n",
        "DATASET_PATH = f\"{DATASET_NAME}.jsonl\"\n",
        "DATASET_URL = f\"https://huggingface.co/datasets/databricks/{DATASET_NAME}/resolve/main/{DATASET_PATH}\"\n",
        "\n",
        "# Finetuned model\n",
        "FINETUNED_MODEL_DIR = f\"./{MODEL_NAME}_{DATASET_NAME}\"\n",
        "FINETUNED_WEIGHTS_PATH = f\"{FINETUNED_MODEL_DIR}/model.weights.h5\"\n",
        "FINETUNED_VOCAB_PATH = f\"{FINETUNED_MODEL_DIR}/vocabulary.spm\"\n",
        "\n",
        "# Converted model\n",
        "HUGGINGFACE_MODEL_DIR = f\"./{MODEL_NAME}_huggingface\"\n",
        "\n",
        "# Deployed model\n",
        "DEPLOYED_MODEL_URI = f\"{BUCKET_URI}/{MODEL_NAME}\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OaIE0BoZyzgE"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "To finetune Gemma, this notebook uses the [databricks-dolly-15k](https://huggingface.co/datasets/databricks/databricks-dolly-15k) test dataset.\n",
        "\n",
        "Download the dataset:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iJSUhANIyqf0"
      },
      "outputs": [],
      "source": [
        "!wget -nv -nc -O $DATASET_PATH $DATASET_URL"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BuOc318GVlB9"
      },
      "source": [
        "## Load Gemma\n",
        "\n",
        "In this step, you will configure Keras precision settings and load Gemma with KerasNLP.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k54KOGzSNrGz"
      },
      "source": [
        "### Keras precision settings\n",
        "\n",
        "When training on NVIDIA GPUs, mixed precision (`keras.mixed_precision.set_global_policy(\"mixed_bfloat16\")`) can be used to speed up training with minimal effect on training quality. In most cases, it is recommended to turn on mixed precision as it saves both memory and time. However, be aware that at small batch sizes, it can inflate memory usage by 1.5x (weights will be loaded twice, at half precision and full precision).\n",
        "\n",
        "For inference, half-precision (`keras.config.set_floatx(\"bfloat16\")`) will work and save memory (while mixed-precision is not applicable).\n",
        "\n",
        "Configure your precision settings:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JodwwVgROfuZ"
      },
      "outputs": [],
      "source": [
        "# Run inferences at half precision\n",
        "keras.config.set_floatx(\"bfloat16\")\n",
        "\n",
        "# Train at mixed precision (enable for large batch sizes)\n",
        "# keras.mixed_precision.set_global_policy(\"mixed_bfloat16\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lj7bzTamp7_f"
      },
      "source": [
        "### Model summary\n",
        "\n",
        "Load the Gemma model using the `GemmaCausalLM.from_preset()` method:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EzHhGJgG5E5T",
        "outputId": "06b4403f-bfdb-4430-8baa-7a525a4bbf64"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Downloading from https://www.kaggle.com/api/v1/models/keras/gemma/keras/gemma_2b_en/1/download/config.json...\n",
            "100%|██████████| 555/555 [00:00<00:00, 634kB/s]\n",
            "Downloading from https://www.kaggle.com/api/v1/models/keras/gemma/keras/gemma_2b_en/1/download/model.weights.h5...\n",
            "100%|██████████| 4.67G/4.67G [02:28<00:00, 33.7MB/s]\n",
            "Downloading from https://www.kaggle.com/api/v1/models/keras/gemma/keras/gemma_2b_en/1/download/tokenizer.json...\n",
            "100%|██████████| 401/401 [00:00<00:00, 554kB/s]\n",
            "Downloading from https://www.kaggle.com/api/v1/models/keras/gemma/keras/gemma_2b_en/1/download/assets/tokenizer/vocabulary.spm...\n",
            "100%|██████████| 4.04M/4.04M [00:00<00:00, 5.27MB/s]\n"
          ]
        }
      ],
      "source": [
        "gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset(MODEL_NAME)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GbXgDqKZ5I8_"
      },
      "source": [
        "Display the model summary:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 385
        },
        "id": "kc77gyOR4xNh",
        "outputId": "8c132d49-feb3-473c-e850-97ebc74e6a71"
      },
      "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,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 (4.67 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> (4.67 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 (4.67 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> (4.67 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": {}
        }
      ],
      "source": [
        "gemma_lm.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LicAd7yjyRiQ"
      },
      "source": [
        "### Test examples\n",
        "\n",
        "Define test examples and functions that will be used to test models before and after finetuning:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LybWJAJeylhB"
      },
      "outputs": [],
      "source": [
        "TEST_EXAMPLES = [\n",
        "    \"What are good activities for a toddler?\",\n",
        "    \"What can we hope to see after rain and sun?\",\n",
        "    \"What's the most famous painting by Monet?\",\n",
        "    \"Who engineered the Statue of Liberty?\",\n",
        "    'Who were \"The Lumières\"?',\n",
        "]\n",
        "\n",
        "# Prompt template for the training data and the finetuning tests\n",
        "PROMPT_TEMPLATE = \"Instruction:\\n{instruction}\\n\\nResponse:\\n{response}\"\n",
        "\n",
        "TEST_PROMPTS = [\n",
        "    PROMPT_TEMPLATE.format(instruction=example, response=\"\")\n",
        "    for example in TEST_EXAMPLES\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hm8bHp160oh6"
      },
      "source": [
        "### Samplers\n",
        "\n",
        "You can control how tokens are generated for `GemmaCausalLM` by calling the `compile()` method with the `sampler` parameter.\n",
        "\n",
        "For example:\n",
        "\n",
        "- `greedy`: picks the next token with the largest probability\n",
        "- `top_k`: randomly picks the next token from the tokens of top K probability\n",
        "\n",
        "To get deterministic outputs in this notebook, make sure you're using the `greedy` sampler:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0Gi17ZmD0-WY"
      },
      "outputs": [],
      "source": [
        "gemma_lm.compile(sampler=\"greedy\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9m5EAkkAsbK5"
      },
      "source": [
        "To learn more about available samplers, see [Samplers](https://keras.io/api/keras_nlp/samplers).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Akqc9HoM8e9X"
      },
      "source": [
        "### Inference before finetuning\n",
        "\n",
        "Check how the model responds to the test examples:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ue-LRnJ_9Mv1",
        "outputId": "e547de93-caec-4b07-8706-ec2385cd7abf"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "What are good activities for a toddler?\n",
            "'\\n\\nWhat are the best activities for a toddler?\\n\\nWhat are the best activities for a toddler?\\n\\nWhat are the best activities for a toddler?\\n\\nWhat are the best activities for a toddler'\n",
            "\n",
            "What can we hope to see after rain and sun?\n",
            "'\\n\\nThe answer is: a lot.\\n\\nThe rain and sun are the two most important elements in the world of photography.\\n\\nThe rain is the most important element because it creates'\n",
            "\n",
            "What's the most famous painting by Monet?\n",
            "\"\\n\\nWhat's the most famous painting by Van Gogh?\\n\\nWhat's the most famous painting by Picasso?\\n\\nWhat's the most famous painting by Dali?\\n\\nWhat'\"\n",
            "\n",
            "Who engineered the Statue of Liberty?\n",
            "'\\n\\nA. George Washington\\nB. Napoleon Bonaparte\\nC. Robert Fulton\\nD. Gustave Eiffel\\n\\nIn the following sentence, underline the correct modifier from the pair given in parentheses. Example 1'\n",
            "\n",
            "Who were \"The Lumières\"?\n",
            "' What did they invent?\\n\\nIn the following sentence, underline the correct modifier from the pair given in parentheses. Example 1. He is (real, $\\\\underline{\\\\text{really}}$) talented'\n",
            "\n"
          ]
        }
      ],
      "source": [
        "for test_example in TEST_EXAMPLES:\n",
        "    response = gemma_lm.generate(test_example, max_length=48)\n",
        "    output = response[len(test_example) :]\n",
        "    print(f\"{test_example}\\n{output!r}\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1zkSfGlQ6ftN"
      },
      "source": [
        "A pretrained model can generate text that deviates from the output you are expecting. Here are some examples:\n",
        "\n",
        "- The output doesn't follow your output requirements.\n",
        "- The output is too generic or not consistent enough.\n",
        "- The output is factually incorrect or outdated.\n",
        "- The output must be aligned with your specific safety policies.\n",
        "\n",
        "More specific inputs (prompt engineering) can fix some of these issues, at the expense of more complex and longer prompts. If the expected output is not part of the model training data, LLMs generate plausible text anyway and produce what is sometimes called hallucinations.\n",
        "\n",
        "You can perform a model finetuning to improve the performance of the model and keep simpler prompts.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oVJ9b5tWOlyn"
      },
      "source": [
        "## Finetune Gemma\n",
        "\n",
        "Finetune your Gemma model to improve its performance in the specific task of answering questions more consistently and more factually.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tk7LUHcCqj1N"
      },
      "source": [
        "### Training data\n",
        "\n",
        "Generate the training examples using the dataset:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "q9YLaaPpvODd",
        "outputId": "e7a3c2c1-f946-432f-b417-b975dac798cb"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Training examples: 1054/10544\n"
          ]
        }
      ],
      "source": [
        "def generate_training_data(training_ratio: int = 100) -> list[str]:\n",
        "    assert 0 < training_ratio <= 100\n",
        "    data = []\n",
        "    with open(DATASET_PATH) as file:\n",
        "        for line in file.readlines():\n",
        "            features = json.loads(line)\n",
        "            # Skip examples with context, for simplicity\n",
        "            if features[\"context\"]:\n",
        "                continue\n",
        "            data.append(PROMPT_TEMPLATE.format(**features))\n",
        "    total_data_count = len(data)\n",
        "    training_data_count = total_data_count * training_ratio // 100\n",
        "    print(f\"Training examples: {training_data_count}/{total_data_count}\")\n",
        "\n",
        "    return data[:training_data_count]\n",
        "\n",
        "\n",
        "# Limit to 10% for test purposes\n",
        "training_data = generate_training_data(training_ratio=10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qm89dWZU52q0"
      },
      "source": [
        "### Low-Rank Adaptation (LoRA)\n",
        "\n",
        "[Low Rank Adaptation](https://arxiv.org/abs/2106.09685) (LoRA) is a finetuning technique which greatly reduces the number of trainable parameters for downstream tasks by freezing the full weights of the model and inserting a smaller number of new trainable weights into the model. This technique makes training much faster and more memory-efficient.\n",
        "\n",
        "Enable LoRA for the model and set the LoRA rank to 4:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iN-I5RhJC1fh"
      },
      "outputs": [],
      "source": [
        "gemma_lm.backbone.enable_lora(rank=4)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DGVreHDoDOkO"
      },
      "source": [
        "Check that the number of trainable parameters is significantly reduced:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 385
        },
        "id": "YCcKAxodDN0S",
        "outputId": "964ab75e-d5da-4d62-b25b-0f38f109e470"
      },
      "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 (4.67 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> (4.67 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 (2.60 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> (2.60 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 (4.67 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> (4.67 GB)\n",
              "</pre>\n"
            ]
          },
          "metadata": {}
        }
      ],
      "source": [
        "gemma_lm.summary()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x8dyzB5MXfo0"
      },
      "source": [
        "The number of trainable parameters decreased from 2.5B down to 1.4M (1,800x less), making it possible to finetune the model with reasonable GPU memory requirements.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bsw_-qIGCZnv"
      },
      "source": [
        "### Finetuning\n",
        "\n",
        "Finetune the model with the training data. This step can take a couple of minutes:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mfk9X11tvPWy",
        "outputId": "cd17f18d-e3ad-422a-df6b-bb2e50cdea97"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[1m1054/1054\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m134s\u001b[0m 77ms/step - loss: 19.3561 - sparse_categorical_accuracy: 0.5872\n"
          ]
        }
      ],
      "source": [
        "def finetune_gemma(model: keras_nlp.models.GemmaCausalLM, data: list[str]):\n",
        "    # Reduce the input sequence length to limit memory usage\n",
        "    model.preprocessor.sequence_length = 128\n",
        "\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",
        "\n",
        "    # Exclude layernorm and bias terms from decay\n",
        "    optimizer.exclude_from_weight_decay(var_names=[\"bias\", \"scale\"])\n",
        "\n",
        "    model.compile(\n",
        "        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "        optimizer=optimizer,\n",
        "        weighted_metrics=[keras.metrics.SparseCategoricalAccuracy()],\n",
        "        sampler=\"greedy\",\n",
        "    )\n",
        "    model.fit(data, epochs=1, batch_size=1)\n",
        "\n",
        "\n",
        "finetune_gemma(gemma_lm, training_data)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "twNI1SKd-4Sp"
      },
      "source": [
        "### Inference after finetuning\n",
        "\n",
        "Test the finetuned model:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "z4zf-UtTVXuG",
        "outputId": "b7a323d2-3863-42ec-965a-c534d4cc1157"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Instruction:\n",
            "What are good activities for a toddler?\n",
            "\n",
            "Response:\n",
            "The best activities for a toddler are those that are fun and engaging.\n",
            "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - \n",
            "Instruction:\n",
            "What can we hope to see after rain and sun?\n",
            "\n",
            "Response:\n",
            "After rain and sun, we can see the rainbow.\n",
            "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - \n",
            "Instruction:\n",
            "What's the most famous painting by Monet?\n",
            "\n",
            "Response:\n",
            "The most famous painting by Monet is \"Impression, Sunrise\".\n",
            "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - \n",
            "Instruction:\n",
            "Who engineered the Statue of Liberty?\n",
            "\n",
            "Response:\n",
            "The Statue of Liberty was designed by a French sculptor, Frederic Auguste Bartholdi\n",
            "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - \n",
            "Instruction:\n",
            "Who were \"The Lumières\"?\n",
            "\n",
            "Response:\n",
            "The Lumières were the inventors of the first motion picture camera. They were\n",
            "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - \n"
          ]
        }
      ],
      "source": [
        "for prompt in TEST_PROMPTS:\n",
        "    output = gemma_lm.generate(prompt, max_length=30)\n",
        "    print(f\"{output}\\n{'- '*40}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XyraRakQrz3i"
      },
      "source": [
        "You should observe that outputs are now structured, more consistent, and more factual.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BqBUxdtO2HLs"
      },
      "source": [
        "## Convert Gemma to Hugging Face Transformers\n",
        "\n",
        "In the next step, the model will be deployed to Vertex AI, served by a [vLLM](https://docs.vllm.ai) container image. vLLM is an optimized LLM serving library which supports Hugging Face [Transformers](https://huggingface.co/docs/transformers). To be loaded by the vLLM service, the finetuned model needs to be converted to the Hugging Face architecture. KerasNLP provides a conversion script for this.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZLmYQmKhZql2"
      },
      "source": [
        "### Checkpoint\n",
        "\n",
        "Save the finetuned model assets:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4NloTkakaIRi"
      },
      "outputs": [],
      "source": [
        "# Make sure the directory exists\n",
        "%mkdir -p $FINETUNED_MODEL_DIR\n",
        "\n",
        "gemma_lm.save_weights(FINETUNED_WEIGHTS_PATH)\n",
        "\n",
        "gemma_lm.preprocessor.tokenizer.save_assets(FINETUNED_MODEL_DIR)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_3124GtDAvJf"
      },
      "source": [
        "List the checkpoint files:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8kJV2NStkclG",
        "outputId": "83d3bb76-4432-4c92-b34a-3da1f7705745"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "4.7G\t./gemma_2b_en_databricks-dolly-15k/model.weights.h5\n",
            "4.1M\t./gemma_2b_en_databricks-dolly-15k/vocabulary.spm\n",
            "4.7G\ttotal\n"
          ]
        }
      ],
      "source": [
        "!du -shc $FINETUNED_MODEL_DIR/*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1D9h3CUaA5am"
      },
      "source": [
        "Release the resources to make sure the GPU is available for the next steps:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I6Uoo52i_QZV"
      },
      "outputs": [],
      "source": [
        "del gemma_lm\n",
        "\n",
        "device = cuda.get_current_device()\n",
        "cuda.select_device(device.id)\n",
        "cuda.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HSaYS78aKCGM"
      },
      "source": [
        "### Model conversion\n",
        "\n",
        "Run the KerasNLP conversion script:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7K4yDhZH2OnQ"
      },
      "outputs": [],
      "source": [
        "# Download the conversion script from KerasNLP tools\n",
        "!wget -nv -nc https://raw.githubusercontent.com/keras-team/keras-nlp/master/tools/gemma/export_gemma_to_hf.py\n",
        "\n",
        "# Run the conversion script\n",
        "# Note: it uses the PyTorch backend of Keras (hence the KERAS_BACKEND env variable)\n",
        "!KERAS_BACKEND=torch python export_gemma_to_hf.py \\\n",
        "    --weights_file $FINETUNED_WEIGHTS_PATH \\\n",
        "    --size $MODEL_SIZE \\\n",
        "    --vocab_path $FINETUNED_VOCAB_PATH \\\n",
        "    --output_dir $HUGGINGFACE_MODEL_DIR"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ma56cdjKH4KW"
      },
      "source": [
        "### Inference with Transformers\n",
        "\n",
        "Before deploying the converted model, test it using the `transformers` library.\n",
        "\n",
        "Load the model and the tokenizer:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "bfd1144bd7ad48e6b3686c5fdf6f1563",
            "fd1dfcfd0cbf4923823736e1789b5112",
            "67a173aa66e94699a157c5f4d7449d56",
            "7413e1164e074cbea6091644a4580e3c",
            "1f3ed8c6c9294b60b13440d435d4d280",
            "6919c81e65e548e09db80694d48d6523",
            "651f75d7f7e54c959364fae5ff6d707b",
            "0a8545ddc700438d8c915fc64a694787",
            "2cea77d31a1a41b5a983693ddbe55fb9",
            "5fb6c6fc55d846d5a8ad639b91ad32ef",
            "5ac0d9be26774111b61a543037e3c7fd"
          ]
        },
        "id": "bMNFa-Tx5E5U",
        "outputId": "91ddd005-46d4-4532-d108-2772426377ba",
        "tags": []
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "bfd1144bd7ad48e6b3686c5fdf6f1563"
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "model = transformers.GemmaForCausalLM.from_pretrained(\n",
        "    HUGGINGFACE_MODEL_DIR,\n",
        "    local_files_only=True,\n",
        "    device_map=\"auto\",  # Library \"accelerate\" to auto-select GPU\n",
        ")\n",
        "tokenizer = transformers.GemmaTokenizer.from_pretrained(\n",
        "    HUGGINGFACE_MODEL_DIR,\n",
        "    local_files_only=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OuiKR6BlzbME"
      },
      "source": [
        "Test the model:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "24196WWF5E5U",
        "outputId": "0f041e93-4308-47fe-a41e-df22f62c555b",
        "tags": []
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Instruction:\n",
            "What are good activities for a toddler?\n",
            "\n",
            "Response:\n",
            "Toddlers are very active and curious. They love to explore and learn\n",
            "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - \n",
            "Instruction:\n",
            "What can we hope to see after rain and sun?\n",
            "\n",
            "Response:\n",
            "After rain and sun, we can see the rainbow.\n",
            "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - \n",
            "Instruction:\n",
            "What's the most famous painting by Monet?\n",
            "\n",
            "Response:\n",
            "The most famous painting by Monet is \"Impression, Sunrise\".\n",
            "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - \n",
            "Instruction:\n",
            "Who engineered the Statue of Liberty?\n",
            "\n",
            "Response:\n",
            "The Statue of Liberty was designed by a French sculptor, Frederic Auguste Bartholdi\n",
            "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - \n",
            "Instruction:\n",
            "Who were \"The Lumières\"?\n",
            "\n",
            "Response:\n",
            "The Lumières were the inventors of the first motion picture camera. They were\n",
            "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - \n"
          ]
        }
      ],
      "source": [
        "def test_transformers_model(\n",
        "    model: transformers.GemmaForCausalLM,\n",
        "    tokenizer: transformers.GemmaTokenizer,\n",
        ") -> None:\n",
        "    for prompt in TEST_PROMPTS:\n",
        "        inputs = tokenizer([prompt], return_tensors=\"pt\").to(model.device)\n",
        "        outputs = model.generate(**inputs, max_length=30)\n",
        "\n",
        "        output = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
        "        print(f\"{output}\\n{'- '*40}\")\n",
        "\n",
        "\n",
        "test_transformers_model(model, tokenizer)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RpR6uXVQ-9Gc"
      },
      "source": [
        "Release the resources:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1xuCTtQp88lX"
      },
      "outputs": [],
      "source": [
        "# Release resources\n",
        "del model, tokenizer\n",
        "\n",
        "# Free GPU RAM\n",
        "torch.cuda.empty_cache()\n",
        "\n",
        "# Restore the default encoding (current issue with the transformers library)\n",
        "locale.getpreferredencoding = lambda: \"UTF-8\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KD0HiIPE4tpI"
      },
      "source": [
        "You're ready to deploy your finetuned model to Vertex AI!\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk,all"
      },
      "source": [
        "## Deploy Gemma to Vertex AI\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RtMu4m_tPQ0u"
      },
      "source": [
        "### Vertex AI initialization\n",
        "\n",
        "Initialize Vertex AI:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tks_NPrJLrWZ"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GlREfDROQFEE"
      },
      "source": [
        "### Model upload\n",
        "\n",
        "Upload the model to the Cloud Storage bucket:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7aSWCofpP_5m"
      },
      "outputs": [],
      "source": [
        "!gcloud storage rsync --recursive --verbosity error $HUGGINGFACE_MODEL_DIR $DEPLOYED_MODEL_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "647rA5KoQ4y6"
      },
      "source": [
        "Check the bucket content:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d3NSSf7TQ4Vz"
      },
      "outputs": [],
      "source": [
        "!gcloud storage du $DEPLOYED_MODEL_URI --readable-sizes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "umrDdBZGRfTK"
      },
      "source": [
        "### Helper functions\n",
        "\n",
        "Define helper functions to deploy the model with a vLLM container:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dZEyVeuf5E5Z",
        "tags": []
      },
      "outputs": [],
      "source": [
        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240220_0936_RC01\"\n",
        "\n",
        "\n",
        "def get_job_name_with_datetime(prefix: str) -> str:\n",
        "    suffix = datetime.datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
        "    return f\"{prefix}{suffix}\"\n",
        "\n",
        "\n",
        "def deploy_model_vllm(\n",
        "    model_name: str,\n",
        "    model_uri: str,\n",
        "    service_account: str,\n",
        "    machine_type: str = \"g2-standard-8\",\n",
        "    accelerator_type: str = \"NVIDIA_L4\",\n",
        "    accelerator_count: int = 1,\n",
        "    max_model_len: int = 8192,\n",
        "    dtype: str = \"bfloat16\",\n",
        ") -> tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    # Upload the model to \"Model Registry\"\n",
        "    job_name = get_job_name_with_datetime(model_name)\n",
        "    vllm_args = [\n",
        "        \"--host=0.0.0.0\",\n",
        "        \"--port=7080\",\n",
        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
        "        \"--swap-space=16\",\n",
        "        \"--gpu-memory-utilization=0.95\",\n",
        "        f\"--max-model-len={max_model_len}\",\n",
        "        f\"--dtype={dtype}\",\n",
        "        \"--disable-log-stats\",\n",
        "    ]\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=job_name,\n",
        "        artifact_uri=model_uri,\n",
        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
        "        serving_container_command=[\"python\", \"-m\", \"vllm.entrypoints.api_server\"],\n",
        "        serving_container_args=vllm_args,\n",
        "        serving_container_ports=[7080],\n",
        "        serving_container_predict_route=\"/generate\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "    )\n",
        "\n",
        "    # Deploy the model to an endpoint to serve \"Online predictions\"\n",
        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
        "    )\n",
        "\n",
        "    return model, endpoint"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OfTTbrx2SwDf"
      },
      "source": [
        "### Model deployment\n",
        "\n",
        "Deploy the model. This step can take 10+ minutes.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7FK3F_eu5E5Z",
        "tags": []
      },
      "outputs": [],
      "source": [
        "MODEL_NAME_VLLM = f\"{MODEL_NAME}-vllm\"\n",
        "\n",
        "# Start with a G2 Series cost-effective configuration\n",
        "match MODEL_SIZE:\n",
        "    case \"2b\":\n",
        "        machine_type = \"g2-standard-8\"\n",
        "        accelerator_type = \"NVIDIA_L4\"\n",
        "        accelerator_count = 1\n",
        "    case \"7b\":\n",
        "        machine_type = \"g2-standard-12\"\n",
        "        accelerator_type = \"NVIDIA_L4\"\n",
        "        accelerator_count = 1\n",
        "    case _:\n",
        "        assert MODEL_SIZE in (\"2b\", \"7b\")\n",
        "\n",
        "# See supported machine/GPU configurations in chosen region:\n",
        "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute\n",
        "\n",
        "# For even more performance, consider V100 and A100 GPUs\n",
        "# > Nvidia Tesla V100\n",
        "# machine_type = \"n1-standard-8\"\n",
        "# accelerator_type = \"NVIDIA_TESLA_V100\"\n",
        "# > Nvidia Tesla A100\n",
        "# machine_type = \"a2-highgpu-1g\"\n",
        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
        "\n",
        "# Larger `max_model_len` values will require more GPU memory\n",
        "max_model_len = 2048\n",
        "\n",
        "model, endpoint = deploy_model_vllm(\n",
        "    MODEL_NAME_VLLM,\n",
        "    DEPLOYED_MODEL_URI,\n",
        "    SERVICE_ACCOUNT,\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
        "    max_model_len=max_model_len,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eg-gVbwnTpYD"
      },
      "source": [
        "### Online inference\n",
        "\n",
        "The model is deployed! Test the endpoint:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "assuAGh7DEUz",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "802c483e-b2ef-4fd7-f1ae-b930f9980e11"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "What are good activities for a toddler?\n",
            "The best activities for a toddler are those that are\n",
            "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - \n",
            "What can we hope to see after rain and sun?\n",
            "After rain and sun, we can see the rainbow\n",
            "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - \n",
            "What's the most famous painting by Monet?\n",
            "The most famous painting by Monet is \"Impression,\n",
            "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - \n",
            "Who engineered the Statue of Liberty?\n",
            "The Statue of Liberty was designed by a French sculptor\n",
            "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - \n",
            "Who were \"The Lumières\"?\n",
            "The Lumières were the inventors of the first motion\n",
            "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - \n"
          ]
        }
      ],
      "source": [
        "def test_vertexai_endpoint(endpoint: aiplatform.Endpoint):\n",
        "    for question, prompt in zip(TEST_EXAMPLES, TEST_PROMPTS):\n",
        "        instance = {\n",
        "            \"prompt\": prompt,\n",
        "            \"max_tokens\": 10,\n",
        "            \"temperature\": 0.0,\n",
        "            \"top_p\": 1.0,\n",
        "            \"top_k\": 1,\n",
        "            \"raw_response\": True,\n",
        "        }\n",
        "        response = endpoint.predict(instances=[instance])\n",
        "        output = response.predictions[0]\n",
        "        print(f\"{question}\\n{output}\\n{'- '*40}\")\n",
        "\n",
        "\n",
        "test_vertexai_endpoint(endpoint)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cH9qQRqyDrkS"
      },
      "source": [
        "> See [vLLM `SamplingParams`](https://github.com/vllm-project/vllm/blob/main/vllm/sampling_params.py) for more details about the sampling parameters supported by vLLM.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TpV-iwP9qw9c"
      },
      "source": [
        "## Clean up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources created in this tutorial:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sx_vKniMq9ZX"
      },
      "outputs": [],
      "source": [
        "delete_model = False\n",
        "delete_objects = False\n",
        "delete_bucket = False\n",
        "\n",
        "if delete_model:\n",
        "    endpoint.delete(force=True)\n",
        "    model.delete()\n",
        "if delete_objects:\n",
        "    !gcloud storage rm --recursive $BUCKET_URI/**\n",
        "if delete_bucket:\n",
        "    !gcloud storage buckets delete $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dzESMIONd6IO"
      },
      "source": [
        "## What's next\n",
        "\n",
        "- Explore the [Vertex AI Model Garden](https://console.cloud.google.com/vertex-ai/model-garden)\n",
        "- See also how to [Serve Gemma open models using GPUs on GKE with vLLM](https://cloud.google.com/kubernetes-engine/docs/tutorials/serve-gemma-gpu-vllm)\n",
        "- Learn more about [KerasLP](https://keras.io/keras_nlp)\n",
        "- Learn more about [vLLM](https://github.com/vllm-project/vllm)\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "environment": {
      "kernel": "conda-root-py",
      "name": "workbench-notebooks.m115",
      "type": "gcloud",
      "uri": "gcr.io/deeplearning-platform-release/workbench-notebooks:m115"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.12.1"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "bfd1144bd7ad48e6b3686c5fdf6f1563": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_fd1dfcfd0cbf4923823736e1789b5112",
              "IPY_MODEL_67a173aa66e94699a157c5f4d7449d56",
              "IPY_MODEL_7413e1164e074cbea6091644a4580e3c"
            ],
            "layout": "IPY_MODEL_1f3ed8c6c9294b60b13440d435d4d280"
          }
        },
        "fd1dfcfd0cbf4923823736e1789b5112": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_6919c81e65e548e09db80694d48d6523",
            "placeholder": "​",
            "style": "IPY_MODEL_651f75d7f7e54c959364fae5ff6d707b",
            "value": "Loading checkpoint shards: 100%"
          }
        },
        "67a173aa66e94699a157c5f4d7449d56": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_0a8545ddc700438d8c915fc64a694787",
            "max": 3,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_2cea77d31a1a41b5a983693ddbe55fb9",
            "value": 3
          }
        },
        "7413e1164e074cbea6091644a4580e3c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_5fb6c6fc55d846d5a8ad639b91ad32ef",
            "placeholder": "​",
            "style": "IPY_MODEL_5ac0d9be26774111b61a543037e3c7fd",
            "value": " 3/3 [00:46&lt;00:00, 11.05s/it]"
          }
        },
        "1f3ed8c6c9294b60b13440d435d4d280": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "6919c81e65e548e09db80694d48d6523": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "651f75d7f7e54c959364fae5ff6d707b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "0a8545ddc700438d8c915fc64a694787": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "2cea77d31a1a41b5a983693ddbe55fb9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "5fb6c6fc55d846d5a8ad639b91ad32ef": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "5ac0d9be26774111b61a543037e3c7fd": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}