Update open-llama notebook by making vLLM the default serving option. (#2237)

notebooks/community/model_garden/model_garden_pytorch_openllama_peft.ipynb

@@ -47,7 +47,7 @@
         "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_openllama_peft.ipynb\">\n",
         "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
         "Open in Vertex AI Workbench\n",
-        "    </a> (A Python-3 CPU notebook is recommended)\n",
+        "    </a> (A Python-3 GPU notebook is recommended)\n",
         "  </td>\n",
         "</table>"
       ]
@@ -60,14 +60,15 @@
       "source": [
         "## Overview\n",
         "\n",
-        "This notebook demonstrates deploying prebuilt OpenLLaMA, and also finetuning and deploying OpenLLaMA with performance efficient finetuning libraries ([PEFT](https://github.com/huggingface/peft)) in Vertex AI.\n",
+        "This notebook demonstrates running local inference with prebuilt OpenLLaMA, deploying prebuilt OpenLLaMA, deploying prebuilt OpenLLaMA with [vLLM](https://github.com/vllm-project/vllm), finetuning and deploying OpenLLaMA with performance efficient finetuning libraries ([PEFT](https://github.com/huggingface/peft)), and evaluating PEFT-finetuned OpenLLaMA in Vertex AI.\n",
         "\n",
         "### Objective\n",
         "\n",
         "- Run local inference with prebuilt OpenLLaMA\n",
         "- Deploy prebuilt OpenLLaMA\n",
-        "- Finetune and deploy OpenLLaMA with PEFT, supporting\n",
-        "- Deploy OpenLLaMA with [vLLM](https://github.com/vllm-project/vllm) to improve serving throughput\n",
+        "- Deploy prebuilt OpenLLaMA with [vLLM](https://github.com/vllm-project/vllm) to improve serving throughput\n",
+        "- Finetune and deploy OpenLLaMA with PEFT\n",
+        "- Evaluate finetuned OpenLLaMA with PEFT\n",
         "\n",
         "| Models | LoRA |\n",
         "| :- | :- |\n",
@@ -93,7 +94,9 @@
       "source": [
         "## Before you begin\n",
         "\n",
-        "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
+        "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands.\n",
+        "\n",
+        "Running local inference with OpenLLaMA requires a GPU."
       ]
     },
     {
@@ -159,7 +162,7 @@
         "id": "6c460088b873"
       },
       "source": [
-        "Fill following variables for experiments environment:"
+        "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\")."
       ]
     },
     {
@@ -177,7 +180,8 @@
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
         "# The Cloud Storage bucket for storing experiments output.\n",
-        "BUCKET_URI = \"\"  # @param {type:\"string\"}\n",
+        "# Start with gs:// prefix, e.g. gs://foo_bucket.\n",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
         "\n",
         "! gcloud config set project $PROJECT_ID\n",
         "\n",
@@ -236,11 +240,17 @@
       "outputs": [],
       "source": [
         "# The pre-built training and serving docker images.\n",
-        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-peft-train\"\n",
-        "PREDICTION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/pytorch-peft-serve\"\n",
+        "TRAIN_DOCKER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train\"\n",
+        ")\n",
+        "PREDICTION_DOCKER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve\"\n",
+        ")\n",
         "VLLM_DOCKER_URI = (\n",
         "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve\"\n",
-        ")"
+        ")\n",
+        "\n",
+        "EVAL_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-lm-evaluation-harness\""
       ]
     },
     {
@@ -260,63 +270,26 @@
       },
       "outputs": [],
       "source": [
-        "import os\n",
         "from datetime import datetime\n",
         "\n",
         "from google.cloud import aiplatform\n",
         "\n",
         "\n",
-        "def get_job_name_with_datetime(prefix: str):\n",
-        "    \"\"\"Gets the job name with date time when triggering training or deployment\n",
+        "def create_name_with_datetime(prefix: str) -> str:\n",
+        "    \"\"\"Creates a name with date time when triggering training or deployment\n",
         "    jobs in Vertex AI.\n",
         "    \"\"\"\n",
         "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
         "\n",
         "\n",
-        "def deploy_model(\n",
-        "    model_name,\n",
-        "    base_model_id,\n",
-        "    finetuned_lora_model_path,\n",
-        "    service_account,\n",
-        "    task,\n",
-        "    machine_type=\"n1-standard-8\",\n",
-        "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
-        "):\n",
-        "    \"\"\"Deploys trained models into Vertex AI.\"\"\"\n",
-        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
-        "    serving_env = {\n",
-        "        \"BASE_MODEL_ID\": base_model_id,\n",
-        "        \"TASK\": task,\n",
-        "    }\n",
-        "    if finetuned_lora_model_path:\n",
-        "        serving_env[\"FINETUNED_LORA_MODEL_PATH\"] = finetuned_lora_model_path\n",
-        "    model = aiplatform.Model.upload(\n",
-        "        display_name=model_name,\n",
-        "        serving_container_image_uri=PREDICTION_DOCKER_URI,\n",
-        "        serving_container_ports=[7080],\n",
-        "        serving_container_predict_route=\"/predictions/peft_serving\",\n",
-        "        serving_container_health_route=\"/ping\",\n",
-        "        serving_container_environment_variables=serving_env,\n",
-        "    )\n",
-        "    model.deploy(\n",
-        "        endpoint=endpoint,\n",
-        "        machine_type=machine_type,\n",
-        "        accelerator_type=accelerator_type,\n",
-        "        accelerator_count=1,\n",
-        "        deploy_request_timeout=1800,\n",
-        "        service_account=service_account,\n",
-        "    )\n",
-        "    return model, endpoint\n",
-        "\n",
-        "\n",
         "def deploy_model_vllm(\n",
-        "    model_name,\n",
-        "    model_id,\n",
-        "    service_account,\n",
-        "    machine_type=\"n1-standard-8\",\n",
-        "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
-        "    accelerator_count=1,\n",
-        "):\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    service_account: str,\n",
+        "    machine_type: str = \"n1-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_TESLA_V100\",\n",
+        "    accelerator_count: int = 1,\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",
@@ -325,7 +298,7 @@
         "        \"--port=7080\",\n",
         "        f\"--model={model_id}\",\n",
         "        f\"--tensor-parallel-size={accelerator_count}\",\n",
-        "        \"--swap-space=16\",\n",
+        "        \"--swap-space=4\",\n",
         "        \"--gpu-memory-utilization=0.95\",\n",
         "        \"--disable-log-stats\",\n",
         "    ]\n",
@@ -390,20 +363,20 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "8neJc8CnDDpu"
+        "id": "V7VOhhHGpUrj"
       },
       "source": [
-        "## Deploy Prebuilt OpenLLaMA\n",
+        "## Deploy Prebuilt OpenLLaMA with vLLM\n",
         "\n",
-        "This section deploys prebuilt OpenLLaMA models on the Endpoint. The model deployment step will take ~15 minutes to complete.\n",
+        "This section deploys prebuilt OpenLLaMA models with [vLLM](https://github.com/vllm-project/vllm) on the Endpoint. The model deployment step will take ~15 minutes to complete.\n",
         "\n",
-        "The peak GPU memory usages for [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b), [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b), and [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) are ~5.3G, ~8.7G and ~15.2G separately with the default settings."
+        "vLLM is a highly optimized LLM serving framework that can significantly increase serving throughput. The higher QPS you have, the more performance benefits you get from using vLLM."
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "2MjaORIIFDVu"
+        "id": "4GTNnnuYqrW_"
       },
       "source": [
         "Set the prebuilt model id."
@@ -413,98 +386,24 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "E8OiHHNNE_wj"
+        "id": "kLsRoc4Kqrkx"
       },
       "outputs": [],
       "source": [
-        "prebuilt_model_id = \"openlm-research/open_llama_3b\"  # @param [\"openlm-research/open_llama_3b\", \"openlm-research/open_llama_7b\", \"openlm-research/open_llama_13b\"]"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "dHFW7yvjaVFV"
-      },
-      "source": [
-        "We use the PEFT serving images to deploy prebuilt OpenLLaMA models, by setting finetuning LoRA model paths as empty."
+        "prebuilt_model_id = \"openlm-research/open_llama_7b\"  # @param [\"openlm-research/open_llama_3b\", \"openlm-research/open_llama_7b\", \"openlm-research/open_llama_13b\"]"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "Uak1pyEeExYM"
-      },
-      "outputs": [],
-      "source": [
-        "model_without_peft, endpoint_without_peft = deploy_model(\n",
-        "    model_name=get_job_name_with_datetime(prefix=\"openllama-serve\"),\n",
-        "    base_model_id=prebuilt_model_id,\n",
-        "    finetuned_lora_model_path=\"\",  # This will avoid override finetuning models.\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
-        "    task=\"causal-language-modeling-lora\",\n",
-        ")\n",
-        "print(\"endpoint_name:\", endpoint_without_peft.name)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sGKIjgmDFRW2"
-      },
-      "source": [
-        "NOTE: The prebuilt model weights will be downloaded on the fly from the orginal location after the deployment succeeds. Thus additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
-        "\n",
-        "Once deployment succeeds, you can send requests to the endpoint with text prompts."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "rDHsCOqvFYBi"
-      },
-      "outputs": [],
-      "source": [
-        "# # Loads an existing endpoint as below.\n",
-        "# endpoint_name = endpoint_without_peft.name\n",
-        "# aip_endpoint_name = (\n",
-        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
-        "# )\n",
-        "# endpoint_without_peft = aiplatform.Endpoint(aip_endpoint_name)\n",
-        "instances = [\n",
-        "    {\"prompt\": \"Hi, Google.\"},\n",
-        "]\n",
-        "response = endpoint_without_peft.predict(instances=instances)\n",
-        "\n",
-        "for prediction in response.predictions[0]:\n",
-        "    print(prediction[\"generated_text\"])"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "8neJc8CnDDpu"
-      },
-      "source": [
-        "## Deploy Prebuilt OpenLLaMA with vLLM\n",
-        "\n",
-        "This section deploys prebuilt OpenLLaMA models with [vLLM](https://github.com/vllm-project/vllm) on the Endpoint. The model deployment step will take ~15 minutes to complete.\n",
-        "\n",
-        "vLLM is an highly optimized LLM serving framework which can increase serving throughput a lot. The higher QPS you have, the more benefits you get using vLLM."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "ee610586eddb"
+        "id": "YI0vaDi6p2fi"
       },
       "outputs": [],
       "source": [
         "model_without_peft_vllm, endpoint_without_peft_vllm = deploy_model_vllm(\n",
-        "    model_name=get_job_name_with_datetime(prefix=\"openllama-serve-vllm\"),\n",
-        "    model_id=\"openlm-research/open_llama_13b\",\n",
+        "    model_name=create_name_with_datetime(prefix=\"openllama-serve-vllm\"),\n",
+        "    model_id=prebuilt_model_id,\n",
         "    service_account=SERVICE_ACCOUNT,\n",
         "    machine_type=\"n1-highmem-8\",\n",
         "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
@@ -515,19 +414,19 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "sGKIjgmDFRW2"
+        "id": "dWYmYWoqqBuZ"
       },
       "source": [
-        "NOTE: The prebuilt model weights will be downloaded on the fly from the orginal location after the deployment succeeds. Thus additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
+        "NOTE: The prebuilt model weights will be downloaded on the fly from the original location after the deployment succeeds. Thus, an additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you can run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
         "\n",
-        "Once deployment succeeds, you can send requests to the endpoint with text prompts."
+        "Once deployment succeeds, you can send requests to the endpoint with text prompts. If you are interested in additional serving parameters, please refer to the vLLM GitHub [examples/api_client.py](https://github.com/vllm-project/vllm/blob/main/examples/api_client.py) for more details."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "63bb8fb78d0b"
+        "id": "fjO4z3qAp3pK"
       },
       "outputs": [],
       "source": [
@@ -548,7 +447,7 @@
       "source": [
         "## Finetune and deploy OpenLLaMA with PEFT\n",
         "\n",
-        "This section demonstrates how to finetune and dpeloy OpenLLaMA with PEFT LoRA."
+        "This section demonstrates how to finetune the OpenLLaMA-7b model, merge the finetuned LoRA adapter with the base model, and serve using vLLM."
       ]
     },
     {
@@ -568,7 +467,8 @@
       },
       "outputs": [],
       "source": [
-        "base_model_id = \"openlm-research/open_llama_3b\"  # @param [\"openlm-research/open_llama_3b\", \"openlm-research/open_llama_7b\", \"openlm-research/open_llama_13b\"]"
+        "# vLLM currently does not support finetuned `open_llama_3b` model yet.\n",
+        "base_model_id = \"openlm-research/open_llama_7b\"  # @param [\"openlm-research/open_llama_7b\", \"openlm-research/open_llama_13b\"]"
       ]
     },
     {
@@ -590,7 +490,9 @@
         "\n",
         "This example uses the dataset [Abirate/english_quotes](https://huggingface.co/datasets/Abirate/english_quotes).\n",
         "\n",
-        "In order to make the finetuning efficiently, we enabled quantization (8bits) when loading pretrained models for finetuning LoRA models. The peak GPU memory usages are ~7G, ~10G and ~16G for finetuning LoRA models for [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b), [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b), and [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) separately with default training parameters and the example dataset. open_llama_3b and open_llama_7b can be finetuned on 1 V100, and open_llama_13b can be finetuned on 1 A100 (40G)."
+        "In order to make the finetuning efficiently, we enabled quantization for loading pretrained models for finetuning LoRA models. Precision options include `\"4bit\"`, `\"8bit\"`, `\"float16\"` (default) and `\"float32\"`, and the precision can be set via `\"--precision_mode\"`. The peak GPU memory usages are ~7G, ~10G and ~16G for finetuning LoRA models for [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b), [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b), and [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) separately with default training parameters and the example dataset. `open_llama_3b` and `open_llama_7b` can be finetuned on **1 V100**, and `open_llama_13b` can be finetuned on **1 A100 (40G)**.\n",
+        "\n",
+        "In this section, the finetuned LoRA adapter will be saved to a GCS bucket specified by the variable `lora_adapter_dir` below; and we merge the LoRa adapter with the base model, and save it to a separate GCS bucket specified by `merged_model_output_dir` below.\n"
       ]
     },
     {
@@ -604,23 +506,32 @@
         "dataset_name = \"Abirate/english_quotes\"  # @param {type:\"string\"}\n",
         "\n",
         "# Worker pool spec.\n",
-        "# Finetunes open_llama_3b and open_llama_7b with 1 V100.\n",
+        "# Finetunes open_llama_7b with 1 V100.\n",
         "machine_type = \"n1-standard-8\"\n",
         "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
-        "# Finetunes and open_llama_13b with 1 A100 (40G).\n",
+        "# Finetunes open_llama_13b with 1 A100 (40G).\n",
         "# machine_type = \"a2-highgpu-1g\"\n",
         "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
         "replica_count = 1\n",
         "accelerator_count = 1\n",
         "\n",
         "# Setup training job.\n",
-        "job_name = get_job_name_with_datetime(\"openllama-lora-train\")\n",
+        "job_name = create_name_with_datetime(\"openllama-lora-train\")\n",
         "train_job = aiplatform.CustomContainerTrainingJob(\n",
         "    display_name=job_name,\n",
         "    container_uri=TRAIN_DOCKER_URI,\n",
         ")\n",
-        "output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
-        "output_dir_gcsfuse = output_dir.replace(\"gs://\", \"/gcs/\")\n",
+        "\n",
+        "# Create a GCS folder to store the LORA adapter.\n",
+        "lora_adapter_dir = create_name_with_datetime(\"openllama-lora-adapter\")\n",
+        "lora_output_dir = os.path.join(MODEL_BUCKET, lora_adapter_dir)\n",
+        "lora_output_dir_gcsfuse = lora_output_dir.replace(\"gs://\", \"/gcs/\")\n",
+        "\n",
+        "# Create a GCS folder to store the merged model with the base model and the\n",
+        "# finetuned LORA adapter.\n",
+        "merged_model_dir = create_name_with_datetime(\"openllama-merged-model\")\n",
+        "merged_model_output_dir = os.path.join(MODEL_BUCKET, merged_model_dir)\n",
+        "merged_model_output_dir_gcsfuse = merged_model_output_dir.replace(\"gs://\", \"/gcs/\")\n",
         "\n",
         "# Pass training arguments and launch job.\n",
         "train_job.run(\n",
@@ -628,7 +539,8 @@
         "        \"--task=causal-language-modeling-lora\",\n",
         "        f\"--pretrained_model_id={base_model_id}\",\n",
         "        f\"--dataset_name={dataset_name}\",\n",
-        "        f\"--output_dir={output_dir_gcsfuse}\",\n",
+        "        f\"--output_dir={lora_output_dir_gcsfuse}\",\n",
+        "        f\"--merge_base_and_lora_output_dir={merged_model_output_dir_gcsfuse}\",\n",
         "        \"--lora_rank=16\",\n",
         "        \"--lora_alpha=32\",\n",
         "        \"--lora_dropout=0.05\",\n",
@@ -643,7 +555,11 @@
         "    boot_disk_size_gb=500,\n",
         ")\n",
         "\n",
-        "print(\"Trained models were saved in: \", output_dir)"
+        "print(\"The finetuned Lora adapter can be found at: \", lora_output_dir)\n",
+        "print(\n",
+        "    \"The finetuned Lora adapter merged with the base model can be found at: \",\n",
+        "    merged_model_output_dir,\n",
+        ")"
       ]
     },
     {
@@ -652,12 +568,14 @@
         "id": "jqmCtkGnhDmp"
       },
       "source": [
-        "### Deploy\n",
-        "This section uploads the model to Model Registry and deploys it on the Endpoint.\n",
+        "### Deploy with vLLM\n",
+        "This section uploads the model to Model Registry and deploys it on the Endpoint. vLLM currently does not support serving finetuned [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b) for now so we will use the [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b) in this example.\n",
         "\n",
         "The model deployment step will take ~15 minutes to complete.\n",
         "\n",
-        "The peak GPU memory usages for [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b), [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b), and [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) with LoRA weights are ~5.3G, ~8.7G and ~15.2G separately with the default settings."
+        "The peak GPU memory usages for [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b), and [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) with LoRA weights are ~8.7G and ~15.2G respectively with the default settings.\n",
+        "\n",
+        "NOTE: vLLM requires a merged model with the base model and the finetuned LoRA adapter. Based on your business need, if you need the base model and the finetuned LoRA weight to be served separately, please consider using the regular Vertex serving instead.\n"
       ]
     },
     {
@@ -668,14 +586,16 @@
       },
       "outputs": [],
       "source": [
-        "model_with_peft, endpoint_with_peft = deploy_model(\n",
-        "    model_name=get_job_name_with_datetime(prefix=\"openllama-peft-serve\"),\n",
-        "    base_model_id=base_model_id,\n",
-        "    finetuned_lora_model_path=output_dir,\n",
+        "model_with_peft_vllm, endpoint_with_peft_vllm = deploy_model_vllm(\n",
+        "    model_name=create_name_with_datetime(prefix=\"openllama-peft-serve-vllm\"),\n",
+        "    model_id=merged_model_output_dir,\n",
         "    service_account=SERVICE_ACCOUNT,\n",
-        "    task=\"causal-language-modeling-lora\",\n",
+        "    machine_type=\"n1-highmem-8\",\n",
+        "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
+        "    accelerator_count=2,\n",
         ")\n",
-        "print(\"endpoint_name:\", endpoint_with_peft.name)"
+        "\n",
+        "print(\"endpoint_name:\", endpoint_with_peft_vllm.name)"
       ]
     },
     {
@@ -684,7 +604,7 @@
         "id": "80b3fd2ace09"
       },
       "source": [
-        "NOTE: After the deployment succeeds, the base model weights will be downloaded one the fly from the original location and LoRA model weights will be downloaded from the GCS bucket used in training above. Thus additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
+        "NOTE: After the deployment succeeds, the base model weights will be downloaded on the fly from the original location and LoRA model weights will be downloaded from the GCS bucket used in training above. Thus, an additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you can run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
         "\n",
         "Once deployment succeeds, you can send requests to the endpoint with text prompts."
       ]
