Gemma 中文指令微调(Kaggle TPU).ipynb

{"metadata":{"colab":{"provenance":[]},"kaggle":{"accelerator":"tpu1vmV38","dataSources":[{"sourceId":10262,"sourceType":"modelInstanceVersion","modelInstanceId":5172},{"sourceId":11371,"sourceType":"modelInstanceVersion","modelInstanceId":5171},{"sourceId":11375,"sourceType":"modelInstanceVersion","modelInstanceId":5172}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false},"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.10.13"},"papermill":{"default_parameters":{},"duration":672.253242,"end_time":"2024-02-21T10:06:37.294427","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2024-02-21T09:55:25.041185","version":"2.5.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"##### Copyright 2024 Google LLC.","metadata":{"id":"60KmTK7o6ppd","papermill":{"duration":0.007422,"end_time":"2024-02-21T09:55:26.786189","exception":false,"start_time":"2024-02-21T09:55:26.778767","status":"completed"},"tags":[]}},{"cell_type":"code","source":"#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.","metadata":{"papermill":{"duration":0.012321,"end_time":"2024-02-21T09:55:26.80406","exception":false,"start_time":"2024-02-21T09:55:26.791739","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-02-25T06:18:28.741196Z","iopub.execute_input":"2024-02-25T06:18:28.741507Z","iopub.status.idle":"2024-02-25T06:18:28.745486Z","shell.execute_reply.started":"2024-02-25T06:18:28.741480Z","shell.execute_reply":"2024-02-25T06:18:28.744891Z"},"trusted":true},"execution_count":1,"outputs":[]},{"cell_type":"markdown","source":"<table class=\"tfo-notebook-buttons\" align=\"left\">\n  <td>\n    <a target=\"_blank\" href=\"https://www.kaggle.com/windmaple/gemma-kaggle-tpu-only\"><img src=\"https://www.kaggle.com/static/images/logos/kaggle-logo-transparent-300.png\" height=\"32\" width=\"70\"/>Run in Kaggle</a>\n  </td>\n</table>","metadata":{"papermill":{"duration":0.005641,"end_time":"2024-02-21T09:55:26.815024","exception":false,"start_time":"2024-02-21T09:55:26.809383","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"# Gemma insutruction tuning on Kaggle TPU using Chinese dataset\n\nThis notebook is an adapted from the official [Gemma distributed tuning tutorial](https://ai.google.dev/gemma/docs/distributed_tuning) and [Gemma Vertex AI tutorial](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_gemma_kerasnlp_to_vertexai.ipynb). It is meant to demonstrate how to instruction tune Gemma 2B (non-instruction tuned variant) on Kaggle TPU so that the finetuned model can better follow Chinese instructions.\n\n(Note that the instruction-tuned variant of Gemma 2B does have some basic capability to follow Chinese instructions but the technique used here can be used to further enhance it.)\n\nThis notebook is also available directly on [Kaggle](https://www.kaggle.com/windmaple/gemma-kaggle-tpu-only).","metadata":{"papermill":{"duration":0.005634,"end_time":"2024-02-21T09:55:26.825916","exception":false,"start_time":"2024-02-21T09:55:26.820282","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"## Overview\n\nGemma is a family of lightweight, state-of-the-art open models built from research and technology used to create Google Gemini models. Gemma can be further finetuned to suit specific needs. But Large Language Models, such as Gemma, can be very large in size and some of them may not fit on a sing accelerator for finetuning. In this case there are two general approaches for finetuning them:\n1. Parameter Efficient Fine-Tuning (PEFT), which seeks to shrink the effective model size by sacrificing some fidelity. LoRA falls in this category and the [Finetune Gemma models in Keras using LoRA](https://ai.google.dev/gemma/docs/lora_tuning) tutorial demonstrates how to finetune the Gemma 2B model `gemma_2b_en` with LoRA using KerasNLP on a single GPU.\n2. Full parameter finetuning with model parallelism. Model parallelism distributes a single model's weights across multiple devices and enables horizontal scaling. You can find out more about distributed training in this [Keras guide](https://keras.io/guides/distribution/).\n\nThis tutorial walks you through using Keras with a JAX backend to finetune the Gemma 7B model with LoRA and model-parallism distributed training on Google's Tensor Processing Unit (TPU). Note that LoRA can be turned off in this tutorial for a slower but more accurate full-parameter tuning.","metadata":{"id":"Tdlq6K0znh3O","papermill":{"duration":0.005633,"end_time":"2024-02-21T09:55:26.83679","exception":false,"start_time":"2024-02-21T09:55:26.831157","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"## Using accelerators\n\nTechnically you can use either TPU or GPU for this tutorial.\n\n### Notes on TPU environments\n\nGoogle has 3 products that provide TPUs:\n* [Colab](https://colab.sandbox.google.com/) provides TPU v2, which is not sufficient for this tutorial.\n* [Kaggle](https://www.kaggle.com/) offers TPU v3 for free and they work for this tutorial.\n* [Cloud TPU](https://cloud.google.com/tpu?hl=en) offers TPU v3 and newer generations. One way to set it up is:\n  1. Create a new [TPU VM](https://cloud.google.com/tpu/docs/managing-tpus-tpu-vm#tpu-vms)\n  2. Set up [SSH port forwarding](https://cloud.google.com/solutions/connecting-securely#port-forwarding-over-ssh) for your intended Jupyter server port\n  3. Install Jupyter and start it on the TPU VM, then connect to Colab through \"Connect to a local runtime\"\n\n### Notes on multi-GPU setup\n\nAlthough this tutorial focuses on the TPU use case, you can easily adapt it for your own needs if you have a multi-GPU machine.\n\nIf you prefer to work through Colab, it's also possible to provision a multi-GPU VM for Colab directly through \"Connect to a custom GCE VM\" in the Colab Connect menu.\n\n\nWe will focus on using the **free TPU from Kaggle** here.","metadata":{"id":"z-jBO5hmDwrc","papermill":{"duration":0.005425,"end_time":"2024-02-21T09:55:26.847611","exception":false,"start_time":"2024-02-21T09:55:26.842186","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"## Before you begin","metadata":{"papermill":{"duration":0.005257,"end_time":"2024-02-21T09:55:26.85841","exception":false,"start_time":"2024-02-21T09:55:26.853153","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"### Gemma setup\n\nTo complete this tutorial, you will first need to complete the setup instructions at [Gemma setup](https://ai.google.dev/gemma/docs/setup). The Gemma setup instructions show you how to do the following:\n\nGemma models are hosted by Kaggle. To use Gemma, request access on Kaggle:\n\n- Sign in or register at [kaggle.com](https://www.kaggle.com)\n- Open the [Gemma model card](https://www.kaggle.com/models/google/gemma) and select _\"Request Access\"_\n- Complete the consent form and accept the terms and conditions\n","metadata":{"id":"aKvTsIkL98BG","papermill":{"duration":0.005259,"end_time":"2024-02-21T09:55:26.869282","exception":false,"start_time":"2024-02-21T09:55:26.864023","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"## Installation\n\nInstall Keras and KerasNLP with the Gemma model.","metadata":{"id":"AO7a1Q4Yyc9Z","papermill":{"duration":0.005175,"end_time":"2024-02-21T09:55:26.880069","exception":false,"start_time":"2024-02-21T09:55:26.874894","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Install Keras 3 last. See https://keras.io/getting_started/ for more details.\n!pip install -q tensorflow-cpu\n!pip install -q -U keras-nlp tensorflow-hub\n!pip install -q -U keras>=3\n!pip install -qU transformers\n!pip install -U sentencepiece","metadata":{"id":"WWEzVJR4Fx9g","papermill":{"duration":37.05282,"end_time":"2024-02-21T09:56:03.93859","exception":false,"start_time":"2024-02-21T09:55:26.88577","status":"completed"},"tags":[],"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-02-25T06:18:28.746585Z","iopub.execute_input":"2024-02-25T06:18:28.746830Z","iopub.status.idle":"2024-02-25T06:19:22.861539Z","shell.execute_reply.started":"2024-02-25T06:18:28.746805Z","shell.execute_reply":"2024-02-25T06:19:22.860334Z"},"trusted":true},"execution_count":2,"outputs":[{"name":"stdout","text":"\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\nkeras-nlp 0.7.0 requires keras-core, which is not installed.\u001b[0m\u001b[31m\n\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n\u001b[0m\n\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n\u001b[0m\n\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\ntensorflow 2.15.0 requires keras<2.16,>=2.15.0, but you have keras 3.0.5 which is incompatible.\ntensorflow-cpu 2.15.0.post1 requires keras<2.16,>=2.15.0, but you have keras 3.0.5 which is incompatible.\u001b[0m\u001b[31m\n\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n\u001b[0m\n\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n\u001b[0m\n\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\nCollecting sentencepiece\n  Downloading sentencepiece-0.2.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.3 MB)\n\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m17.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m0:01\u001b[0m\n\u001b[?25hInstalling collected packages: sentencepiece\nSuccessfully installed sentencepiece-0.2.0\n\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n\u001b[0m\n\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n","output_type":"stream"}]},{"cell_type":"markdown","source":"### Set up Keras JAX backend","metadata":{"id":"fr9VnPm7FoMf","papermill":{"duration":0.006513,"end_time":"2024-02-21T09:56:03.951855","exception":false,"start_time":"2024-02-21T09:56:03.945342","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"Import JAX and run a sanity check on TPU. Kaggle offers TPUv3-8 devices which have 8 TPU cores with 16GB of memory each.","metadata":{"id":"lbZsUvfhwL2D","papermill":{"duration":0.00599,"end_time":"2024-02-21T09:56:03.964473","exception":false,"start_time":"2024-02-21T09:56:03.958483","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import jax\n\njax.devices()","metadata":{"id":"BK4MpHLKGujb","outputId":"a60376b8-0937-45fc-809b-33eaa92cbc6c","papermill":{"duration":8.126711,"end_time":"2024-02-21T09:56:12.097265","exception":false,"start_time":"2024-02-21T09:56:03.970554","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-02-25T06:19:22.863499Z","iopub.execute_input":"2024-02-25T06:19:22.863840Z","iopub.status.idle":"2024-02-25T06:19:31.217628Z","shell.execute_reply.started":"2024-02-25T06:19:22.863812Z","shell.execute_reply":"2024-02-25T06:19:31.216818Z"},"trusted":true},"execution_count":3,"outputs":[{"execution_count":3,"output_type":"execute_result","data":{"text/plain":"[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),\n TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1),\n TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0),\n TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1),\n TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0),\n TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1),\n TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0),\n TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]"},"metadata":{}}]},{"cell_type":"code","source":"import os\n\n# The Keras 3 distribution API is only implemented for the JAX backend for now\nos.environ[\"KERAS_BACKEND\"] = \"jax\"\n# Pre-allocate 90% of TPU memory to minimize memory fragmentation and allocation\n# overhead\nos.environ[\"XLA_PYTHON_CLIENT_MEM_FRACTION\"] = \"0.9\"","metadata":{"id":"WEgg_OVIL2HY","papermill":{"duration":0.01342,"end_time":"2024-02-21T09:56:12.117529","exception":false,"start_time":"2024-02-21T09:56:12.104109","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-02-25T06:19:31.218594Z","iopub.execute_input":"2024-02-25T06:19:31.218950Z","iopub.status.idle":"2024-02-25T06:19:31.222794Z","shell.execute_reply.started":"2024-02-25T06:19:31.218920Z","shell.execute_reply":"2024-02-25T06:19:31.222086Z"},"trusted":true},"execution_count":4,"outputs":[]},{"cell_type":"markdown","source":"## Load model","metadata":{"id":"wo1xkzr62hXN","papermill":{"duration":0.00603,"end_time":"2024-02-21T09:56:12.129531","exception":false,"start_time":"2024-02-21T09:56:12.123501","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import keras\nimport keras_nlp","metadata":{"id":"kFCmWEKdMA_Y","outputId":"57359fb1-ea3e-482a-aaf3-dcffc265a5ec","papermill":{"duration":8.153944,"end_time":"2024-02-21T09:56:20.315559","exception":false,"start_time":"2024-02-21T09:56:12.161615","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-02-25T06:19:31.224528Z","iopub.execute_input":"2024-02-25T06:19:31.224776Z","iopub.status.idle":"2024-02-25T06:19:38.352286Z","shell.execute_reply.started":"2024-02-25T06:19:31.224750Z","shell.execute_reply":"2024-02-25T06:19:38.351050Z"},"trusted":true},"execution_count":5,"outputs":[]},{"cell_type":"markdown","source":"### Notes on mixed precision training on NVIDIA GPUs\n\nWhen 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 case, 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\nFor inference, half-precision (`keras.config.set_floatx(\"bfloat16\")`) will work and save memory while mixed-precision is not applicable.","metadata":{"id":"bx3m8f1dB7nk","papermill":{"duration":0.005945,"end_time":"2024-02-21T09:56:20.327808","exception":false,"start_time":"2024-02-21T09:56:20.321863","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Uncomment the line below if you want to enable mixed precision training on GPUs\n# keras.mixed_precision.set_global_policy('mixed_bfloat16')","metadata":{"id":"T0lHxEDX03gp","papermill":{"duration":0.0122,"end_time":"2024-02-21T09:56:20.345701","exception":false,"start_time":"2024-02-21T09:56:20.333501","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-02-25T06:19:38.353483Z","iopub.execute_input":"2024-02-25T06:19:38.354011Z","iopub.status.idle":"2024-02-25T06:19:38.357466Z","shell.execute_reply.started":"2024-02-25T06:19:38.353977Z","shell.execute_reply":"2024-02-25T06:19:38.356742Z"},"trusted":true},"execution_count":6,"outputs":[]},{"cell_type":"markdown","source":"To load the model with the weights and tensors distributed across TPUs, first create a new `DeviceMesh`. `DeviceMesh` represents a collection of hardware devices configured for distributed computation and was introduced in Keras 3 as part of the unified distribution API.\n\nThe distribution API enables data and model parallelism, allowing for efficient scaling of deep learning models on multiple accelerators and hosts. It leverages the underlying framework (e.g. JAX) to distribute the program and tensors according to the sharding directives through a procedure called single program, multiple data (SPMD) expansion. Check out more details in the new [Keras 3 distribution API guide](https://keras.io/guides/distribution/).","metadata":{"id":"xrR8TpVS6uPs","papermill":{"duration":0.005852,"end_time":"2024-02-21T09:56:20.357594","exception":false,"start_time":"2024-02-21T09:56:20.351742","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Create a device mesh with (1, 8) shape so that the weights are sharded across\n# all 8 TPUs.\ndevice_mesh = keras.distribution.DeviceMesh(\n    (1, 8),\n    [\"batch\", \"model\"],\n    devices=keras.distribution.list_devices())","metadata":{"id":"7gxEkpUiP1Qf","papermill":{"duration":0.012454,"end_time":"2024-02-21T09:56:20.375812","exception":false,"start_time":"2024-02-21T09:56:20.363358","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-02-25T06:19:38.358347Z","iopub.execute_input":"2024-02-25T06:19:38.358589Z","iopub.status.idle":"2024-02-25T06:19:38.370232Z","shell.execute_reply.started":"2024-02-25T06:19:38.358565Z","shell.execute_reply":"2024-02-25T06:19:38.369483Z"},"trusted":true},"execution_count":7,"outputs":[]},{"cell_type":"markdown","source":"`LayoutMap` from the distribution API specifies how the weights and tensors should be sharded or replicated, using the string keys, for example, `token_embedding/embeddings` below, which are treated like regex to match tensor paths. Matched tensors are sharded with model dimensions (8 TPUs); others will be fully replicated.","metadata":{"id":"gTSJUwkC-7c6","papermill":{"duration":0.00582,"end_time":"2024-02-21T09:56:20.387661","exception":false,"start_time":"2024-02-21T09:56:20.381841","status":"completed"},"tags":[]}},{"cell_type":"code","source":"model_dim = \"model\"\n\nlayout_map = keras.distribution.LayoutMap(device_mesh)\n\n# Weights that match 'token_embedding/embeddings' will be sharded on 8 TPUs\nlayout_map[\"token_embedding/embeddings\"] = (None, model_dim)\n# Regex to match against the query, key and value matrices in the decoder\n# attention layers\nlayout_map[\"decoder_block.*attention.*(query|key|value).*kernel\"] = (\n    None, model_dim, None)\n\nlayout_map[\"decoder_block.*attention_output.*kernel\"] = (\n    None, None, model_dim)\nlayout_map[\"decoder_block.*ffw_gating.*kernel\"] = (model_dim, None)\nlayout_map[\"decoder_block.*ffw_linear.*kernel\"] = (None, model_dim)","metadata":{"id":"8Wgh8h0qQCcu","papermill":{"duration":0.013198,"end_time":"2024-02-21T09:56:20.406598","exception":false,"start_time":"2024-02-21T09:56:20.3934","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-02-25T06:19:38.371156Z","iopub.execute_input":"2024-02-25T06:19:38.371436Z","iopub.status.idle":"2024-02-25T06:19:38.380554Z","shell.execute_reply.started":"2024-02-25T06:19:38.371412Z","shell.execute_reply":"2024-02-25T06:19:38.379812Z"},"trusted":true},"execution_count":8,"outputs":[]},{"cell_type":"markdown","source":"`ModelParallel` allows you to shard model weights or activation tensors across all devcies on the `DeviceMesh`. In this case, some of the Gemma 7B model weights are sharded across 8 TPU chips according the `layout_map` defined above. Now load the model in the distributed way.","metadata":{"id":"6n4Zlvk9ALhZ","papermill":{"duration":0.005938,"end_time":"2024-02-21T09:56:20.418485","exception":false,"start_time":"2024-02-21T09:56:20.412547","status":"completed"},"tags":[]}},{"cell_type":"code","source":"model_parallel = keras.distribution.ModelParallel(\n    device_mesh, layout_map, batch_dim_name=\"batch\")\n\nkeras.distribution.set_distribution(model_parallel)\ngemma_lm = keras_nlp.models.GemmaCausalLM.from_preset(\"gemma_2b_en\")","metadata":{"id":"bu48vUnbQj0p","outputId":"fd216acb-852c-46f4-e8ac-2d8f91362d24","papermill":{"duration":145.668669,"end_time":"2024-02-21T09:58:46.092826","exception":false,"start_time":"2024-02-21T09:56:20.424157","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-02-25T06:19:38.381408Z","iopub.execute_input":"2024-02-25T06:19:38.381672Z","iopub.status.idle":"2024-02-25T06:20:24.977280Z","shell.execute_reply.started":"2024-02-25T06:19:38.381648Z","shell.execute_reply":"2024-02-25T06:20:24.976133Z"},"trusted":true},"execution_count":9,"outputs":[{"name":"stderr","text":"Attaching 'config.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\nAttaching 'config.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\nAttaching 'model.weights.h5' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\nAttaching 'tokenizer.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\nAttaching 'assets/tokenizer/vocabulary.spm' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\nnormalizer.cc(51) LOG(INFO) precompiled_charsmap is empty. use identity normalization.\n","output_type":"stream"}]},{"cell_type":"markdown","source":"Now verify that the model has been partitioned correctly. Let's take `decoder_block_1` as an example.","metadata":{"id":"ORCOIawAvpZ1","papermill":{"duration":0.006357,"end_time":"2024-02-21T09:58:46.105898","exception":false,"start_time":"2024-02-21T09:58:46.099541","status":"completed"},"tags":[]}},{"cell_type":"code","source":"decoder_block_1 = gemma_lm.backbone.get_layer('decoder_block_1')\nprint(type(decoder_block_1))\nfor variable in decoder_block_1.weights:\n  print(f'{variable.path:<58}  {str(variable.shape):<16}  {str(variable.value.sharding.spec)}')","metadata":{"id":"DqT7TRHKvoMK","papermill":{"duration":0.014281,"end_time":"2024-02-21T09:58:46.126669","exception":false,"start_time":"2024-02-21T09:58:46.112388","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-02-25T06:20:24.978518Z","iopub.execute_input":"2024-02-25T06:20:24.978782Z","iopub.status.idle":"2024-02-25T06:20:24.983981Z","shell.execute_reply.started":"2024-02-25T06:20:24.978755Z","shell.execute_reply":"2024-02-25T06:20:24.983219Z"},"trusted":true},"execution_count":10,"outputs":[{"name":"stdout","text":"<class 'keras_nlp.src.models.gemma.gemma_decoder_block.GemmaDecoderBlock'>\ndecoder_block_1/pre_attention_norm/scale                    (2048,)           PartitionSpec(None,)\ndecoder_block_1/attention/query/kernel                      (8, 2048, 256)    PartitionSpec(None, 'model', None)\ndecoder_block_1/attention/key/kernel                        (1, 2048, 256)    PartitionSpec(None, 'model', None)\ndecoder_block_1/attention/value/kernel                      (1, 2048, 256)    PartitionSpec(None, 'model', None)\ndecoder_block_1/attention/attention_output/kernel           (8, 256, 2048)    PartitionSpec(None, None, 'model')\ndecoder_block_1/pre_ffw_norm/scale                          (2048,)           PartitionSpec(None,)\ndecoder_block_1/ffw_gating/kernel                           (2048, 16384)     PartitionSpec('model', None)\ndecoder_block_1/ffw_gating_2/kernel                         (2048, 16384)     PartitionSpec('model', None)\ndecoder_block_1/ffw_linear/kernel                           (16384, 2048)     PartitionSpec(None, 'model')\n","output_type":"stream"}]},{"cell_type":"markdown","source":"## Load instruction dataset","metadata":{}},{"cell_type":"code","source":"!wget -O baike.jsonl https://huggingface.co/datasets/Hello-SimpleAI/HC3-Chinese/raw/main/baike.jsonl","metadata":{"execution":{"iopub.status.busy":"2024-02-25T06:20:24.986255Z","iopub.execute_input":"2024-02-25T06:20:24.986531Z","iopub.status.idle":"2024-02-25T06:20:25.769667Z","shell.execute_reply.started":"2024-02-25T06:20:24.986504Z","shell.execute_reply":"2024-02-25T06:20:25.768302Z"},"trusted":true},"execution_count":11,"outputs":[{"name":"stdout","text":"--2024-02-25 06:20:25--  https://huggingface.co/datasets/Hello-SimpleAI/HC3-Chinese/raw/main/baike.jsonl\nResolving huggingface.co (huggingface.co)... 18.244.202.118, 18.244.202.60, 18.244.202.68, ...\nConnecting to huggingface.co (huggingface.co)|18.244.202.118|:443... connected.\nHTTP request sent, awaiting response... 200 OK\nLength: 5005244 (4.8M) [text/plain]\nSaving to: ‘baike.jsonl’\n\nbaike.jsonl         100%[===================>]   4.77M  14.6MB/s    in 0.3s    \n\n2024-02-25 06:20:25 (14.6 MB/s) - ‘baike.jsonl’ saved [5005244/5005244]\n\n","output_type":"stream"}]},{"cell_type":"code","source":"import re\nimport json\ndata = []\ncontext = \"你是一个知识丰富的人工智能助手，用户将用中文向你提问，你将根据你的知识用中文来如实回答问题。\\n\"\nwith open(\"baike.jsonl\") as file:\n    for line in file:\n        features = json.loads(line)        \n        template = context + \"问题：\\n{question}\\n答案：\\n{human_answers[0]}\"\n        data.append(template.format(**features))\n\n# Manually construct a test case; \n# Already made sure the finetuning dataset contains nothing about zsh\ntest_prompt = context + \"问题：\\n我有一个信息科学相关的问题，请用中文回答，什么是 zsh\\n答案：\\n\"\n# 4616 in total in baike split\ntrain_data = data[:4600]","metadata":{"_kg_hide-output":true,"id":"6MVJlsuSXCcf","outputId":"fcfda7bb-899e-4606-88e1-42eed9e70fc0","papermill":{"duration":43.816972,"end_time":"2024-02-21T09:59:56.217135","exception":false,"start_time":"2024-02-21T09:59:12.400163","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-02-25T06:20:25.771273Z","iopub.execute_input":"2024-02-25T06:20:25.771562Z","iopub.status.idle":"2024-02-25T06:20:25.832424Z","shell.execute_reply.started":"2024-02-25T06:20:25.771535Z","shell.execute_reply":"2024-02-25T06:20:25.831561Z"},"trusted":true},"execution_count":12,"outputs":[]},{"cell_type":"markdown","source":"## Inference before finetuning","metadata":{"id":"jc0ZzYIW0TSN","papermill":{"duration":0.005944,"end_time":"2024-02-21T09:58:46.139044","exception":false,"start_time":"2024-02-21T09:58:46.1331","status":"completed"},"tags":[]}},{"cell_type":"code","source":"gemma_lm.generate(test_prompt, max_length=200)","metadata":{"id":"ClaTyBp3Tgr4","outputId":"f2cf1ac7-469a-4c91-adff-74c2e3a1de89","papermill":{"duration":26.22201,"end_time":"2024-02-21T09:59:12.367352","exception":false,"start_time":"2024-02-21T09:58:46.145342","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-02-25T06:20:25.833549Z","iopub.execute_input":"2024-02-25T06:20:25.833821Z","iopub.status.idle":"2024-02-25T06:20:42.381527Z","shell.execute_reply.started":"2024-02-25T06:20:25.833795Z","shell.execute_reply":"2024-02-25T06:20:42.380383Z"},"trusted":true},"execution_count":13,"outputs":[{"execution_count":13,"output_type":"execute_result","data":{"text/plain":"'你是一个知识丰富的人工智能助手，用户将用中文向你提问，你将根据你的知识用中文来如实回答问题。\\n问题：\\n我有一个信息科学相关的问题，请用中文回答，什么是 zsh\\n答案：\\nzsh 是一个命令行界面（CLI）的 shell，它支持许多命令行工具，包括 bash， fish， ksh， mksh， pdksh， tcsh， zsh， 和 yash。\\nzsh 是一个命令行界面（CLI）的 shell，它支持许多命令行工具，包括 bash， fish， ksh， mksh， pdksh， tcsh， zsh， 和 yash。\\nzsh 是一个命令行界面（CLI）的 shell，它支持许多命令行工具，包括 bash， fish， ksh， mksh， pdksh， tcsh， zsh， 和 yash。\\nzsh 是一个命令行界面（CLI）的 shell，它支持'"},"metadata":{}}]},{"cell_type":"markdown","source":"The model starts to repeat itself after a few sentences, which is not good.","metadata":{}},{"cell_type":"markdown","source":"## Finetune","metadata":{"id":"IcPCXCwvXC7t","papermill":{"duration":0.006625,"end_time":"2024-02-21T09:59:12.393943","exception":false,"start_time":"2024-02-21T09:59:12.387318","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"Perform finetuning using [Low Rank Adaptation](https://arxiv.org/abs/2106.09685) (LoRA). LoRA is a fine-tuning 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. Basically LoRA reparameterizes the larger full weight matrices by 2 smaller low-rank matrices AxB to train and this technique makes training much faster and more memory-efficient.","metadata":{"id":"xiLW0SpI1PfC","papermill":{"duration":0.032799,"end_time":"2024-02-21T09:59:56.345251","exception":false,"start_time":"2024-02-21T09:59:56.312452","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Enable LoRA for the model and set the LoRA rank to 4.\ngemma_lm.backbone.enable_lora(rank=4)","metadata":{"id":"3o_Gi3v_jp7s","papermill":{"duration":0.527031,"end_time":"2024-02-21T09:59:56.902475","exception":false,"start_time":"2024-02-21T09:59:56.375444","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-02-25T06:20:42.382794Z","iopub.execute_input":"2024-02-25T06:20:42.383078Z","iopub.status.idle":"2024-02-25T06:20:43.032532Z","shell.execute_reply.started":"2024-02-25T06:20:42.383050Z","shell.execute_reply":"2024-02-25T06:20:43.031405Z"},"trusted":true},"execution_count":14,"outputs":[]},{"cell_type":"code","source":"# Fine-tune on the IMDb movie reviews dataset.\n\n# Limit the input sequence length to 128 to control memory usage.\ngemma_lm.preprocessor.sequence_length = 128\n# Use AdamW (a common optimizer for transformer models).\noptimizer = keras.optimizers.AdamW(\n    learning_rate=5e-5,\n    weight_decay=0.01,\n)\n# Exclude layernorm and bias terms from decay.\noptimizer.exclude_from_weight_decay(var_names=[\"bias\", \"scale\"])\n\ngemma_lm.compile(\n    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n    optimizer=optimizer,\n    weighted_metrics=[keras.metrics.SparseCategoricalAccuracy()],\n)\ngemma_lm.summary()\ngemma_lm.fit(train_data, epochs=5, batch_size=32)","metadata":{"_kg_hide-output":true,"id":"1-hQFy7hXWRl","outputId":"eb21b839-1e23-4be5-afcb-90a7bb0e4167","papermill":{"duration":362.718929,"end_time":"2024-02-21T10:05:59.651432","exception":false,"start_time":"2024-02-21T09:59:56.932503","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-02-25T06:20:43.033681Z","iopub.execute_input":"2024-02-25T06:20:43.033981Z","iopub.status.idle":"2024-02-25T06:29:05.524103Z","shell.execute_reply.started":"2024-02-25T06:20:43.033953Z","shell.execute_reply":"2024-02-25T06:29:05.523286Z"},"trusted":true},"execution_count":15,"outputs":[{"output_type":"display_data","data":{"text/plain":"\u001b[1mPreprocessor: \"gemma_causal_lm_preprocessor\"\u001b[0m\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Preprocessor: \"gemma_causal_lm_preprocessor\"</span>\n</pre>\n"},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n┃\u001b[1m \u001b[0m\u001b[1mTokenizer (type)                                  \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m                                            Vocab #\u001b[0m\u001b[1m \u001b[0m┃\n┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n│ gemma_tokenizer (\u001b[38;5;33mGemmaTokenizer\u001b[0m)                   │                                             \u001b[38;5;34m256,000\u001b[0m │\n└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n┃<span style=\"font-weight: bold\"> Tokenizer (type)                                   </span>┃<span style=\"font-weight: bold\">                                             Vocab # </span>┃\n┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n│ gemma_tokenizer (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">GemmaTokenizer</span>)                   │                                             <span style=\"color: #00af00; text-decoration-color: #00af00\">256,000</span> │\n└────────────────────────────────────────────────────┴─────────────────────────────────────────────────────┘\n</pre>\n"},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"\u001b[1mModel: \"gemma_causal_lm\"\u001b[0m\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\">Model: \"gemma_causal_lm\"</span>\n</pre>\n"},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n┃\u001b[1m \u001b[0m\u001b[1mLayer (type)                 \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mOutput Shape             \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m        Param #\u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1mConnected to              \u001b[0m\u001b[1m \u001b[0m┃\n┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n│ padding_mask (\u001b[38;5;33mInputLayer\u001b[0m)     │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n│ token_ids (\u001b[38;5;33mInputLayer\u001b[0m)        │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m)              │               \u001b[38;5;34m0\u001b[0m │ -                          │\n├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n│ gemma_backbone                │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m2048\u001b[0m)        │   \u001b[38;5;34m2,507,536,384\u001b[0m │ padding_mask[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m],        │\n│ (\u001b[38;5;33mGemmaBackbone\u001b[0m)               │                           │                 │ token_ids[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m]            │\n├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n│ token_embedding               │ (\u001b[38;5;45mNone\u001b[0m, \u001b[38;5;45mNone\u001b[0m, \u001b[38;5;34m256000\u001b[0m)      │     \u001b[38;5;34m524,288,000\u001b[0m │ gemma_backbone[\u001b[38;5;34m0\u001b[0m][\u001b[38;5;34m0\u001b[0m]       │\n│ (\u001b[38;5;33mReversibleEmbedding\u001b[0m)         │                           │                 │                            │\n└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n┃<span style=\"font-weight: bold\"> Layer (type)                  </span>┃<span style=\"font-weight: bold\"> Output Shape              </span>┃<span style=\"font-weight: bold\">         Param # </span>┃<span style=\"font-weight: bold\"> Connected to               </span>┃\n┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n│ padding_mask (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)     │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n│ token_ids (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">InputLayer</span>)        │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>)              │               <span style=\"color: #00af00; text-decoration-color: #00af00\">0</span> │ -                          │\n├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n│ gemma_backbone                │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">2048</span>)        │   <span style=\"color: #00af00; text-decoration-color: #00af00\">2,507,536,384</span> │ padding_mask[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>],        │\n│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">GemmaBackbone</span>)               │                           │                 │ token_ids[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>]            │\n├───────────────────────────────┼───────────────────────────┼─────────────────┼────────────────────────────┤\n│ token_embedding               │ (<span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00d7ff; text-decoration-color: #00d7ff\">None</span>, <span style=\"color: #00af00; text-decoration-color: #00af00\">256000</span>)      │     <span style=\"color: #00af00; text-decoration-color: #00af00\">524,288,000</span> │ gemma_backbone[<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>][<span style=\"color: #00af00; text-decoration-color: #00af00\">0</span>]       │\n│ (<span style=\"color: #0087ff; text-decoration-color: #0087ff\">ReversibleEmbedding</span>)         │                           │                 │                            │\n└───────────────────────────────┴───────────────────────────┴─────────────────┴────────────────────────────┘\n</pre>\n"},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"\u001b[1m Total params: \u001b[0m\u001b[38;5;34m2,507,536,384\u001b[0m (9.34 GB)\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Total params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">2,507,536,384</span> (9.34 GB)\n</pre>\n"},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"\u001b[1m Trainable params: \u001b[0m\u001b[38;5;34m1,363,968\u001b[0m (5.20 MB)\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">1,363,968</span> (5.20 MB)\n</pre>\n"},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"\u001b[1m Non-trainable params: \u001b[0m\u001b[38;5;34m2,506,172,416\u001b[0m (9.34 GB)\n","text/html":"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"font-weight: bold\"> Non-trainable params: </span><span style=\"color: #00af00; text-decoration-color: #00af00\">2,506,172,416</span> (9.34 GB)\n</pre>\n"},"metadata":{}},{"name":"stdout","text":"Epoch 1/5\n","output_type":"stream"},{"name":"stderr","text":"/usr/local/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py:756: UserWarning: Some donated buffers were not usable: ShapedArray(float32[256000,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]).\nSee an explanation at https://jax.readthedocs.io/en/latest/faq.html#buffer-donation.\n  warnings.warn(\"Some donated buffers were not usable:\"\n","output_type":"stream"},{"name":"stdout","text":"\u001b[1m143/144\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m━\u001b[0m \u001b[1m0s\u001b[0m 638ms/step - loss: 2.9938 - sparse_categorical_accuracy: 0.4366","output_type":"stream"},{"name":"stderr","text":"/usr/local/lib/python3.10/site-packages/jax/_src/interpreters/mlir.py:756: UserWarning: Some donated buffers were not usable: ShapedArray(float32[256000,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[1,256,256]), ShapedArray(float32[8,256,256]), ShapedArray(float32[256,16384]), ShapedArray(float32[256,16384]), ShapedArray(float32[16384,256]).\nSee an explanation at https://jax.readthedocs.io/en/latest/faq.html#buffer-donation.\n  warnings.warn(\"Some donated buffers were not usable:\"\n","output_type":"stream"},{"name":"stdout","text":"\u001b[1m144/144\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m138s\u001b[0m 800ms/step - loss: 2.9890 - sparse_categorical_accuracy: 0.4374\nEpoch 2/5\n\u001b[1m144/144\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m91s\u001b[0m 629ms/step - loss: 1.9191 - sparse_categorical_accuracy: 0.6265\nEpoch 3/5\n\u001b[1m144/144\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m91s\u001b[0m 630ms/step - loss: 1.8129 - sparse_categorical_accuracy: 0.6410\nEpoch 4/5\n\u001b[1m144/144\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m91s\u001b[0m 629ms/step - loss: 1.7858 - sparse_categorical_accuracy: 0.6431\nEpoch 5/5\n\u001b[1m144/144\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m91s\u001b[0m 629ms/step - loss: 1.7715 - sparse_categorical_accuracy: 0.6446\n","output_type":"stream"},{"execution_count":15,"output_type":"execute_result","data":{"text/plain":"<keras.src.callbacks.history.History at 0x78d960646b90>"},"metadata":{}}]},{"cell_type":"markdown","source":"Note that enabling LoRA reduces the number of trainable parameters significantly, from 7 billion to only 11 million.\n\nIn total it took <10 mininutes.","metadata":{"id":"CnpeavB4fZ7Y","papermill":{"duration":0.137841,"end_time":"2024-02-21T10:05:59.925998","exception":false,"start_time":"2024-02-21T10:05:59.788157","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"## Inference after finetuning","metadata":{"id":"lBiOKlAy2MAe","papermill":{"duration":0.137064,"end_time":"2024-02-21T10:06:00.201917","exception":false,"start_time":"2024-02-21T10:06:00.064853","status":"completed"},"tags":[]}},{"cell_type":"code","source":"gemma_lm.generate(test_prompt, max_length=200)","metadata":{"id":"9yNyJ8CLXfw0","outputId":"f0a1c1e9-2221-4e3a-83df-829a47488b1c","papermill":{"duration":31.150786,"end_time":"2024-02-21T10:06:31.488004","exception":false,"start_time":"2024-02-21T10:06:00.337218","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-02-25T06:29:05.525139Z","iopub.execute_input":"2024-02-25T06:29:05.525413Z","iopub.status.idle":"2024-02-25T06:29:23.400801Z","shell.execute_reply.started":"2024-02-25T06:29:05.525387Z","shell.execute_reply":"2024-02-25T06:29:23.399898Z"},"trusted":true},"execution_count":16,"outputs":[{"execution_count":16,"output_type":"execute_result","data":{"text/plain":"'你是一个知识丰富的人工智能助手，用户将用中文向你提问，你将根据你的知识用中文来如实回答问题。\\n问题：\\n我有一个信息科学相关的问题，请用中文回答，什么是 zsh\\n答案：\\nzsh（Z Shell）是一种UNIX/Linux操作系统中的一个内部命令。 \\nzsh是一种高效率的交互式终端用户命令语言。它在命令行中提供了一个类似于Unix Bourne或Shell（BASH）的shell。 \\nzsh的优点是：\\n支持命令行历史记录（command history）'"},"metadata":{}}]},{"cell_type":"markdown","source":"Now it gives much better answer in Chinese than the pretrained variant.","metadata":{"id":"inqB1e_v0xP5","papermill":{"duration":0.136614,"end_time":"2024-02-21T10:06:31.761087","exception":false,"start_time":"2024-02-21T10:06:31.624473","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"## Convert to Hugging Face\n\nMany ppl prefer to use Hugging Face than Keras for whatever reason. It's easy to make the conversion.","metadata":{}},{"cell_type":"code","source":"# Finetuned model\nFINETUNED_MODEL_DIR = f\"./finetuned_gemma\"\nFINETUNED_WEIGHTS_PATH = f\"{FINETUNED_MODEL_DIR}/model.weights.h5\"\nFINETUNED_VOCAB_PATH = f\"{FINETUNED_MODEL_DIR}/vocabulary.spm\"\n\n# Converted model\nHUGGINGFACE_MODEL_DIR = f\"./gemma_huggingface\"\n\nMODEL_NAME = \"gemma_2b_en\"\n\n# Deduce model size from name format: \"gemma[_instruct]_{2b,7b}_en\"\nMODEL_SIZE = MODEL_NAME.split(\"_\")[-2]","metadata":{"execution":{"iopub.status.busy":"2024-02-25T06:29:23.402121Z","iopub.execute_input":"2024-02-25T06:29:23.402476Z","iopub.status.idle":"2024-02-25T06:29:23.407294Z","shell.execute_reply.started":"2024-02-25T06:29:23.402441Z","shell.execute_reply":"2024-02-25T06:29:23.406550Z"},"trusted":true},"execution_count":17,"outputs":[]},{"cell_type":"code","source":"# Make sure the directory exists\n%mkdir -p $FINETUNED_MODEL_DIR\n\ngemma_lm.save_weights(FINETUNED_WEIGHTS_PATH)\n\ngemma_lm.preprocessor.tokenizer.save_assets(FINETUNED_MODEL_DIR)","metadata":{"execution":{"iopub.status.busy":"2024-02-25T06:29:23.408213Z","iopub.execute_input":"2024-02-25T06:29:23.408476Z","iopub.status.idle":"2024-02-25T06:29:40.504590Z","shell.execute_reply.started":"2024-02-25T06:29:23.408449Z","shell.execute_reply":"2024-02-25T06:29:40.503242Z"},"trusted":true},"execution_count":18,"outputs":[]},{"cell_type":"code","source":"!du -shc $FINETUNED_MODEL_DIR/*","metadata":{"execution":{"iopub.status.busy":"2024-02-25T06:29:40.505917Z","iopub.execute_input":"2024-02-25T06:29:40.506190Z","iopub.status.idle":"2024-02-25T06:29:41.167062Z","shell.execute_reply.started":"2024-02-25T06:29:40.506162Z","shell.execute_reply":"2024-02-25T06:29:41.165605Z"},"trusted":true},"execution_count":19,"outputs":[{"name":"stdout","text":"9.4G\t./finetuned_gemma/model.weights.h5\n4.1M\t./finetuned_gemma/vocabulary.spm\n9.4G\ttotal\n","output_type":"stream"}]},{"cell_type":"code","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","metadata":{"execution":{"iopub.status.busy":"2024-02-25T06:29:41.168641Z","iopub.execute_input":"2024-02-25T06:29:41.168960Z","iopub.status.idle":"2024-02-25T06:34:15.362493Z","shell.execute_reply.started":"2024-02-25T06:29:41.168930Z","shell.execute_reply":"2024-02-25T06:34:15.361128Z"},"trusted":true},"execution_count":20,"outputs":[{"name":"stdout","text":"2024-02-25 06:31:03 URL:https://raw.githubusercontent.com/keras-team/keras-nlp/master/tools/gemma/export_gemma_to_hf.py [11761/11761] -> \"export_gemma_to_hf.py\" [1]\n/usr/local/lib/python3.10/site-packages/jax/_src/cloud_tpu_init.py:75: UserWarning: JAX_USE_PJRT_C_API_ON_TPU no longer has an effect (the new TPU runtime is always enabled now). Unset the environment variable to disable this warning.\n  warnings.warn(\n\n-> Loading Keras weights from file `./finetuned_gemma/model.weights.h5`...\nAttaching 'config.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\nAttaching 'config.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\nAttaching 'model.weights.h5' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\nAttaching 'tokenizer.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\nAttaching 'assets/tokenizer/vocabulary.spm' from model 'keras/gemma/keras/gemma_2b_en/2' to your Kaggle notebook...\nnormalizer.cc(51) LOG(INFO) precompiled_charsmap is empty. use identity normalization.\n/usr/local/lib/python3.10/site-packages/keras/src/saving/saving_lib.py:396: UserWarning: Skipping variable loading for optimizer 'adam', because it has 2 variables whereas the saved optimizer has 146 variables. \n  trackable.load_own_variables(weights_store.get(inner_path))\n\n-> Loading HuggingFace Gemma `2B` model...\n\n✅ Model loading complete.\n\n-> Converting weights from KerasNLP Gemma to HuggingFace Gemma...\n\n✅ Weights converted successfully.\n\n-> Saving HuggingFace model to `./gemma_huggingface`...\n\n✅ Saving complete. Model saved at `./gemma_huggingface`.\n\n-> Saving HuggingFace Gemma tokenizer to `./gemma_huggingface`...\n\n✅ Saving complete. Tokenizer saved at `./gemma_huggingface`.\n","output_type":"stream"}]},{"cell_type":"code","source":"import transformers\nmodel = transformers.GemmaForCausalLM.from_pretrained(\n    HUGGINGFACE_MODEL_DIR,\n    local_files_only=True,\n    device_map=\"auto\",  # Library \"accelerate\" to auto-select GPU\n)\ntokenizer = transformers.GemmaTokenizer.from_pretrained(\n    HUGGINGFACE_MODEL_DIR,\n    local_files_only=True,\n)","metadata":{"execution":{"iopub.status.busy":"2024-02-25T06:34:15.364068Z","iopub.execute_input":"2024-02-25T06:34:15.364391Z","iopub.status.idle":"2024-02-25T06:34:20.462510Z","shell.execute_reply.started":"2024-02-25T06:34:15.364361Z","shell.execute_reply":"2024-02-25T06:34:20.461326Z"},"trusted":true},"execution_count":21,"outputs":[{"name":"stderr","text":"/usr/local/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n  from .autonotebook import tqdm as notebook_tqdm\nLoading checkpoint shards: 100%|██████████| 3/3 [00:01<00:00,  2.69it/s]\n","output_type":"stream"}]},{"cell_type":"code","source":"def test_transformers_model(\n    model: transformers.GemmaForCausalLM,\n    tokenizer: transformers.GemmaTokenizer,\n) -> None:   \n    inputs = tokenizer([test_prompt], return_tensors=\"pt\").to(model.device)\n    outputs = model.generate(**inputs, max_length=200)\n\n    output = tokenizer.decode(outputs[0], skip_special_tokens=True)\n    print(f\"{output}\\n{'- '*40}\")\n\n# This run on CPU so it is a bit slow\ntest_transformers_model(model, tokenizer)","metadata":{"execution":{"iopub.status.busy":"2024-02-25T06:34:20.463591Z","iopub.execute_input":"2024-02-25T06:34:20.464099Z","iopub.status.idle":"2024-02-25T06:34:59.576527Z","shell.execute_reply.started":"2024-02-25T06:34:20.464069Z","shell.execute_reply":"2024-02-25T06:34:59.575424Z"},"trusted":true},"execution_count":22,"outputs":[{"name":"stdout","text":"你是一个知识丰富的人工智能助手，用户将用中文向你提问，你将根据你的知识用中文来如实回答问题。\n问题：\n我有一个信息科学相关的问题，请用中文回答，什么是 zsh\n答案：\nzsh（Z Shell）是一个POSIX兼容的shell，它在BSD/OS和Linux系统上被广泛使用。 \nzsh是Z shell的缩写，Z shell是Unix shell的一种，它继承了Bourne shell的特性，并增加了许多新的特性。 \nzsh的特性包括： \n1.支持多级目录 \n2.支持命令别名 \n3.支持命令补全 \n4.支持命令历史 \n5.支持命令行参数\n- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - \n","output_type":"stream"}]},{"cell_type":"markdown","source":"Tis is very much similar to the KerasNLP output we had before, so I think our HF conversion worked.","metadata":{}},{"cell_type":"markdown","source":"# Final note\n\n* Here we used Gemma 2B. Technically you can use Gemma 7B, but sadly Kaggle only offers 20G of hard drive disk space, so you can't easily store the converted HF file.\n* TPU v3 is so much faster than the free T4 GPU on Google Colab. ","metadata":{}}]}