Simon-LLM-Application-Gemma-2b-LORA-Fine-Tuning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "G3MMAcssHTML"
   },
   "source": [
    "<link rel=\"stylesheet\" href=\"/site-assets/css/gemma.css\">\n",
    "<link rel=\"stylesheet\" href=\"https://fonts.googleapis.com/css2?family=Google+Symbols:opsz,wght,FILL,GRAD@20..48,100..700,0..1,-50..200\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SDEExiAk4fLb"
   },
   "source": [
    "# Fine-tune Gemma models in Keras using LoRA\n",
    "\n",
    "改編自 Google Cloud 官方範例：https://www.apache.org/licenses/LICENSE-2.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lSGRSsRPgkzK"
   },
   "source": [
    "## Overview\n",
    "\n",
    "Gemma 是⼀系列輕量級的開放式 LLM 模型，建構於與打造 Gemini 模型相同的研發技術之上。\n",
    "\n",
    "大型語言模型 (LLMs) 像 Gemma 已被證明能有效執行各種自然語言處理 (NLP) 任務。LLM 首先會透過自我監督的方式，在大量文本資料集上進行預訓練。預訓練可幫助 LLM 學習通用知識，例如詞語之間的統計關係。然後，LLM 可以使用特定領域的資料進行微調，以執行下游任務 (例如情緒分析)。\n",
    "\n",
    "LLMs 的模型非常龐大。對於大多數應用程式來說，完整微調 (更新模型中的所有參數) 並非必要，因為典型的微調資料集通常遠遠小於預訓練資料集。\n",
    "\n",
    "低秩適應 (LoRA) 是一種微調技術，可以大幅減少下游任務的可訓練參數量。該技術透過凍結模型權重並引入少量新權重的方式來實現這一目標。LoRA 微調速度更快、記憶體消耗更少，生成的模型權重更小 (幾百 MB)，同時還能維持模型輸出质量。\n",
    "\n",
    "這個 Colab ，我將引導您使用 KerasNLP 執行 LoRA 微調，使用 Gemma 2B 模型以及 Databricks Dolly 15k 資料集 。此資料集包含 15,000 個由人類生成的高品質提示/回覆配對，專為微調 LLM 而設計。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "w1q6-W_mKIT-"
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lyhHCMfoRZ_v"
   },
   "source": [
    "### 取得 Gemma 訪問權限\n",
    "\n",
    "要完成本教學，您首先需要完成 [Gemma 設定](https://ai.google.dev/gemma/docs/setup) 中的設定指示。Gemma 設定指示將指導您完成以下步驟：\n",
    "\n",
    "* 在 [kaggle.com](https://kaggle.com) 上獲得 Gemma 的訪問權限。\n",
    "* 選擇一個具有足夠資源的 Colab 運行時環境，以運行 Gemma 2B 模型。\n",
    "* 生成並配置您的 Kaggle 使用者名稱和 API 金鑰。\n",
    "\n",
    "完成 Gemma 設定後，請進入下一個部分，為您的 Colab 環境設置環境變數。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AZ5Qo0fxRZ1V"
   },
   "source": [
    "### 選擇運行時環境\n",
    "\n",
    "要完成本教學，您需要選擇一個具有足夠資源的 Colab 運行時環境來執行 Gemma 模型。在這裡，您可以使用 T4 GPU：\n",
    "\n",
    "1. 在 Colab 視窗的右上角，選擇 &#9662;（**額外的連線選項**）。\n",
    "2. 選擇 **變更運行時類型**。\n",
    "3. 在 **硬體加速器** 下，選擇 **T4 GPU**。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hsPC0HRkJl0K"
   },
   "source": [
    "### 配置您的 API 金鑰\n",
    "\n",
    "要使用 Gemma，您必須提供您的 Kaggle 使用者名稱和 Kaggle API 金鑰。\n",
    "\n",
    "要生成 Kaggle API 金鑰，請前往 Kaggle 使用者檔案的 **帳戶** 標籤頁並選擇 **Create New Token**。這將會下載一個包含您的 API 認證的 `kaggle.json` 檔案。\n",
    "\n",
    "在 Colab 中，選擇左側面板中的 **Secrets**（🔑），然後新增您的 Kaggle 使用者名稱和 Kaggle API 金鑰。將您的使用者名稱儲存為 `KAGGLE_USERNAME`，並將您的 API 金鑰儲存為 `KAGGLE_KEY`。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7iOF6Yo-wUEC"
   },
   "source": [
    "### Set environment variables\n",
    "\n",
    "Set environment variables for `KAGGLE_USERNAME` and `KAGGLE_KEY`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "executionInfo": {
     "elapsed": 2836,
     "status": "ok",
     "timestamp": 1725684340358,
     "user": {
      "displayName": "劉育維",
      "userId": "07932650701621055368"
     },
     "user_tz": -480
    },
    "id": "0_EdOg9DPK6Q"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from google.colab import userdata\n",
    "\n",
    "# Note: `userdata.get` is a Colab API. If you're not using Colab, set the env\n",
    "# vars as appropriate for your system.\n",
    "\n",
    "os.environ[\"KAGGLE_USERNAME\"] = userdata.get('KAGGLE_USERNAME')\n",
    "os.environ[\"KAGGLE_KEY\"] = userdata.get('KAGGLE_KEY')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CuEUAKJW1QkQ"
   },
   "source": [
    "### Install dependencies\n",
    "\n",
    "Install Keras, KerasNLP, and other dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 22473,
     "status": "ok",
     "timestamp": 1725684362827,
     "user": {
      "displayName": "劉育維",
      "userId": "07932650701621055368"
     },
     "user_tz": -480
    },
    "id": "1eeBtYqJsZPG",
    "outputId": "a3f22767-ffdc-4c85-f690-f252db7cd8f2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m572.2/572.2 kB\u001b[0m \u001b[31m6.5 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[31m31.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.1/1.1 MB\u001b[0m \u001b[31m21.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h"
     ]
    }
   ],
   "source": [
    "# Install Keras 3 last. See https://keras.io/getting_started/ for more details.\n",
    "!pip install -q -U keras-nlp\n",
    "!pip install -q -U \"keras>=3\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rGLS-l5TxIR4"
   },
   "source": [
    "### Select a backend\n",
    "\n",
    "Keras 是一個高階、多框架的深度學習 API，旨在簡化使用並提升使用便利性。使用 Keras 3，您可以在以下三種後端之一上運行工作流程：TensorFlow、JAX 或 PyTorch。\n",
    "\n",
    "在本教學中，請將後端配置為 JAX。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "executionInfo": {
     "elapsed": 5,
     "status": "ok",
     "timestamp": 1725684362828,
     "user": {
      "displayName": "劉育維",
      "userId": "07932650701621055368"
     },
     "user_tz": -480
    },
    "id": "yn5uy8X8sdD0"
   },
   "outputs": [],
   "source": [
    "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\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hZs8XXqUKRmi"
   },
   "source": [
    "### Import packages\n",
    "\n",
    "Import Keras and KerasNLP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "executionInfo": {
     "elapsed": 8664,
     "status": "ok",
     "timestamp": 1725684371488,
     "user": {
      "displayName": "劉育維",
      "userId": "07932650701621055368"
     },
     "user_tz": -480
    },
    "id": "FYHyPUA9hKTf"
   },
   "outputs": [],
   "source": [
    "import keras\n",
    "import keras_nlp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9T7xe_jzslv4"
   },
   "source": [
    "## Load Dataset\n",
    "\n",
    "本次，我們使用的文件是：[erhwenkuo/dolly-15k-chinese-zhtw](https://huggingface.co/datasets/erhwenkuo/dolly-15k-chinese-zhtw)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 4371,
     "status": "ok",
     "timestamp": 1725684375850,
     "user": {
      "displayName": "劉育維",
      "userId": "07932650701621055368"
     },
     "user_tz": -480
    },
    "id": "xRaNCPUXKoa7",
    "outputId": "693d6010-8c9d-436e-bd74-a7f13c98d384"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-09-07 04:46:10--  https://huggingface.co/datasets/erhwenkuo/dolly-15k-chinese-zhtw/resolve/main/data/train-00000-of-00001-839cf763a52639ec.parquet\n",
      "Resolving huggingface.co (huggingface.co)... 3.165.160.61, 3.165.160.12, 3.165.160.59, ...\n",
      "Connecting to huggingface.co (huggingface.co)|3.165.160.61|:443... connected.\n",
      "HTTP request sent, awaiting response... 302 Found\n",
      "Location: https://cdn-lfs.huggingface.co/repos/7f/a4/7fa4bedfecc28e6c287d997e1cf54c95a43f10cd85204e4590a9e456dfe93acf/54fcd90b0dfa518bac827f1016031be9a53684e9594d74afbd4b2e51b28c44ab?response-content-disposition=inline%3B+filename*%3DUTF-8%27%27train-00000-of-00001-839cf763a52639ec.parquet%3B+filename%3D%22train-00000-of-00001-839cf763a52639ec.parquet%22%3B&Expires=1725943571&Policy=eyJTdGF0ZW1lbnQiOlt7IkNvbmRpdGlvbiI6eyJEYXRlTGVzc1RoYW4iOnsiQVdTOkVwb2NoVGltZSI6MTcyNTk0MzU3MX19LCJSZXNvdXJjZSI6Imh0dHBzOi8vY2RuLWxmcy5odWdnaW5nZmFjZS5jby9yZXBvcy83Zi9hNC83ZmE0YmVkZmVjYzI4ZTZjMjg3ZDk5N2UxY2Y1NGM5NWE0M2YxMGNkODUyMDRlNDU5MGE5ZTQ1NmRmZTkzYWNmLzU0ZmNkOTBiMGRmYTUxOGJhYzgyN2YxMDE2MDMxYmU5YTUzNjg0ZTk1OTRkNzRhZmJkNGIyZTUxYjI4YzQ0YWI%7EcmVzcG9uc2UtY29udGVudC1kaXNwb3NpdGlvbj0qIn1dfQ__&Signature=jRRnPpkcpFhAhfBlqSOGFpj1PHvbWAd3IPAhhvHohyQRHsVU0pQ7xNraDSgb%7E%7E8L41ejAQ4oxOlsgJIfDUTnotlVS0%7EDskFXUFSy1VgrFsykOOwie4WxSDpEu1xzNXpAb4f44R6j%7Es2ZD3BXWVg%7EAHqgV9jUjra%7ErhNX8mKtvSsC2cUR9zZaPKsNH3rPPffOlpGxvG%7EI%7ENAVzFOG05e%7Evr57oeroZ5UgRmQ4ajX--AcTOTWXYXTSJnOfH49gJbz7iBFsy%7ES40e9kz7xGP0Py5EO3GEn4-Vbji4PP-EPTU38-V1rVCq%7EIUTxm-Vhr5rk5N9bhWujqZMSckrBNWOBGSw__&Key-Pair-Id=K3ESJI6DHPFC7 [following]\n",
      "--2024-09-07 04:46:11--  https://cdn-lfs.huggingface.co/repos/7f/a4/7fa4bedfecc28e6c287d997e1cf54c95a43f10cd85204e4590a9e456dfe93acf/54fcd90b0dfa518bac827f1016031be9a53684e9594d74afbd4b2e51b28c44ab?response-content-disposition=inline%3B+filename*%3DUTF-8%27%27train-00000-of-00001-839cf763a52639ec.parquet%3B+filename%3D%22train-00000-of-00001-839cf763a52639ec.parquet%22%3B&Expires=1725943571&Policy=eyJTdGF0ZW1lbnQiOlt7IkNvbmRpdGlvbiI6eyJEYXRlTGVzc1RoYW4iOnsiQVdTOkVwb2NoVGltZSI6MTcyNTk0MzU3MX19LCJSZXNvdXJjZSI6Imh0dHBzOi8vY2RuLWxmcy5odWdnaW5nZmFjZS5jby9yZXBvcy83Zi9hNC83ZmE0YmVkZmVjYzI4ZTZjMjg3ZDk5N2UxY2Y1NGM5NWE0M2YxMGNkODUyMDRlNDU5MGE5ZTQ1NmRmZTkzYWNmLzU0ZmNkOTBiMGRmYTUxOGJhYzgyN2YxMDE2MDMxYmU5YTUzNjg0ZTk1OTRkNzRhZmJkNGIyZTUxYjI4YzQ0YWI%7EcmVzcG9uc2UtY29udGVudC1kaXNwb3NpdGlvbj0qIn1dfQ__&Signature=jRRnPpkcpFhAhfBlqSOGFpj1PHvbWAd3IPAhhvHohyQRHsVU0pQ7xNraDSgb%7E%7E8L41ejAQ4oxOlsgJIfDUTnotlVS0%7EDskFXUFSy1VgrFsykOOwie4WxSDpEu1xzNXpAb4f44R6j%7Es2ZD3BXWVg%7EAHqgV9jUjra%7ErhNX8mKtvSsC2cUR9zZaPKsNH3rPPffOlpGxvG%7EI%7ENAVzFOG05e%7Evr57oeroZ5UgRmQ4ajX--AcTOTWXYXTSJnOfH49gJbz7iBFsy%7ES40e9kz7xGP0Py5EO3GEn4-Vbji4PP-EPTU38-V1rVCq%7EIUTxm-Vhr5rk5N9bhWujqZMSckrBNWOBGSw__&Key-Pair-Id=K3ESJI6DHPFC7\n",
      "Resolving cdn-lfs.huggingface.co (cdn-lfs.huggingface.co)... 108.138.94.14, 108.138.94.122, 108.138.94.25, ...\n",
      "Connecting to cdn-lfs.huggingface.co (cdn-lfs.huggingface.co)|108.138.94.14|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 7492947 (7.1M) [binary/octet-stream]\n",
      "Saving to: ‘databricks-dolly-15k-zhtw.parquet’\n",
      "\n",
      "databricks-dolly-15 100%[===================>]   7.15M  2.07MB/s    in 3.5s    \n",
      "\n",
      "2024-09-07 04:46:15 (2.07 MB/s) - ‘databricks-dolly-15k-zhtw.parquet’ saved [7492947/7492947]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget -O databricks-dolly-15k-zhtw.parquet https://huggingface.co/datasets/erhwenkuo/dolly-15k-chinese-zhtw/resolve/main/data/train-00000-of-00001-839cf763a52639ec.parquet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "45UpBDfBgf0I"
   },
   "source": [
    "預處理資料。本教學使用 1000 個訓練範例的子集來加快 notebook 的執行速度。若需更高品質的微調，請考慮使用更多的訓練資料。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "executionInfo": {
     "elapsed": 4289,
     "status": "ok",
     "timestamp": 1725684380127,
     "user": {
      "displayName": "劉育維",
      "userId": "07932650701621055368"
     },
     "user_tz": -480
    },
    "id": "ZiS-KU9osh_N"
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "# Read the Parquet file into a DataFrame\n",
    "df = pd.read_parquet(\"databricks-dolly-15k-zhtw.parquet\")\n",
    "\n",
    "data = []\n",
    "\n",
    "# Iterate over each row in the DataFrame\n",
    "for _, row in df.iterrows():\n",
    "    # Check if there is context, and skip if true\n",
    "    if row[\"context\"]:\n",
    "        continue\n",
    "    # Format the entire example as a single string\n",
    "    template = \"Instruction:\\n{instruction}\\n\\nResponse:\\n{response}\"\n",
    "    data.append(template.format(instruction=row[\"instruction\"], response=row[\"response\"]))\n",
    "\n",
    "# Only use 1000 training examples, to keep it fast\n",
    "data = data[:1000]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7RCE3fdGhDE5"
   },
   "source": [
    "## 載入模型\n",
    "\n",
    "KerasNLP 提供了許多熱門[模型架構](https://keras.io/api/keras_nlp/models/)的實作。在本教學中，您將使用 `GemmaCausalLM` 建立一個模型，這是一個用於因果語言建模的端對端 Gemma 模型。因果語言模型會根據先前的標記預測下一個標記。\n",
    "\n",
    "使用 `from_preset` 方法來建立模型：\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 385
    },
    "executionInfo": {
     "elapsed": 55171,
     "status": "ok",
     "timestamp": 1725684435282,
     "user": {
      "displayName": "劉育維",
      "userId": "07932650701621055368"
     },
     "user_tz": -480
    },
    "id": "vz5zLEyLstfn",
    "outputId": "1729630d-270d-4e75-bda1-f22897ec2a49"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Preprocessor: \"gemma_causal_lm_preprocessor\"</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1mPreprocessor: \"gemma_causal_lm_preprocessor\"\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\"> Tokenizer (type)                                   </span>┃<span style=\"font-weight: bold\">                                             Vocab # </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│ gemma_tokenizer (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">GemmaTokenizer</span>)                   │                                             <span style=\"color: #00af00; text-decoration-color: #00af00\">256,000</span> │\n",
       "└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1mTokenizer (type)                                  \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m                                            Vocab #\u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│ gemma_tokenizer (\u001b[38;5;33mGemmaTokenizer\u001b[0m)                   │                                             \u001b[38;5;34m256,000\u001b[0m │\n",
       "└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Model: \"gemma_causal_lm\"</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1mModel: \"gemma_causal_lm\"\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\"> Layer (type)                  </span>┃<span style=\"font-weight: bold\"> Output Shape              </span>┃<span style=\"font-weight: bold\">         Param # </span>┃<span style=\"font-weight: bold\"> Connected to               </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│ padding_mask (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)     │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ token_ids (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)        │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ gemma_backbone                │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">2304</span>)        │   <span style=\"color: #00af00; text-decoration-color: #00af00\">2,614,341,888</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\">589,824,000</span> │ gemma_backbone[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>]       │\n",
       "│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">ReversibleEmbedding</span>)         │                           │                 │                            │\n",
       "└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1mLayer (type)                 \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mOutput Shape             \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m        Param #\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mConnected to              \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│ padding_mask (\u001b[38;5;33mInputLayer\u001b[0m)     │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ token_ids (\u001b[38;5;33mInputLayer\u001b[0m)        │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ gemma_backbone                │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m2304\u001b[0m)        │   \u001b[38;5;34m2,614,341,888\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;34m589,824,000\u001b[0m │ gemma_backbone[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m]       │\n",
       "│ (\u001b[38;5;33mReversibleEmbedding\u001b[0m)         │                           │                 │                            │\n",
       "└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Total params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">2,614,341,888</span> (9.74 GB)\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m Total params: \u001b[0m\u001b[38;5;34m2,614,341,888\u001b[0m (9.74 GB)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">2,614,341,888</span> (9.74 GB)\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m Trainable params: \u001b[0m\u001b[38;5;34m2,614,341,888\u001b[0m (9.74 GB)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Non-trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> (0.00 B)\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m Non-trainable params: \u001b[0m\u001b[38;5;34m0\u001b[0m (0.00 B)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "gemma_lm = keras_nlp.models.GemmaCausalLM.from_preset(\"gemma2_2b_en\")\n",
    "gemma_lm.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Nl4lvPy5zA26"
   },
   "source": [
    "`from_preset` 方法會從預設的架構和權重來實例化模型。在上述程式碼中，字串 \"gemma2_2b_en\" 指定了預設的架構，即擁有 20 億個參數的 Gemma 模型。\n",
    "\n",
    "注意：還有一個具有 70 億參數的 Gemma 模型可供使用。若要在 Colab 中運行更大的模型，您需要付費方案中提供的高級 GPU 訪問權限。或者，您可以在 Kaggle 或 Google Cloud 上進行 [Gemma 7B 模型的分佈式調優](https://ai.google.dev/gemma/docs/distributed_tuning)。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "G_L6A5J-1QgC"
   },
   "source": [
    "## 微調前的推理\n",
    "\n",
    "在本節中，您將使用各種提示來查詢模型，觀察它的回應。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 25333,
     "status": "ok",
     "timestamp": 1725684460605,
     "user": {
      "displayName": "劉育維",
      "userId": "07932650701621055368"
     },
     "user_tz": -480
    },
    "id": "ZwQz3xxxKciD",
    "outputId": "b7a3d40a-c39a-4995-95e0-09f82c5db9f5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Instruction:\n",
      "愛麗絲的父母有三個女兒：艾米、傑西，第三個女兒叫什麼名字？\t\n",
      "\n",
      "Response:\n",
      "艾莉絲的母親有一個女兒叫傑西。\n",
      "\n",
      "Instruction:\n",
      "艾米和傑西是愛麗絲的兩個姐姐。\t\n",
      "\n",
      "Question:\n",
      "艾米和傑西分別是愛麗絲的大姐姐和小姐姐。\n",
      "\n",
      "Instruction:\n",
      "傑西是愛麗絲的妹妹。\t\n",
      "\n",
      "Response:\n",
      "艾米和傑西都是愛麗絲的姐姐，但是艾米是大姐姐，傑西是小姐姐。\n"
     ]
    }
   ],
   "source": [
    "prompt = template.format(\n",
    "    instruction=\"愛麗絲的父母有三個女兒：艾米、傑西，第三個女兒叫什麼名字？\t\",\n",
    "    response=\"\",\n",
    ")\n",
    "sampler = keras_nlp.samplers.TopKSampler(k=5, seed=2)\n",
    "gemma_lm.compile(sampler=sampler)\n",
    "print(gemma_lm.generate(prompt, max_length=256))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Pt7Nr6a7tItO"
   },
   "source": [
    "## LoRA 微調\n",
    "\n",
    "為了讓模型提供更好的回應，使用 Databricks Dolly 15k 資料集進行低秩適應（LoRA）微調。\n",
    "\n",
    "LoRA 的秩（rank）決定了可訓練矩陣的維度，這些矩陣會被添加到大型語言模型（LLM）的原始權重中。它控制了微調調整的表現力和精確性。\n",
    "\n",
    "較高的秩意味著可以進行更詳細的更改，但也會增加可訓練參數的數量。較低的秩則意味著較少的計算負擔，但可能導致適應精度較低。\n",
    "\n",
    "本教學使用的 LoRA 秩為 4。在實務中，建議從相對較小的秩開始（例如 4、8、16），這樣在實驗中計算更為高效。訓練模型並評估其在任務上的性能改進，然後在隨後的試驗中逐漸增加秩，看是否進一步提升性能。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 385
    },
    "executionInfo": {
     "elapsed": 27,
     "status": "ok",
     "timestamp": 1725684460606,
     "user": {
      "displayName": "劉育維",
      "userId": "07932650701621055368"
     },
     "user_tz": -480
    },
    "id": "RCucu6oHz53G",
    "outputId": "62028177-1ab7-4c4d-d0c5-785d4a396873"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Preprocessor: \"gemma_causal_lm_preprocessor\"</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1mPreprocessor: \"gemma_causal_lm_preprocessor\"\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\"> Tokenizer (type)                                   </span>┃<span style=\"font-weight: bold\">                                             Vocab # </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│ gemma_tokenizer (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">GemmaTokenizer</span>)                   │                                             <span style=\"color: #00af00; text-decoration-color: #00af00\">256,000</span> │\n",
       "└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1mTokenizer (type)                                  \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m                                            Vocab #\u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│ gemma_tokenizer (\u001b[38;5;33mGemmaTokenizer\u001b[0m)                   │                                             \u001b[38;5;34m256,000\u001b[0m │\n",
       "└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Model: \"gemma_causal_lm\"</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1mModel: \"gemma_causal_lm\"\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃<span style=\"font-weight: bold\"> Layer (type)                  </span>┃<span style=\"font-weight: bold\"> Output Shape              </span>┃<span style=\"font-weight: bold\">         Param # </span>┃<span style=\"font-weight: bold\"> Connected to               </span>┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│ padding_mask (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)     │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ token_ids (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)        │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ gemma_backbone                │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">2304</span>)        │   <span style=\"color: #00af00; text-decoration-color: #00af00\">2,617,270,528</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\">589,824,000</span> │ gemma_backbone[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>]       │\n",
       "│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">ReversibleEmbedding</span>)         │                           │                 │                            │\n",
       "└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n",
       "</pre>\n"
      ],
      "text/plain": [
       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
       "┃\u001b[1m \u001b[0m\u001b[1mLayer (type)                 \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mOutput Shape             \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m        Param #\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mConnected to              \u001b[0m\u001b[1m \u001b[0m┃\n",
       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
       "│ padding_mask (\u001b[38;5;33mInputLayer\u001b[0m)     │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ token_ids (\u001b[38;5;33mInputLayer\u001b[0m)        │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n",
       "├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n",
       "│ gemma_backbone                │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m2304\u001b[0m)        │   \u001b[38;5;34m2,617,270,528\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;34m589,824,000\u001b[0m │ gemma_backbone[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m]       │\n",
       "│ (\u001b[38;5;33mReversibleEmbedding\u001b[0m)         │                           │                 │                            │\n",
       "└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Total params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">2,617,270,528</span> (9.75 GB)\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m Total params: \u001b[0m\u001b[38;5;34m2,617,270,528\u001b[0m (9.75 GB)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">2,928,640</span> (11.17 MB)\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m Trainable params: \u001b[0m\u001b[38;5;34m2,928,640\u001b[0m (11.17 MB)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Non-trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">2,614,341,888</span> (9.74 GB)\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1m Non-trainable params: \u001b[0m\u001b[38;5;34m2,614,341,888\u001b[0m (9.74 GB)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "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()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hQQ47kcdpbZ9"
   },
   "source": [
    "請注意，啟用 LoRA 會大幅減少可訓練參數的數量（從 26 億減少到 290 萬）。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 1909922,
     "status": "ok",
     "timestamp": 1725686370505,
     "user": {
      "displayName": "劉育維",
      "userId": "07932650701621055368"
     },
     "user_tz": -480
    },
    "id": "_Peq7TnLtHse",
    "outputId": "bf8af61d-3c6c-4485-fc56-d36dc9fe2950"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/2\n",
      "\u001b[1m1000/1000\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m983s\u001b[0m 941ms/step - loss: 1.1238 - sparse_categorical_accuracy: 0.4821\n",
      "Epoch 2/2\n",
      "\u001b[1m1000/1000\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m923s\u001b[0m 896ms/step - loss: 0.9916 - sparse_categorical_accuracy: 0.5133\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.src.callbacks.history.History at 0x7f68801f73d0>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Limit the input sequence length to 256 (to control memory usage).\n",
    "gemma_lm.preprocessor.sequence_length = 256\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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bx3m8f1dB7nk"
   },
   "source": [
    "### 關於在 NVIDIA GPU 上進行混合精度微調的注意事項\n",
    "\n",
    "建議在微調時使用全精度。在 NVIDIA GPU 上進行微調時，請注意，您可以使用混合精度（`keras.mixed_precision.set_global_policy('mixed_bfloat16')`）來加快訓練速度，同時對訓練品質的影響最小。混合精度微調會消耗更多的記憶體，因此僅適用於較大的 GPU。\n",
    "\n",
    "在推理時，使用半精度（`keras.config.set_floatx(\"bfloat16\")`）即可節省記憶體，而混合精度則不適用。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "executionInfo": {
     "elapsed": 19,
     "status": "ok",
     "timestamp": 1725686370506,
     "user": {
      "displayName": "劉育維",
      "userId": "07932650701621055368"
     },
     "user_tz": -480
    },
    "id": "T0lHxEDX03gp"
   },
   "outputs": [],
   "source": [
    "# Uncomment the line below if you want to enable mixed precision training on GPUs\n",
    "# keras.mixed_precision.set_global_policy('mixed_bfloat16')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4yd-1cNw1dTn"
   },
   "source": [
    "## 微調後的推理\n",
    "\n",
    "微調後的回應將遵循提示中提供的指令。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "H55JYJ1a1Kos"
   },
   "source": [
    "### Europe Trip Prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 19369,
     "status": "ok",
     "timestamp": 1725687364198,
     "user": {
      "displayName": "劉育維",
      "userId": "07932650701621055368"
     },
     "user_tz": -480
    },
    "id": "Y7cDJHy8WfCB",
    "outputId": "64ffba1b-386b-4ca2-def4-9e7954961750"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Instruction:\n",
      "哪一種是魚類？Tope還是Rope？\t\n",
      "\n",
      "Response:\n",
      "rope\n"
     ]
    }
   ],
   "source": [
    "prompt = template.format(\n",
    "    instruction=\"哪一種是魚類？Tope還是Rope？\t\",\n",
    "    response=\"\",\n",
    ")\n",
    "sampler = keras_nlp.samplers.TopKSampler(k=5, seed=2)\n",
    "gemma_lm.compile(sampler=sampler)\n",
    "print(gemma_lm.generate(prompt, max_length=256))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I8kFG12l0mVe"
   },
   "source": [
    "請注意，為了示範用途，本教學僅在資料集的一小部分上進行一次 epoch 的微調，且使用較低的 LoRA 秩值。若要從微調後的模型中獲得更好的回應，您可以嘗試：\n",
    "\n",
    "1. 增加微調資料集的大小\n",
    "2. 增加訓練步數（epochs）\n",
    "3. 設定更高的 LoRA 秩值\n",
    "4. 修改超參數值，例如 `learning_rate`（學習率）和 `weight_decay`（權重衰減）。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gSsRdeiof_rJ"
   },
   "source": [
    "## 總結與下一步\n",
    "\n",
    "本教學介紹了如何使用 KerasNLP 對 Gemma 模型進行 LoRA 微調。接下來可以查看以下文件：\n",
    "\n",
    "* 學習如何 [使用 Gemma 模型生成文本](https://ai.google.dev/gemma/docs/get_started)。\n",
    "* 學習如何進行 [Gemma 模型的分佈式微調和推理](https://ai.google.dev/gemma/docs/distributed_tuning)。\n",
    "* 學習如何 [使用 Vertex AI 配合 Gemma 開放模型](https://cloud.google.com/vertex-ai/docs/generative-ai/open-models/use-gemma)。\n",
    "* 學習如何 [使用 KerasNLP 微調 Gemma 並部署到 Vertex AI](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_gemma_kerasnlp_to_vertexai.ipynb)。\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}