@@ -697,19 +617,153 @@
       },
       "outputs": [],
       "source": [
-        "# # Loads an existing endpoint as below.\n",
-        "# endpoint_name = endpoint_with_peft.name\n",
-        "# aip_endpoint_name = (\n",
-        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
-        "# )\n",
-        "# endpoint_with_peft = aiplatform.Endpoint(aip_endpoint_name)\n",
-        "instances = [\n",
-        "    {\"prompt\": \"Hi, Google.\"},\n",
-        "]\n",
-        "response = endpoint_with_peft.predict(instances=instances)\n",
+        "instance = {\n",
+        "    \"prompt\": \"Hi, Google. How are you doing?\",\n",
+        "    \"n\": 1,\n",
+        "    \"max_tokens\": 32,\n",
+        "}\n",
+        "response = endpoint_with_peft_vllm.predict(instances=[instance])\n",
+        "print(response.predictions[0])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "JmuUk3l1DoEo"
+      },
+      "source": [
+        "## Evaluate PEFT-finetuned OpenLLaMA\n",
         "\n",
-        "for prediction in response.predictions[0]:\n",
-        "    print(prediction[\"generated_text\"])"
+        "This section demonstrates how to evaluate the OpenLLaMA model fintuned with PEFT LoRA using EleutherAI's [Language Model Evaluation Harness (lm-evaluation-harness)](https://github.com/EleutherAI/lm-evaluation-harness) with Vertex CustomJob.\n",
+        "\n",
+        "This example uses the dataset [HellaSwag](https://allenai.org/data/hellaswag). All supported tasks are listed in [this task table](https://github.com/EleutherAI/lm-evaluation-harness/blob/master/docs/task_table.md)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "gM4SXaquDoEo"
+      },
+      "outputs": [],
+      "source": [
+        "eval_dataset = \"hellaswag\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "# Evaluates open_llama_3b and open_llama_7b with 1 V100.\n",
+        "machine_type = \"n1-standard-8\"\n",
+        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
+        "# Evaluates open_llama_13b with 1 A100 (40G).\n",
+        "# machine_type = \"a2-highgpu-1g\"\n",
+        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "replica_count = 1\n",
+        "accelerator_count = 1\n",
+        "\n",
+        "# Setup evaluation job.\n",
+        "job_name = create_name_with_datetime(prefix=\"openllama-peft-eval\")\n",
+        "eval_output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
+        "eval_output_dir_gcsfuse = eval_output_dir.replace(\"gs://\", \"/gcs/\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Yt4Tth8hDoEo"
+      },
+      "outputs": [],
+      "source": [
+        "# Prepare evaluation script that runs the evaluation harness.\n",
+        "script_path = \"./eval_script.py\"  # @param {type:\"string\"}\n",
+        "\n",
+        "eval_command = f\"\"\"import subprocess\n",
+        "\n",
+        "\n",
+        "subprocess.call([\n",
+        "    'python',\n",
+        "    'main.py',\n",
+        "    '--model',\n",
+        "    'hf-causal-experimental',\n",
+        "    '--model_args',\n",
+        "    'pretrained={base_model_id},peft={output_dir_gcsfuse}',\n",
+        "    '--tasks',\n",
+        "    '{eval_dataset}',\n",
+        "    '--output_path',\n",
+        "    '{eval_output_dir_gcsfuse}',\n",
+        "])\n",
+        "\"\"\"\n",
+        "\n",
+        "with open(script_path, \"w\") as fp:\n",
+        "    fp.write(eval_command)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ItWB0WS__CX-"
+      },
+      "source": [
+        "### Submit evaluation CustomJob"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "BbyIk99bDoEo"
+      },
+      "outputs": [],
+      "source": [
+        "# Pass evaluation arguments and launch job.\n",
+        "eval_job = aiplatform.CustomJob.from_local_script(\n",
+        "    display_name=job_name,\n",
+        "    script_path=script_path,\n",
+        "    container_uri=EVAL_DOCKER_URI,\n",
+        "    replica_count=replica_count,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    base_output_dir=eval_output_dir,\n",
+        ")\n",
+        "\n",
+        "eval_job.run()\n",
+        "\n",
+        "print(\"Evaluation results were saved in:\", eval_output_dir)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kN0lE2iu_NXN"
+      },
+      "source": [
+        "### Fetch and print evaluation results"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "927oRxoADoEp"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "\n",
+        "from google.cloud import storage\n",
+        "\n",
+        "# Fetch evaluation results.\n",
+        "storage_client = storage.Client()\n",
+        "BUCKET_NAME = BUCKET_URI.split(\"gs://\")[1]\n",
+        "bucket = storage_client.get_bucket(BUCKET_NAME)\n",
+        "RESULT_FILE_PATH = eval_output_dir[len(BUCKET_URI) + 1 :]\n",
+        "blob = bucket.blob(RESULT_FILE_PATH)\n",
+        "raw_result = blob.download_as_string()\n",
+        "\n",
+        "# Print evaluation results.\n",
+        "result = json.loads(raw_result)\n",
+        "result_formatted = json.dumps(result, indent=2)\n",
+        "print(f\"Evaluation result:\\n{result_formatted}\")"
       ]
     },
     {
@@ -729,18 +783,17 @@
       },
       "outputs": [],
       "source": [
-        "# Delete custom train jobs.\n",
+        "# Delete custom train and evaluation jobs.\n",
         "train_job.delete()\n",
+        "eval_job.delete()\n",
         "\n",
-        "# Undeploy model and delete endpoint.\n",
-        "endpoint_without_peft.delete(force=True)\n",
-        "endpoint_with_peft.delete(force=True)\n",
+        "# Undeploy models and delete endpoints.\n",
         "endpoint_without_peft_vllm.delete(force=True)\n",
+        "endpoint_with_peft_vllm.delete(force=True)\n",
         "\n",
         "# Delete models.\n",
-        "model_without_peft.delete()\n",
-        "model_with_peft.delete()\n",
-        "model_without_peft_vllm.delete()"
+        "model_without_peft_vllm.delete()\n",
+        "model_with_peft_vllm.delete()"
       ]
     }
   ],

