part_llama3_1_finetuning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "cellView": "form",
    "id": "7d9bbf86da5e"
   },
   "outputs": [],
   "source": [
    "# Copyright 2024 Google LLC\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "99c1c3fc2ca5"
   },
   "source": [
    "# Vertex AI Model Garden - Llama 3.1 Finetuning\n",
    "\n",
    "<table><tbody><tr>\n",
    "  <td style=\"text-align: center\">\n",
    "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_llama3_1_finetuning.ipynb\">\n",
    "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
    "    </a>\n",
    "  </td>\n",
    "  <td style=\"text-align: center\">\n",
    "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_llama3_1_finetuning.ipynb\">\n",
    "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
    "    </a>\n",
    "  </td>\n",
    "</tr></tbody></table>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3de7470326a2"
   },
   "source": [
    "## Overview\n",
    "\n",
    "This notebook demonstrates finetuning and deploying Llama 3.1 models with Vertex AI. All of the examples in this notebook use parameter efficient finetuning methods [PEFT (LoRA)](https://github.com/huggingface/peft) to reduce training and storage costs. LoRA (Low-Rank Adaptation) is one approach of Parameter Efficient FineTuning (PEFT), where pretrained model weights are frozen and rank decomposition matrices representing the change in model weights are trained during finetuning.\n",
    "After finetuning, we can deploy models on Vertex with GPU.\n",
    "\n",
    "\n",
    "### Objective\n",
    "\n",
    "- Finetune Llama 3.1 models with Vertex AI Custom Training Jobs.\n",
    "- Deploy finetuned Llama 3.1 models on Vertex AI Prediction.\n",
    "- Send prediction requests to your finetuned Llama 3.1 models.\n",
    "\n",
    "\n",
    "### Costs\n",
    "\n",
    "This tutorial uses billable components of Google Cloud:\n",
    "\n",
    "* Vertex AI\n",
    "* Cloud Storage\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "264c07757582"
   },
   "source": [
    "## Before you begin"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Installing google-cloud-aiplatform\n",
      "Installing tensorflow\n",
      "Requirement already satisfied: transformers in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (4.44.2)\n",
      "Requirement already satisfied: filelock in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from transformers) (3.15.4)\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.23.2 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from transformers) (0.25.1)\n",
      "Requirement already satisfied: numpy>=1.17 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from transformers) (1.26.4)\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from transformers) (24.1)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from transformers) (6.0.1)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from transformers) (2024.7.24)\n",
      "Requirement already satisfied: requests in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from transformers) (2.32.3)\n",
      "Requirement already satisfied: safetensors>=0.4.1 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from transformers) (0.4.5)\n",
      "Requirement already satisfied: tokenizers<0.20,>=0.19 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from transformers) (0.19.1)\n",
      "Requirement already satisfied: tqdm>=4.27 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from transformers) (4.66.4)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from huggingface-hub<1.0,>=0.23.2->transformers) (2024.6.1)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from huggingface-hub<1.0,>=0.23.2->transformers) (4.12.2)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from requests->transformers) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from requests->transformers) (3.7)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from requests->transformers) (1.26.19)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from requests->transformers) (2024.7.4)\n",
      "Requirement already satisfied: vllm in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (0.6.1.post2)\n",
      "Requirement already satisfied: psutil in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (5.9.3)\n",
      "Requirement already satisfied: sentencepiece in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (0.2.0)\n",
      "Requirement already satisfied: numpy<2.0.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (1.26.4)\n",
      "Requirement already satisfied: requests in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (2.32.3)\n",
      "Requirement already satisfied: tqdm in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (4.66.4)\n",
      "Requirement already satisfied: py-cpuinfo in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (9.0.0)\n",
      "Requirement already satisfied: transformers>=4.43.2 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (4.44.2)\n",
      "Requirement already satisfied: tokenizers>=0.19.1 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (0.19.1)\n",
      "Requirement already satisfied: protobuf in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (3.19.6)\n",
      "Requirement already satisfied: aiohttp in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (3.10.0)\n",
      "Requirement already satisfied: openai>=1.40.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (1.47.0)\n",
      "Requirement already satisfied: uvicorn[standard] in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (0.30.5)\n",
      "Requirement already satisfied: pydantic>=2.9 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (2.9.2)\n",
      "Requirement already satisfied: pillow in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (10.4.0)\n",
      "Requirement already satisfied: prometheus-client>=0.18.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (0.20.0)\n",
      "Requirement already satisfied: prometheus-fastapi-instrumentator>=7.0.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (7.0.0)\n",
      "Requirement already satisfied: tiktoken>=0.6.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (0.7.0)\n",
      "Requirement already satisfied: lm-format-enforcer==0.10.6 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (0.10.6)\n",
      "Requirement already satisfied: outlines<0.1,>=0.0.43 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (0.0.46)\n",
      "Requirement already satisfied: typing-extensions>=4.10 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (4.12.2)\n",
      "Requirement already satisfied: filelock>=3.10.4 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (3.15.4)\n",
      "Requirement already satisfied: partial-json-parser in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (0.2.1.1.post4)\n",
      "Requirement already satisfied: pyzmq in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (26.0.3)\n",
      "Requirement already satisfied: msgspec in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (0.18.6)\n",
      "Requirement already satisfied: gguf==0.9.1 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (0.9.1)\n",
      "Requirement already satisfied: importlib-metadata in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (8.0.0)\n",
      "Requirement already satisfied: mistral-common>=1.4.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (1.4.3)\n",
      "Requirement already satisfied: pyyaml in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (6.0.1)\n",
      "Requirement already satisfied: einops in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (0.8.0)\n",
      "Requirement already satisfied: ray>=2.9 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (2.34.0)\n",
      "Requirement already satisfied: nvidia-ml-py in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (11.495.46)\n",
      "Requirement already satisfied: torch==2.4.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (2.4.0)\n",
      "Requirement already satisfied: torchvision==0.19 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (0.19.0)\n",
      "Requirement already satisfied: xformers==0.0.27.post2 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (0.0.27.post2)\n",
      "Requirement already satisfied: vllm-flash-attn==2.6.1 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (2.6.1)\n",
      "Requirement already satisfied: fastapi>=0.114.1 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from vllm) (0.115.0)\n",
      "Requirement already satisfied: interegular>=0.3.2 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from lm-format-enforcer==0.10.6->vllm) (0.3.3)\n",
      "Requirement already satisfied: packaging in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from lm-format-enforcer==0.10.6->vllm) (24.1)\n",
      "Requirement already satisfied: sympy in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from torch==2.4.0->vllm) (1.13.3)\n",
      "Requirement already satisfied: networkx in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from torch==2.4.0->vllm) (3.3)\n",
      "Requirement already satisfied: jinja2 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from torch==2.4.0->vllm) (3.1.4)\n",
      "Requirement already satisfied: fsspec in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from torch==2.4.0->vllm) (2024.6.1)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from torch==2.4.0->vllm) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from torch==2.4.0->vllm) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from torch==2.4.0->vllm) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cudnn-cu12==9.1.0.70 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from torch==2.4.0->vllm) (9.1.0.70)\n",
      "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from torch==2.4.0->vllm) (12.1.3.1)\n",
      "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from torch==2.4.0->vllm) (11.0.2.54)\n",
      "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from torch==2.4.0->vllm) (10.3.2.106)\n",
      "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from torch==2.4.0->vllm) (11.4.5.107)\n",
      "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from torch==2.4.0->vllm) (12.1.0.106)\n",
      "Requirement already satisfied: nvidia-nccl-cu12==2.20.5 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from torch==2.4.0->vllm) (2.20.5)\n",
      "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from torch==2.4.0->vllm) (12.1.105)\n",
      "Requirement already satisfied: triton==3.0.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from torch==2.4.0->vllm) (3.0.0)\n",
      "Requirement already satisfied: nvidia-nvjitlink-cu12 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from nvidia-cusolver-cu12==11.4.5.107->torch==2.4.0->vllm) (12.6.68)\n",
      "Requirement already satisfied: starlette<0.39.0,>=0.37.2 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from fastapi>=0.114.1->vllm) (0.37.2)\n",
      "Requirement already satisfied: jsonschema<5.0.0,>=4.21.1 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from mistral-common>=1.4.0->vllm) (4.23.0)\n",
      "Requirement already satisfied: anyio<5,>=3.5.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from openai>=1.40.0->vllm) (3.7.1)\n",
      "Requirement already satisfied: distro<2,>=1.7.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from openai>=1.40.0->vllm) (1.9.0)\n",
      "Requirement already satisfied: httpx<1,>=0.23.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from openai>=1.40.0->vllm) (0.27.0)\n",
      "Requirement already satisfied: jiter<1,>=0.4.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from openai>=1.40.0->vllm) (0.5.0)\n",
      "Requirement already satisfied: sniffio in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from openai>=1.40.0->vllm) (1.3.1)\n",
      "Requirement already satisfied: lark in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from outlines<0.1,>=0.0.43->vllm) (1.2.2)\n",
      "Requirement already satisfied: nest-asyncio in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from outlines<0.1,>=0.0.43->vllm) (1.6.0)\n",
      "Requirement already satisfied: cloudpickle in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from outlines<0.1,>=0.0.43->vllm) (2.2.1)\n",
      "Requirement already satisfied: diskcache in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from outlines<0.1,>=0.0.43->vllm) (5.6.3)\n",
      "Requirement already satisfied: numba in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from outlines<0.1,>=0.0.43->vllm) (0.60.0)\n",
      "Requirement already satisfied: referencing in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from outlines<0.1,>=0.0.43->vllm) (0.35.1)\n",
      "Requirement already satisfied: datasets in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from outlines<0.1,>=0.0.43->vllm) (3.0.0)\n",
      "Requirement already satisfied: pycountry in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from outlines<0.1,>=0.0.43->vllm) (24.6.1)\n",
      "Requirement already satisfied: pyairports in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from outlines<0.1,>=0.0.43->vllm) (2.1.1)\n",
      "Requirement already satisfied: annotated-types>=0.6.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from pydantic>=2.9->vllm) (0.7.0)\n",
      "Requirement already satisfied: pydantic-core==2.23.4 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from pydantic>=2.9->vllm) (2.23.4)\n",
      "Requirement already satisfied: click>=7.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from ray>=2.9->vllm) (8.1.7)\n",
      "Requirement already satisfied: msgpack<2.0.0,>=1.0.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from ray>=2.9->vllm) (1.0.8)\n",
      "Requirement already satisfied: aiosignal in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from ray>=2.9->vllm) (1.3.1)\n",
      "Requirement already satisfied: frozenlist in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from ray>=2.9->vllm) (1.4.1)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from requests->vllm) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from requests->vllm) (3.7)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from requests->vllm) (1.26.19)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from requests->vllm) (2024.7.4)\n",
      "Requirement already satisfied: regex>=2022.1.18 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from tiktoken>=0.6.0->vllm) (2024.7.24)\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.16.4 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from tokenizers>=0.19.1->vllm) (0.25.1)\n",
      "Requirement already satisfied: safetensors>=0.4.1 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from transformers>=4.43.2->vllm) (0.4.5)\n",
      "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from aiohttp->vllm) (2.3.4)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from aiohttp->vllm) (23.2.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from aiohttp->vllm) (6.0.5)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from aiohttp->vllm) (1.9.4)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from aiohttp->vllm) (4.0.3)\n",
      "Requirement already satisfied: zipp>=0.5 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from importlib-metadata->vllm) (3.19.2)\n",
      "Requirement already satisfied: h11>=0.8 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from uvicorn[standard]->vllm) (0.14.0)\n",
      "Requirement already satisfied: httptools>=0.5.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from uvicorn[standard]->vllm) (0.6.1)\n",
      "Requirement already satisfied: python-dotenv>=0.13 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from uvicorn[standard]->vllm) (1.0.1)\n",
      "Requirement already satisfied: uvloop!=0.15.0,!=0.15.1,>=0.14.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from uvicorn[standard]->vllm) (0.19.0)\n",
      "Requirement already satisfied: watchfiles>=0.13 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from uvicorn[standard]->vllm) (0.22.0)\n",
      "Requirement already satisfied: websockets>=10.4 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from uvicorn[standard]->vllm) (12.0)\n",
      "Requirement already satisfied: exceptiongroup in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from anyio<5,>=3.5.0->openai>=1.40.0->vllm) (1.2.2)\n",
      "Requirement already satisfied: httpcore==1.* in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from httpx<1,>=0.23.0->openai>=1.40.0->vllm) (1.0.5)\n",
      "Requirement already satisfied: jsonschema-specifications>=2023.03.6 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from jsonschema<5.0.0,>=4.21.1->mistral-common>=1.4.0->vllm) (2023.12.1)\n",
      "Requirement already satisfied: rpds-py>=0.7.1 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from jsonschema<5.0.0,>=4.21.1->mistral-common>=1.4.0->vllm) (0.19.1)\n",
      "Requirement already satisfied: pyarrow>=15.0.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from datasets->outlines<0.1,>=0.0.43->vllm) (17.0.0)\n",
      "Requirement already satisfied: dill<0.3.9,>=0.3.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from datasets->outlines<0.1,>=0.0.43->vllm) (0.3.8)\n",
      "Requirement already satisfied: pandas in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from datasets->outlines<0.1,>=0.0.43->vllm) (2.2.2)\n",
      "Requirement already satisfied: xxhash in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from datasets->outlines<0.1,>=0.0.43->vllm) (3.5.0)\n",
      "Requirement already satisfied: multiprocess in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from datasets->outlines<0.1,>=0.0.43->vllm) (0.70.16)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from jinja2->torch==2.4.0->vllm) (2.0.1)\n",
      "Requirement already satisfied: llvmlite<0.44,>=0.43.0dev0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from numba->outlines<0.1,>=0.0.43->vllm) (0.43.0)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from sympy->torch==2.4.0->vllm) (1.3.0)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from pandas->datasets->outlines<0.1,>=0.0.43->vllm) (2.9.0)\n",
      "Requirement already satisfied: pytz>=2020.1 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from pandas->datasets->outlines<0.1,>=0.0.43->vllm) (2024.1)\n",
      "Requirement already satisfied: tzdata>=2022.7 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from pandas->datasets->outlines<0.1,>=0.0.43->vllm) (2024.1)\n",
      "Requirement already satisfied: six>=1.5 in /opt/conda/envs/tensorflow/lib/python3.10/site-packages (from python-dateutil>=2.8.2->pandas->datasets->outlines<0.1,>=0.0.43->vllm) (1.16.0)\n",
      "Installing dependencies complete\n"
     ]
    }
   ],
   "source": [
    "print(\"Installing google-cloud-aiplatform\")\n",
    "! pip3 install -q google-cloud-aiplatform\n",
    "print(\"Installing tensorflow\")\n",
    "!pip3 install -q tensorflow\n",
    "!pip3 install -q transformers\n",
    "!pip3 install --upgrade transformers\n",
    "!pip3 install --upgrade vllm\n",
    "\n",
    "print(\"Installing dependencies complete\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup Google Cloud project\n",
    "\n",
    "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
    "\n",
    "2. [Make sure that you have GPU quota for Vertex Training (finetuing) and Vertex Prediction (serving)](https://cloud.google.com/docs/quotas/view-manage). The quota name for Vertex Training is \"Custom model training your-gpu-type per region\" and the quota name for Vertex Prediction is \"Custom model serving your-gpu-type per region\" such as `Custom model training Nvidia L4 GPUs per region` and `Custom model serving Nvidia L4 GPUs per region` for L4 GPUs. [Submit a quota increase request](https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota) if additional quota is needed. At minimum, running this notebook requires 4 L4s for finetuning and 1 L4 for serving. More GPUs may be needed for larger models and different finetuning configurations. To secure GPUs for larger models, ask your customer engineer to get you allowlisted for a Shared Reservation or a Dynamic Workload Scheduler.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import the necessary packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fatal: destination path 'vertex-ai-samples' already exists and is not an empty directory.\n"
     ]
    }
   ],
   "source": [
    "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define environment variables\n",
    "\n",
    "Set the following variables for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the specified region (`REGION`). Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\")."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "cellView": "form",
    "id": "855d6b96f291",
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-09-24 06:12:27.748793: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2024-09-24 06:12:27.885403: I tensorflow/core/util/port.cc:104] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2024-09-24 06:12:28.731018: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/cuda/lib64:/usr/local/nccl2/lib:/usr/local/cuda/extras/CUPTI/lib64:/usr/lib/x86_64-linux-gnu/:/opt/conda/lib\n",
      "2024-09-24 06:12:28.731173: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/cuda/lib64:/usr/local/nccl2/lib:/usr/local/cuda/extras/CUPTI/lib64:/usr/lib/x86_64-linux-gnu/:/opt/conda/lib\n",
      "2024-09-24 06:12:28.731185: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "region is not ASIA Southeast\n",
      "Enabling Vertex AI API and Compute Engine API.\n",
      "Operation \"operations/acat.p2-390782612502-a4773bae-3116-45d5-9f60-a712c1f71534\" finished successfully.\n",
      "Using this GCS Bucket: gs://ai-sandbox-company-71-llama31-us-central1\n",
      "Updated property [core/project].\n",
      "Creating gs://ai-sandbox-company-71-llama31-us-central1/...\n",
      "\u001b[1;31mERROR:\u001b[0m (gcloud.storage.buckets.create) HTTPError 409: Your previous request to create the named bucket succeeded and you already own it.\n"
     ]
    }
   ],
   "source": [
    "import importlib\n",
    "import os\n",
    "import uuid\n",
    "from datetime import datetime\n",
    "from typing import Tuple\n",
    "\n",
    "from google.cloud import aiplatform\n",
    "import tensorflow\n",
    "\n",
    "import socket\n",
    "import re\n",
    "\n",
    "common_util = importlib.import_module(\n",
    "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
    ")\n",
    "models, endpoints = {}, {}\n",
    "\n",
    "UNIQUE_PREFIX = socket.gethostname()\n",
    "UNIQUE_PREFIX = re.sub('[^A-Za-z0-9]+', '', UNIQUE_PREFIX)\n",
    "\n",
    "# Cloud project id.\n",
    "PROJECT_IDS = !(gcloud config get-value core/project)\n",
    "PROJECT_ID = PROJECT_IDS[0]  # @param {type:\"string\"}\n",
    "\n",
    "# The region you want to launch jobs in.\n",
    "PREFIX_NUM_ONLY=int(str(re.search(r'\\d+', UNIQUE_PREFIX).group()))\n",
    "REGION_ALLOCATE=PREFIX_NUM_ONLY%3\n",
    "if REGION_ALLOCATE == 0:\n",
    "    REGION = \"asia-southeast1\"\n",
    "elif REGION_ALLOCATE == 1:\n",
    "    REGION = \"us-central1\"\n",
    "    print(\"region is not ASIA Southeast\")\n",
    "else:\n",
    "    REGION = \"europe-west4\"\n",
    "\n",
    "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
    "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
    "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
    "\n",
    "# The service account looks like:\n",
    "# '@.iam.gserviceaccount.com'\n",
    "# Please go to https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console\n",
    "# and create service account with `Vertex AI User` and `Storage Object Admin` roles.\n",
    "# The service account for deploying fine tuned model.\n",
    "SERVICE_ACCOUNT = !(gcloud config get-value core/account)  # @param {type:\"string\"}\n",
    "SERVICE_ACCOUNT = SERVICE_ACCOUNT[0]  # @param {type:\"string\"}\n",
    "\n",
    "# The Cloud Storage bucket for storing experiments output.\n",
    "# Remove prefix gs://, e.g. foo_bucket.\n",
    "BUCKET_NAME = f\"{PROJECT_ID}-{UNIQUE_PREFIX}-{REGION}\"\n",
    "BUCKET_URI = f\"gs://{BUCKET_NAME}\"  # @param {type:\"string\"}\n",
    "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
    "\n",
    "! gcloud config set project $PROJECT_ID\n",
    "! gcloud storage buckets create {BUCKET_URI} --project={PROJECT_ID} --location={REGION}\n",
    "\n",
    "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
    "MODEL_BUCKET = os.path.join(BUCKET_URI, \"llama3_1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize Vertex AI API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initializing Vertex AI API.\n"
     ]
    }
   ],
   "source": [
    "print(\"Initializing Vertex AI API.\")\n",
    "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The pre-built serving docker image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240721_0916_RC00\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Access Llama 3.1 models on Vertex AI for GPU based serving\n",
    "The original models from Meta are converted into the Hugging Face format for serving in Vertex AI.\n",
    "Accept the model agreement to access the models:\n",
    "1. Open the [Llama 3.1 model card](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama3_1) from [Vertex AI Model Garden](https://cloud.google.com/model-garden).\n",
    "2. Review and accept the agreement in the pop-up window on the model card page. If you have previously accepted the model agreement, there will not be a pop-up window on the model card page and this step is not needed.\n",
    "3. After accepting the agreement of Llama 3.1, a `gs://` URI containing Llama 3.1 pretrained and finetuned models will be shared.\n",
    "4. Paste the URI in the `VERTEX_AI_MODEL_GARDEN_LLAMA3_1` field below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Copying gs://vertex-model-garden-public-us/llama3.1/Meta-Llama-3.1-8B/LICENSE [Content-Type=application/octet-stream]...\n",
      "Copying gs://vertex-model-garden-public-us/llama3.1/Meta-Llama-3.1-8B/README.md [Content-Type=text/markdown]...\n",
      "Copying gs://vertex-model-garden-public-us/llama3.1/Meta-Llama-3.1-8B/USE_POLICY.md [Content-Type=text/markdown]...\n",
      "Copying gs://vertex-model-garden-public-us/llama3.1/Meta-Llama-3.1-8B/config.json [Content-Type=application/json]...\n",
      "Copying gs://vertex-model-garden-public-us/llama3.1/Meta-Llama-3.1-8B/generation_config.json [Content-Type=application/json]...\n",
      "Copying gs://vertex-model-garden-public-us/llama3.1/Meta-Llama-3.1-8B/tokenizer_config.json [Content-Type=application/json]...\n",
      "Copying gs://vertex-model-garden-public-us/llama3.1/Meta-Llama-3.1-8B/model-00002-of-00004.safetensors [Content-Type=application/octet-stream]...\n",
      "Copying gs://vertex-model-garden-public-us/llama3.1/Meta-Llama-3.1-8B/model-00003-of-00004.safetensors [Content-Type=application/octet-stream]...\n",
      "Copying gs://vertex-model-garden-public-us/llama3.1/Meta-Llama-3.1-8B/special_tokens_map.json [Content-Type=application/json]...\n",
      "Copying gs://vertex-model-garden-public-us/llama3.1/Meta-Llama-3.1-8B/model-00001-of-00004.safetensors [Content-Type=application/octet-stream]...\n",
      "Copying gs://vertex-model-garden-public-us/llama3.1/Meta-Llama-3.1-8B/model.safetensors.index.json [Content-Type=application/json]...\n",
      "Copying gs://vertex-model-garden-public-us/llama3.1/Meta-Llama-3.1-8B/model-00004-of-00004.safetensors [Content-Type=application/octet-stream]...\n",
      "Copying gs://vertex-model-garden-public-us/llama3.1/Meta-Llama-3.1-8B/tokenizer.json [Content-Type=application/json]...\n",
      "/ [13/13 files][ 15.0 GiB/ 15.0 GiB] 100% Done 307.7 MiB/s ETA 00:00:00         \n",
      "Operation completed over 13 objects/15.0 GiB.                                    \n"
     ]
    }
   ],
   "source": [
    "# The Llama 3.1 base model.\n",
    "base_model_name = \"Meta-Llama-3.1-8B\"  # @param [\"meta-llama/Meta-Llama-3.1-8B\", \"meta-llama/Meta-Llama-3.1-8B-Instruct\", \"meta-llama/Meta-Llama-3.1-70B\", \"meta-llama/Meta-Llama-3.1-70B-Instruct\"] {isTemplate:true}\n",
    "! gsutil -m cp -R gs://vertex-model-garden-public-us/llama3.1/{base_model_name} {MODEL_BUCKET}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define common functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def deploy_model_vllm(\n",
    "    model_name: str,\n",
    "    model_id: str,\n",
    "    service_account: str,\n",
    "    machine_type: str = \"g2-standard-8\",\n",
    "    accelerator_type: str = \"NVIDIA_L4\",\n",
    "    accelerator_count: int = 1,\n",
    "    gpu_memory_utilization: float = 0.9,\n",
    "    max_model_len: int = 8192,\n",
    "    max_loras: int = 1,\n",
    "    max_cpu_loras: int = 16,\n",
    ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
    "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
    "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
    "\n",
    "    vllm_args = [\n",
    "        \"python\",\n",
    "        \"-m\",\n",
    "        \"vllm.entrypoints.api_server\",\n",
    "        \"--host=0.0.0.0\",\n",
    "        \"--port=7080\",\n",
    "        f\"--model={model_id}\",\n",
    "        f\"--tensor-parallel-size={accelerator_count}\",\n",
    "        \"--swap-space=16\",\n",
    "        f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
    "        f\"--max-model-len={max_model_len}\",\n",
    "        f\"--max-loras={max_loras}\",\n",
    "        f\"--max-cpu-loras={max_cpu_loras}\",\n",
    "        \"--disable-log-stats\",\n",
    "    ]\n",
    "\n",
    "    env_vars = {\n",
    "        \"MODEL_ID\": model_id,\n",
    "        \"DEPLOY_SOURCE\": \"notebook\",\n",
    "    }\n",
    "\n",
    "    model = aiplatform.Model.upload(\n",
    "        display_name=model_name,\n",
    "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
    "        serving_container_args=vllm_args,\n",
    "        serving_container_ports=[7080],\n",
    "        serving_container_predict_route=\"/generate\",\n",
    "        serving_container_health_route=\"/ping\",\n",
    "        serving_container_environment_variables=env_vars,\n",
    "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
    "        serving_container_deployment_timeout=7200,\n",
    "    )\n",
    "    print(\n",
    "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
    "    )\n",
    "    model.deploy(\n",
    "        endpoint=endpoint,\n",
    "        machine_type=machine_type,\n",
    "        accelerator_type=accelerator_type,\n",
    "        accelerator_count=accelerator_count,\n",
    "        deploy_request_timeout=7200,\n",
    "        service_account=service_account,\n",
    "    )\n",
    "    print(\"endpoint_name:\", endpoint.name)\n",
    "\n",
    "    return model, endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cb56d402e84a"
   },
   "source": [
    "## Finetune with HuggingFace PEFT and deploy with vLLM on GPUs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set dataset\n",
    "\n",
    "Use the Vertex AI SDK to create and run the custom training jobs.\n",
    "\n",
    "This notebook uses [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) dataset as an example.\n",
    "You can set `dataset_name` to any existing [Hugging Face dataset](https://huggingface.co/datasets) name, and set `instruct_column_in_dataset` to the name of the dataset column containing training data. The [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) has only one column `text`, and therefore we set `instruct_column_in_dataset` to `text` in this notebook.\n",
    "\n",
    "### (Optional) Prepare a custom JSONL dataset for finetuning\n",
    "\n",
    "You can prepare a JSONL file where each line is a valid JSON string as your custom training dataset. For example, here is one line from the [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) dataset:\n",
    "```\n",
    "{\"text\": \"### Human: Hola### Assistant: \\u00a1Hola! \\u00bfEn qu\\u00e9 puedo ayudarte hoy?\"}\n",
    "```\n",
    "\n",
    "The JSON object has a key `text`, which should match `instruct_column_in_dataset`; The value should be one training data point, i.e. a string. After you prepared your JSONL file, you can either upload it to [Hugging Face datasets](https://huggingface.co/datasets) or [Google Cloud Storage](https://cloud.google.com/storage).\n",
    "\n",
    "- To upload a JSONL dataset to [Hugging Face datasets](https://huggingface.co/datasets), follow the instructions on [Uploading Datasets](https://huggingface.co/docs/hub/en/datasets-adding). Then, set `dataset_name` to the name of your newly created dataset on Hugging Face.\n",
    "\n",
    "- To upload a JSONL dataset to [Google Cloud Storage](https://cloud.google.com/storage), follow the instructions on [Upload objects from a filesystem](https://cloud.google.com/storage/docs/uploading-objects). Then, set `dataset_name` to the `gs://` URI to your JSONL file. For example: `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`.\n",
    "\n",
    "Optionally update the `instruct_column_in_dataset` field below if your JSON objects use a key other than the default `text`.\n",
    "\n",
    "### (Optional) Format your data with custom JSON template\n",
    "\n",
    "Sometimes, your dataset might have multiple text columns and you want to construct the training data with a template. You can prepare a JSON template in the following format:\n",
    "\n",
    "```\n",
    "{\n",
    "  \"description\": \"Template used by Llama 3.1, accepting text-bison format.\",\n",
    "  \"source\": \"https://cloud.google.com/vertex-ai/generative-ai/docs/models/tune-text-models-supervised#dataset-format\",\n",
    "  \"prompt_input\": \"<|start_header_id|>user<|end_header_id|>\\n\\n{input_text}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\\n\\n{output_text}<|eot_id|>\",\n",
    "  \"instruction_separator\": \"<|start_header_id|>user<|end_header_id|>\\n\\n\",\n",
    "  \"response_separator\": \"<|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
    "}\n",
    "```\n",
    "\n",
    "As an example, the template above can be used to format the following training data (this line comes from `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`):\n",
    "\n",
    "```\n",
    "{\"input_text\":\"TRANSCRIPT: \\nREASON FOR EVALUATION:,\\n\\n LABEL:\",\"output_text\":\"Chiropractic\"}\n",
    "```\n",
    "\n",
    "This example template simply concatenates `input_text` with `output_text` with some special tokens in between.\n",
    "\n",
    "To try such custom dataset, you can make the following changes:\n",
    "1. Set `template` to `llama3-text-bison`\n",
    "1. Set `train_dataset_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`\n",
    "1. Set `train_split_name` to `train`\n",
    "1. Set `eval_dataset_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_eval_sample.jsonl`\n",
    "1. Set `eval_split_name` to `train` (**NOT** `test`)\n",
    "1. Set `instruct_column_in_dataset` as `input_text`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "cellView": "form",
    "id": "KwAW99YZHTdy"
   },
   "outputs": [],
   "source": [
    "# Template name or gs:// URI to a custom template.\n",
    "template = \"openassistant-guanaco\"  # @param {type:\"string\"}\n",
    "\n",
    "# Hugging Face dataset name or gs:// URI to a custom JSONL dataset.\n",
    "train_dataset_name = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
    "train_split_name = \"train\"  # @param {type:\"string\"}\n",
    "eval_dataset_name = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
    "eval_split_name = \"test\"  # @param {type:\"string\"}\n",
    "\n",
    "# Name of the dataset column containing training text input.\n",
    "instruct_column_in_dataset = \"text\"  # @param {type:\"string\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finetune\n",
    "Use the Vertex AI SDK to create and run the custom training jobs.\n",
    "\n",
    "**Note**:\n",
    "1. We recommend setting `finetuning_precision_mode` to `4bit` because it enables using fewer hardware resources for finetuning.\n",
    "1. We recommend using NVIDIA_L4 for 8B models and NVIDIA_A100_80GB for 70B models.\n",
    "1. If `max_steps>0`, it will precedence over `epochs`. One can set a small `max_steps` value to quickly check the pipeline.\n",
    "1. With the default setting, training takes between 1.5 ~ 2 hours."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "cellView": "form",
    "id": "ivVGS9dHXPOz"
   },
   "outputs": [],
   "source": [
    "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20240724_0936_RC00\"\n",
    "\n",
    "# The Llama 3.1 base model.\n",
    "MODEL_ID = os.path.join(MODEL_BUCKET, base_model_name)\n",
    "# The accelerator to use.\n",
    "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_A100_80GB\"]\n",
    "accelerator_count = 4\n",
    "\n",
    "# Batch size for finetuning.\n",
    "per_device_train_batch_size = 1  # @param{type:\"integer\"}\n",
    "gradient_accumulation_steps = 8  # @param{type:\"integer\"}\n",
    "# Maximum sequence length.\n",
    "max_seq_length = 4096  # @param{type:\"integer\"}\n",
    "# Setting a positive `max_steps` here will override `num_epochs`\n",
    "max_steps = -1  # @param{type:\"integer\"}\n",
    "num_epochs = 1.0  # @param{type:\"number\"}\n",
    "# Precision mode for finetuning.\n",
    "finetuning_precision_mode = \"4bit\"  # @param [\"4bit\", \"8bit\", \"float16\"]\n",
    "# Learning rate.\n",
    "learning_rate = 5e-5  # @param{type:\"number\"}\n",
    "lr_scheduler_type = \"cosine\"  # @param{type:\"string\"}\n",
    "# LoRA parameters.\n",
    "lora_rank = 16  # @param{type:\"integer\"}\n",
    "lora_alpha = 32  # @param{type:\"integer\"}\n",
    "lora_dropout = 0.05  # @param{type:\"number\"}\n",
    "enable_gradient_checkpointing = True\n",
    "attn_implementation = \"flash_attention_2\"\n",
    "optimizer = \"paged_adamw_32bit\"\n",
    "warmup_ratio = \"0.01\"\n",
    "report_to = \"tensorboard\"\n",
    "save_steps = 10\n",
    "logging_steps = save_steps\n",
    "\n",
    "replica_count = 1\n",
    "\n",
    "# common_util.check_quota(\n",
    "#     project_id=PROJECT_ID,\n",
    "#     region=REGION,\n",
    "#     accelerator_type=accelerator_type,\n",
    "#     accelerator_count=accelerator_count,\n",
    "#     is_for_training=True,\n",
    "# )\n",
    "\n",
    "job_name = common_util.get_job_name_with_datetime(\"llama3_1-lora-train\").replace(\n",
    "    \"_\", \"-\"\n",
    ")\n",
    "\n",
    "base_output_dir = os.path.join(STAGING_BUCKET, job_name)\n",
    "# Create a GCS folder to store the LORA adapter.\n",
    "lora_output_dir = os.path.join(base_output_dir, \"adapter\")\n",
    "# Create a GCS folder to store the merged model with the base model and the\n",
    "# finetuned LORA adapter.\n",
    "merged_model_output_dir = os.path.join(base_output_dir, \"merged-model\")\n",
    "\n",
    "eval_args = [\n",
    "    f\"--eval_dataset_path={eval_dataset_name}\",\n",
    "    f\"--eval_column={instruct_column_in_dataset}\",\n",
    "    f\"--eval_template={template}\",\n",
    "    f\"--eval_split={eval_split_name}\",\n",
    "    f\"--eval_steps={save_steps}\",\n",
    "    \"--eval_tasks=builtin_eval\",\n",
    "    \"--eval_metric_name=loss\",\n",
    "]\n",
    "\n",
    "train_job_args = [\n",
    "    \"--config_file=vertex_vision_model_garden_peft/deepspeed_zero2_4gpu.yaml\",\n",
    "    \"--task=instruct-lora\",\n",
    "    \"--completion_only=True\",\n",
    "    f\"--pretrained_model_id={MODEL_ID}\",\n",
    "    f\"--dataset_name={train_dataset_name}\",\n",
    "    f\"--train_split_name={train_split_name}\",\n",
    "    f\"--instruct_column_in_dataset={instruct_column_in_dataset}\",\n",
    "    f\"--output_dir={lora_output_dir}\",\n",
    "    f\"--merge_base_and_lora_output_dir={merged_model_output_dir}\",\n",
    "    f\"--per_device_train_batch_size={per_device_train_batch_size}\",\n",
    "    f\"--gradient_accumulation_steps={gradient_accumulation_steps}\",\n",
    "    f\"--lora_rank={lora_rank}\",\n",
    "    f\"--lora_alpha={lora_alpha}\",\n",
    "    f\"--lora_dropout={lora_dropout}\",\n",
    "    f\"--max_steps={max_steps}\",\n",
    "    f\"--max_seq_length={max_seq_length}\",\n",
    "    f\"--learning_rate={learning_rate}\",\n",
    "    f\"--lr_scheduler_type={lr_scheduler_type}\",\n",
    "    f\"--precision_mode={finetuning_precision_mode}\",\n",
    "    f\"--enable_gradient_checkpointing={enable_gradient_checkpointing}\",\n",
    "    f\"--num_epochs={num_epochs}\",\n",
    "    f\"--attn_implementation={attn_implementation}\",\n",
    "    f\"--optimizer={optimizer}\",\n",
    "    f\"--warmup_ratio={warmup_ratio}\",\n",
    "    f\"--report_to={report_to}\",\n",
    "    f\"--logging_output_dir={base_output_dir}\",\n",
    "    f\"--save_steps={save_steps}\",\n",
    "    f\"--logging_steps={logging_steps}\",\n",
    "    f\"--template={template}\",\n",
    "] + eval_args"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create TensorBoard"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating Tensorboard\n",
      "Create Tensorboard backing LRO: projects/390782612502/locations/us-central1/tensorboards/1602142373297520640/operations/3256761309852598272\n",
      "Tensorboard created. Resource name: projects/390782612502/locations/us-central1/tensorboards/1602142373297520640\n",
      "To use this Tensorboard in another session:\n",
      "tb = aiplatform.Tensorboard('projects/390782612502/locations/us-central1/tensorboards/1602142373297520640')\n",
      "Creating TensorboardExperiment\n",
      "TensorboardExperiment created. Resource name: projects/390782612502/locations/us-central1/tensorboards/1602142373297520640/experiments/llama3-1-lora-train-20240924-061434\n",
      "To use this TensorboardExperiment in another session:\n",
      "tb experiment = aiplatform.TensorboardExperiment('projects/390782612502/locations/us-central1/tensorboards/1602142373297520640/experiments/llama3-1-lora-train-20240924-061434')\n",
      "Training Output directory:\n",
      "gs://ai-sandbox-company-71-llama31-us-central1/temporal/llama3-1-lora-train-20240924-061434 \n",
      "View Training:\n",
      "https://console.cloud.google.com/ai/platform/locations/us-central1/training/3441773863715733504?project=390782612502\n",
      "View backing custom job:\n",
      "https://console.cloud.google.com/ai/platform/locations/us-central1/training/5382748337298472960?project=390782612502\n",
      "View tensorboard:\n",
      "https://us-central1.tensorboard.googleusercontent.com/experiment/projects+390782612502+locations+us-central1+tensorboards+1602142373297520640+experiments+5382748337298472960\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob run completed. Resource name: projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504\n",
      "Training did not produce a Managed Model returning None. Training Pipeline projects/390782612502/locations/us-central1/trainingPipelines/3441773863715733504 is not configured to upload a Model. Create the Training Pipeline with model_serving_container_image_uri and model_display_name passed in. Ensure that your training script saves to model to os.environ['AIP_MODEL_DIR'].\n",
      "LoRA adapter was saved in:  gs://ai-sandbox-company-71-llama31-us-central1/temporal/llama3-1-lora-train-20240924-061434/adapter\n",
      "Trained and merged models were saved in:  gs://ai-sandbox-company-71-llama31-us-central1/temporal/llama3-1-lora-train-20240924-061434/merged-model\n"
     ]
    }
   ],
   "source": [
    "tensorboard = aiplatform.Tensorboard.create(job_name)\n",
    "exp = aiplatform.TensorboardExperiment.create(\n",
    "    tensorboard_experiment_id=job_name, tensorboard_name=tensorboard.name\n",
    ")\n",
    "\n",
    "# Pass training arguments and launch job.\n",
    "train_job = aiplatform.CustomContainerTrainingJob(\n",
    "    display_name=job_name,\n",
    "    container_uri=TRAIN_DOCKER_URI,\n",
    ")\n",
    "# Worker pool spec.\n",
    "try:\n",
    "    accelerator_type = \"NVIDIA_L4\"\n",
    "    machine_type = \"g2-standard-48\"\n",
    "    accelerator_count = 4\n",
    "    train_job.run(\n",
    "        args=train_job_args,\n",
    "        environment_variables={\"WANDB_DISABLED\": True},\n",
    "        replica_count=replica_count,\n",
    "        machine_type=machine_type,\n",
    "        accelerator_type=accelerator_type,\n",
    "        accelerator_count=accelerator_count,\n",
    "        boot_disk_size_gb=500,\n",
    "        service_account=SERVICE_ACCOUNT,\n",
    "        tensorboard=tensorboard.resource_name,\n",
    "        base_output_dir=base_output_dir,\n",
    "    )\n",
    "except Exception as e:\n",
    "    print(f\"Error: {e}, NVIDIA_L4 failed, trying NVIDIA_A100_80GB / g2-standard-96\t\")\n",
    "    accelerator_type = \"NVIDIA_TESLA_A100\" #\"NVIDIA_A100_80GB\"\n",
    "    machine_type = \"a2-highgpu-1g\" #\"a2-ultragpu-4g\"\n",
    "    accelerator_count = 1    \n",
    "    train_job.run(\n",
    "        args=train_job_args,\n",
    "        environment_variables={\"WANDB_DISABLED\": True},\n",
    "        replica_count=replica_count,\n",
    "        machine_type=machine_type,\n",
    "        accelerator_type=accelerator_type,\n",
    "        accelerator_count=accelerator_count,\n",
    "        boot_disk_size_gb=500,\n",
    "        service_account=SERVICE_ACCOUNT,\n",
    "        tensorboard=tensorboard.resource_name,\n",
    "        base_output_dir=base_output_dir,\n",
    "    )\n",
    "\n",
    "print(\"LoRA adapter was saved in: \", lora_output_dir)\n",
    "print(\"Trained and merged models were saved in: \", merged_model_output_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deploy\n",
    "This section uploads the model to Model Registry and deploys it on the Endpoint. It takes 15 minutes to 1 hour to finish."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Update Config file\n",
    "from google.cloud import storage\n",
    "import json\n",
    "\n",
    "config = '''{\n",
    "  \"_name_or_path\": \"/tmp/base_model_dir/Meta-Llama-3.1-8B\",\n",
    "  \"architectures\": [\n",
    "    \"LlamaForCausalLM\"\n",
    "  ],\n",
    "  \"attention_bias\": false,\n",
    "  \"attention_dropout\": 0.0,\n",
    "  \"bos_token_id\": 128000,\n",
    "  \"eos_token_id\": 128001,\n",
    "  \"hidden_act\": \"silu\",\n",
    "  \"hidden_size\": 4096,\n",
    "  \"initializer_range\": 0.02,\n",
    "  \"intermediate_size\": 14336,\n",
    "  \"max_position_embeddings\": 131072,\n",
    "  \"mlp_bias\": false,\n",
    "  \"model_type\": \"llama\",\n",
    "  \"num_attention_heads\": 32,\n",
    "  \"num_hidden_layers\": 32,\n",
    "  \"num_key_value_heads\": 8,\n",
    "  \"pretraining_tp\": 1,\n",
    "  \"rms_norm_eps\": 1e-05,\n",
    "  \"rope_scaling\": {\n",
    "    \"factor\": 8.0,\n",
    "    \"type\": \"dynamic\"\n",
    "  },\n",
    "  \"rope_theta\": 500000.0,\n",
    "  \"tie_word_embeddings\": false,\n",
    "  \"torch_dtype\": \"float16\",\n",
    "  \"transformers_version\": \"4.43.1\",\n",
    "  \"use_cache\": true,\n",
    "  \"vocab_size\": 128257\n",
    "}'''\n",
    "\n",
    "data = json.loads(config)\n",
    "\n",
    "storage_client = storage.Client()\n",
    "bucket = storage_client.bucket(BUCKET_NAME)\n",
    "config_path = merged_model_output_dir.split(f'{BUCKET_NAME}/',1)\n",
    "blob = bucket.blob(f'{config_path[1]}/config.json')\n",
    "blob.upload_from_string(data=json.dumps(data),content_type='application/json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "qmHW6m8xG_4U"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Deploying models in:  gs://ai-sandbox-company-71-llama31-us-central1/temporal/llama3-1-lora-train-20240924-061434/merged-model\n",
      "Deploying models in:  gs://ai-sandbox-company-71-llama31-us-central1/temporal/llama3-1-lora-train-20240924-061434/merged-model\n",
      "Creating Endpoint\n",
      "Create Endpoint backing LRO: projects/390782612502/locations/us-central1/endpoints/5603921595216166912/operations/4440082111944196096\n",
      "Endpoint created. Resource name: projects/390782612502/locations/us-central1/endpoints/5603921595216166912\n",
      "To use this Endpoint in another session:\n",
      "endpoint = aiplatform.Endpoint('projects/390782612502/locations/us-central1/endpoints/5603921595216166912')\n",
      "Creating Model\n",
      "Create Model backing LRO: projects/390782612502/locations/us-central1/models/2504275171213312000/operations/5205694048597180416\n",
      "Model created. Resource name: projects/390782612502/locations/us-central1/models/2504275171213312000@1\n",
      "To use this Model in another session:\n",
      "model = aiplatform.Model('projects/390782612502/locations/us-central1/models/2504275171213312000@1')\n",
      "Deploying llama3-1-vllm-serve-20240924-075329 on g2-standard-12 with 1 NVIDIA_L4 GPU(s).\n",
      "Deploying model to Endpoint : projects/390782612502/locations/us-central1/endpoints/5603921595216166912\n",
      "Deploy Endpoint model backing LRO: projects/390782612502/locations/us-central1/endpoints/5603921595216166912/operations/2899851039383486464\n"
     ]
    }
   ],
   "source": [
    "print(\"Deploying models in: \", merged_model_output_dir)\n",
    "\n",
    "# The pre-built serving docker image for vLLM.\n",
    "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240721_0916_RC00\"\n",
    "\n",
    "# Find Vertex AI prediction supported accelerators and regions in [here](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute).\n",
    "if \"8b\" in MODEL_ID.lower():\n",
    "    machine_type = \"g2-standard-12\"\n",
    "    accelerator_type = \"NVIDIA_L4\"\n",
    "    accelerator_count = 1\n",
    "else:\n",
    "    machine_type = \"g2-standard-96\"\n",
    "    accelerator_type = \"NVIDIA_L4\"\n",
    "    accelerator_count = 8\n",
    "\n",
    "common_util.check_quota(\n",
    "    project_id=PROJECT_ID,\n",
    "    region=REGION,\n",
    "    accelerator_type=accelerator_type,\n",
    "    accelerator_count=accelerator_count,\n",
    "    is_for_training=False,\n",
    ")\n",
    "print(\"Deploying models in: \", merged_model_output_dir)\n",
    "\n",
    "gpu_memory_utilization = 0.85\n",
    "max_model_len = 8192  # Maximum context length.\n",
    "\n",
    "# Ensure max_model_len does not exceed the limit\n",
    "if max_model_len > 8192:\n",
    "    raise ValueError(\"max_model_len cannot exceed 8192\")\n",
    "\n",
    "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
    "    model_name=common_util.get_job_name_with_datetime(prefix=\"llama3_1-vllm-serve\"),\n",
    "    model_id=merged_model_output_dir,\n",
    "    service_account=SERVICE_ACCOUNT,\n",
    "    machine_type=machine_type,\n",
    "    accelerator_type=accelerator_type,\n",
    "    accelerator_count=accelerator_count,\n",
    "    gpu_memory_utilization=gpu_memory_utilization,\n",
    "    max_model_len=max_model_len,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predict\n",
    "\n",
    "Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://docs.vllm.ai/en/latest/dev/sampling_params.html).\n",
    "\n",
    "Example:\n",
    "\n",
    "```\n",
    "Human: What is a car?\n",
    "Assistant:  A car, or a motor car, is a road-connected human-transportation system used to move people or goods from one place to another. The term also encompasses a wide range of vehicles, including motorboats, trains, and aircrafts. Cars typically have four wheels, a cabin for passengers, and an engine or motor. They have been around since the early 19th century and are now one of the most popular forms of transportation, used for daily commuting, shopping, and other purposes.\n",
    "```\n",
    "Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loads an existing endpoint instance using the endpoint name:\n",
    "- Using `endpoint_name = endpoint.name` allows us to get the\n",
    "  endpoint name of the endpoint `endpoint` created in the cell\n",
    "  above.\n",
    "- Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
    "  an existing endpoint with the ID 1234567890123456789.\n",
    "You may uncomment the code below to load an existing endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
    "# aip_endpoint_name = (\n",
    "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
    "# )\n",
    "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
    "# print(\"endpoint\",endpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"What is a car?\"  # @param {type: \"string\"}\n",
    "max_tokens = 50  # @param {type:\"integer\"}\n",
    "temperature = 1.0  # @param {type:\"number\"}\n",
    "top_p = 1.0  # @param {type:\"number\"}\n",
    "top_k = 1  # @param {type:\"integer\"}\n",
    "raw_response = False  # @param {type:\"boolean\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Overrides parameters for inferences.\n",
    "If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`,\n",
    "you can reduce the maximum number of output tokens, such as set max_tokens as 20."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "2UYUNn60G_4U"
   },
   "outputs": [],
   "source": [
    "instances = [\n",
    "    {\n",
    "        \"prompt\": prompt,\n",
    "        \"max_tokens\": max_tokens,\n",
    "        \"temperature\": temperature,\n",
    "        \"top_p\": top_p,\n",
    "        \"top_k\": top_k,\n",
    "        \"raw_response\": raw_response,\n",
    "    },\n",
    "]\n",
    "response = endpoints[\"vllm_gpu\"].predict(instances=instances)\n",
    "\n",
    "for prediction in response.predictions:\n",
    "    print(prediction)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "af21a3cff1e0"
   },
   "source": [
    "## Clean up resources\n",
    "### Delete the model and endpoint\n",
    "Delete the experiment models and endpoints to recycle the resources and avoid unnecessary continouous charges that may incur."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# train_job.delete()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Undeploy model and delete endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# for endpoint in endpoints.values():\n",
    "#     endpoint.delete(force=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Delete models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "911406c1561e",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# for model in models.values():\n",
    "#     model.delete()\n",
    "\n",
    "# delete_bucket = False  # @param {type:\"boolean\"}\n",
    "# if delete_bucket:\n",
    "#     ! gsutil -m rm -r $BUCKET_NAME"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "model_garden_pytorch_llama3_1_finetuning.ipynb",
   "toc_visible": false
  },
  "environment": {
   "kernel": "conda-env-tensorflow-tensorflow",
   "name": "workbench-notebooks.m124",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m124"
  },
  "kernelspec": {
   "display_name": "TensorFlow 2-11 (Local)",
   "language": "python",
   "name": "conda-env-tensorflow-tensorflow"
  },
  "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}