notebooks/official/gen_ai/tune_peft.ipynb

@@ -0,0 +1,696 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9f0d0f32-23b4-41a6-b364-579da297c326"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Copyright & License (click to expand)\n",
+        "# Copyright 2023 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": "dd53d60c-97eb-4c72-91ea-f274a753ab34"
+      },
+      "source": [
+        "# Vertex AI Tuning a PEFT model\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/gen_ai/tune_peft.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/gen_ai/tune_peft.ipynb\">\n",
+        "        <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "    <td>\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/gen_ai/tune_peft.ipynb\">\n",
+        "       <img src=\"https://www.gstatic.com/cloud/images/navigation/vertex-ai.svg\" alt=\"Vertex AI logo\">Open in Vertex AI Workbench\n",
+        "    </a>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9ef820fb-1203-4cab-965f-17093a4ba25e"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This tutorial demonstrates how to use Vertex AI to tune a PEFT large-language model (LLM) and make a prediction. This workflow improves a model's accuracy by fine-tuning a base model with a training dataset.\n",
+        "\n",
+        "Learn more about [Vertex AI Generative AI Models](https://cloud.google.com/vertex-ai/docs/generative-ai/learn/models#model_naming_scheme)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "74b00940-376c-4056-90fb-d22c1ce6eedf"
+      },
+      "source": [
+        "### Objective\n",
+        "\n",
+        "In this tutorial, you learn to use `Vertex AI LLM` to tune and deploy a PEFT large language model model.\n",
+        "\n",
+        "\n",
+        "This tutorial uses the following Google Cloud ML services:\n",
+        "\n",
+        "- `Vertex AI LLM`\n",
+        "- `Vertex AI Model Garden`\n",
+        "- `Vertex AI Prediction`\n",
+        "\n",
+        "\n",
+        "The steps performed include:\n",
+        "\n",
+        "- Get the Vertex AI LLM model.\n",
+        "- Tune the model.\n",
+        "  - This will automatically create a Vertex AI endpoint and deploy the model to it.\n",
+        "- Make a prediction using `Vertex AI LLM`.\n",
+        "- Make a prediction using `Vertex AI Prediction`"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c6b43693-b20a-41bd-b5b8-5ad414517162"
+      },
+      "source": [
+        "### Model\n",
+        "\n",
+        "The pre-trained LLM model is a text-bison (Decoder only) model for text generation."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6d7b5435-e947-49bb-9ce3-aa8a42c30118"
+      },
+      "source": [
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI\n",
+        "pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage\n",
+        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
+        "Calculator](https://cloud.google.com/products/calculator/)\n",
+        "to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0cbf01f0-5f6e-4bcd-903f-84ccaad5332c"
+      },
+      "source": [
+        "## Installation\n",
+        "\n",
+        "Install the following packages required to execute this notebook."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "d1f49723-9cbc-4fac-8e6d-48f33db68255"
+      },
+      "outputs": [],
+      "source": [
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform  \"shapely<2.0.0\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ff7f7e74-60d7-4032-b335-216564f43e35"
+      },
+      "source": [
+        "### Colab only: Uncomment the following cell to restart the kernel"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "5c10aa63-17a3-4162-a397-fe570a962cb3"
+      },
+      "outputs": [],
+      "source": [
+        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
+        "# import IPython\n",
+        "\n",
+        "# app = IPython.Application.instance()\n",
+        "# app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "b37d4259-7e39-417b-8879-24f7575732c8"
+      },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "### Set your project ID\n",
+        "\n",
+        "**If you don't know your project ID**, try the following:\n",
+        "* Run `gcloud config list`.\n",
+        "* Run `gcloud projects list`.\n",
+        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "caaf0d7e-c6cb-4e56-af5c-553db5180e00"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "054d794d-cd2e-4280-95ac-859b264ea2d6"
+      },
+      "source": [
+        "#### Region\n",
+        "\n",
+        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0121bf60-1acd-4272-afaf-aa54b4ded263"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"us-central1\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "eac9e842-d225-4876-836f-afdb1937d800"
+      },
+      "source": [
+        "### Authenticate your Google Cloud account\n",
+        "\n",
+        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below.\n",
+        "\n",
+        "**1. Vertex AI Workbench**\n",
+        "* Do nothing as you are already authenticated.\n",
+        "\n",
+        "**2. Local JupyterLab instance, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "23082eec-b1bd-4594-b5b5-56fe2b74db6f"
+      },
+      "outputs": [],
+      "source": [
+        "# ! gcloud auth login"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3c20f923-3c46-4d6d-80d2-d7cb22b1a8da"
+      },
+      "source": [
+        "**3. Colab, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "60302a3f-fad9-452c-8998-a9c9822d2732"
+      },
+      "outputs": [],
+      "source": [
+        "# from google.colab import auth\n",
+        "# auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ac33116d-b079-46cb-9614-86326c211e00"
+      },
+      "source": [
+        "**4. Service account or other**\n",
+        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "52c100d7-172f-4578-a3bf-f6e6d193ee6b"
+      },
+      "source": [
+        "### Create a Cloud Storage bucket\n",
+        "\n",
+        "Create a storage bucket to store intermediate artifacts such as datasets."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "638142db-edea-47c2-a4be-f43e6ff4c6f0"
+      },
+      "outputs": [],
+      "source": [
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bf01f385-3c69-45ca-b72b-84fb45b15f25"
+      },
+      "source": [
+        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "a0eabc6a-c7cd-4964-8d27-a020a526b0d9"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "42eb59a2-a959-4cf2-b312-77c726baa361"
+      },
+      "source": [
+        "#### Service Account\n",
+        "\n",
+        "You use a service account to create Vertex AI Pipeline jobs."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b357b4b6-eb6b-404e-b5a8-2771519ce569"
+      },
+      "outputs": [],
+      "source": [
+        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0f85a0f2-4ffe-4679-9a23-bc3c755e896a"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "IS_COLAB = \"google.colab\" in sys.modules\n",
+        "if (\n",
+        "    SERVICE_ACCOUNT == \"\"\n",
+        "    or SERVICE_ACCOUNT is None\n",
+        "    or SERVICE_ACCOUNT == \"[your-service-account]\"\n",
+        "):\n",
+        "    # Get your service account from gcloud\n",
+        "    if not IS_COLAB:\n",
+        "        shell_output = !gcloud auth list 2>/dev/null\n",
+        "        SERVICE_ACCOUNT = shell_output[2].replace(\"*\", \"\").strip()\n",
+        "\n",
+        "    else:  # IS_COLAB:\n",
+        "        shell_output = ! gcloud projects describe $PROJECT_ID --format=\"value(projectNumber)\"\n",
+        "        project_number = shell_output[0]\n",
+        "        SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "\n",
+        "    print(\"Service Account:\", SERVICE_ACCOUNT)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f64a1440-c23c-4b7b-9b8a-e8a4ce600265"
+      },
+      "source": [
+        "#### Set service account access for Vertex AI Pipelines\n",
+        "\n",
+        "Run the following commands to grant your service account access to read and write pipeline artifacts in the bucket that you created in the previous step. You only need to run this step once per service account."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "072a6ea1-c58a-4f9b-9569-6d68631dfed5"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectCreator $BUCKET_URI\n",
+        "\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectViewer $BUCKET_URI"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e6a924d0-a034-4e53-b240-03d356c7b7a6"
+      },
+      "source": [
+        "### Import libraries and define constants"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "463729ba-ec3c-4302-95bf-80207b0f9e2d"
+      },
+      "outputs": [],
+      "source": [
+        "import google.cloud.aiplatform as aiplatform\n",
+        "from vertexai.preview.language_models import (TextGenerationModel,\n",
+        "                                              TuningEvaluationSpec)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "a522acfe-d0b6-4b4e-b201-0a4ccf59b133"
+      },
+      "source": [
+        "## Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c845aca6-4f72-4d3b-b9ed-de4a18fcbbf8"
+      },
+      "outputs": [],
+      "source": [
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "7ec1bf44-dc64-47e0-9bd7-c2d5fc3d0851"
+      },
+      "source": [
+        "### Load pretrained model\n",
+        "\n",
+        "Load the pretrained text-bison model from Vertex AI LLM Model Garden."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "22c4fec6-9c10-4f08-80b1-b6e457453103"
+      },
+      "outputs": [],
+      "source": [
+        "model = TextGenerationModel.from_pretrained(\"google/text-bison@001\")\n",
+        "\n",
+        "model.list_tuned_model_names()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8f9a48e3-8464-41de-93d7-c451824d0ece"
+      },
+      "source": [
+        "### Tune the model\n",
+        "\n",
+        "Next, you tune the model using the `tune_model()` method, with the following parameters:\n",
+        "\n",
+        "`training_data`: A pandas Dataframe or Cloud Storage location of the training data for tuning the model.<br>\n",
+        "`learning_rate_multiplier`: A multiplier to apply to the recommended learning rate. To use the recommended learning rate, use a multiple of 1.0. <br>\n",
+        "`train_steps`: The number of steps to run for model tuning. The batch size varies by tuning location:<br>\n",
+        "- us-central1 has a batch size of 8.\n",
+        "- europe-west4 has a batch size of 24.<br>\n",
+        "\n",
+        "If there are 240 examples in a training dataset, in europe-west4, it takes 240 / 24 = 10 steps to process the entire dataset once. In us-central1, it takes 240 / 8 = 30 steps to process the entire dataset once. The default value is 300.<br>\n",
+        "\n",
+        "`tuning_job_location`: The region where the tuning job should be run. Supported regions are: `us-central1` and `europe-west4`.<br>\n",
+        "`tuned_model_location`: The region where the tuned model should be deployed. Only us-central1 is currently supported"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "8ab391f7-9229-491e-9eff-6e80e9e1d2a7"
+      },
+      "outputs": [],
+      "source": [
+        "tuning_evaluation_spec = TuningEvaluationSpec(\n",
+        "    evaluation_data=\"gs://cloud-samples-data/vertex-ai/model-evaluation/peft_eval_sample.jsonl\",\n",
+        "    evaluation_interval=20,\n",
+        "    enable_early_stopping=True,\n",
+        ")\n",
+        "model.tune_model(\n",
+        "    training_data=\"gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl\",\n",
+        "    # LR multiplier set to 1.0 and training steps set to 100 for fast iteration and demo purposes.\n",
+        "    # Generally speaking, avoid large training steps which may lead to catastrophic forgetting.\n",
+        "    train_steps=100,\n",
+        "    learning_rate_multiplier=1.0,\n",
+        "    tuning_job_location=\"europe-west4\",\n",
+        "    tuned_model_location=REGION,\n",
+        "    model_display_name=\"test_model\",\n",
+        "    tuning_evaluation_spec=tuning_evaluation_spec,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5680557f-67bd-4e8c-a383-02ab655246c5"
+      },
+      "source": [
+        "### Make a prediction with Vertex AI LLM\n",
+        "\n",
+        "Now, make a prediction using the `predict()` method from the Vertex AI LLM interface."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "r_0HHwrj96f8"
+      },
+      "outputs": [],
+      "source": [
+        "prompt = \"TRANSCRIPT: \\nPROCEDURE PERFORMED: , Umbilical hernia repair.,PROCEDURE:,  After informed consent was obtained, the patient was brought to the operative suite and placed supine on the operating table.  The patient was sedated, and an adequate local anesthetic was administered using 1% lidocaine without epinephrine.  The patient was prepped and draped in the usual sterile manner.,A standard curvilinear umbilical incision was made, and dissection was carried down to the hernia sac using a combination of Metzenbaum scissors and Bovie electrocautery.  The sac was cleared of overlying adherent tissue, and the fascial defect was delineated.  The fascia was cleared of any adherent tissue for a distance of 1.5 cm from the defect.  The sac was then placed into the abdominal cavity and the defect was closed primarily using simple interrupted 0 Vicryl sutures.  The umbilicus was then re-formed using 4-0 Vicryl to tack the umbilical skin to the fascia.,The wound was then irrigated using sterile saline, and hemostasis was obtained using Bovie electrocautery.  The skin was approximated with 4-0 Vicryl in a subcuticular fashion.  The skin was prepped with benzoin, and Steri-Strips were applied.  A dressing was then applied.  All surgical counts were reported as correct.,Having tolerated the procedure well, the patient was subsequently taken to the recovery room in good and stable condition.\\n\\n LABEL: \""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b8c48d2b-bca0-44f3-96ab-0fdde26dd2a9"
+      },
+      "outputs": [],
+      "source": [
+        "print(model.predict(prompt))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ae39e95e-9553-4cd8-98cc-602a0d70e940"
+      },
+      "source": [
+        "### Get the deployed Vertex AI Endpoint resource\n",
+        "\n",
+        "Next, get the Vertex AI Endpoint resource that the model was automatically deployed to."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ade86880-a392-4db6-8152-8e58b50d376b"
+      },
+      "outputs": [],
+      "source": [
+        "endpoint = aiplatform.Endpoint(model._endpoint.resource_name)\n",
+        "print(endpoint)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9795902e-e124-48f3-9951-02deb6f85fff"
+      },
+      "source": [
+        "### Make a prediction using Vertex AI Prediction\n",
+        "\n",
+        "Now, make a prediction using the `predict()` method from the Vertex AI Prediction interface, with the following parameters:\n",
+        "\n",
+        "- `instances`: A list of one or more instances for prediction. Each instance has the format:\n",
+        "    - { \"content\": the_text_input }\n",
+        "- `parameters`: Parameters passed to the model for the model's predict method. The corresponding examples is default values."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "35615fa2-c923-4bfd-b194-5265bbe81ce1"
+      },
+      "outputs": [],
+      "source": [
+        "endpoint.predict(\n",
+        "    instances=[{\"prompt\": prompt}],\n",
+        "    parameters={\n",
+        "        \"temperature\": 0.0,\n",
+        "        \"maxDecodeSteps\": 128,\n",
+        "        \"topP\": 0.95,\n",
+        "        \"topK\": 40,\n",
+        "    },\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0xlMGasOaBJ6"
+      },
+      "source": [
+        "### Run Post-Tuning Evaluation\n",
+        "\n",
+        "Note that the format of `ground_truth_data` should be a JSONL file where each line is a json of the following format:\n",
+        "\n",
+        "```\n",
+        "{\n",
+        "  \"prompt\": \"your input/prompt text\",\n",
+        "  \"ground_truth\": \"your ground truth output text\"\n",
+        "}\n",
+        "```\n",
+        "\n",
+        "* \"prompt\" corresponds to the \"input_text\" in the train dataset. This is needed for batch prediction\n",
+        "* \"ground_truth\" corresponds to the \"output_text\" in the train dataset. This is needed for evaluation.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "G55d36wwaI6X"
+      },
+      "outputs": [],
+      "source": [
+        "from vertexai.preview.language_models import EvaluationTextGenerationSpec\n",
+        "\n",
+        "tuned_model = model\n",
+        "\n",
+        "# Uncomment the following to load a tuned model if the tuning session is broken\n",
+        "# tuned_model = TextGenerationModel.from_pretrained(\"google/text-bison@001\")\n",
+        "# tuned_model.get_tuned_model(f'projects/{PROJECT_ID}/locations/us-central1/models/3890975937629519872')\n",
+        "\n",
+        "# Text generation example\n",
+        "evaluation_task_spec = EvaluationTextGenerationSpec(\n",
+        "    ground_truth_data=[\n",
+        "        \"gs://cloud-samples-data/vertex-ai/model-evaluation/peft_test_sample.jsonl\"\n",
+        "    ]\n",
+        ")\n",
+        "\n",
+        "tuned_model.evaluate(task_spec=evaluation_task_spec)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "99c3c746-4f85-4fd9-8467-d5017477c012"
+      },
+      "source": [
+        "## Cleaning up\n",
+        "\n",
+        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
+        "\n",
+        "Otherwise, you can delete the individual resources you created in this tutorial."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0f467b0a-07c5-4c49-b6cd-7588c9e3985b"
+      },
+      "outputs": [],
+      "source": [
+        "delete_bucket = False\n",
+        "\n",
+        "endpoint.undeploy_all()\n",
+        "endpoint.delete()\n",
+        "\n",
+        "vertex_model.delete()\n",
+        "\n",
+        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "    ! gsutil rm -rf {BUCKET_URI}"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "tune_peft.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}