Update run_linter.sh

Fix spelling error `notebooked` → `notebooks`

.cloud-build/execute_changed_notebooks_helper.py

@@ -238,7 +238,7 @@ def _get_notebook_python_version(notebook_path: str) -> str:
 
             # Look for the python version specification pattern
             re_match = re.search(
-                "python version = (\d\.\d)", markdown, flags=re.IGNORECASE
+                "python version = (\d+\.\d+)", markdown, flags=re.IGNORECASE
             )
             if re_match:
                 # get the version number

.github/workflows/linter/requirements.txt

@@ -2,9 +2,9 @@ git+https://github.com/tensorflow/docs
 ipython
 jupyter
 nbconvert
-black==23.12.1
-pyupgrade==3.15.0
+black==24.4.2
+pyupgrade==3.16.0
 isort==5.13.2
-flake8==7.0.0
-nbqa==1.7.1
+flake8==7.1.0
+nbqa==1.8.5
 

.github/workflows/linter/run_linter.sh

@@ -58,7 +58,7 @@ done
 # Only check notebooks in test folders modified in this pull request.
 # Note: Use process substitution to persist the data in the array
 if [ ${#notebooks[@]} -eq 0 ]; then
-    echo "Checking for changed notebooked using git"
+    echo "Checking for changed notebooks using git"
     while read -r file || [ -n "$line" ]; do
         notebooks+=("$file")
     done < <(git diff --name-only main... | grep '\.ipynb$')

README.md

@@ -1,37 +1,176 @@
-# Google Cloud Vertex AI Samples
+# ![Google Cloud](https://avatars.githubusercontent.com/u/2810941?s=60&v=4) Google Cloud Vertex AI Samples
 
-[![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](LICENSE)
-
-Welcome to the Google Cloud [Vertex AI](https://cloud.google.com/vertex-ai/docs/) sample repository.
+This repository contains notebooks, code samples, sample apps, and other resources that demonstrate how to use, develop and manage machine learning and generative AI workflows using Google Cloud Vertex AI.
 
 ## Overview
 
-The repository contains [notebooks](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/master/notebooks) and [community content](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/master/community-content) that demonstrate how to develop and manage ML workflows using Google Cloud Vertex AI.
+[Vertex AI](https://cloud.google.com/vertex-ai) is a fully-managed, unified AI development platform for building and using generative AI. This repository is designed to help you get started with Vertex AI. Whether you're new to Vertex AI or an experienced ML practitioner, you'll find valuable resources here.
+
+For more Vertex AI Generative AI notebook samples, please visit the Vertex AI [Generative AI](https://github.com/GoogleCloudPlatform/generative-ai) GitHub repository.
+
+## Explore, learn and contribute
+
+You can explore, learn, and contribute to this repository to unleash the full potential of machine learning on Vertex AI! 
+
+### Explore and learn
+
+Explore this repository, follow the links in the header section of each of the notebooks to -
+
+![Colab](https://cloud.google.com/ml-engine/images/colab-logo-32px.png)  Open and run the notebook in [Colab](https://colab.google/)\
+![Colab Enterprise](https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png)  Open and run the notebook in [Colab Enterprise](https://cloud.google.com/colab/docs/introduction)\
+![Workbench](https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32)  Open and run the notebook in [Vertex AI Workbench](https://cloud.google.com/vertex-ai/docs/workbench/introduction)\
+![Github](https://cloud.google.com/ml-engine/images/github-logo-32px.png)  View the notebook on Github
+
+### Contribute
+
+See the [Contributing Guide](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/CONTRIBUTING.md).
+
+## Get started
+
+To get started using Vertex AI, you must have a Google Cloud project. 
+
+- If you don't have a Google Cloud project, you can learn and build on GCP for free using [Free Trail](https://cloud.google.com/free).
+- Once you have a Google Cloud project, you can learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment).
+
 
 ## Repository structure
 
 ```bash
-├── community-content - Sample code and tutorials contributed by the community
 ├── notebooks
-│   ├── community - Notebooks contributed by the community
 │   ├── official - Notebooks demonstrating use of each Vertex AI service
 │   │   ├── automl
 │   │   ├── custom
 │   │   ├── ...
+│   ├── community - Notebooks contributed by the community
+│   │   ├── model_garden
+│   │   ├── ...
+├── community-content - Sample code and tutorials contributed by the community
+
 ```
+## Examples
 
-## Contributing
+<!-- markdownlint-disable MD033 -->
+<table>
 
-Contributions welcome! See the [Contributing Guide](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/CONTRIBUTING.md).
+  <tr>
+    <th style="text-align: center;">Category</th>
+    <th style="text-align: center;">Product</th>
+    <th style="text-align: center;">Description</th>
+  </tr>
+  <tr>
+    <td>Model</td>
+    <td>
+      <a href="notebooks/community/model_garden"><code>Model Garden/</code></a>
+    </td>
+    <td>
+      Curated collection of first-party, open-source, and third-party models available on Vertex AI including Gemini, Gemma, Llama 3, Claude 3 and many more.
+    </td>
+  </tr>
+  <tr>
+    <td>Data</td>
+    <td>
+      <a href="notebooks/official/feature_store"><code>Feature Store/</code></a>
+    </td>
+    <td>
+      Set up and manage online serving using Vertex AI Feature Store.
+    </td>
+  </tr>
+  <tr>
+    <td></td>
+    <td>
+      <a href="notebooks/official/datasets"><code>datasets/</code></a>
+    </td>
+    <td>
+      Use BigQuery and Data Labeling service with Vertex AI.
+    </td>
+  </tr>
+  <tr>
+    <td>Model development</td>
+    <td>
+      <a href="notebooks/official/automl"><code>automl/</code></a>
+    </td>
+    <td>
+      Train and make predictions on AutoML models
+    </td>
+  </tr>
+  <tr>
+    <td></td>
+    <td>
+      <a href="notebooks/official/custom"><code>custom/</code></a>
+    </td>
+    <td>
+      Create, deploy and serve custom  models on Vertex AI
+    </td>
+  </tr>
+  <tr>
+    <td></td>
+    <td>
+      <a href="notebooks/official/ray_on_vertex_ai"><code>ray_on_vertex_ai/</code></a>
+    </td>
+    <td>
+      Use Colab Enterprise and Vertex AI SDK for Python to connect to the Ray Cluster.
+    </td>
+  </tr>
+  <tr>
+    <td>Deploy and use</td>
+    <td>
+      <a href="notebooks/official/prediction"><code>prediction/</code></a>
+    </td>
+    <td>
+      Build, train and deploy models using prebuilt containers for custom training and prediction.
+    </td>
+  </tr>
+  <tr>
+    <td></td>
+    <td>
+      <a href="notebooks/official/model_registry"><code>model_registry/</code></a>
+    </td>
+    <td>
+      Use Model Registry to create and register a model.
+    </td>
+  </tr>
+  <tr>
+    <td></td>
+    <td>
+      <a href="notebooks/official/explainable_ai"><code>Explainable AI/</code></a>
+    </td>
+    <td>
+      Use Vertex Explainable AI's feature-based and example-based explanations to explain how or why a model produced a specific prediction.
+    </td>
+  </tr>
+  <tr>
+    <td></td>
+    <td>
+      <a href="notebooks/official/ml_metadata"><code>ml_metadata/</code></a>
+    </td>
+    <td>
+      Record the metadata and artifacts and query that metadata to help analyze, debug, and audit the performance of your ML system.
+    </td>
+  </tr>
+  <tr>
+    <td>Tools</td>
+    <td>
+      <a href="notebooks/official/pipelines"><code>Pipelines/</code></a>
+    </td>
+    <td>
+      Use `Vertex AI Pipelines` and `Google Cloud Pipeline Components` to build, tune, or deploy a custom model.
+    </td>
+  </tr>
+</table>
+<!-- markdownlint-enable MD033 -->
 
-## Getting help
 
-Please use the [issues page](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues) to provide feedback or submit a bug report.
+## Get help
+
+Please use the [Issues page](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues) to provide feedback or submit a bug report.
 
 ## Disclaimer
 
 This is not an officially supported Google product. The code in this repository is for demonstrative purposes only.
 
-## Feedback
 
-Please feel free to fill out our [survey](https://bit.ly/vertex-ai-samples-survey) to give us feedback on the repo and its content.
+## References
+- [Vertex AI Jupyter Notebook tutorials](https://cloud.google.com/vertex-ai/docs/tutorials/jupyter-notebooks)
+- Vertex AI [Generative AI](https://github.com/GoogleCloudPlatform/generative-ai) GitHub repository
+- [Vertex AI documentaton](https://cloud.google.com/vertex-ai/docs)
+  

community-content/CODEOWNERS

@@ -10,6 +10,7 @@
 /pipeline_components @Ark-kun
 /pipeline_components/image_ml_model_training @lakeyk
 /prediction_featurestore_integration @googleapis/vertex-prediction-team
+/vertex_model_garden/model_oss/notebook_util @minwoo33park
 /vertex_model_garden/model_oss/util @weigary
 /vertex_model_garden/model_oss/diffusers @weigary
 /vertex_model_garden/model_oss/keras @dstnluong-google
@@ -23,6 +24,8 @@
 /vertex_model_garden/model_oss/tfvision @dstnluong-google
 /vertex_model_garden/model_oss/fvlm @minwoo33park
 /vertex_model_garden/model_oss/imagebind @kathyyu-google
+/vertex_model_garden/model_oss/llava @py4
+/vertex_model_garden/model_oss/vllm @kathyyu-google
 /vertex_model_garden/benchmarking_reports @lavraicse
-
+/vertex_model_garden/model_oss/autogluon @lavraicse
 

community-content/cpr-examples/timm_serving/requirements.txt

@@ -1,5 +1,5 @@
 absl-py==1.1.0
-fastapi==0.75.2
+fastapi==0.109.1
 uvicorn==0.18.2
 timm==0.5.4
 smart_open==6.0.0

community-content/pytorch_image_classification_single_gpu_with_vertex_sdk_and_torchserve/trainer/requirements.txt

@@ -1,3 +1,3 @@
-torch==1.13.1
+torch==2.2.0
 torchvision==0.9.1
 tensorboard==2.5.0

community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/logger/requirements.txt

@@ -1,4 +1,4 @@
 google-cloud-bigquery==2.20.0
 tensorflow==2.7.2
-pillow==10.0.1
+pillow==10.3.0
 tf-agents==0.8.0

community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/prediction_container/requirements.txt

@@ -1,4 +1,4 @@
 google-cloud-pubsub==2.5.0
-pillow==10.0.1
+pillow==10.3.0
 tf-agents==0.8.0
 tensorflow==2.7.2

community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/simulator/requirements.txt

@@ -1,5 +1,5 @@
 dataclasses==0.6
 google-cloud-aiplatform==1.8.1
 tensorflow==2.7.2
-pillow==10.0.1
+pillow==10.3.0
 tf-agents==0.8.0

community-content/vertex_cpr_samples/README.md

@@ -0,0 +1,15 @@
+# Vertex AI custom prediction routines samples
+
+## Overview
+Vertex Custom Prediction Routines(CPR) simplify the process of building custom containers 
+and make local model testing easy. Here are the sameple codes for different libraries.
+
+
+### Objectives
+The objective is to provide various samples for Vertex Custom Prediction Routine(CPR).
+
+
+### Supporting libraries
+* torch
+* sklearn
+* xgboost

community-content/vertex_cpr_samples/torch/predictor_transformers.py

@@ -0,0 +1,73 @@
+import ast
+import json
+import os
+import pickle
+import torch
+
+from google.cloud.aiplatform.utils import prediction_utils
+from google.cloud.aiplatform.prediction.predictor import Predictor
+from transformers import AutoModelForQuestionAnswering
+from typing import Dict, List
+
+class TorchTransformersPredictor(Predictor):
+
+    def __init__(self):
+        return
+
+    def load(self, artifacts_uri: str) -> None:
+        prediction_utils.download_model_artifacts(artifacts_uri)
+        
+        if os.path.isfile("setup_config.json"):
+            with open("setup_config.json") as setup_config_file:
+                self.setup_config = json.load(setup_config_file)
+
+        if os.path.exists("model.pt"):
+            self.model = AutoModelForQuestionAnswering.from_pretrained("model.pt")
+            self.model.eval()
+        else:
+            raise ValueError("One of the following model files must be provided: model.pt.")
+
+    def preprocess(self, prediction_input: dict) -> torch.Tensor:
+        max_length = self.setup_config["max_length"]
+        instances = prediction_input["instances"]
+        question_context = ast.literal_eval(instances)
+        question = question_context["question"]
+        context = question_context["context"]
+        inputs = self.tokenizer.encode_plus(
+            question,
+            context,
+            max_length=int(max_length),
+            pad_to_max_length=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        input_ids = inputs["input_ids"]
+        attention_mask = inputs["attention_mask"]
+        return torch.Tensor(input_ids, attention_mask)
+    
+    @torch.inference_mode()
+    def predict(self, instances: torch.Tensor) -> List[str]:
+        input_ids, attention_mask = instances
+        outputs = self._model(input_ids, attention_mask)
+        answer_start_scores = outputs.start_logits
+        answer_end_scores = outputs.end_logits
+        
+        num_rows, num_cols = answer_start_scores.shape
+        inferences = []
+        for i in range(num_rows):
+            answer_start_scores_one_seq = answer_start_scores[i].unsqueeze(0)
+            answer_start = torch.argmax(answer_start_scores_one_seq)
+            answer_end_scores_one_seq = answer_end_scores[i].unsqueeze(0)
+            answer_end = torch.argmax(answer_end_scores_one_seq) + 1
+            prediction = self.tokenizer.convert_tokens_to_string(
+                self.tokenizer.convert_ids_to_tokens(
+                    input_ids[i].tolist()[answer_start:answer_end]
+                )
+            )
+            inferences.append(prediction)
+        return inferences
+
+    def postprocess(self, prediction_results: List[str]) -> Dict:
+        return {"predictions": prediction_results}
+     
+

community-content/vertex_model_garden/benchmarking_reports/pytorch_openllama_7b_finetune_benchmark_report.md

@@ -1,227 +0,0 @@
-# Benchmark report on fine tuning the OpenLLaMA 7B model on Google Cloud Vertex Model Garden
-
-Gary Wei, Software Engineer, Google Cloud
-Dustin Luong, Software Engineer, Google Cloud
-Changyu Zhu, Software Engineer, Google Cloud
-Genquan Duan, Software Engineer, Google Cloud
-
-## Introduction
-
-Fine-tuning of LLMs can be non-trivial to find an optimal configuration of
-machine types, training parameters, and other hyperparameters that achieves a
-good balance between cost efficiency and model performance. To facilitate users
-in conducting tuning experiments, this report benchmarks OpenLLaMA 7B
-fine-tuning on Google Cloud Vertex Model Garden, demonstrating both efficiency
-and effectiveness. The observations are general and can be applied to other LLM
-models.
-
-We benchmarked fine tuning algorithms [LoRA](https://arxiv.org/abs/2106.09685)
-and [QLoRA](https://arxiv.org/abs/2305.14314) supported by
-[huggingface PEFT libraries](https://github.com/huggingface/peft). LoRA, short
-for Low-Rank Adaptation of Large Language Models, is an improved fine tuning
-method where instead of fine tuning all the weights that constitute the weight
-matrix of the pre-trained large language model, two smaller matrices that
-approximate this larger matrix are fine-tuned. QLoRA is an even more
-memory-efficient version of LoRA, where the pretrained model is loaded to GPU
-memory as quantized 4-bit weights, while preserving similar effectiveness to
-LoRA. We also provide simple scripts and parameter settings to reproduce the
-results reported in this report.
-
-In general, there are many factors that affect the performance of fine-tuning
-experiments, such as hardware settings, parameters, cost, and accuracy. It is
-impractical to obtain benchmarks for all possible combinations of these factors.
-Instead, we focus on tuning a subset of related parameters and evaluating their
-impact on a set of chosen metrics. The evaluation metrics are GPU memory usage,
-percentage of parameters tuned, tuning speed, cost, and accuracy. The tuning
-parameters are batch size, lora rank, maximum sequence length, and maximum
-training steps.
-
-## Key takeaways
-
--   **Use QLoRA to minimize the peak GPU requirements**: The QLoRA can
-    significantly reduce the peak GPU memory usage by ~75% compared to LoRA. For
-    OpenLLaMA7b, the peak memory is ~28G for LoRA and ~7G for QLoRA.
--   **Use LoRA to maximize the tuning speed and minimize the tuning cost**: LoRA
-    is ~66% faster than QLoRA in fine tuning speed. LoRA/QLoRA tuning cost is
-    low generally, while LoRA is even ~40% cheaper than QLoRA with the same
-    parameters. Suggest to use QLoRA for limited GPU memories, and LoRA for
-    limited training budgets. For OpenLLaMA7b, the tuning speed for LoRA/QLoRA
-    ~5 samples / 3 samples per second, and the tuning cost for LoRA/QLoRA in 500
-    steps is ~$1/$1.7 on `a2-highgpu-1g` with 1 A100 40G GPU. The tuning cost
-    for QLoRA in 500 steps is $6.75 on n1-standard-8 with 1 V100 GPU, while LoRA
-    could not run because of OOM.
--   **Use QLoRA to tune models with large sequence lengths**. For OpenLLaMA7b,
-    the max sequence length for QLoRA can be 2048 when consuming 16.3G GPU,
-    while the max sequence length for LoRA is 512 when consuming 28.2G GPU, and
-    encounter OOM when max sequence length is 1024.
--   **Both LoRA and QLoRA give similar accuracy improvement after fine tuning.**
-    For OpenLLaMA7b, both LoRA/QLoRA can improve the average accuracy by ~4%
-    evaluating on 3 typical tasks (ARC challenge, HellaSwag and TruthfulQA),
-    after training 1875 steps on dataset
-    [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco).
--   **Use a big batch size if GPU memory is not a constraint**. For OpenLLaMA7b
-    with other default parameters, we suggest using a batch size as 24 for
-    QLoRA, but 2 for LoRA when tuning with 1 A100 40G. We also suggest using a
-    batch size as 8 for QLoRA when tuning with 1 V100. Tuning with LoRA and
-    batch size as 1 got OOM and we don't recommend tuning LoRA with 1 V100.
-
-## Benchmark Details
-
-### Experiment Setup
-
-The benchmark dataset is
-[timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco).
-The training dataset is directly downloaded from hugging face to the VM, before
-every experiment.
-
-The default tuning parameters during benchmark are:
-
--   Host VM: a2-highgpu-1g
--   Accelerator type: 1 A100 40G
--   batch size: 2
--   lora_rank: 16
--   max_seq_length: 512
--   precision_mode: float16
--   max_train_steps: 500
-
-For simplicity, we set the precision mode to `float16` when tuning LoRA models,
-and set the precision to `4bit` for QLoRA.
-
-Sample script to start fine tuning dockers in a VM on GCP.
-
-```shell
-IMAGE_TAG=us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:latest
-docker run  --runtime=nvidia  -e NVIDIA_VISIBLE_DEVICES=0 \
---rm --name "test_gpu"  -it --pull=always ${IMAGE_TAG} \
---task=instruct-lora \
---pretrained_model_id=openlm-research/open_llama_7b \
---dataset_name="timdettmers/openassistant-guanaco"  \
---instruct_column_in_dataset="text" \
---precision_mode="float16" \
---output_dir=<OUTPUT DIR>  \
---lora_rank=2 \
---max_sequence_length=512 \
---learning_rate=2e-4 \
---max_steps=50
-```
-
-### GPU Memory
-
-In this benchmark, we investigated the impact of batch size, lora rank, and
-maximum sequence length on GPU memory, and then made recommendations on the
-maximum batch size for different GPUs.
-
-#### Peak GPU memory by batch size (GB)
-
-<img src="images/openllama_7b_fine_tune_benchmark_report/openllama-7b-peak-gpu-vs-batch-size.png" width="600">
-
--   The QLoRA can significantly reduce the peak GPU memory usage by ~75%
-    compared to LoRA. The peak GPU memory is ~28G for LoRA and ~7G for QLoRA
-    when batch size is 2.
--   QLoRA can support much larger batch sizes than LoRA
-    -   We can use a batch size as 32 for QLoRA, but only 2 for LoRA on 1 A100
-        40G.
-    -   We can use a batch size of 8 for QLoRA on 1 V100 GPU. LoRA will fail
-        with OOM even with a batch size of 1.
-
-#### Peak GPU memory by LoRA rank (GB)
-
-<img src="images/openllama_7b_fine_tune_benchmark_report/openllama-7b-peak-gpu-vs-lora-rank.png" width="600">
-
--   Peak GPU memories are quite similar for different LoRA ranks for both
-    LoRA/QLoRA.
--   The peak GPU memory increasing percentages are very small generally when
-    LoRA rank increases.
-    -   The peak GPU memory increases from 28G with LoRA rank 4 to 29.09G with
-        LoRA rank 64, and the increasing percentage is only ~3.9%.
-
-#### Peak GPU memory by max sequence length for LoRA/QLoRA (GB)
-
-<img src="images/openllama_7b_fine_tune_benchmark_report/openllama-7b-peak-gpu-vs-max-seq-length.png" width="600">
-
--   The peak GPU increases quickly when max sequence length increases for both
-    LoRA/QLoRA, and the increasing rate of LoRA is much faster than QLoRA.
-    -   For LoRA tuning, the GPU memory increased from 20.5G (max sequence
-        length=256) to 28.2G (max sequence length=512), an increase of ~37%.
-    -   For QLoRA tuning, the GPU memory increased from 6.94G (max sequence
-        length=256) to 7.57G (max sequence length=512), an increase of ~9%.
--   The max sequence length for QLoRA can be 2048 when consuming 16.3G GPU,
-    while the max sequence length for LoRA is 512 when consuming 28.2G GPU, and
-    encounter OOM when max sequence length is 1024.
-
-### Fine Tuning Parameters
-
-This section shows the number/percentage of trainable parameters, and the sizes
-of the fine tuned models. LoRA and QLoRA differ only in how they represent the
-precision of their parameters. The total number of parameters and the number of
-trainable parameters are the same for both methods.
-
-| LoRA Rank | Finetuned parameters | Total parameters | Trainable Parameter Percentage | Fine tuned model size (MB) |
-| --------- | -------------------- | ---------------- | ------------------------------ | -------------------------- |
-| 8         | 2.00E+07             | 6.76E+09         | 0.3%                           | 76.4                       |
-| 16        | 4.00E+07             | 6.78E+09         | 0.6%                           | 152.65                     |
-| 32        | 8.00E+07             | 6.82E+09         | 1.2%                           | 305.15                     |
-| 64        | 1.60E+08             | 6.90E+09         | 2.3%                           | 610.15                     |
-
-
-LoRA/QLoRA tunes quite a small fraction (only 0.3% with LoRA rank=8) of all
-parameters, and the tuned models are very small (only 76.4MB with LoRA rank=8).
-
-### Fine Tuning Speed And Costs
-
-The fine-tuning speed and cost are affected by various factors, such as the
-GPUs, LoRA ranks, and max sequence lengths.
-
--   LoRA is ~66% faster than QLoRA in fine tuning speed. The tuning speed for
-    LoRA/QLoRA ~5 samples / 3 samples per second on 1 A100 40G GPU
--   Higher LoRA ranks, slower tuning speed for both LoRA/QLoRA.
-    -   LoRA tuning speed reduces from ~5 samples per second with LoRA rank as 8
-        to ~4 samples per second with LoRA rank as 64, slowed down by 20%.
-    -   QLoRA tuning speed reduces from ~3 samples per second with LoRA rank as
-        8 to ~2.5 samples per second with LoRA rank as 64, slowed down by 17%.
-
-<img src="images/openllama_7b_fine_tune_benchmark_report/openllama-7b-tune-speed-vs-lora-rank.png" width="600">
-
--   Longer sequence lengths, slower tuning speed.
-    -   LoRA tuning speed reduces from ~5.56 samples per second with max
-        sequence length as 256 to ~4.84 samples per second with max sequence
-        length as 512 slowed down by 13%.
-    -   LoRA tuning speed reduces from ~2.95 samples per second with max
-        sequence length as 256 to ~2.88 samples per second with max sequence
-        length as 512 slowed down by ~2.4%.
-
-<img src="images/openllama_7b_fine_tune_benchmark_report/openllama-7b-tune-speed-lora-qlora.png" width="600">
-
--   LoRA/QLoRA tuning cost is low generally, while LoRA is even ~40% cheaper
-    than QLoRA with the same parameters.
-    -   The LoRA/QLoRA fine tuning cost for 500 steps is ~$1/$1.7 on 1 A100 40G.
-    -   The tuning cost for QLoRA in 500 steps is $6.75 on n1-standard-8 with 1
-        V100 GPU, while LoRA could not run because of OOM.
-
-<img src="images/openllama_7b_fine_tune_benchmark_report/openllama-7b-tune-cost-lora-qlora.png" width="600">
-
-### Accuracy
-
-We fine tuned Open Llama 7B model with
-[timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco),
-and report accuracy similar to the
-[HuggingFace leaderboard](https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard)
-using
-[Eleuther AI Language Model Evaluation Harness](https://github.com/EleutherAI/lm-evaluation-harness).
-[HuggingFace leaderboard](https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard)
-mainly compares models on ARC, HellaSwag, MMLU, and TruthfulQA. The authors did
-not publish OpenLLaMA 7B on MMLU
-([link](https://huggingface.co/openlm-research/open_llama_7b)). Therefore, we
-only benchmark accuracies on ARC, HellaSwag, and TruthfulQA.
-
-|                                                              | Mean | ARC  | HellaSwag | TruthfulQA | Tuning Parameters                                            |
-| ------------------------------------------------------------ | ---- | ---- | --------- | ---------- | ------------------------------------------------------------ |
-| OpenLLaMA7B ([Original Report](https://huggingface.co/openlm-research/open_llama_7b)) | 0.49 | 0.41 | 0.73      | 0.34       | n/a                                                          |
-| OpenLLaMA7B ([Re-run with lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness)) | 0.51 | 0.47 | 0.72      | 0.35       | n/a                                                          |
-| OpenLLaMA7B+LoRA                                             | 0.56 | 0.48 | 0.74      | 0.45       | LoRA Rank=16; Max Sequence Length=512;Learning Rate=1e-4; Train steps=1875 |
-| OpenLLaMA7B+QLoRA                                            | 0.53 | 0.45 | 0.73      | 0.42       | LoRA Rank=16; Max Sequence Length=512; Learning Rate=1e-4; Train steps=1875 |
-
--   The base OpenLLaMA7B model gets better performance (2%) when using the
-    [Eleuther AI Language Model Evaluation Harness](https://github.com/EleutherAI/lm-evaluation-harness).
--   LoRA/QLoRA can improve the performance by ~2-4% when trained for 1875 steps
-    with learning rate 1e-4.

community-content/vertex_model_garden/model_oss/autogluon/dockerfile/autogluon_serve.Dockerfile

@@ -0,0 +1,50 @@
+# Dockerfile for serving dockers with AutoGluon.
+#
+# To build:
+# docker build -f model_oss/autogluon/dockerfile/serve.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+FROM pytorch/pytorch:2.1.2-cuda11.8-cudnn8-runtime
+
+USER root
+
+# AutoGluon might require libgomp for some dependencies.
+RUN apt-get update && apt-get -y upgrade && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim \
+        libgomp1
+
+# Install AutoGluon and other dependencies.
+RUN pip install --upgrade pip
+RUN pip install autogluon==1.0.0
+RUN pip install flask==3.0.0
+
+# Dependencies needed to work with GCS.
+RUN pip install absl-py==2.0.0
+RUN pip install google-cloud-storage==2.7.0
+
+# Copy scripts into the container.
+COPY model_oss/autogluon /autogluon
+COPY model_oss/util /autogluon/util
+WORKDIR /autogluon
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+RUN wget https://github.com/pallets/flask/blob/main/LICENSE.rst
+
+# Expose the port the app runs on.
+EXPOSE 8501
+
+# Set the working directory to a specific path for consistency.
+WORKDIR /autogluon
+
+# Change to a non-root user for security purposes.
+RUN useradd -m autogluonuser
+USER autogluonuser
+
+# Run Flask application.
+CMD ["python", "serve.py"]

community-content/vertex_model_garden/model_oss/autogluon/dockerfile/autogluon_train.Dockerfile

@@ -0,0 +1,36 @@
+# Dockerfile for training dockers with Autogluon.
+#
+# To build:
+# docker build -f model_oss/autogluon/dockerfile/train.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+FROM pytorch/pytorch:2.1.2-cuda11.8-cudnn8-runtime
+
+# Install tools.
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PIP_ROOT_USER_ACTION=ignore
+
+RUN apt-get update && apt-get -y upgrade && apt-get install -y --no-install-recommends \
+apt-utils \
+curl \
+wget \
+git \
+jq \
+gnupg \
+build-essential \
+tesseract-ocr \
+vim
+
+# Install libraries.
+RUN pip install autogluon==1.0.0
+
+COPY model_oss/autogluon /autogluon
+WORKDIR /autogluon
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+ENTRYPOINT ["python", "train.py"]

community-content/vertex_model_garden/model_oss/autogluon/serve.py

@@ -0,0 +1,87 @@
+r"""AutoGluon serving binary.
+
+This module sets up a Flask web server for serving predictions from a
+trained AutoGluon model. The server exposes two endpoints:
+
+1.  `/ping`: A health check endpoint that returns "pong" to
+    indicate that the server is running.
+2.  `/predict`: An endpoint that accepts POST requests with JSON content.
+    Each request should contain one or more instances for which the
+    predictions are desired. The endpoint returns the predictions and
+    associated probabilities in a JSON response.
+
+The server expects an environment variable `model_path` that points to
+the directory where the AutoGluon model artifacts are
+stored. If `model_path` is not provided, it defaults to '/autogluon/models'.
+"""
+
+import json
+import logging
+import os
+
+from autogluon.tabular import TabularPredictor
+import flask
+import pandas as pd
+
+from util import constants
+from util import fileutils
+
+_SUCCESS_STATUS = 200
+_ERROR_STATUS = 500
+_PORT = 8501
+
+app = flask.Flask(__name__)
+# Check the environment variables.
+model_dir = os.getenv('model_path', '/autogluon/models')
+logging.info('Model directory passed by the user is: %s', model_dir)
+# If the model is on GCS then copy it to a local folder first.
+if model_dir.startswith(constants.GCS_URI_PREFIX):
+  gcs_path = model_dir[len(constants.GCS_URI_PREFIX) :]
+  local_model_dir = os.path.join(constants.LOCAL_MODEL_DIR, gcs_path)
+  logging.info('Download %s to %s', model_dir, local_model_dir)
+  fileutils.download_gcs_dir_to_local(model_dir, local_model_dir)
+  model_dir = local_model_dir
+  logging.info('Local model directory is: %s', model_dir)
+
+
+# Load the predictor at startup.
+predictor = TabularPredictor.load(model_dir)
+
+
+@app.route('/ping', methods=['GET'])
+def ping() -> flask.Response:
+  """Health check route."""
+  return flask.Response('pong', status=_SUCCESS_STATUS)
+
+
+@app.route('/predict', methods=['POST'])
+def predict() -> flask.Response:
+  """Prediction route."""
+  try:
+    # Extract JSON content from the POST request.
+    data = flask.request.get_json(force=True)
+    instances = data.get('instances', [])
+
+    # Convert instances to DataFrame.
+    df_to_predict = pd.DataFrame(instances)
+
+    # Perform prediction.
+    predictions = predictor.predict(df_to_predict).tolist()
+    response = {'predictions': predictions}
+
+    return flask.Response(
+        json.dumps(response),
+        status=_SUCCESS_STATUS,
+        mimetype='application/json',
+    )
+
+  except Exception as e:  # pylint: disable=broad-exception-caught
+    return flask.Response(
+        json.dumps({'error': str(e)}),
+        status=_ERROR_STATUS,
+        mimetype='application/json',
+    )
+
+
+if __name__ == '__main__':
+  app.run(host='0.0.0.0', port=_PORT)

community-content/vertex_model_garden/model_oss/autogluon/train.py

@@ -0,0 +1,144 @@
+"""AutoGluon training binary. """
+
+import argparse
+import json
+from typing import Any
+
+from autogluon.tabular import TabularPredictor
+import pandas as pd
+
+
+class BaseConfig:
+
+  def to_dict(self) -> dict[str, Any]:
+    return {
+        key: value for key, value in self.__dict__.items() if value is not None
+    }
+
+
+class DataConfig(BaseConfig):
+
+  def __init__(self, train_data_path: Any) -> None:
+    self.train_data_path = train_data_path
+
+
+class ProblemConfig(BaseConfig):
+
+  def __init__(self, label: Any, problem_type: Any) -> None:
+    self.label = label
+    self.problem_type = problem_type
+
+
+class EvaluationConfig(BaseConfig):
+
+  def __init__(self, eval_metric: Any) -> None:
+    self.eval_metric = eval_metric
+
+
+class TrainingConfig(BaseConfig):
+  """Config for training."""
+
+  def __init__(
+      self,
+      time_limit: Any,
+      presets: Any,
+      hyperparameters: Any,
+      model_save_path: str,
+  ) -> None:
+    self.time_limit = time_limit
+    self.hyperparameters = hyperparameters
+    self.presets = presets
+    self.model_save_path = model_save_path
+
+
+def parse_args() -> (
+    tuple[DataConfig, ProblemConfig, EvaluationConfig, TrainingConfig]
+):
+  """Parse command line arguments."""
+  parser = argparse.ArgumentParser(description="AutoGluon Tabular Predictor")
+  # Add arguments for each config class
+  parser.add_argument(
+      "--train_data_path",
+      type=str,
+      required=True,
+      help="Path to the input data CSV file.",
+  )
+  parser.add_argument(
+      "--label", type=str, required=True, help="Target variable column name."
+  )
+  parser.add_argument(
+      "--problem_type",
+      type=str,
+      choices=["binary", "multiclass", "regression", "quantile"],
+      default=None,
+      help="Problem type.",
+  )
+  parser.add_argument(
+      "--eval_metric", type=str, default=None, help="Evaluation metric to use."
+  )
+  # Add arguments for TrainingConfig if needed
+  parser.add_argument(
+      "--time_limit",
+      type=int,
+      default=None,
+      help="Time limit in seconds for training.",
+  )
+  parser.add_argument(
+      "--presets",
+      type=str,
+      default="medium_quality",
+      help="Presets used for training ",
+  )
+  parser.add_argument(
+      "--hyperparameters",
+      type=json.loads,
+      default=None,
+      help="Hyperparameter dictionary in JSON format.",
+  )
+  parser.add_argument(
+      "--model_save_path",
+      type=str,
+      default=None,
+      help="Path to save the trained model.",
+  )
+  args = parser.parse_args()
+
+  data_config = DataConfig(train_data_path=args.train_data_path)
+  problem_config = ProblemConfig(
+      label=args.label, problem_type=args.problem_type
+  )
+  eval_config = EvaluationConfig(eval_metric=args.eval_metric)
+  training_config = TrainingConfig(
+      time_limit=args.time_limit,
+      presets=args.presets,
+      hyperparameters=args.hyperparameters,
+      model_save_path=args.model_save_path,
+  )
+
+  return data_config, problem_config, eval_config, training_config
+
+
+def main() -> None:
+  data_config, problem_config, eval_config, training_config = parse_args()
+
+  # Load the training data.
+  data = pd.read_csv(data_config.train_data_path)
+
+  # Create a TabularPredictor.
+  predictor = TabularPredictor(
+      label=problem_config.label,
+      eval_metric=eval_config.eval_metric,
+      path=training_config.model_save_path,
+  )
+
+  # Fit the model
+  predictor.fit(
+      data,
+      presets=training_config.presets,
+      time_limit=training_config.time_limit,
+      hyperparameters=training_config.hyperparameters,
+  )
+
+
+if __name__ == "__main__":
+  main()

community-content/vertex_model_garden/model_oss/cloudnerf/configs/360.gin

@@ -0,0 +1,25 @@
+# The provided content is a configuration file for the ZipNeRF
+# PyTorch implementation.
+
+# Sets the name of the experiment to 'test'.
+Config.exp_name = 'test'
+# Specifies the dataset loader, in this case, 'llff' for light field.
+Config.dataset_loader = 'llff'
+# Defines the near and far clipping planes for the camera view.
+Config.near = 0.2
+Config.far = 1e6
+# Image downsampling.
+Config.factor = 4
+
+# For the model configurations.
+Model.raydist_fn = 'power_transformation'
+Model.opaque_background = True
+
+# Disables the computation of density normals and RGB values, and sets
+# the grid level dimension to 1 for PropMLP.
+PropMLP.disable_density_normals = True
+PropMLP.disable_rgb = True
+PropMLP.grid_level_dim = 1
+
+# Disable density normals for NerfMLP
+NerfMLP.disable_density_normals = True

community-content/vertex_model_garden/model_oss/cloudnerf/configs/360_glo.gin

@@ -0,0 +1,21 @@
+# The provided content is a configuration file for Generative
+# Latent Optimization (GLO) vectors in the Pytorch implemnetation of ZipNeRF.
+
+# Specifies the dataset loader, in this case, 'llff' for light field.
+Config.dataset_loader = 'llff'
+# Defines the near and far clipping planes for the camera view.
+Config.near = 0.2
+Config.far = 1e6
+# Image downsampling.
+Config.factor = 4
+
+# For the model configurations.
+Model.raydist_fn = 'power_transformation'
+Model.num_glo_features = 128
+Model.opaque_background = True
+
+PropMLP.disable_density_normals = True
+PropMLP.disable_rgb = True
+PropMLP.grid_level_dim = 1
+
+NerfMLP.disable_density_normals = True

community-content/vertex_model_garden/model_oss/cloudnerf/configs/accelerate_config.yaml

@@ -0,0 +1,18 @@
+# The provided content is a configuration file running ZipNeRF
+# training on 8 gpu machine.
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+gpu_ids: all
+machine_rank: 0
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

community-content/vertex_model_garden/model_oss/cloudnerf/dockerfile/pytorch_cloudnerf_base.Dockerfile

@@ -0,0 +1,120 @@
+# Dockerfile for ZipNeRF base image.
+#
+# To build:
+# docker build -f model_oss/cloudnerf/dockerfile/pytorch_cloudnerf_base.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+FROM pytorch/pytorch:2.1.0-cuda11.8-cudnn8-devel
+
+USER root
+
+ARG COLMAP_GIT_COMMIT=main
+ARG CUDA_ARCHITECTURES=60;70;75;80;86
+
+# Prevent stop building ubuntu at time zone selection.
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update -y --allow-releaseinfo-change && apt-get -y upgrade && apt-get install -y --no-install-recommends \
+        curl \
+        g++ \
+        wget \
+        vim \
+        bash \
+        cmake \
+        imagemagick \
+        ninja-build \
+        build-essential \
+        libboost-program-options-dev \
+        libboost-filesystem-dev \
+        libboost-graph-dev \
+        libboost-system-dev \
+        libeigen3-dev \
+        libflann-dev \
+        libfreeimage-dev \
+        libmetis-dev \
+        libgoogle-glog-dev \
+        libgtest-dev \
+        libsqlite3-dev \
+        libglew-dev \
+        qtbase5-dev \
+        libqt5opengl5-dev \
+        libcgal-dev \
+        libceres-dev \
+        git \
+        git-lfs \
+        python3-cffi \
+        python3-cryptography \
+        libffi-dev \
+        python-dev
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Install google cloud CLI.
+RUN wget -q https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-cli-430.0.0-linux-x86.tar.gz
+RUN tar xzf google-cloud-cli-430.0.0-linux-x86.tar.gz
+RUN ./google-cloud-sdk/install.sh -q
+# Make sure gsutil will use the default service account.
+RUN echo '[GoogleCompute]\nservice_account = default' > /etc/boto.cfg
+
+# Install deps and install gsutil.
+RUN pip install gsutil==5.27
+
+# When building colmap in colab, the link error "undefined reference.
+# to '_glapi_tls_Current'" happens. A solution is to install "libglvnd"
+# as described in this page https://github.com/colmap/colmap/issues/1271.
+RUN git clone --depth 1 --branch v1.7.0 https://github.com/NVIDIA/libglvnd && \
+    apt-get install -y libxext-dev libx11-dev x11proto-gl-dev && \
+    cd libglvnd/  && \
+    apt-get install -y autoconf automake libtool && \
+    apt-get install -y libffi-dev && \
+    ./autogen.sh && \
+    ./configure && \
+    make  -j4 && \
+    make install
+
+RUN apt remove nvidia-cuda-toolkit -y \
+    nvidia-cuda-toolkit \
+    nvidia-cuda-toolkit-gcc
+
+# Install libraries.
+ENV PIP_ROOT_USER_ACTION=ignore
+RUN python3 -m pip install --upgrade pip
+
+ENV CUDA_HOME=/usr/local/cuda
+
+RUN git clone --branch main https://github.com/SuLvXiangXin/zipnerf-pytorch.git
+# Set current directory to the downloaded 'zipnerf-pytorch' repository.
+WORKDIR ./zipnerf-pytorch
+# Using git reset command to pin it down to a specific version.
+RUN git reset --hard 4de3d21ebb9e15412d36951b56e2d713fddd812b
+COPY model_oss/cloudnerf/requirements.txt requirements.txt
+RUN pip install -r requirements.txt
+
+# Install gridencoder extensions and nvdiffrast (for textured mesh).
+RUN cd .. && \
+    TORCH_CUDA_ARCH_LIST="6.0 7.0 7.5 8.0 8.6+PTX" CXX=g++ pip install ./zipnerf-pytorch/gridencoder
+
+# Install cuda version of torch_scatter.
+RUN pip install torch-scatter==2.1.2 -f https://data.pyg.org/whl/torch-2.0.1+cu118.html
+RUN pip install google-cloud-aiplatform==1.25.0
+RUN pip install google-cloud-storage==2.9.0
+
+# Build and install COLMAP.
+RUN git clone  --depth 1 --branch 3.8 https://github.com/colmap/colmap.git
+RUN cd colmap && \
+    git fetch https://github.com/colmap/colmap.git ${COLMAP_GIT_COMMIT} && \
+    mkdir build && \
+    cd build && \
+    cmake .. -GNinja -DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHITECTURES} && \
+    ninja && \
+    ninja install && \
+    cd .. && rm -rf colmap
+
+RUN git clone  --depth 1 --branch v1.0.2 https://github.com/dranjan/python-plyfile.git
+
+RUN sed -i "20 i\sys.path.append('/workspace/zipnerf-pytorch/internal/pycolmap')" /workspace/zipnerf-pytorch/internal/datasets.py
+RUN sed -i "21 i\sys.path.append('/workspace/zipnerf-pytorch/internal/pycolmap/pycolmap')" /workspace/zipnerf-pytorch/internal/datasets.py

community-content/vertex_model_garden/model_oss/cloudnerf/dockerfile/pytorch_cloudnerf_calibrate.Dockerfile

@@ -0,0 +1,16 @@
+# Dockerfile for ZipNeRF COLMAP image calibration.
+#
+# To build:
+# docker build -f model_oss/cloudnerf/dockerfile/cloudnerf_pytorch_calibrate.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+FROM us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-cloudnerf-base:20231206_0923_RC00
+
+COPY model_oss/cloudnerf/local_colmap_and_resize.sh /workspace/zipnerf-pytorch/scripts/local_colmap_and_resize.sh
+
+WORKDIR /workspace/zipnerf-pytorch/
+
+ENTRYPOINT ["bash","scripts/local_colmap_and_resize.sh"]

community-content/vertex_model_garden/model_oss/cloudnerf/dockerfile/pytorch_cloudnerf_render.Dockerfile

@@ -0,0 +1,22 @@
+# Dockerfile for ZipNeRF rendering.
+#
+# To build:
+# docker build -f model_oss/cloudnerf/dockerfile/pytorch_cloudnerf_render.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+FROM us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-cloudnerf-base:20231206_0923_RC00
+
+COPY model_oss/cloudnerf/render.sh /workspace/zipnerf-pytorch/scripts/render.sh
+COPY model_oss/cloudnerf/configs/360.gin /workspace/zipnerf-pytorch/configs/360.gin
+COPY model_oss/cloudnerf/configs/360_glo.gin /workspace/zipnerf-pytorch/configs/360_glo.gin
+COPY model_oss/cloudnerf/configs/accelerate_config.yaml /root/.cache/huggingface/accelerate/default_config.yaml
+RUN sed -i '324s/.*/            keyframe_names = fp.read().splitlines()/' /workspace/zipnerf-pytorch/internal/camera_utils.py
+
+ENV PYTHONPATH "${PYTHONPATH}:/workspace/zipnerf-pytorch/util"
+
+WORKDIR /workspace/zipnerf-pytorch/
+
+ENTRYPOINT ["bash", "scripts/render.sh"]

community-content/vertex_model_garden/model_oss/cloudnerf/dockerfile/pytorch_cloudnerf_train.Dockerfile

@@ -0,0 +1,21 @@
+# Dockerfile for ZipNeRF training.
+#
+# To build:
+# docker build -f model_oss/cloudnerf/dockerfile/pytorch_cloudnerf_train.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+FROM us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-cloudnerf-base:20231206_0923_RC00
+
+COPY model_oss/cloudnerf/train.sh /workspace/zipnerf-pytorch/scripts/train.sh
+COPY model_oss/cloudnerf/configs/360.gin /workspace/zipnerf-pytorch/configs/360.gin
+COPY model_oss/cloudnerf/configs/360_glo.gin /workspace/zipnerf-pytorch/configs/360_glo.gin
+COPY model_oss/cloudnerf/configs/accelerate_config.yaml /root/.cache/huggingface/accelerate/default_config.yaml
+
+ENV PYTHONPATH "${PYTHONPATH}:/workspace/zipnerf-pytorch/util"
+
+WORKDIR /workspace/zipnerf-pytorch/
+
+ENTRYPOINT ["bash", "scripts/train.sh"]

community-content/vertex_model_garden/model_oss/cloudnerf/local_colmap_and_resize.sh

@@ -0,0 +1,144 @@
+#!/bin/bash
+# This script runs colmap for scale invariant feature (SIFT) extraction and
+# matching to map camera extrinsics and intrinsics values for ZipNeRF,
+# given a folder of images and videos
+# from a GCS bucket. It uses ffmepg to extract an image from a video at
+# 1fps. The folder can contain images or videos. If both images and videos
+# are present, the extracted frames from the videos is added to the images
+# to create the final combined image dataset.
+# vv-docker:google3-begin(internal)
+# TODO(b/314042136): Specify cloudnerf colmap fps.
+# vv-docker:google3-end
+
+# Initialize variables.
+use_gpu=1  # Default to 1 (assuming the docker is run on a machine with GPU)
+gcs_dataset_path=""
+gcs_experiment_path=""
+camera=""
+
+# This loop processes command-line arguments for configuring the container.
+# It supports arguments for GPU usage, dataset and experiment paths,
+# and camera type.
+while [[ $# -gt 0 ]]; do
+  case $1 in
+    -use_gpu)
+      use_gpu="$2"
+      if ! [[ $use_gpu =~ ^[0-9]+$ ]]; then
+        echo "Error: -use_gpu must be an integer."
+        exit 1
+      fi
+      shift # past argument
+      shift # past value
+      ;;
+    -gcs_dataset_path)
+      gcs_dataset_path="$2"
+      shift # past argument
+      shift # past value
+      ;;
+    -gcs_experiment_path)
+      gcs_experiment_path="$2"
+      shift # past argument
+      shift # past value
+      ;;
+    -camera)
+      camera="$2"
+      if [[ $camera != "OPENCV" && $camera != "OPENCV_FISHEYE" ]]; then
+        echo "Error: -camera must be either 'OPENCV' or 'OPENCV_FISHEYE'."
+        exit 1
+      fi
+      shift # past argument
+      shift # past value
+      ;;
+    *)    # unknown option
+      echo "Unknown option: $1" >&2
+      exit 1
+      ;;
+  esac
+done
+
+local_folder="dataset_content"
+images_folder="dataset_images"
+images_subfolder="images"
+output_folder="$images_folder/$images_subfolder"
+
+# Create the local folder if it doesn't exist
+mkdir -p "$local_folder"
+mkdir -p "$output_folder"
+
+# Download the content from the GCS URI
+gsutil -m cp -r "$gcs_dataset_path"/* "$local_folder/"
+
+# Process files in the local folder
+for file in "$local_folder"/*; do
+  if [[ -f "$file" ]]; then
+    # Check if the file is an image (e.g., jpg, png, etc.)
+    if file --mime-type "$file" | grep -q "image"; then
+      # Copy the image to the "images" subfolder within the "dataset_images" folder
+      cp "$file" "$output_folder/$(basename "$file")"
+    elif file --mime-type "$file" | grep -q "video"; then
+      # Use FFmpeg to extract an image every 30 frames from the video
+      ffmpeg -i "$file" -vf "select='not(mod(n,30))'" "$output_folder/$(basename "$file" ."${file##*.}")_%03d.jpg"
+    else
+      echo "Skipping unsupported file: $file"
+    fi
+  fi
+done
+
+# Run COLMAP Feature extraction
+colmap feature_extractor \
+    --database_path "$local_folder"/database.db \
+    --image_path "$output_folder" \
+    --ImageReader.single_camera 1 \
+    --ImageReader.camera_model "$camera" \
+    --SiftExtraction.use_gpu "$use_gpu"
+
+# Run COLMAP Feature matching
+colmap exhaustive_matcher \
+    --database_path "$local_folder"/database.db \
+    --SiftMatching.use_gpu "$use_gpu"
+
+# Bundle adjustment. The default Mapper tolerance is unnecessarily large,
+# decreasing it speeds up bundle adjustment steps.
+mkdir -p "$local_folder"/sparse
+colmap mapper \
+    --database_path "$local_folder"/database.db \
+    --image_path "$output_folder" \
+    --output_path "$local_folder"/sparse \
+    --Mapper.ba_global_function_tolerance=0.000001
+
+# Downsample images at 1/2, 1/4, 1/8 scales. Save feature matching to
+# sqlite database.
+# All input and output images:
+#   $gcs_dataset_path
+#   $gcs_experiment_path/data/images
+# Downsampled output images:
+#   $gcs_experiment_path/data/images_2/
+#   $gcs_experiment_path/data/images_4/
+#   $gcs_experiment_path/data/images_8/
+# COLMAP sparse reconstruction files: project.ini, images.bin,
+# cameras.bin, points3D.bin
+#   $gcs_experiment_path/data/sparse/0/
+cp -r "$output_folder" "$images_folder"/images_2
+pushd "$images_folder"/images_2
+ls | xargs -P 8 -I {} mogrify -resize 50% {}
+popd
+gsutil -m cp -r "$images_folder"/images_2/*  "$gcs_experiment_path"/data/images_2
+
+cp -r "$output_folder" "$images_folder"/images_4
+pushd "$images_folder"/images_4
+ls | xargs -P 8 -I {} mogrify -resize 25% {}
+popd
+gsutil -m cp -r "$images_folder"/images_4/*  "$gcs_experiment_path"/data/images_4
+
+cp -r "$output_folder" "$images_folder"/images_8
+pushd "$images_folder"/images_8
+ls | xargs -P 8 -I {} mogrify -resize 12.5% {}
+popd
+gsutil -m cp "$images_folder"/images_8/*  "$gcs_experiment_path"/data/images_8
+
+# Copy images and sparse reconstruction files to gcs experiment folder.
+gsutil -m cp "$images_folder"/images/*  "$gcs_experiment_path"/data/images
+gsutil -m cp -r "$local_folder"/sparse  "$gcs_experiment_path"/data
+gsutil -m cp "$local_folder"/database.db  "$gcs_experiment_path"/data
+
+echo "Processing complete."

community-content/vertex_model_garden/model_oss/cloudnerf/render.sh

@@ -0,0 +1,117 @@
+#!/bin/bash
+# This script runs rendering for ZipNeRF given an experiment folder
+# from a GCS bucket with colmap dataset.
+
+# Initialize associative array for arguments.
+declare -A args
+
+# vv-docker:google3-begin(internal)
+# TODO(b/311468174): Pass gin config file from gcs bucket.
+# vv-docker:google3-end
+# Function to parse named arguments.
+parse_args() {
+  while [[ $# -gt 0 ]]; do
+    key="$1"
+    case $key in
+      -gcs_experiment_path|-gin_config_file|-gcs_keyframes_file)
+        args[$key]="$2"
+        shift # past argument
+        shift # past value
+        ;;
+      -training_job_name)
+        training_job_name="$2"
+        shift # past argument
+        shift # past value
+        ;;
+      -rendering_job_name)
+        rendering_job_name="$2"
+        shift # past argument
+        shift # past value
+        ;;
+      -render_path_frames|-factor|-render_video_fps)
+        args[$key]="$2"
+        if ! [[ ${args[$key]} =~ ^[0-9]+$ ]]; then
+          echo "Error: $key must be an integer."
+          exit 1
+        fi
+        shift # past argument
+        shift # past value
+        ;;
+      *)
+        echo "Unknown option: $1" >&2
+        exit 1
+        ;;
+    esac
+  done
+}
+
+# Function to create a directory if it doesn't exist.
+create_dir_if_not_exists() {
+  local dir_path=$1
+  if [[ ! -d "$dir_path" ]]; then
+    echo "Creating folder: $dir_path"
+    mkdir "$dir_path"
+  else
+    echo "Folder $dir_path already exists."
+  fi
+}
+
+# Function to launch rendering.
+launch_rendering() {
+  local keyframes_file=$1
+  local render_bindings=(
+    "--gin_configs=${args[-gin_config_file]}"
+    "--gin_bindings=Config.data_dir='${DATASET_PATH}'"
+    "--gin_bindings=Config.exp_name='${EXPERIMENT}'"
+    "--gin_bindings=Config.render_path=True"
+    "--gin_bindings=Config.render_path_frames=${args[-render_path_frames]}"
+    "--gin_bindings=Config.render_video_fps=${args[-render_video_fps]}"
+    "--gin_bindings=Config.factor=${args[-factor]}"
+  )
+
+  if [[ -n $keyframes_file ]]; then
+    render_bindings+=("--gin_bindings=Config.render_spline_keyframes='${keyframes_file}'")
+  fi
+
+  accelerate launch render.py "${render_bindings[@]}"
+}
+
+# Parse arguments.
+parse_args "$@"
+
+# Extract folder names and paths.
+scene_folder_name=$(basename "${args[-gcs_experiment_path]}")
+local_dataset_path="local_dataset"
+local_experiment_path="exp"
+exp_folder_name=$(basename "${args[-gcs_experiment_path]}")
+DATASET_PATH="$local_experiment_path/$exp_folder_name/data"
+CHECKPOINTS_PATH="$local_experiment_path/$exp_folder_name/checkpoints"
+OUTPUT_RENDER_PATH="$local_experiment_path/$scene_folder_name/render"
+EXPERIMENT=$exp_folder_name
+
+# Create necessary directories.
+create_dir_if_not_exists "$local_dataset_path"
+create_dir_if_not_exists "$local_experiment_path"
+create_dir_if_not_exists "$local_experiment_path/$exp_folder_name"
+create_dir_if_not_exists "$CHECKPOINTS_PATH"
+
+# Create the file log_render.txt in the exp folder.
+touch "$local_experiment_path/$exp_folder_name/log_render.txt"
+
+# Copy experiment from GCS bucket to local
+gsutil -m cp -r "${args[-gcs_experiment_path]}/data" "$local_experiment_path/$exp_folder_name" || exit 1
+gsutil -m cp -r "${args[-gcs_experiment_path]}/checkpoints/${training_job_name}/*" "$CHECKPOINTS_PATH" || exit 1
+
+# Check and copy keyframes file.
+if [[ -n ${args[-gcs_keyframes_file]} ]]; then
+  keyframes_file_basename=$(basename "${args[-gcs_keyframes_file]}")
+  local_keyframes_file="$local_dataset_path/$keyframes_file_basename"
+  gsutil cp "${args[-gcs_keyframes_file]}" "$local_keyframes_file" || exit 1
+  echo "Local keyframe file: $local_keyframes_file"
+  launch_rendering "$local_keyframes_file"
+else
+  launch_rendering ""
+fi
+
+# Copy rendered data back to GCS.
+gsutil -m cp -r "$OUTPUT_RENDER_PATH" "${args[-gcs_experiment_path]}/render/${rendering_job_name}"

community-content/vertex_model_garden/model_oss/cloudnerf/requirements.txt

@@ -0,0 +1,24 @@
+--find-links https://download.pytorch.org/whl/torch_stable.html
+
+torch==2.0.1+cu118
+numpy==1.26.1
+absl_py==2.0.0
+accelerate==0.24.0
+gin_config==0.5.0
+imageio==2.31.6
+imageio-ffmpeg==0.4.9
+matplotlib==3.8.0
+mediapy==1.1.9
+ninja==1.11.1.1
+opencv_contrib_python==4.8.1.78
+opencv_python==4.8.1.78
+Pillow==10.3.0
+rawpy==0.18.1
+scipy==1.11.3
+scikit-image==0.22.0
+scikit-learn==1.5.0
+tensorboard==2.15.0
+tensorboardX==2.6.2.2
+tqdm==4.66.3
+trimesh==4.0.1
+xatlas==0.0.8

community-content/vertex_model_garden/model_oss/cloudnerf/train.sh

@@ -0,0 +1,94 @@
+#!/bin/bash
+
+# Initialize variables.
+training_job_name=""
+gcs_experiment_path=""
+gin_config_file="configs/360.gin"
+factor=4
+max_training_steps=25000
+
+# Parse named arguments.
+while [[ $# -gt 0 ]]; do
+  case $1 in
+    -training_job_name)
+      training_job_name="$2"
+      shift # past argument
+      shift # past value
+      ;;
+    -gcs_experiment_path)
+      gcs_experiment_path="$2"
+      shift # past argument
+      shift # past value
+      ;;
+    -gin_config_file)
+      gin_config_file="$2"
+      shift # past argument
+      shift # past value
+      ;;
+    -factor)
+      factor="$2"
+      if ! [[ $factor =~ ^[0-9]+$ ]]; then
+        echo "Error: -factor must be an integer."
+        exit 1
+      fi
+      shift # past argument
+      shift # past value
+      ;;
+    -max_training_steps)
+      max_training_steps="$2"
+      if ! [[ $max_training_steps =~ ^[0-9]+$ ]]; then
+        echo "Error: -max_training_steps must be an integer."
+        exit 1
+      fi
+      shift # past argument
+      shift # past value
+      ;;
+    *)    # unknown option
+      echo "Unknown option: $1" >&2
+      exit 1
+      ;;
+  esac
+done
+
+# Function to create a directory if it doesn't exist.
+create_dir_if_not_exists() {
+  local dir_path=$1
+  if [[ ! -d "$dir_path" ]]; then
+    echo "Creating folder: $dir_path"
+    mkdir "$dir_path"
+  else
+    echo "Folder $dir_path already exists."
+  fi
+}
+
+# Extract folder names and paths.
+scene_folder_name=$(basename "${gcs_experiment_path}")
+local_dataset_path="local_dataset"
+local_experiment_path="exp"
+DATASET_PATH="$local_experiment_path/$scene_folder_name/data"
+EXPERIMENT=$scene_folder_name
+
+# Create necessary directories.
+create_dir_if_not_exists "$local_dataset_path"
+create_dir_if_not_exists "$local_experiment_path"
+create_dir_if_not_exists "$local_experiment_path/$scene_folder_name"
+
+# Copy experiment from GCS bucket to local.
+gsutil -m cp -r "${gcs_experiment_path}/data" "$local_experiment_path/$scene_folder_name" || exit 1
+
+echo "GCS Experiment: $gcs_experiment_path"
+echo "Gin Config File: $gin_config_file"
+echo "Factor: $factor"
+echo "Scene: $scene_folder_name"
+echo "Local Dataset: $DATASET_PATH"
+echo "Local Experiment: $EXPERIMENT"
+
+accelerate launch train.py --gin_configs="$gin_config_file" \
+  --gin_bindings="Config.data_dir = '${DATASET_PATH}'" \
+  --gin_bindings="Config.exp_name = '${EXPERIMENT}'" \
+  --gin_bindings="Config.factor = ${factor}" \
+  --gin_bindings="Config.max_steps = ${max_training_steps}"
+
+gsutil -m rm -r "${gcs_experiment_path}/checkpoints/${training_job_name}"
+gsutil -m cp -r "$local_experiment_path/$scene_folder_name/config.gin" "${gcs_experiment_path}/${training_job_name}_config.gin"
+gsutil -m cp -r "$local_experiment_path/$scene_folder_name/checkpoints/*/*" "${gcs_experiment_path}/checkpoints/${training_job_name}"

community-content/vertex_model_garden/model_oss/fvlm/dockerfile/jax_fvlm_model_conversion.Dockerfile

@@ -0,0 +1,83 @@
+# This Dockerfile converts JAX vision transformer model to
+# tensorflow saved model format.
+# Here is an example to build this dockerfile:
+# PROJECT="your gcp project"
+# IMAGE_TAG="jax-f-vlm-model-conversion:${USER}-test"
+# docker build -f model_oss/fvlm/dockerfile/jax_fvlm_model_conversion.Dockerfile . -t "${IMAGE_TAG}"
+# docker tag "${IMAGE_TAG}" "gcr.io/${PROJECT}/${IMAGE_TAG}"
+# docker push "gcr.io/${PROJECT}/${IMAGE_TAG}"
+
+
+# See https://cloud.google.com/tensorflow-enterprise/docs/overview for details.
+FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-12.py310:m110
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Install basic libs
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        cmake \
+        curl \
+        wget \
+        sudo \
+        gnupg \
+        libsm6 \
+        libxext6 \
+        libxrender-dev \
+        lsb-release \
+        ca-certificates \
+        build-essential \
+        git \
+        libgl1
+
+# Copy Apache license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Install required libs
+RUN pip install --upgrade pip
+
+# Using the commit 6712c224985c694001ba8ee68697bbf4dcb32edb on Jan 4th, 2024.
+ARG COMMIT_ID=6712c224985c694001ba8ee68697bbf4dcb32edb
+RUN git clone -c \
+remote.origin.fetch=+${COMMIT_ID}:refs/remotes/origin/${COMMIT_ID} \
+https://github.com/google-research/google-research --no-checkout --progress \
+--depth 1
+WORKDIR ./google-research
+RUN git sparse-checkout init --cone
+RUN git sparse-checkout set fvlm
+RUN git checkout ${COMMIT_ID}
+
+# The following pip installs are pinned down versions satisfying
+# fvlm/requirements.txt file.
+# NOTE: Using `no-deps` flag to avoid overwriting of dependent library
+# versions. For example, both `chex` and `jax` can overwrite each other's
+# `jax-lib` version.
+# Note: The following libraries are pinned down versions of:
+# https://github.com/google-research/google-research/blob/master/fvlm/requirements.txt
+RUN pip install --no-cache-dir tensorflow==2.12.0
+RUN pip install --no-cache-dir tensorflow-datasets==4.9.2
+RUN pip install --no-cache-dir numpy==1.23.5
+RUN pip install --no-cache-dir torch==2.0.1
+RUN pip install --no-cache-dir torchvision==0.15.2
+RUN pip install --no-cache-dir opencv-python==4.7.0.72
+RUN pip install --no-cache-dir tqdm==4.65.0
+RUN pip install --no-cache-dir git+https://github.com/openai/CLIP.git@a1d071733d7111c9c014f024669f959182114e33
+RUN pip install --no-cache-dir Pillow==9.5.0
+RUN pip install --no-cache-dir orbax-checkpoint==0.3.3
+RUN pip install --no-cache-dir gin-config==0.5.0
+RUN pip install --no-cache-dir pycocotools==2.0.6
+RUN pip install --no-cache-dir contextlib2==21.6.0
+RUN pip install --no-cache-dir ml-collections==0.1.1
+RUN pip install --no-cache-dir chex==0.1.7
+RUN pip install --no-cache-dir optax==0.1.5
+# Dependencies already included. Use no-deps to not update numpy.
+RUN pip install --no-cache-dir --no-deps flax==0.7.2
+RUN pip install --no-cache-dir --no-deps clu==0.0.9
+RUN pip install --no-cache-dir jax[cuda11_cudnn86]==0.4.9 \
+--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+RUN pip install --no-cache-dir ml-dtypes==0.2.0
+RUN pip install --no-cache-dir tensorflow_text==2.12.0
+
+WORKDIR ./fvlm
+ENV PYTHONPATH ./
+
+ENTRYPOINT ["python", "export_saved_model.py"]

community-content/vertex_model_garden/model_oss/fvlm/dockerfile/jax_fvlm_train_gpu.Dockerfile

@@ -0,0 +1,78 @@
+# This Dockerfile trains the F-VLM model on GPU.
+# Here is an example to build this dockerfile:
+# PROJECT="your gcp project"
+# IMAGE_TAG="jax-f-vlm-train:${USER}-test"
+# docker build -f model_oss/fvlm/dockerfile/jax_fvlm_train_gpu.Dockerfile . -t "${IMAGE_TAG}"
+# docker tag "${IMAGE_TAG}" "gcr.io/${PROJECT}/${IMAGE_TAG}"
+# docker push "gcr.io/${PROJECT}/${IMAGE_TAG}"
+
+# See https://cloud.google.com/tensorflow-enterprise/docs/overview for details.
+FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-12.py310:m110
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Install basic libs
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        cmake \
+        curl \
+        wget \
+        sudo \
+        gnupg \
+        libsm6 \
+        libxext6 \
+        libxrender-dev \
+        lsb-release \
+        ca-certificates \
+        build-essential \
+        git
+
+# Copy Apache license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Install required libs
+RUN pip install --upgrade pip
+# The following pip installs are pinned down versions satisfying
+# fvlm/requirements.txt file.
+# Get F-VLM repository by using git sparse-checkout to avoid downloading entire
+# google-research repository.
+# Using the commit 6712c224985c694001ba8ee68697bbf4dcb32edb on Jan 4th, 2024.
+ARG COMMIT_ID=6712c224985c694001ba8ee68697bbf4dcb32edb
+RUN git clone -c \
+remote.origin.fetch=+${COMMIT_ID}:refs/remotes/origin/${COMMIT_ID} \
+https://github.com/google-research/google-research --no-checkout --progress \
+--depth 1
+WORKDIR ./google-research
+RUN git sparse-checkout init --cone
+RUN git sparse-checkout set fvlm
+RUN git checkout ${COMMIT_ID}
+
+# Note: The following libraries are pinned down versions of:
+# https://github.com/google-research/google-research/blob/master/fvlm/requirements.txt
+RUN pip install --no-cache-dir tensorflow==2.12.0
+RUN pip install --no-cache-dir tensorflow-datasets==4.9.2
+RUN pip install --no-cache-dir numpy==1.23.5
+RUN pip install --no-cache-dir torch==2.0.1
+RUN pip install --no-cache-dir torchvision==0.15.2
+RUN pip install --no-cache-dir opencv-python==4.7.0.72
+RUN pip install --no-cache-dir tqdm==4.65.0
+RUN pip install --no-cache-dir git+https://github.com/openai/CLIP.git@a1d071733d7111c9c014f024669f959182114e33
+RUN pip install --no-cache-dir Pillow==9.5.0
+RUN pip install --no-cache-dir orbax-checkpoint==0.3.3
+RUN pip install --no-cache-dir gin-config==0.5.0
+RUN pip install --no-cache-dir pycocotools==2.0.6
+RUN pip install --no-cache-dir contextlib2==21.6.0
+RUN pip install --no-cache-dir ml-collections==0.1.1
+RUN pip install --no-cache-dir chex==0.1.7
+RUN pip install --no-cache-dir optax==0.1.5
+# Dependencies already included. Use no-deps to not update numpy.
+RUN pip install --no-cache-dir --no-deps flax==0.7.2
+RUN pip install --no-cache-dir --no-deps clu==0.0.9
+# Installing jax at the very end with GPU support.
+# NOTE: Not using `no-deps` flag here because we need CUDA support.
+RUN pip install --no-cache-dir jax[cuda11_cudnn86]==0.4.9 \
+--find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+
+WORKDIR ./fvlm
+ENV PYTHONPATH ./
+
+ENTRYPOINT ["python", "train_and_eval.py"]

community-content/vertex_model_garden/model_oss/fvlm/dockerfile/jax_fvlm_train_tpu.Dockerfile

@@ -0,0 +1,138 @@
+# This Dockerfile trains the F-VLM model on TPU.
+# Here is an example to build this dockerfile:
+# PROJECT="your gcp project"
+# IMAGE_TAG="jax-f-vlm-train-tpu:${USER}-test"
+# docker build -f model_oss/fvlm/dockerfile/jax_fvlm_train_tpu.Dockerfile . -t "${IMAGE_TAG}"
+# docker tag "${IMAGE_TAG}" "gcr.io/${PROJECT}/${IMAGE_TAG}"
+# docker push "gcr.io/${PROJECT}/${IMAGE_TAG}"
+
+FROM python:3.11
+
+# Get libtpu shared library. See go/what-is-libtpu.
+RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.0/libtpu.so -o /lib/libtpu.so
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Install basic libs
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        cmake \
+        curl \
+        wget \
+        sudo \
+        gnupg \
+        libsm6 \
+        libxext6 \
+        libxrender-dev \
+        lsb-release \
+        ca-certificates \
+        build-essential \
+        git \
+        libgl1
+
+# Copy Apache license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Install required libs
+RUN pip install --upgrade pip
+
+# Get F-VLM repository by using git sparse-checkout to avoid downloading entire
+# google-research repository.
+# Using the commit 05ece4b1c97285b48b51fa44321ccb2cb347406a on Dec 11th, 2023.
+ARG COMMIT_ID=05ece4b1c97285b48b51fa44321ccb2cb347406a
+RUN git clone -c \
+remote.origin.fetch=+${COMMIT_ID}:refs/remotes/origin/${COMMIT_ID} \
+https://github.com/google-research/google-research --no-checkout --progress \
+--depth 1
+WORKDIR ./google-research
+RUN git sparse-checkout init --cone
+RUN git sparse-checkout set fvlm
+RUN git checkout ${COMMIT_ID}
+
+# Note: The following libraries are pinned down versions of:
+# https://github.com/google-research/google-research/blob/master/fvlm/requirements.txt
+RUN pip install --no-cache-dir ml_dtypes==0.3.1
+RUN pip install --no-cache-dir tensorstore==0.1.51
+RUN pip install --no-cache-dir MarkupSafe==2.1.3
+RUN pip install --no-cache-dir Pillow==9.5.0
+RUN pip install --no-cache-dir PyYAML==6.0.1
+RUN pip install --no-cache-dir absl_py==1.4.0
+RUN pip install --no-cache-dir array_record==0.4.1
+RUN pip install --no-cache-dir astunparse==1.6.3
+RUN pip install --no-cache-dir cachetools==5.3.1
+RUN pip install --no-cache-dir certifi==2023.7.22
+RUN pip install --no-cache-dir charset_normalizer==3.3.0
+RUN pip install --no-cache-dir chex==0.1.83
+RUN pip install --no-cache-dir click==8.1.7
+RUN pip install --no-cache-dir clip==0.2.0
+RUN pip install --no-cache-dir clu==0.0.9
+RUN pip install --no-cache-dir contourpy==1.1.1
+RUN pip install --no-cache-dir cycler==0.12.1
+RUN pip install --no-cache-dir dm_tree==0.1.8
+RUN pip install --no-cache-dir etils==1.5.1
+RUN pip install --no-cache-dir filelock==3.12.4
+RUN pip install --no-cache-dir flatbuffers==23.5.26
+RUN pip install --no-cache-dir flax==0.7.4
+RUN pip install --no-cache-dir fonttools==4.43.1
+RUN pip install --no-cache-dir fsspec==2023.9.2
+RUN pip install --no-cache-dir ftfy==6.1.1
+RUN pip install --no-cache-dir gast==0.5.4
+RUN pip install --no-cache-dir gin_config==0.5.0
+RUN pip install --no-cache-dir google_auth==2.23.3
+RUN pip install --no-cache-dir google_auth_oauthlib==1.0.0
+RUN pip install --no-cache-dir google_pasta==0.2.0
+RUN pip install --no-cache-dir googleapis_common_protos==1.61.0
+RUN pip install --no-cache-dir grpcio==1.59.0
+RUN pip install --no-cache-dir h5py==3.10.0
+RUN pip install --no-cache-dir importlib_resources==6.1.0
+RUN pip install --no-cache-dir 'jax[tpu]==0.4.18' \
+-f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+RUN pip install --no-cache-dir jaxlib==0.4.18
+RUN pip install --no-cache-dir jinja2==3.1.2
+RUN pip install --no-cache-dir keras==2.14.0
+RUN pip install --no-cache-dir kiwisolver==1.4.5
+RUN pip install --no-cache-dir libclang==16.0.6
+RUN pip install --no-cache-dir markdown==3.5
+RUN pip install --no-cache-dir matplotlib==3.8.0
+RUN pip install --no-cache-dir mpmath==1.3.0
+RUN pip install --no-cache-dir networkx==3.1
+RUN pip install --no-cache-dir numpy==1.26.0
+RUN pip install --no-cache-dir nvidia_cublas_cu12==12.1.3.1
+RUN pip install --no-cache-dir nvidia_cuda_cupti_cu12==12.1.105
+RUN pip install --no-cache-dir nvidia_cuda_nvrtc_cu12==12.1.105
+RUN pip install --no-cache-dir nvidia_cuda_runtime_cu12==12.1.105
+RUN pip install --no-cache-dir nvidia_cudnn_cu12==8.9.2.26
+RUN pip install --no-cache-dir nvidia_cufft_cu12==11.0.2.54
+RUN pip install --no-cache-dir nvidia_curand_cu12==10.3.2.106
+RUN pip install --no-cache-dir nvidia_cusolver_cu12==11.4.5.107
+RUN pip install --no-cache-dir nvidia_cusparse_cu12==12.1.0.106
+RUN pip install --no-cache-dir nvidia_nccl_cu12==2.18.1
+RUN pip install --no-cache-dir nvidia_nvjitlink_cu12==12.2.140
+RUN pip install --no-cache-dir nvidia_nvtx_cu12==12.1.105
+RUN pip install --no-cache-dir opencv_python==4.8.1.78
+RUN pip install --no-cache-dir orbax_checkpoint==0.4.1
+RUN pip install --no-cache-dir promise==2.3
+RUN pip install --no-cache-dir protobuf==3.20.3
+RUN pip install --no-cache-dir psutil==5.9.5
+RUN pip install --no-cache-dir pyasn1==0.5.0
+RUN pip install --no-cache-dir pycocotools==2.0.7
+RUN pip install --no-cache-dir pygments==2.16.1
+RUN pip install --no-cache-dir regex==2023.10.3
+RUN pip install --no-cache-dir rich==13.6.0
+RUN pip install --no-cache-dir scipy==1.11.3
+RUN pip install --no-cache-dir sympy==1.12
+RUN pip install --no-cache-dir tensorboard==2.14.1
+RUN pip install --no-cache-dir tensorboard_data_server==0.7.1
+RUN pip install --no-cache-dir tensorflow==2.14.0
+RUN pip install --no-cache-dir tensorflow_datasets==4.9.3
+RUN pip install --no-cache-dir torch==2.1.0
+RUN pip install --no-cache-dir torchvision==0.16.0
+RUN pip install --no-cache-dir urllib3==2.0.6
+RUN pip install --no-cache-dir wcwidth==0.2.8
+RUN pip install --no-cache-dir werkzeug==3.0.0
+RUN pip install --no-cache-dir wheel==0.41.2
+RUN pip install --no-cache-dir tensorflow_text==2.14.0
+
+WORKDIR ./fvlm
+ENV PYTHONPATH ./
+
+ENTRYPOINT ["python", "train_and_eval.py"]

community-content/vertex_model_garden/model_oss/llava/dockerfile/serve.Dockerfile

@@ -0,0 +1,71 @@
+FROM pytorch/torchserve:0.9.0-gpu
+
+USER root
+
+# Install tools.
+RUN apt-get update -y --allow-releaseinfo-change && apt-get -y upgrade && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim \
+        git
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+ENV INFER_PORT=7080
+ENV MNG_PORT=7081
+ENV MODEL_NAME="llava_serving"
+ENV PATH="/home/model-server/:${PATH}"
+ENV PATH="/usr/local/cuda-12.1/bin:${PATH}"
+ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda-12.1/lib64:$LD_LIBRARY_PATH
+ENV NVIDIA_VISIBLE_DEVICES=all
+
+# Get 'LLaVA' repository from github.
+RUN git clone https://github.com/haotian-liu/LLaVA /home/model-server/LLaVA
+WORKDIR /home/model-server/LLaVA
+# Using git reset command to pin it down to a specific version.
+RUN git reset --hard 7775b12d6b20cd69089be7a18ea02615a59621cd
+
+# Install the package.
+RUN python3 -m pip install --upgrade pip
+RUN pip install google-cloud-storage==2.13.0
+RUN pip install absl-py==2.0.0
+RUN pip install -e .
+
+# Copy model artifacts.
+COPY model_oss/llava/handler.py /home/model-server/handler.py
+COPY model_oss/llava/model_handler_setup.py /home/model-server/model_handler_setup.py
+COPY model_oss/util/ /home/model-server/util/
+ENV PYTHONPATH /home/model-server
+WORKDIR /home/model-server
+
+# Create torchserve configuration file.
+RUN echo \
+    "default_response_timeout=1800\n" \
+    "service_envelope=json\n" \
+    "inference_address=http://0.0.0.0:${INFER_PORT}\n"  \
+    "management_address=http://0.0.0.0:${MNG_PORT}\n"  \
+    "default_workers_per_model=DEFAULT_WORKERS_PER_MODEL" >> /home/model-server/config.properties
+
+# Expose ports.
+EXPOSE ${INFER_PORT}
+EXPOSE ${MNG_PORT}
+
+# Archive model artifacts and dependencies.
+# Do not set --model-file and --serialized-file because model and checkpoint will be dynamically loaded in handler.py.
+RUN torch-model-archiver \
+--model-name=${MODEL_NAME} \
+--version=1.0 \
+--handler=/home/model-server/handler.py \
+--runtime=python3 \
+--export-path=/home/model-server/model-store \
+--archive-format=default \
+--force
+
+# Run Torchserve HTTP serve to respond to prediction requests.
+# Use $NUM_GPU workers unless overriden by $TS_NUM_WORKERS
+CMD ["TOTAL=$(nvidia-smi", "--list-gpus","|","wc","-l)","&&", "TS_NUM_WORKERS=${TS_NUM_WORKERS:-$TOTAL}","&&", "sed","-i","\"s/DEFAULT_WORKERS_PER_MODEL/$TS_NUM_WORKERS/g\"","/home/model-server/config.properties", "&&", \
+    "torchserve", "--start", \
+     "--ts-config", "/home/model-server/config.properties", \
+     "--models", "${MODEL_NAME}=${MODEL_NAME}.mar", \
+     "--model-store", "/home/model-server/model-store"]

community-content/vertex_model_garden/model_oss/llava/handler.py

@@ -0,0 +1,174 @@
+"""Customer handler for LLava 1.5 OSS model.
+
+The code is based on here: https://github.com/haotian-liu/LLaVA
+handler based on:
+https://github.com/haotian-liu/LLaVA/blob/main/llava/eval/run_llava.py
+There are two supported variant:
+1. liuhaotian/llava-v1.5-13b: 13B params
+2. liuhaotian/llava-v1.5-7b: 7B params
+"""
+
+import os
+import re
+from typing import Any, Dict, List
+
+from llava import constants as llava_constants
+from llava import conversation
+from llava import mm_utils
+from llava.model import builder
+import model_handler_setup
+import torch
+from ts.torch_handler import base_handler
+
+from util import constants
+from util import image_format_converter
+
+
+DEFAULT_MODEL_ID = "liuhaotian/llava-v1.5-7b"
+
+
+class LlavaHandler(base_handler.BaseHandler):
+  """Custom handler for LLava model."""
+
+  def initialize(self, context: Any):
+    """Initializes model, tokenizer, and other components."""
+    self.map_location = model_handler_setup.get_map_location(context=context)
+    self.device = model_handler_setup.get_model_device(
+        map_location=self.map_location, context=context
+    )
+    self.manifest = context.manifest
+    self.model_id = model_handler_setup.get_model_id(
+        default_model_id=DEFAULT_MODEL_ID
+    )
+
+    # Allows 4bit and 8bit quantiziation using BnB nf4.
+    precision = os.environ.get("PRECISION_MODE")
+    load_8bit = precision == constants.PRECISION_MODE_8
+    load_4bit = precision == constants.PRECISION_MODE_4
+
+    self.tokenizer, self.model, self.image_processor, self.context_len = (
+        builder.load_pretrained_model(
+            model_path=self.model_id,
+            model_base=None,
+            model_name=mm_utils.get_model_name_from_path(self.model_id),
+            load_8bit=load_8bit,
+            load_4bit=load_4bit,
+        )
+    )
+
+  def preprocess(self, data: List[Dict[str, Any]]) -> Any:
+    """Runs the preprocessing to tokenize image and the prompt."""
+    if len(data) > 1:
+      raise ValueError(
+          "LLava original repo currently does not support batch inference."
+          " https://github.com/haotian-liu/LLaVA/issues/754"
+      )
+    data = data[0]
+    prompt, base64_image = data["prompt"], data["base64_image"]
+
+    # Adds proper image token to the prompt.
+    image_token_se = (
+        llava_constants.DEFAULT_IM_START_TOKEN
+        + llava_constants.DEFAULT_IMAGE_TOKEN
+        + llava_constants.DEFAULT_IM_END_TOKEN
+    )
+    if llava_constants.IMAGE_PLACEHOLDER in prompt:
+      if self.model.config.mm_use_im_start_end:
+        prompt = re.sub(
+            llava_constants.IMAGE_PLACEHOLDER, image_token_se, prompt
+        )
+      else:
+        prompt = re.sub(
+            llava_constants.IMAGE_PLACEHOLDER,
+            llava_constants.DEFAULT_IMAGE_TOKEN,
+            prompt,
+        )
+    else:
+      if self.model.config.mm_use_im_start_end:
+        prompt = image_token_se + "\n" + prompt
+      else:
+        prompt = llava_constants.DEFAULT_IMAGE_TOKEN + "\n" + prompt
+
+    # Formats the prompt as a conversation to be fed to the model.
+    conv = conversation.conv_llava_v1.copy()
+    conv.append_message(role=conv.roles[0], message=prompt)
+    conv.append_message(role=conv.roles[1], message=None)
+    prompt = conv.get_prompt()
+
+    # Tokenizes the prompt that includes special image token as well.
+    input_ids = (
+        mm_utils.tokenizer_image_token(
+            prompt=prompt,
+            tokenizer=self.tokenizer,
+            image_token_index=llava_constants.IMAGE_TOKEN_INDEX,
+            return_tensors="pt",
+        )
+        .unsqueeze(0)
+        .to(self.device)
+    )
+
+    images = [
+        image_format_converter.base64_to_image(image_str=base64_image).convert(
+            "RGB"
+        )
+    ]
+    # Gets the image embedding.
+    images_tensor = mm_utils.process_images(
+        images=images,
+        image_processor=self.image_processor,
+        model_cfg=self.model.config,
+    ).to(self.device, dtype=torch.float16)
+
+    self.stop_str = conversation.conv_llava_v1.sep2
+    self.keywords = [self.stop_str]
+
+    return input_ids, images_tensor
+
+  def inference(
+      self, input_ids: List[torch.Tensor], images_tensor: torch.Tensor
+  ) -> List[torch.Tensor]:
+    """Runs the inference."""
+    stopping_criteria = mm_utils.KeywordsStoppingCriteria(
+        keywords=self.keywords, tokenizer=self.tokenizer, input_ids=input_ids
+    )
+
+    with torch.inference_mode():
+      output_ids = self.model.generate(
+          input_ids=input_ids,
+          images=images_tensor,
+          do_sample=False,
+          temperature=0,
+          top_p=None,
+          num_beams=1,
+          max_new_tokens=512,
+          use_cache=True,
+          stopping_criteria=[stopping_criteria],
+      )
+
+    return output_ids
+
+  def postprocess(
+      self, output_ids: List[torch.Tensor], input_token_len: int
+  ) -> List[str]:
+    """Runs the postprocessing to convert token ids to string."""
+    outputs = self.tokenizer.batch_decode(
+        output_ids[:, input_token_len:], skip_special_tokens=True
+    )[0]
+    outputs = outputs.strip()
+    if outputs.endswith(self.stop_str):
+      outputs = outputs[: -len(self.stop_str)]
+    outputs = outputs.strip()
+
+    return [outputs]
+
+  def handle(self, data: List[Dict[str, Any]], context: Any) -> List[str]:
+    """Handles an incoming request by passing it through `preprocess`, `inference`, and `postprocess`."""
+    input_ids, images_tensor = self.preprocess(data=data)
+    model_output = self.inference(
+        input_ids=input_ids, images_tensor=images_tensor
+    )
+
+    input_token_len = input_ids.shape[1]
+    return self.postprocess(
+        output_ids=model_output, input_token_len=input_token_len
+    )

community-content/vertex_model_garden/model_oss/llava/model_handler_setup.py

@@ -0,0 +1,77 @@
+"""Common utility functions for setting up and initializing the model and the handler."""
+
+import logging
+import os
+from typing import Any
+
+import torch
+
+from util import constants
+from util import fileutils
+
+
+def get_model_id(default_model_id: str) -> str:
+  """Gets a model id or a local model path.
+
+  Args:
+    default_model_id: Default model id for the corresponding model set in the
+      handler.
+
+  Returns:
+    str: model id or a local model path.
+  """
+  # The model id can be either:
+  # 1) a huggingface model card id, like "Salesforce/blip", or
+  # 2) a GCS path to the model files, like "gs://foo/bar".
+  # If it's a model card id, the model will be loaded from huggingface.
+  model_id = (
+      default_model_id
+      if os.environ.get("MODEL_ID") is None
+      else os.environ["MODEL_ID"]
+  )
+
+  # Else it will be downloaded from GCS to local first.
+  # Since the transformers from_pretrained API can't read from GCS.
+  if model_id.startswith(constants.GCS_URI_PREFIX):
+    gcs_path = model_id[len(constants.GCS_URI_PREFIX) :]
+    local_model_dir = os.path.join(constants.LOCAL_MODEL_DIR, gcs_path)
+    logging.info("Download %s to %s", model_id, local_model_dir)
+    fileutils.download_gcs_dir_to_local(model_id, local_model_dir)
+    model_id = local_model_dir
+
+  return model_id
+
+
+def get_map_location(context: Any) -> str:
+  """Gets model map location.
+
+  Args:
+    context: Torchserve worker context.
+
+  Returns:
+    str: Mapping location.
+  """
+  properties = context.system_properties
+  return (
+      "cuda"
+      if torch.cuda.is_available() and properties.get("gpu_id") is not None
+      else "cpu"
+  )
+
+
+def get_model_device(map_location: str, context: Any) -> torch.device:
+  """Gets model accelerator device.
+
+  Args:
+    map_location: Model map location.
+    context: TorchServe worker context.
+
+  Returns:
+    torch.Device: Device to load the model into.
+  """
+  properties = context.system_properties
+  return torch.device(
+      map_location + ":" + str(properties.get("gpu_id"))
+      if torch.cuda.is_available() and properties.get("gpu_id") is not None
+      else map_location
+  )

community-content/vertex_model_garden/model_oss/notebook_util/common_util.py

@@ -0,0 +1,447 @@
+"""Common util functions for notebook."""
+
+import base64
+import datetime
+import io
+import json
+import os
+import subprocess
+from typing import Any, Dict, Sequence
+
+from google.cloud import storage
+import matplotlib.pyplot as plt
+import numpy as np
+from PIL import Image
+import requests
+import tensorflow as tf
+import yaml
+
+
+GCS_URI_PREFIX = "gs://"
+CHECKPOINT_BUCKET = "gs://model_garden_checkpoints"
+
+
+def convert_numpy_array_to_byte_string_via_tf_tensor(
+    np_array: np.ndarray,
+) -> str:
+  """Serializes a numpy array to tensor bytes.
+
+  Args:
+    np_array: A numpy array.
+
+  Returns:
+    A tensor bytes.
+  """
+  tensor_array = tf.convert_to_tensor(np_array)
+  tensor_byte_string = tf.io.serialize_tensor(tensor_array)
+  return tensor_byte_string.numpy()
+
+
+def get_jpeg_bytes(local_image_path: str, new_width: int = -1) -> bytes:
+  """Returns jpeg bytes given an image path and resizes if required.
+
+  Args:
+    local_image_path: A string of local image path.
+    new_width: An integer of new image width.
+
+  Returns:
+    A jpeg bytes.
+  """
+  image = Image.open(local_image_path)
+  if new_width <= 0:
+    new_image = image
+  else:
+    width, height = image.size
+    print("original input image size: ", width, " , ", height)
+    new_height = int(height * new_width / width)
+    print("new input image size: ", new_width, " , ", new_height)
+    new_image = image.resize((new_width, new_height))
+  buffered = io.BytesIO()
+  new_image.save(buffered, format="JPEG")
+  return buffered.getvalue()
+
+
+def gcs_fuse_path(path: str) -> str:
+  """Try to convert path to gcsfuse path if it starts with gs:// else do not modify it.
+
+  Args:
+    path: A string of path.
+
+  Returns:
+    A gcsfuse path.
+  """
+  path = path.strip()
+  if path.startswith("gs://"):
+    return "/gcs/" + path[5:]
+  return path
+
+
+def get_job_name_with_datetime(prefix: str) -> str:
+  """Gets a job name by adding current time to prefix.
+
+  Args:
+    prefix: A string of job name prefix.
+
+  Returns:
+    A job name.
+  """
+  return prefix + datetime.datetime.now().strftime("_%Y%m%d_%H%M%S")
+
+
+def create_job_name(prefix: str) -> str:
+  """Creates a job name.
+
+  Args:
+    prefix: A string of job name prefix.
+
+  Returns:
+    A job name.
+  """
+  user = os.environ.get("USER")
+  now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+  job_name = f"{prefix}-{user}-{now}"
+  return job_name
+
+
+def save_subset_annotation(
+    input_annotation_path: str, output_annotation_path: str
+):
+  """Saves a subset of COCO annotation json file with CCA 4.0 license.
+
+  Args:
+    input_annotation_path: A string of input annotation path.
+    output_annotation_path: A string of output annotation path.
+  """
+
+  with open(input_annotation_path) as f:
+    coco_json = json.load(f)
+
+  img_ids = set()
+  images = []
+  annotations = []
+
+  for img in coco_json["images"]:
+    if img["license"] in [4, 5]:  # CCA 4.0 license.
+      img_ids.add(img["id"])
+      images.append(img)
+
+  for ann in coco_json["annotations"]:
+    if ann["image_id"] in img_ids:
+      annotations.append(ann)
+
+  new_json = {
+      "info": coco_json["info"],
+      "licenses": coco_json["licenses"],
+      "images": images,
+      "annotations": annotations,
+      "categories": coco_json["categories"],
+  }
+
+  with open(output_annotation_path, "w") as f:
+    json.dump(new_json, f)
+
+
+def image_to_base64(image: Any, image_format: str = "JPEG") -> str:
+  """Converts an image to base64.
+
+  Args:
+    image: A PIL.Image instance.
+    image_format: A string of image format.
+
+  Returns:
+    A base64 string.
+  """
+  buffer = io.BytesIO()
+  image.save(buffer, format=image_format)
+  image_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
+  return image_str
+
+
+def base64_to_image(image_str: str) -> Any:
+  """Convert base64 encoded string to an image.
+
+  Args:
+    image_str: A string of base64 encoded image.
+
+  Returns:
+    A PIL.Image instance.
+  """
+  image = Image.open(io.BytesIO(base64.b64decode(image_str)))
+  return image
+
+
+def image_grid(imgs: Sequence[Any], rows: int = 2, cols: int = 2) -> Any:
+  """Creates an image grid.
+
+  Args:
+    imgs: A list of PIL.Image instances.
+    rows: An integer of number of rows.
+    cols: An integer of number of columns.
+
+  Returns:
+    A PIL.Image instance.
+  """
+  w, h = imgs[0].size
+  grid = Image.new(
+      mode="RGB", size=(cols * w + 10 * cols, rows * h), color=(255, 255, 255)
+  )
+  for i, img in enumerate(imgs):
+    grid.paste(img, box=(i % cols * w + 10 * i, i // cols * h))
+  return grid
+
+
+def display_image(image: Any):
+  """Displays an image.
+
+  Args:
+    image: A PIL.Image instance.
+  """
+  _ = plt.figure(figsize=(20, 15))
+  plt.grid(False)
+  plt.imshow(image)
+
+
+def download_gcs_file_to_local(gcs_uri: str, local_path: str):
+  """Download a gcs file to a local path.
+
+  Args:
+    gcs_uri: A string of file path on GCS.
+    local_path: A string of local file path.
+  """
+  if not gcs_uri.startswith(GCS_URI_PREFIX):
+    raise ValueError(
+        f"{gcs_uri} is not a GCS path starting with {GCS_URI_PREFIX}."
+    )
+  client = storage.Client()
+  os.makedirs(os.path.dirname(local_path), exist_ok=True)
+  with open(local_path, "wb") as f:
+    client.download_blob_to_file(gcs_uri, f)
+
+
+def download_image(url: str) -> str:
+  """Downloads an image from the given URL.
+
+  Args:
+    url: The URL of the image to download.
+
+  Returns:
+    base64 encoded image.
+  """
+  response = requests.get(url)
+  return Image.open(io.BytesIO(response.content))
+
+
+def resize_image(image: Any, new_width: int = 1000) -> Any:
+    """Resizes an image to a certain width.
+
+    Args:
+      image: The image which has to be resized.
+      new_width: New width of the image.
+
+    Returns:
+      New resized image.
+    """
+    width, height = image.size
+    new_height = int(height * new_width / width)
+    new_img = image.resize((new_width, new_height))
+    return new_img
+
+
+def load_img(path: str) -> Any:
+  """Reads image from path and return PIL.Image instance.
+
+  Args:
+    path: A string of image path.
+
+  Returns:
+    A PIL.Image instance.
+  """
+  img = tf.io.read_file(path)
+  img = tf.image.decode_jpeg(img, channels=3)
+  return Image.fromarray(np.uint8(img)).convert("RGB")
+
+
+def decode_image(
+    image_str_tensor: tf.string, new_height: int, new_width: int
+) -> tf.float32:
+  """Converts and resizes image bytes to image tensor.
+
+  Args:
+    image_str_tensor: A string of image bytes.
+    new_height: An integer of new image height.
+    new_width: An integer of new image width.
+
+  Returns:
+    An image tensor.
+  """
+  image = tf.io.decode_image(image_str_tensor, 3, expand_animations=False)
+  image = tf.image.resize(image, (new_height, new_width))
+  return image
+
+
+def get_label_map(label_map_yaml_filepath: str) -> Dict[int, str]:
+  """Returns class id to label mapping given a filepath to the label map.
+
+  Args:
+    label_map_yaml_filepath: A string of label map yaml file path.
+
+  Returns:
+    A dictionary of class id to label mapping.
+  """
+  with tf.io.gfile.GFile(label_map_yaml_filepath, "rb") as input_file:
+    label_map = yaml.safe_load(input_file.read())["label_map"]
+  return label_map
+
+
+def get_prediction_instances(test_filepath: str, new_width: int = -1) -> Any:
+  """Generate instance from image path to pass to Vertex AI Endpoint for prediction.
+
+  Args:
+    test_filepath: A string of test image path.
+    new_width: An integer of new image width.
+
+  Returns:
+    A list of instances.
+  """
+  if new_width <= 0:
+    with tf.io.gfile.GFile(test_filepath, "rb") as input_file:
+      encoded_string = base64.b64encode(input_file.read()).decode("utf-8")
+  else:
+    img = load_img(test_filepath)
+    width, height = img.size
+    print("original input image size: ", width, " , ", height)
+    new_height = int(height * new_width / width)
+    new_img = img.resize((new_width, new_height))
+    print("resized input image size: ", new_width, " , ", new_height)
+    buffered = io.BytesIO()
+    new_img.save(buffered, format="JPEG")
+    encoded_string = base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+  instances = [{
+      "encoded_image": {"b64": encoded_string},
+  }]
+  return instances
+
+
+def get_quota(project_id: str, region: str, resource_id: str) -> int:
+  """Returns the quota for a resource in a region.
+
+  Args:
+    project_id: The project id.
+    region: The region.
+    resource_id: The resource id.
+
+  Returns:
+    The quota for the resource in the region. Returns -1 if can not figure out
+    the quota.
+
+  Raises:
+    RuntimeError: If the command to get quota fails.
+  """
+  service_endpoint = "aiplatform.googleapis.com"
+
+  command = (
+      "gcloud alpha services quota list"
+      f" --service={service_endpoint} --consumer=projects/{project_id}"
+      f" --filter='{service_endpoint}/{resource_id}' --format=json"
+  )
+  process = subprocess.run(
+      command, shell=True, capture_output=True, text=True, check=True
+  )
+  if process.returncode == 0:
+    quota_data = json.loads(process.stdout)
+  else:
+    raise RuntimeError(f"Error fetching quota data: {process.stderr}")
+
+  if not quota_data or "consumerQuotaLimits" not in quota_data[0]:
+    return -1
+  if (
+      not quota_data[0]["consumerQuotaLimits"]
+      or "quotaBuckets" not in quota_data[0]["consumerQuotaLimits"][0]
+  ):
+    return -1
+  all_regions_data = quota_data[0]["consumerQuotaLimits"][0]["quotaBuckets"]
+  for region_data in all_regions_data:
+    if (
+        region_data.get("dimensions")
+        and region_data["dimensions"]["region"] == region
+    ):
+      if "effectiveLimit" in region_data:
+        return int(region_data["effectiveLimit"])
+      else:
+        return 0
+  return -1
+
+
+def get_resource_id(accelerator_type: str, is_for_training: bool) -> str:
+  """Returns the resource id for a given accelerator type and the use case.
+
+  Args:
+    accelerator_type: The accelerator type.
+    is_for_training: Whether the resource is used for training. Set false for
+      serving use case.
+
+  Returns:
+    The resource id.
+  """
+  training_accelerator_map = {
+      "NVIDIA_TESLA_V100": "custom_model_training_nvidia_v100_gpus",
+      "NVIDIA_L4": "custom_model_training_nvidia_l4_gpus",
+      "NVIDIA_TESLA_A100": "custom_model_training_nvidia_a100_gpus",
+      "NVIDIA_A100_80GB": "custom_model_training_nvidia_a100_80gb_gpus",
+      "NVIDIA_TESLA_T4": "custom_model_training_nvidia_t4_gpus",
+      "TPU_V5e": "custom_model_training_tpu_v5e",
+      "TPU_V3": "custom_model_training_tpu_v3",
+  }
+  serving_accelerator_map = {
+      "NVIDIA_TESLA_V100": "custom_model_serving_nvidia_v100_gpus",
+      "NVIDIA_L4": "custom_model_serving_nvidia_l4_gpus",
+      "NVIDIA_TESLA_A100": "custom_model_serving_nvidia_a100_gpus",
+      "NVIDIA_A100_80GB": "custom_model_serving_nvidia_a100_80gb_gpus",
+      "NVIDIA_TESLA_T4": "custom_model_serving_nvidia_t4_gpus",
+      "TPU_V5e": "custom_model_serving_tpu_v5e",
+  }
+  if is_for_training:
+    if accelerator_type in training_accelerator_map:
+      return training_accelerator_map[accelerator_type]
+    else:
+      raise ValueError(
+          f"Could not find accelerator type: {accelerator_type} for training."
+      )
+  else:
+    if accelerator_type in serving_accelerator_map:
+      return serving_accelerator_map[accelerator_type]
+    else:
+      raise ValueError(
+          f"Could not find accelerator type: {accelerator_type} for serving."
+      )
+
+
+def check_quota(
+    project_id: str,
+    region: str,
+    accelerator_type: str,
+    accelerator_count: int,
+    is_for_training: bool,
+):
+  """Checks if the project and the region has the required quota."""
+  resource_id = get_resource_id(accelerator_type, is_for_training)
+  quota = get_quota(project_id, region, resource_id)
+  quota_request_instruction = (
+      "Either use "
+      "a different region or request additional quota. Follow "
+      "instructions here "
+      "https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota"
+      " to check quota in a region or request additional quota for "
+      "your project."
+  )
+  if quota == -1:
+    raise ValueError(
+        f"Quota not found for: {resource_id} in {region}."
+        f" {quota_request_instruction}"
+    )
+  if quota < accelerator_count:
+    raise ValueError(
+        f"Quota not enough for {resource_id} in {region}: {quota} <"
+        f" {accelerator_count}. {quota_request_instruction}"
+    )

community-content/vertex_model_garden/model_oss/open_clip/handler.py

@@ -10,28 +10,15 @@ import open_clip
 import torch
 from ts.torch_handler.base_handler import BaseHandler
 
-from google3.cloud.ml.applications.vision.model_garden.model_oss.util import constants
-from google3.cloud.ml.applications.vision.model_garden.model_oss.util import fileutils
-from google3.cloud.ml.applications.vision.model_garden.model_oss.util import image_format_converter
-
-
-@enum.unique
-class Precision(enum.Enum):
-  AMP = "amp"
-  AMP_BF16 = "amp_bf16"
-  AMP_BFLOAT16 = "amp_bfloat16"
-  BF16 = "bf16"
-  FP16 = "fp16"
-  PURE_BF16 = "pure_bf16"
-  PURE_FP16 = "pure_fp16"
-  FP32 = "fp32"
+from util import constants
+from util import fileutils
+from util import image_format_converter
 
 
 # Supported checkpoint&model pairs:
 # https://github.com/mlfoundations/open_clip#pretrained-model-interface
-_DEFAULT_CHECKPOINT = "openai"
 _DEFAULT_MODEL = "RN50"
-_DEFAULT_PRECISION = Precision.AMP
+_BIOMED_CLIP_MODEL = "microsoft/BiomedCLIP"
 _ZERO_CLASSIFICATION = "zero-shot-image-classification"
 _FEATURE_EMBEDDING = "feature-embedding"
 _VALID_TASKS = frozenset([_ZERO_CLASSIFICATION, _FEATURE_EMBEDDING])
@@ -45,6 +32,21 @@ _TEXT_FEATURES_KEY = "text_features"
 class OpenclipHandler(BaseHandler):
   """Custom handler for OpenCLIP."""
 
+  @enum.unique
+  class Precision(enum.Enum):
+    AMP = "amp"
+    AMP_BF16 = "amp_bf16"
+    AMP_BFLOAT16 = "amp_bfloat16"
+    # For the difference between floating points and "pure" floating points, see
+    # https://github.com/mlfoundations/open_clip/blob/0142d279298a4ca0138316286f775fe9d7bdbb94/src/open_clip/factory.py#L232C58-L232C58
+    BF16 = "bf16"
+    FP16 = "fp16"
+    PURE_BF16 = "pure_bf16"
+    PURE_FP16 = "pure_fp16"
+    FP32 = "fp32"
+
+  _DEFAULT_PRECISION = Precision.AMP
+
   def initialize(self, context: Any):
     """Custom initialize."""
 
@@ -61,29 +63,32 @@ class OpenclipHandler(BaseHandler):
     )
     self.manifest = context.manifest
 
-    model_name = os.environ.get("MODEL", _DEFAULT_MODEL)
-    precision = os.environ.get("PRECISION", _DEFAULT_PRECISION)
-    checkpoint = os.environ.get("CHECKPOINT", _DEFAULT_CHECKPOINT)
+    self.model_name = os.environ.get("MODEL", None)
+    if not self.model_name:
+      self.model_name = os.environ.get("MODEL_ID", _DEFAULT_MODEL)
+    precision = os.environ.get("PRECISION", self._DEFAULT_PRECISION)
+    checkpoint = os.environ.get("CHECKPOINT")
     self.task = os.environ.get("TASK", _FEATURE_EMBEDDING)
     if self.task not in _VALID_TASKS:
       raise ValueError(f"Invalid task: {self.task}.")
     logging.info(
         "Handler initializing task:%s, model:%s, precision:%s, checkpoint:%s",
         self.task,
-        model_name,
+        self.model_name,
         precision,
         checkpoint,
     )
 
-    if checkpoint != _DEFAULT_CHECKPOINT:
+    if fileutils.is_gcs_path(checkpoint):
       local_fname = os.path.join(constants.LOCAL_MODEL_DIR, "model.pt")
       fileutils.download_gcs_file_to_local(checkpoint, local_fname)
       checkpoint = local_fname
 
-    self.model, _, self.preprocessor = open_clip.create_model_and_transforms(
-        model_name, pretrained=checkpoint, precision=precision
+    self.model, self.preprocessor = open_clip.create_model_from_pretrained(
+        self.model_name, pretrained=checkpoint, precision=precision
     )
-    self.tokenizer = open_clip.get_tokenizer(model_name)
+    self.model.to(self.device)
+    self.tokenizer = open_clip.get_tokenizer(self.model_name)
 
     self.initialized = True
 
@@ -102,9 +107,32 @@ class OpenclipHandler(BaseHandler):
       processed_list.append(sample)
     return processed_list
 
+  def _biomedclip_inference(
+      self, data: List[Dict[str, Any]], *args, **kwargs
+  ) -> List[List[float]]:
+    """Inference for BiomedCLIP model."""
+    texts = torch.stack(
+        [item[_TEXT_KEY][0] for item in data if _TEXT_KEY in item]
+    ).to(self.map_location)
+    images = torch.stack(
+        [item[_IMAGE_KEY][0] for item in data if _IMAGE_KEY in item]
+    ).to(self.map_location)
+    if texts.shape[0] == 0 or images.shape[0] == 0:
+      return []
+    with torch.no_grad():
+      image_features, text_features, logit_scale = self.model(images, texts)
+      logits = (
+          (logit_scale * image_features @ text_features.t())
+          .detach()
+          .softmax(dim=-1)
+      )
+      return logits.cpu().numpy().tolist()
+
   def inference(
       self, data: List[Dict[str, Any]], *args, **kwargs
   ) -> List[Dict[str, Any]]:
+    if _BIOMED_CLIP_MODEL in self.model_name:
+      return self._biomedclip_inference(data)
     feature_list = []
     with torch.no_grad(), torch.cuda.amp.autocast():
       for item in data:
@@ -120,6 +148,8 @@ class OpenclipHandler(BaseHandler):
 
   def postprocess(self, features: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     """Postprocess the image/text featreus for downstream task."""
+    if _BIOMED_CLIP_MODEL in self.model_name:
+      return features
     preds = []
     if self.task == _FEATURE_EMBEDDING:
       for item in features:
@@ -139,4 +169,4 @@ class OpenclipHandler(BaseHandler):
         text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)
         preds.append(text_probs.tolist())
 
-    return preds
+    return preds

community-content/vertex_model_garden/model_oss/peft/dockerfile/serve.Dockerfile

@@ -83,7 +83,7 @@ EXPOSE ${mng_port}
 
 # Set environments.
 ENV TASK "causal-language-modeling-lora"
-ENV BASE_MODEL_ID "openlm-research/open_llama_7b"
+ENV MODEL_ID "openlm-research/open_llama_7b"
 ENV PRECISION_LOADING_MODE "float16"
 ENV FINETUNED_LORA_MODEL_PATH ""
 

community-content/vertex_model_garden/model_oss/peft/handler.py

@@ -3,15 +3,17 @@
 # pylint: disable=g-importing-member
 # pylint: disable=logging-fstring-interpolation
 
-import logging
 import os
-from typing import Any, List
+import time
+from typing import Any, List, Tuple
 
 from absl import logging
+from awq import AutoAWQForCausalLM
 from diffusers import DPMSolverMultistepScheduler
 from diffusers import StableDiffusionPipeline
 from peft import PeftModel
 from PIL import Image
+import psutil
 import torch
 import transformers
 from transformers import AutoModelForCausalLM
@@ -24,6 +26,12 @@ from util import constants
 from util import fileutils
 from util import image_format_converter
 
+if os.path.exists(constants.SHARED_MEM_DIR):
+  logging.info(
+      "SharedMemorySizeMb: %s",
+      psutil.disk_usage(constants.SHARED_MEM_DIR).free / 1e6,
+  )
+
 # Tasks
 TEXT_TO_IMAGE_LORA = "text-to-image-lora"
 SEQUENCE_CLASSIFICATION_LORA = "sequence-classification-lora"
@@ -33,8 +41,13 @@ INSTRUCT_LORA = "instruct-lora"
 # Inference parameters.
 _NUM_INFERENCE_STEPS = 25
 _MAX_LENGTH_DEFAULT = 200
+_MAX_TOKENS_DEFAULT = None
+_TEMPERATURE_DEFAULT = 1.0
+_TOP_P_DEFAULT = 1.0
 _TOP_K_DEFAULT = 10
 
+logging.set_verbosity(os.environ.get("LOG_LEVEL", logging.INFO))
+
 
 class PeftHandler(BaseHandler):
   """Custom handler for Peft models."""
@@ -59,16 +72,24 @@ class PeftHandler(BaseHandler):
         "PRECISION_LOADING_MODE", constants.PRECISION_MODE_16
     )
     self.task = os.environ.get("TASK", CAUSAL_LANGUAGE_MODELING_LORA)
-    self.base_model_id = os.environ.get(
-        "BASE_MODEL_ID", "openlm-research/open_llama_7b"
-    )
-    if fileutils.is_gcs_path(self.base_model_id):
+    self.base_model_id = os.environ.get("BASE_MODEL_ID", None)
+    self.model_id = self.base_model_id
+    if not self.base_model_id:
+      self.model_id = os.environ.get("MODEL_ID", "")
+    self.quantization = os.environ.get("QUANTIZATION", None)
+    logging.info(f"Load base model id from MODEL_ID:{self.model_id}.")
+    if not self.model_id:
+      self.model_id = os.environ.get("AIP_STORAGE_URI", "")
+      logging.info(f"Load base model id from AIP_STORAGE_URI: {self.model_id}.")
+    if not self.model_id:
+      raise ValueError("Base model id is must be set.")
+    if fileutils.is_gcs_path(self.model_id):
       fileutils.download_gcs_dir_to_local(
-          self.base_model_id,
+          self.model_id,
           constants.LOCAL_BASE_MODEL_DIR,
           skip_hf_model_bin=True,
       )
-      self.base_model_id = constants.LOCAL_BASE_MODEL_DIR
+      self.model_id = constants.LOCAL_BASE_MODEL_DIR
     self.finetuned_lora_model_path = os.environ.get(
         "FINETUNED_LORA_MODEL_PATH", ""
     )
@@ -79,7 +100,7 @@ class PeftHandler(BaseHandler):
       self.finetuned_lora_model_path = constants.LOCAL_MODEL_DIR
 
     logging.info(
-        f"Using task:{self.task}, base model:{self.base_model_id}, lora model:"
+        f"Using task:{self.task}, base model:{self.model_id}, lora model:"
         f" {self.finetuned_lora_model_path}, and precision"
         f" {self.precision_mode}."
     )
@@ -87,30 +108,32 @@ class PeftHandler(BaseHandler):
     self.pipeline = None
     self.model = None
     self.tokenizer = None
+    start_time = time.perf_counter()
+    logging.info("Started PEFT handler initialization at: %s", start_time)
     if self.task == TEXT_TO_IMAGE_LORA:
       pipeline = StableDiffusionPipeline.from_pretrained(
-          self.base_model_id, torch_dtype=torch.float16
+          self.model_id, torch_dtype=torch.float16
       )
       logging.debug("Initialized the base model for text to image.")
       pipeline.scheduler = DPMSolverMultistepScheduler.from_config(
           pipeline.scheduler.config
       )
       logging.debug("Initialized the scheduler for text to image.")
-      if self.finetuned_lora_model_path:
-        pipeline.unet.load_attn_procs(self.finetuned_lora_model_path)
-      logging.debug("Initialized the LoRA model for text to image.")
       # This is to reduce GPU memory requirements.
       pipeline.enable_xformers_memory_efficient_attention()
       pipeline = pipeline.to(self.map_location)
       # Reduces memory footprint.
       pipeline.enable_attention_slicing()
+      if self.finetuned_lora_model_path:
+        pipeline.load_lora_weights(self.finetuned_lora_model_path)
+      logging.debug("Initialized the LoRA model for text to image.")
       self.pipeline = pipeline
       logging.info("Initialized the text to image pipelines.")
     elif self.task == SEQUENCE_CLASSIFICATION_LORA:
-      tokenizer = AutoTokenizer.from_pretrained(self.base_model_id)
+      tokenizer = AutoTokenizer.from_pretrained(self.model_id)
       logging.debug("Initialized the tokenizer for sequence classification.")
       model = AutoModelForSequenceClassification.from_pretrained(
-          self.base_model_id, torch_dtype=torch.float16
+          self.model_id, torch_dtype=torch.float16
       )
       logging.debug("Initialized the base model for sequence classification.")
       if self.finetuned_lora_model_path:
@@ -122,55 +145,74 @@ class PeftHandler(BaseHandler):
     elif (
         self.task == CAUSAL_LANGUAGE_MODELING_LORA or self.task == INSTRUCT_LORA
     ):
-      tokenizer = AutoTokenizer.from_pretrained(self.base_model_id)
+      tokenizer = AutoTokenizer.from_pretrained(self.model_id)
       logging.debug("Initialized the tokenizer.")
       if self.task == CAUSAL_LANGUAGE_MODELING_LORA:
-        if self.precision_mode == constants.PRECISION_MODE_32:
-          model = AutoModelForCausalLM.from_pretrained(
-              self.base_model_id,
-              return_dict=True,
-              torch_dtype=torch.float32,
-              device_map="auto",
-          )
-        elif self.precision_mode == constants.PRECISION_MODE_16:
-          model = AutoModelForCausalLM.from_pretrained(
-              self.base_model_id,
-              return_dict=True,
-              torch_dtype=torch.bfloat16,
-              device_map="auto",
-          )
-        elif self.precision_mode == constants.PRECISION_MODE_8:
-          quantization_config = BitsAndBytesConfig(
-              load_in_8bit=True, int8_threshold=0
-          )
-          model = AutoModelForCausalLM.from_pretrained(
-              self.base_model_id,
-              return_dict=True,
-              torch_dtype=torch.float16,
-              device_map="auto",
-              quantization_config=quantization_config,
-          )
+        if self.quantization == constants.AWQ:
+          model = AutoAWQForCausalLM.from_quantized(self.model_id)
+        elif self.quantization == constants.GPTQ or not self.quantization:
+          if self.precision_mode == constants.PRECISION_MODE_32:
+            model = AutoModelForCausalLM.from_pretrained(
+                self.model_id,
+                return_dict=True,
+                torch_dtype=torch.float32,
+                device_map="auto",
+            )
+          elif self.precision_mode == constants.PRECISION_MODE_16B:
+            model = AutoModelForCausalLM.from_pretrained(
+                self.model_id,
+                return_dict=True,
+                torch_dtype=torch.bfloat16,
+                device_map="auto",
+            )
+          elif self.precision_mode == constants.PRECISION_MODE_16:
+            model = AutoModelForCausalLM.from_pretrained(
+                self.model_id,
+                return_dict=True,
+                torch_dtype=torch.float16,
+                device_map="auto",
+            )
+          elif self.precision_mode == constants.PRECISION_MODE_8:
+            quantization_config = BitsAndBytesConfig(
+                load_in_8bit=True, int8_threshold=0
+            )
+            model = AutoModelForCausalLM.from_pretrained(
+                self.model_id,
+                return_dict=True,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                quantization_config=quantization_config,
+            )
+          else:
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_compute_dtype=torch.bfloat16,
+            )
+            model = AutoModelForCausalLM.from_pretrained(
+                self.model_id,
+                return_dict=True,
+                device_map="auto",
+                torch_dtype=torch.bfloat16,
+                quantization_config=quantization_config,
+            )
         else:
-          quantization_config = BitsAndBytesConfig(
-              load_in_4bit=True,
-              bnb_4bit_quant_type="nf4",
-              bnb_4bit_compute_dtype=torch.bfloat16,
-          )
-          model = AutoModelForCausalLM.from_pretrained(
-              self.base_model_id,
-              return_dict=True,
-              device_map="auto",
-              torch_dtype=torch.bfloat16,
-              quantization_config=quantization_config,
-          )
+          raise ValueError(f"Invalid QUANTIZATION value: {self.quantization}")
       else:
-        model = AutoModelForCausalLM.from_pretrained(
-            self.base_model_id,
-            return_dict=True,
-            torch_dtype=torch.bfloat16,
-            trust_remote_code=True,
-            device_map="auto",
-        )
+        try:
+          model = AutoModelForCausalLM.from_pretrained(
+              self.model_id,
+              torch_dtype=torch.bfloat16,
+              trust_remote_code=True,
+              device_map="auto",
+          )
+        except:  # pylint: disable=bare-except
+          model = AutoModelForCausalLM.from_pretrained(
+              self.model_id,
+              torch_dtype=torch.bfloat16,
+              trust_remote_code=True,
+              device_map="auto",
+          )
       logging.debug("Initialized the base model.")
       if self.finetuned_lora_model_path:
         model = PeftModel.from_pretrained(model, self.finetuned_lora_model_path)
@@ -186,35 +228,57 @@ class PeftHandler(BaseHandler):
       raise ValueError(f"Invalid TASK: {self.task}")
 
     self.initialized = True
-    logging.info("The PEFT handler was initialized.")
+    end_time = time.perf_counter()
+    logging.info("The PEFT handler was initialize at: %s", end_time)
+    logging.info("Handler initiation took %s seconds", end_time - start_time)
 
   def preprocess(self, data: Any) -> Any:
     """Preprocesses input data."""
     # Assumes that the parameters are same in one request. We parse the
     # parameters from the first instance for all instances in one request.
+    # For generation length: `max_length` defines the maximum length of the
+    # sequence to be generated, including both input and output tokens.
+    # `max_length` is overridden by `max_new_tokens` if also set.
+    # `max_new_tokens` defines the maximum number of new tokens to generate,
+    # ignoring the current number of tokens.
+    # Reference:
+    # https://github.com/huggingface/transformers/blob/574a5384557b1aaf98ddb13ea9eb0a0ee8ff2cb2/src/transformers/generation/configuration_utils.py#L69-L73
     max_length = _MAX_LENGTH_DEFAULT
+    max_tokens = _MAX_TOKENS_DEFAULT
+    temperature = _TEMPERATURE_DEFAULT
+    top_p = _TOP_P_DEFAULT
     top_k = _TOP_K_DEFAULT
 
     prompts = [item["prompt"] for item in data]
     if "max_length" in data[0]:
       max_length = data[0]["max_length"]
+    if "max_tokens" in data[0]:
+      max_tokens = data[0]["max_tokens"]
+    if "temperature" in data[0]:
+      temperature = data[0]["temperature"]
+    if "top_p" in data[0]:
+      top_p = data[0]["top_p"]
     if "top_k" in data[0]:
       top_k = data[0]["top_k"]
 
-    return prompts, max_length, top_k
+    return prompts, max_length, max_tokens, temperature, top_p, top_k
 
-  def inference(self, data: Any, *args, **kwargs) -> List[Image.Image]:
+  def inference(
+      self, data: Any, *args, **kwargs
+  ) -> Tuple[List[str], List[Image.Image]]:
     """Runs the inference."""
-    prompts, max_length, top_k = data
+    prompts, max_length, max_tokens, temperature, top_p, top_k = data
     logging.debug(
-        f"Inference prompts={prompts}, max_length={max_length}, top_k={top_k}."
+        f"Inference prompts={prompts}, max_length={max_length},"
+        f" max_tokens={max_tokens}, temperature={temperature}, top_p={top_p},"
+        f" top_k={top_k}."
     )
     if self.task == TEXT_TO_IMAGE_LORA:
       predicted_results = self.pipeline(
           prompt=prompts, num_inference_steps=_NUM_INFERENCE_STEPS
       ).images
     elif self.task == SEQUENCE_CLASSIFICATION_LORA:
-      encoded_input = self.tokenizer(prompts, return_tensors="pt")
+      encoded_input = self.tokenizer(prompts, return_tensors="pt", padding=True)
       encoded_input.to(self.map_location)
       with torch.no_grad():
         outputs = self.model(**encoded_input)
@@ -226,25 +290,43 @@ class PeftHandler(BaseHandler):
       predicted_results = self.pipeline(
           prompts,
           max_length=max_length,
+          max_new_tokens=max_tokens,
           do_sample=True,
+          temperature=temperature,
+          top_p=top_p,
           top_k=top_k,
           num_return_sequences=1,
           eos_token_id=self.tokenizer.eos_token_id,
+          return_full_text=False,
       )
     else:
       raise ValueError(f"Invalid TASK: {self.task}")
-    return predicted_results
+    return prompts, predicted_results
 
   def postprocess(self, data: Any) -> List[str]:
     """Postprocesses output data."""
+    prompts, predicted_results = data
     if self.task == TEXT_TO_IMAGE_LORA:
       # Converts the images to base64 string.
       outputs = [
-          image_format_converter.image_to_base64(image) for image in data
+          image_format_converter.image_to_base64(image)
+          for image in predicted_results
       ]
+    elif self.task == SEQUENCE_CLASSIFICATION_LORA:
+      outputs = predicted_results
     else:
-      outputs = data
+      outputs = []
+      for prompt, predicted_result in zip(prompts, predicted_results):
+        formatted_output = self._format_text_generation_output(
+            prompt=prompt, output=predicted_result[0]["generated_text"]
+        )
+        outputs.append(formatted_output)
     return outputs
 
+  def _format_text_generation_output(self, prompt: str, output: str) -> str:
+    """Formats text generation output."""
+    output = output.strip("\n")
+    return f"Prompt:\n{prompt.strip()}\nOutput:\n{output}"
 
-# pylint: enable=logging-fstring-interpolation
+
+# pylint: enable=logging-fstring-interpolation

community-content/vertex_model_garden/model_oss/pic2word/dockerfile/serve.Dockerfile

@@ -6,6 +6,7 @@ ENV infer_port=7080
 ENV mng_port=7081
 ENV model_name="pic2word"
 ENV PATH="/home/model-server/:${PATH}"
+ENV PYTHONPATH="$PYTHONPATH:/home/model-server/composed_image_retrieval:/home/model-server/composed_image_retrieval/src:/home/model-server"
 
 # Copy license.
 RUN apt-get update && apt-get install -y --no-install-recommends \
@@ -84,6 +85,7 @@ RUN pip uninstall dataclasses -y
 
 # Copy model artifacts.
 COPY model_oss/pic2word/handler.py /home/model-server/handler.py
+COPY model_oss/util/ /home/model-server/util/
 
 # Create torchserve configuration file.
 RUN echo \

community-content/vertex_model_garden/model_oss/pic2word/handler.py

@@ -2,7 +2,7 @@
 
 from argparse import Namespace  # pylint: disable=g-importing-member
 import os
-from typing import Any
+from typing import Any, List
 
 from absl import logging
 from data import CustomFolder
@@ -25,7 +25,7 @@ _COCO_DATASET_NAME = "coco"
 _MODEL_NAME = "ViT-L/14"
 _LOCAL_QUERY_PATH = "./query/"
 _IMAGE_OUTPUT_LOCAL_DIR = "demo_out/images"
-_OUTPUT_LOCAL_DIR = "/demo_out/"
+_OUTPUT_LOCAL_DIR = "./demo_out/"
 _DATA_DIR = "data"
 _CHECKPOINT_DIR = "checkpoint/pic2word_model.pt"
 _REQUEST_PROMPTS = "prompts"
@@ -33,6 +33,7 @@ _REQUEST_OUTPUT_STORAGE_DIR = "output_storage_dir"
 _REQUEST_IMAGE_PATH = "image_path"
 _REQUEST_IMAGE_FILE_NAME = "image_file_name"
 _RESPONSE_MSG = "Successfully retrieved images."
+_PICKLE_DIR_PATH = "gs://pic2word-bucket/pickle/"
 
 
 class ModelHandler(BaseHandler):
@@ -49,6 +50,8 @@ class ModelHandler(BaseHandler):
   def initialize(self, context: Any):
     """Initialize."""
     logging.info("Initializing pic2word.")
+    # Download pickle file for COCO
+    fileutils.download_gcs_dir_to_local(_PICKLE_DIR_PATH, "./data")
 
     # Download COCO dataset. The model looks for this folder specifically
     # during image retrieval to generate a response for each request.
@@ -157,11 +160,11 @@ class ModelHandler(BaseHandler):
         _IMAGE_OUTPUT_LOCAL_DIR, self.output_storage_dir
     )
 
-  def handle(self, data: Any, context: Any) -> str:  # pylint: disable=unused-argument
+  def handle(self, data: Any, context: Any) -> List[str]:  # pylint: disable=unused-argument
     """Runs preprocess, inference, and post-processing."""
     logging.info("Received Pic2Word inference request")
     model_input = self.preprocess(data)
     self.inference(model_input)
     self.postprocess()
     logging.info("Done handling input.")
-    return _RESPONSE_MSG
+    return [_RESPONSE_MSG]

community-content/vertex_model_garden/model_oss/transformers/biogpt_handler.py

@@ -0,0 +1,98 @@
+"""Custom handler for huggingface/biogpt models."""
+
+import os
+from typing import Any, List
+
+from absl import logging
+import torch
+from transformers import BioGptForCausalLM, BioGptTokenizer
+from transformers import pipeline
+from ts.torch_handler.base_handler import BaseHandler
+
+from util import constants
+from util import fileutils
+
+# Tasks
+TEXT_GENERATION = "text-generation"
+
+# prompt specific parameters
+MAX_LENGTH = 200
+NUM_RETURN_SEQUENCES = 10
+
+# Default Model ID
+DEFAULT_MODEL_ID = "microsoft/biogpt"
+
+logging.set_verbosity(os.environ.get("LOG_LEVEL", logging.INFO))
+
+
+class BioGPTHandler(BaseHandler):
+  """Custom handler for BioGPT models."""
+
+  def initialize(self, context: Any):
+    """Initializes the handler."""
+    logging.info("Start to initialize the BioGPT handler.")
+    properties = context.system_properties
+    self.map_location = (
+        "cuda"
+        if torch.cuda.is_available() and properties.get("gpu_id") is not None
+        else "cpu"
+    )
+
+    self.device = torch.device(
+        self.map_location + ":" + str(properties.get("gpu_id"))
+        if torch.cuda.is_available() and properties.get("gpu_id") is not None
+        else self.map_location
+    )
+    self.manifest = context.manifest
+    self.max_length = int(os.environ.get("MAX_LENGTH", MAX_LENGTH))
+    self.num_return_sequences = int(
+        os.environ.get("NUM_RETURN_SEQUENCES", NUM_RETURN_SEQUENCES)
+    )
+    self.base_model_id = os.environ.get("BASE_MODEL_ID", None)
+    self.model_id = self.base_model_id
+    if not self.base_model_id:
+      self.model_id = os.environ.get("MODEL_ID", "microsoft/biogpt")
+    if fileutils.is_gcs_path(self.model_id):
+      fileutils.download_gcs_dir_to_local(
+          self.model_id,
+          constants.LOCAL_BASE_MODEL_DIR,
+          skip_hf_model_bin=True,
+      )
+      self.model_id = constants.LOCAL_BASE_MODEL_DIR
+
+    logging.info(f"Using base model:{self.model_id}")
+
+    self.pipeline = None
+    self.tokenizer = None
+
+    self.tokenizer = BioGptTokenizer.from_pretrained(self.model_id)
+    logging.debug("Initialized the BioGPT tokenizer.")
+    model = BioGptForCausalLM.from_pretrained(self.model_id)
+    logging.debug("Initialized the base model.")
+    self.pipeline = pipeline(
+        TEXT_GENERATION, model=model, tokenizer=self.tokenizer
+    )
+
+    self.initialized = True
+    logging.info("The BioGPT handler was initialized.")
+
+  def preprocess(self, data: Any) -> str:
+    """Preprocess input data."""
+    prompt = data[0]["prompt"]
+    return prompt
+
+  def inference(self, data: Any, *args, **kwargs) -> str:
+    """Run the inference."""
+    logging.debug(f"Inference prompts={data}")
+    predicted_results = self.pipeline(
+        data,
+        max_length=self.max_length,
+        num_return_sequences=self.num_return_sequences,
+        do_sample=True,
+    )[0]["generated_text"]
+    return predicted_results
+
+  def postprocess(self, data: Any) -> List[str]:
+    """Postprocesses output data."""
+    output = data.replace("<|endoftext|></s>", "")
+    return [output]

community-content/vertex_model_garden/model_oss/util/fileutils.py

@@ -232,10 +232,13 @@ def download_video_from_gcs_to_local(video_file_path: str) -> Tuple[str, str]:
   """
   _, local_video_file_name = os.path.split(video_file_path)
   file_extension = os.path.splitext(video_file_path)[1]
-  remote_video_file_name = local_video_file_name.replace(
-      file_extension, '_overlay.mp4'
-  )
-  local_file_path = generate_tmp_path(os.path.splitext(video_file_path)[1])
+  if file_extension:
+    remote_video_file_name = local_video_file_name.replace(
+        file_extension, '_overlay.mp4'
+    )
+  else:
+    remote_video_file_name = local_video_file_name + '_overlay.mp4'
+  local_file_path = generate_tmp_path(file_extension)
   logging.info('Downloading %s to %s...', video_file_path, local_file_path)
   download_gcs_file_to_local(video_file_path, local_file_path)
   return local_file_path, remote_video_file_name
@@ -251,7 +254,10 @@ def get_output_video_file(video_output_file_path: str) -> str:
     str: Local video output file path.
   """
   file_extension = os.path.splitext(video_output_file_path)[1]
-  out_local_video_file_name = video_output_file_path.replace(
-      file_extension, '_overlay' + file_extension
-  )
+  if file_extension:
+    out_local_video_file_name = video_output_file_path.replace(
+        file_extension, '_overlay' + file_extension
+    )
+  else:
+    out_local_video_file_name = video_output_file_path + '_overlay'
   return out_local_video_file_name

community-content/vertex_model_garden/model_oss/vllm/dockerfile/serve.Dockerfile

@@ -1,5 +1,5 @@
 # Dockerfile for vLLM serving.
-#
+# It requires at least an n1-highmem-16 machine to build.
 # To build:
 # docker build -f model_oss/vllm/dockerfile/serve.Dockerfile . -t ${YOUR_IMAGE_TAG}
 #
@@ -7,67 +7,63 @@
 # docker tag ${YOUR_IMAGE_TAG} gcr.io/{YOUR_PROJECT}/${YOUR_IMAGE_TAG}
 # docker push gcr.io/{YOUR_PROJECT}/${YOUR_IMAGE_TAG}
 
-# The base image is required by vllm
+# The base image is required by vllm:
 # https://vllm.readthedocs.io/en/latest/getting_started/installation.html
+# Refer to the nvcr docker hub for the full list:
+# https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch/tags
 FROM nvcr.io/nvidia/pytorch:22.12-py3
 
 USER root
 
 # Install tools.
 ENV DEBIAN_FRONTEND=noninteractive
-RUN apt-get update
+RUN apt-get update && apt-get -y upgrade
 RUN apt-get install -y --no-install-recommends apt-utils
 RUN apt-get install -y --no-install-recommends curl
 RUN apt-get install -y --no-install-recommends wget
 RUN apt-get install -y --no-install-recommends git
-RUN apt-get install -y --no-install-recommends jq
-RUN apt-get install -y --no-install-recommends gnupg
 
 # Copy license.
 RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
 
 # Install libraries.
 ENV PIP_ROOT_USER_ACTION=ignore
-RUN python3 -m pip install --upgrade pip
+RUN python -m pip install --upgrade pip
 RUN pip install google-cloud-storage==2.7.0
 RUN pip install absl-py==1.4.0
-
-# Install pytorch
-RUN pip install --upgrade torch==2.0.1
+RUN pip install boto3==1.26.9
 
 # Install vllm deps.
-RUN pip install xformers==0.0.20
 RUN pip install ninja==1.11.1
 RUN pip install psutil==5.9.5
-RUN pip install ray==2.6.2
+RUN pip install ray==2.7.0
 RUN pip install sentencepiece==0.1.99
 RUN pip install fastapi==0.100.1
-RUN pip install uvicorn==0.23.2
+RUN pip install uvicorn[standard]==0.23.2
 RUN pip install pydantic==1.10.12
-
-# Install transformers from source.
-WORKDIR /workspace
-RUN git clone https://github.com/huggingface/transformers.git
-WORKDIR transformers
-# Pin the commit to add-code-llama at 08/25/2023
-RUN git reset --hard 015f8e110d270a0ad42de4ae5b98198d69eb1964
-RUN pip install -e .
-WORKDIR /workspace
+RUN pip install --upgrade torch==2.1.1 --index-url https://download.pytorch.org/whl/cu118
+RUN pip install --upgrade xformers==0.0.23 --index-url https://download.pytorch.org/whl/cu118
+RUN pip install transformers==4.34.0
+RUN pip install packaging==23.2
 
 # Install vllm from source.
 RUN git clone https://github.com/vllm-project/vllm.git
 WORKDIR vllm
-# Pin the version to a fixed git commit on 08/16/2023.
-RUN git reset --hard d1744376ae9fdbfa6a2dc763e1c67309e138fa3d
+# Pin the version to a fixed git commit on 12/20/2023.
+# https://github.com/vllm-project/vllm/tree/bd29cf3d3ad3dd06105f1a4bb9023bb23bdfd5ed
+RUN git reset --hard bd29cf3d3ad3dd06105f1a4bb9023bb23bdfd5ed
 # Apply a patch to vllm source:
 # 1) For models on Huggingface hub: if the model has multiple bin files, each
 #    bin file is downloaded separately and gets deleted after loading to GPU
 # 2) For models on GCS bucket: each model bin files is download separately
 #    and gets deleted after loading to GPU.
-# 3) Support code-llama model loading.
+
+
 COPY model_oss/vllm/vllm.patch /tmp/vllm.patch
 RUN git apply /tmp/vllm.patch
-RUN pip install -e .
+RUN pip install -e . -v
+
+COPY model_oss/vllm/vllm_startup_prober.sh /model_garden/scripts/vllm_startup_prober.sh
 
 # Expose port 7080 for host serving.
-EXPOSE 7080
+EXPOSE 7080

community-content/vertex_model_garden/model_oss/vllm/vllm.patch

@@ -1,20 +1,135 @@
+diff --git a/pyproject.toml b/pyproject.toml
+deleted file mode 100644
+index b197256..0000000
+--- a/pyproject.toml
++++ /dev/null
+@@ -1,34 +0,0 @@
+-[build-system]
+-# Should be mirrored in requirements-build.txt
+-requires = [
+-    "ninja",
+-    "packaging",
+-    "setuptools >= 49.4.0",
+-    "torch == 2.1.2",
+-    "wheel",
+-]
+-build-backend = "setuptools.build_meta"
+-
+-[tool.ruff.lint]
+-select = [
+-    # pycodestyle
+-    "E",
+-    # Pyflakes
+-    "F",
+-    # pyupgrade
+-    # "UP",
+-    # flake8-bugbear
+-    "B",
+-    # flake8-simplify
+-    "SIM",
+-    # isort
+-    # "I",
+-]
+-ignore = [
+-    # star imports
+-    "F405", "F403",
+-    # lambda expression assignment
+-    "E731",
+-    # line too long, handled by black formatting
+-    "E501",
+-]
+diff --git a/requirements.txt b/requirements.txt
+index 92ba0a7..3506a73 100644
+--- a/requirements.txt
++++ b/requirements.txt
+@@ -5,9 +5,7 @@ pandas  # Required for Ray data.
+ pyarrow  # Required for Ray data.
+ sentencepiece  # Required for LLaMA tokenizer.
+ numpy
+-torch == 2.1.2
+ transformers >= 4.36.0  # Required for Mixtral.
+-xformers == 0.0.23.post1  # Required for CUDA 12.1.
+ fastapi
+ uvicorn[standard]
+ pydantic == 1.10.13  # Required for OpenAI server.
+diff --git a/setup.py b/setup.py
+index 811d494..6e0ac70 100644
+--- a/setup.py
++++ b/setup.py
+@@ -12,7 +12,7 @@ from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME,
+ 
+ ROOT_DIR = os.path.dirname(__file__)
+ 
+-MAIN_CUDA_VERSION = "12.1"
++MAIN_CUDA_VERSION = "11.8"
+ 
+ # Supported NVIDIA GPU architectures.
+ NVIDIA_SUPPORTED_ARCHS = {"7.0", "7.5", "8.0", "8.6", "8.9", "9.0"}
+@@ -123,10 +123,11 @@ def get_torch_arch_list() -> Set[str]:
+     arch_list = torch_arch_list.intersection(valid_archs)
+     # If none of the specified architectures are valid, raise an error.
+     if not arch_list:
+-        raise RuntimeError(
++        print(
+             "None of the CUDA/ROCM architectures in `TORCH_CUDA_ARCH_LIST` env "
+             f"variable ({env_arch_list}) is supported. "
+             f"Supported CUDA/ROCM architectures are: {valid_archs}.")
++        return None
+     invalid_arch_list = torch_arch_list - valid_archs
+     if invalid_arch_list:
+         warnings.warn(
 diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py
-index 99fe593..e11246b 100644
+index 7e58069..6791248 100644
 --- a/vllm/engine/arg_utils.py
 +++ b/vllm/engine/arg_utils.py
-@@ -1,12 +1,43 @@
- import argparse
- import dataclasses
- from dataclasses import dataclass
-+import os
- from typing import Optional, Tuple
+@@ -5,6 +5,78 @@ from typing import Optional, Tuple
  
-+from google.cloud import storage
  from vllm.config import (CacheConfig, ModelConfig, ParallelConfig,
                           SchedulerConfig)
- 
- 
++from vllm.logger import init_logger
++import os
++from google.cloud import storage
++import boto3
++
++logger = init_logger(__name__)
 +GCS_PREFIX = "gs://"
++S3_PREFIX = "s3://"
++
++
++def is_s3_path(input_path: str) -> bool:
++    return input_path.startswith(S3_PREFIX)
++
++
++def download_s3_dir_to_local(s3_dir: str, local_dir: str):
++    if os.path.isdir(local_dir):
++        return
++    # s3://bucket_name/dir
++    bucket_name = s3_dir.split('/')[2]
++    prefix = s3_dir[len(S3_PREFIX + bucket_name) :].strip('/')
++
++    access_key_id = os.environ['AWS_ACCESS_KEY_ID']
++    secret_key = os.environ['AWS_SECRET_ACCESS_KEY']
++    client = boto3.client(
++        's3',
++        aws_access_key_id=access_key_id,
++        aws_secret_access_key=secret_key,
++    )
++    blobs = client.list_objects_v2(Bucket=bucket_name, Prefix=prefix)
++    if not blobs:
++        raise ValueError(f"No blobs found in {s3_dir}")
++    for blob in blobs['Contents']:
++        name = blob['Key']
++        if name[-1] == '/':
++            continue
++        file_path = name[len(prefix) :].strip('/')
++        local_file_path = os.path.join(local_dir, file_path)
++        os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
++        if file_path.endswith(".bin") or file_path.endswith(".safetensors"):
++            with open(local_file_path, 'w') as f:
++                f.write(f'{S3_PREFIX}{bucket_name}/{prefix}/{file_path}')
++        else:
++            print(f"==> Download {s3_dir}/{file_path} to {local_file_path}")
++            client.download_file(bucket_name, name, local_file_path)
 +
 +
 +def is_gcs_path(input_path: str) -> bool:
@@ -29,283 +144,651 @@ index 99fe593..e11246b 100644
 +    prefix = gcs_dir[len(GCS_PREFIX + bucket_name) :].strip('/')
 +    client = storage.Client()
 +    blobs = client.list_blobs(bucket_name, prefix=prefix)
++    if not blobs:
++        raise ValueError(f"No blobs found in {gcs_dir}")
 +    for blob in blobs:
 +        if blob.name[-1] == '/':
 +            continue
 +        file_path = blob.name[len(prefix) :].strip('/')
 +        local_file_path = os.path.join(local_dir, file_path)
 +        os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
-+        if file_path.endswith(".bin"):
++        if file_path.endswith(".bin") or file_path.endswith(".safetensors"):
 +            with open(local_file_path, 'w') as f:
 +                f.write(f'{GCS_PREFIX}{bucket_name}/{prefix}/{file_path}')
 +        else:
 +            print(f"==> Download {gcs_dir}/{file_path} to {local_file_path}")
 +            blob.download_to_filename(local_file_path)
-+
-+
+ 
+ 
  @dataclass
- class EngineArgs:
-     """Arguments for vLLM engine."""
-@@ -143,6 +174,19 @@ class EngineArgs:
-     def create_engine_configs(
-         self,
-     ) -> Tuple[ModelConfig, CacheConfig, ParallelConfig, SchedulerConfig]:
-+        # Preprocess GCS paths.
+@@ -37,6 +109,14 @@ class EngineArgs:
+     max_context_len_to_capture: int = 8192
+ 
+     def __post_init__(self):
++        if not self.model:
++            self.model = os.environ.get("AIP_STORAGE_URI", "")
++            logger.info(
++                f"Load self.model from AIP_STORAGE_URI: {self.model}."
++            )
++        if not self.model:
++            raise ValueError("self.model is must be set.")
++
+         if self.tokenizer is None:
+             self.tokenizer = self.model
+ 
+@@ -52,7 +132,7 @@ class EngineArgs:
+         parser.add_argument(
+             '--model',
+             type=str,
+-            default='facebook/opt-125m',
++            default=None,
+             help='name or path of the huggingface model to use')
+         parser.add_argument(
+             '--tokenizer',
+@@ -212,9 +292,39 @@ class EngineArgs:
+         engine_args = cls(**{attr: getattr(args, attr) for attr in attrs})
+         return engine_args
+ 
++    def process_gcs(self):
++        # Download GCS tokenizer.
 +        if is_gcs_path(self.tokenizer) and self.tokenizer != self.model:
 +            local_dir = "/tmp/gcs_tokenizer"
 +            download_gcs_dir_to_local(self.tokenizer, local_dir)
 +            self.tokenizer = local_dir
++        # Download GCS model without bin files.
 +        if is_gcs_path(self.model):
-+            # Download GCS model without bin files.
 +            local_dir = "/tmp/gcs_model"
 +            download_gcs_dir_to_local(self.model, local_dir)
 +            if self.tokenizer == self.model:
 +                self.tokenizer = local_dir
 +            self.model = local_dir
 +
-         # Initialize the configs.
++    def process_s3(self):
++        # Download S3 tokenizer.
++        if is_s3_path(self.tokenizer) and self.tokenizer != self.model:
++            local_dir = "/tmp/s3_tokenizer"
++            download_s3_dir_to_local(self.tokenizer, local_dir)
++            self.tokenizer = local_dir
++        # Download S3 model without bin files.
++        if is_s3_path(self.model):
++            local_dir = "/tmp/s3_model"
++            download_s3_dir_to_local(self.model, local_dir)
++            if self.tokenizer == self.model:
++                self.tokenizer = local_dir
++            self.model = local_dir
++
+     def create_engine_configs(
+         self,
+     ) -> Tuple[ModelConfig, CacheConfig, ParallelConfig, SchedulerConfig]:
++        self.process_gcs()
++        self.process_s3()
          model_config = ModelConfig(self.model, self.tokenizer,
                                     self.tokenizer_mode, self.trust_remote_code,
+                                    self.download_dir, self.load_format,
+diff --git a/vllm/engine/async_llm_engine.py b/vllm/engine/async_llm_engine.py
+index d854a20..158b5fe 100644
+--- a/vllm/engine/async_llm_engine.py
++++ b/vllm/engine/async_llm_engine.py
+@@ -377,9 +377,8 @@ class AsyncLLMEngine:
+                     shortened_token_ids = shortened_token_ids[:self.
+                                                               max_log_len]
+             logger.info(f"Received request {request_id}: "
+-                        f"prompt: {shortened_prompt!r}, "
+-                        f"sampling params: {sampling_params}, "
+-                        f"prompt token ids: {shortened_token_ids}.")
++                        f"prompt len: {len(shortened_prompt)}, "
++                        f"sampling params: {sampling_params}.")
+ 
+         if not self.is_running:
+             if self.start_engine_loop:
 diff --git a/vllm/entrypoints/api_server.py b/vllm/entrypoints/api_server.py
-index 58ea2e2..350e209 100644
+index 6910b32..d5bbe34 100644
 --- a/vllm/entrypoints/api_server.py
 +++ b/vllm/entrypoints/api_server.py
-@@ -15,6 +15,10 @@ TIMEOUT_KEEP_ALIVE = 5  # seconds.
+@@ -1,4 +1,5 @@
+ import argparse
++import copy
+ import json
+ from typing import AsyncGenerator
+ 
+@@ -10,13 +11,27 @@ from vllm.engine.arg_utils import AsyncEngineArgs
+ from vllm.engine.async_llm_engine import AsyncLLMEngine
+ from vllm.sampling_params import SamplingParams
+ from vllm.utils import random_uuid
++from vllm.entrypoints.openai.api_server import init_openai_api_server, create_chat_completion
++from vllm.entrypoints.openai.protocol import ChatCompletionRequest
++from vllm.logger import init_logger
+ 
++logger = init_logger(__name__)
+ TIMEOUT_KEEP_ALIVE = 5  # seconds.
  TIMEOUT_TO_PREVENT_DEADLOCK = 1  # seconds.
  app = FastAPI()
+ engine = None
+ 
  
 +# Required by Vertex deployment.
 +@app.get("/ping")
 +async def ping() -> Response:
 +    return Response(status_code=200)
- 
- @app.post("/generate")
- async def generate(request: Request) -> Response:
-@@ -26,6 +30,9 @@ async def generate(request: Request) -> Response:
++
++def format_output(prompt: str, output: str):
++    output = output.strip("\n")
++    return f"Prompt:\n{prompt.strip()}\nOutput:\n{output}"
++
++
+ @app.get("/health")
+ async def health() -> Response:
+     """Health check."""
+@@ -33,8 +48,16 @@ async def generate(request: Request) -> Response:
      - other fields: the sampling parameters (See `SamplingParams` for details).
      """
      request_dict = await request.json()
++    is_chat_completion = request_dict.get("@requestFormat", "") == "chatCompletions"
++    if is_chat_completion:
++        chat_completion_request = ChatCompletionRequest(**request_dict)
++        return await create_chat_completion(chat_completion_request, request)
 +    is_on_vertex = "instances" in request_dict
 +    if is_on_vertex:
 +        request_dict = request_dict["instances"][0]
      prompt = request_dict.pop("prompt")
      stream = request_dict.pop("stream", False)
++    raw_response = request_dict.pop("raw_response", False)
      sampling_params = SamplingParams(**request_dict)
-@@ -63,7 +70,10 @@ async def generate(request: Request) -> Response:
+     request_id = random_uuid()
+ 
+@@ -42,12 +65,33 @@ async def generate(request: Request) -> Response:
+ 
+     # Streaming case
+     async def stream_results() -> AsyncGenerator[bytes, None]:
++        prior_request_output = None
+         async for request_output in results_generator:
+             prompt = request_output.prompt
+-            text_outputs = [
+-                prompt + output.text for output in request_output.outputs
+-            ]
+-            ret = {"text": text_outputs}
++            text_outputs = []
++            for i, output in enumerate(request_output.outputs):
++                if prior_request_output is not None:
++                    prior_output = prior_request_output.outputs[i]
++                    text_output = output.text[len(prior_output.text):]
++                else:
++                    text_output = output.text
++                text_outputs.append(text_output)
++            ret = {"predictions": text_outputs}
++            if raw_response:
++                output_token_counts = []
++                for i, output in enumerate(request_output.outputs):
++                    if prior_request_output is not None:
++                        prior_output = prior_request_output.outputs[i]
++                        output_token_count = len(output.token_ids) - len(prior_output.token_ids)
++                    else:
++                        output_token_count = len(output.token_ids)
++                    output_token_counts.append(output_token_count)
++                cumulative_logprobs = [output.cumulative_logprob for output in request_output.outputs]
++                ret.update({
++                    "output_token_counts": output_token_counts,
++                    "cumulative_logprobs": cumulative_logprobs
++                })
++            prior_request_output = copy.deepcopy(request_output)
+             yield (json.dumps(ret) + "\0").encode("utf-8")
+ 
+     if stream:
+@@ -63,24 +107,40 @@ async def generate(request: Request) -> Response:
+         final_output = request_output
+ 
      assert final_output is not None
-     prompt = final_output.prompt
-     text_outputs = [prompt + output.text for output in final_output.outputs]
+-    prompt = final_output.prompt
+-    text_outputs = [prompt + output.text for output in final_output.outputs]
 -    ret = {"text": text_outputs}
-+    if is_on_vertex:
-+        ret = {"predictions": text_outputs}
++    if raw_response:
++        text_outputs = [output.text for output in final_output.outputs]
++        output_token_counts = [len(output.token_ids) for output in final_output.outputs]
++        cumulative_logprobs = [output.cumulative_logprob for output in final_output.outputs]
++        ret = {
++            "predictions": text_outputs,
++            "output_token_counts": output_token_counts,
++            "cumulative_logprobs": cumulative_logprobs
++        }
 +    else:
-+        ret = {"text": text_outputs}
++        prompt = final_output.prompt
++        text_outputs = [format_output(prompt, output.text) for output in final_output.outputs]
++        ret = {"predictions": text_outputs}
      return JSONResponse(ret)
  
  
-diff --git a/vllm/model_executor/models/llama.py b/vllm/model_executor/models/llama.py
-index 93ab499..eca1b89 100644
---- a/vllm/model_executor/models/llama.py
-+++ b/vllm/model_executor/models/llama.py
-@@ -85,6 +85,7 @@ class LlamaAttention(nn.Module):
-         hidden_size: int,
-         num_heads: int,
-         num_kv_heads: int,
-+        rope_theta: float = 10000,
-     ):
-         super().__init__()
-         self.hidden_size = hidden_size
-@@ -99,6 +100,7 @@ class LlamaAttention(nn.Module):
-         self.q_size = self.num_heads * self.head_dim
-         self.kv_size = self.num_kv_heads * self.head_dim
-         self.scaling = self.head_dim**-0.5
-+        self.rope_theta = rope_theta
+ if __name__ == "__main__":
++    logger.info("Starting API server...")
+     parser = argparse.ArgumentParser()
+     parser.add_argument("--host", type=str, default=None)
+     parser.add_argument("--port", type=int, default=8000)
+     parser.add_argument("--ssl-keyfile", type=str, default=None)
+     parser.add_argument("--ssl-certfile", type=str, default=None)
++    parser.add_argument("--chat-template", type=str, default=None)
++    parser.add_argument("--response-role", type=str, default="assistant")
+     parser = AsyncEngineArgs.add_cli_args(parser)
+     args = parser.parse_args()
  
-         self.qkv_proj = ColumnParallelLinear(
-             hidden_size,
-@@ -118,6 +120,7 @@ class LlamaAttention(nn.Module):
-         self.attn = PagedAttentionWithRoPE(self.num_heads,
-                                            self.head_dim,
-                                            self.scaling,
-+                                           base=self.rope_theta,
-                                            rotary_dim=self.head_dim,
-                                            num_kv_heads=self.num_kv_heads)
+     engine_args = AsyncEngineArgs.from_cli_args(args)
+     engine = AsyncLLMEngine.from_engine_args(engine_args)
  
-@@ -143,10 +146,15 @@ class LlamaDecoderLayer(nn.Module):
-     def __init__(self, config: LlamaConfig):
-         super().__init__()
-         self.hidden_size = config.hidden_size
-+        try:
-+            rope_theta = config.rope_theta
-+        except AttributeError:
-+            rope_theta = 10000
-         self.self_attn = LlamaAttention(
-             hidden_size=self.hidden_size,
-             num_heads=config.num_attention_heads,
-             num_kv_heads=config.num_key_value_heads,
-+            rope_theta=rope_theta,
-         )
-         self.mlp = LlamaMLP(
-             hidden_size=self.hidden_size,
++    logger.info("Initializing OpenAI API server...")
++    init_openai_api_server(args, engine)
++
+     uvicorn.run(app,
+                 host=args.host,
+                 port=args.port,
+diff --git a/vllm/entrypoints/openai/api_server.py b/vllm/entrypoints/openai/api_server.py
+index be5f419..0cda03e 100644
+--- a/vllm/entrypoints/openai/api_server.py
++++ b/vllm/entrypoints/openai/api_server.py
+@@ -37,65 +37,10 @@ from vllm.utils import random_uuid
+ TIMEOUT_KEEP_ALIVE = 5  # seconds
+ 
+ logger = init_logger(__name__)
+-served_model = None
+-app = fastapi.FastAPI()
+ engine = None
+ response_role = None
+-
+-
+-def parse_args():
+-    parser = argparse.ArgumentParser(
+-        description="vLLM OpenAI-Compatible RESTful API server.")
+-    parser.add_argument("--host", type=str, default=None, help="host name")
+-    parser.add_argument("--port", type=int, default=8000, help="port number")
+-    parser.add_argument("--allow-credentials",
+-                        action="store_true",
+-                        help="allow credentials")
+-    parser.add_argument("--allowed-origins",
+-                        type=json.loads,
+-                        default=["*"],
+-                        help="allowed origins")
+-    parser.add_argument("--allowed-methods",
+-                        type=json.loads,
+-                        default=["*"],
+-                        help="allowed methods")
+-    parser.add_argument("--allowed-headers",
+-                        type=json.loads,
+-                        default=["*"],
+-                        help="allowed headers")
+-    parser.add_argument("--served-model-name",
+-                        type=str,
+-                        default=None,
+-                        help="The model name used in the API. If not "
+-                        "specified, the model name will be the same as "
+-                        "the huggingface name.")
+-    parser.add_argument("--chat-template",
+-                        type=str,
+-                        default=None,
+-                        help="The file path to the chat template, "
+-                        "or the template in single-line form "
+-                        "for the specified model")
+-    parser.add_argument("--response-role",
+-                        type=str,
+-                        default="assistant",
+-                        help="The role name to return if "
+-                        "`request.add_generation_prompt=true`.")
+-    parser.add_argument("--ssl-keyfile",
+-                        type=str,
+-                        default=None,
+-                        help="The file path to the SSL key file")
+-    parser.add_argument("--ssl-certfile",
+-                        type=str,
+-                        default=None,
+-                        help="The file path to the SSL cert file")
+-
+-    parser = AsyncEngineArgs.add_cli_args(parser)
+-    return parser.parse_args()
+-
+-
+-app.add_middleware(MetricsMiddleware)  # Trace HTTP server metrics
+-app.add_route("/metrics", metrics)  # Exposes HTTP metrics
+-
++max_model_len = None
++tokenizer = None
+ 
+ def create_error_response(status_code: HTTPStatus,
+                           message: str) -> JSONResponse:
+@@ -123,21 +68,10 @@ def load_chat_template(args, tokenizer):
+         logger.warning("No chat template provided. Chat API will not work.")
+ 
+ 
+-@app.exception_handler(RequestValidationError)
+ async def validation_exception_handler(_, exc):
+     return create_error_response(HTTPStatus.BAD_REQUEST, str(exc))
+ 
+ 
+-async def check_model(request) -> Optional[JSONResponse]:
+-    if request.model == served_model:
+-        return
+-    ret = create_error_response(
+-        HTTPStatus.NOT_FOUND,
+-        f"The model `{request.model}` does not exist.",
+-    )
+-    return ret
+-
+-
+ async def check_length(
+     request: Union[ChatCompletionRequest, CompletionRequest],
+     prompt: Optional[str] = None,
+@@ -165,23 +99,11 @@ async def check_length(
+         return input_ids, None
+ 
+ 
+-@app.get("/health")
+ async def health() -> Response:
+     """Health check."""
+     return Response(status_code=200)
+ 
+ 
+-@app.get("/v1/models")
+-async def show_available_models():
+-    """Show available models. Right now we only have one model."""
+-    model_cards = [
+-        ModelCard(id=served_model,
+-                  root=served_model,
+-                  permission=[ModelPermission()])
+-    ]
+-    return ModelList(data=model_cards)
+-
+-
+ def create_logprobs(
+     token_ids: List[int],
+     top_logprobs: Optional[List[Optional[Dict[int, float]]]] = None,
+@@ -217,7 +139,6 @@ def create_logprobs(
+     return logprobs
+ 
+ 
+-@app.post("/v1/chat/completions")
+ async def create_chat_completion(request: ChatCompletionRequest,
+                                  raw_request: Request):
+     """Completion API similar to OpenAI's API.
+@@ -229,12 +150,8 @@ async def create_chat_completion(request: ChatCompletionRequest,
+         - function_call (Users should implement this by themselves)
+         - logit_bias (to be supported by vLLM engine)
+     """
+-    error_check_ret = await check_model(request)
+-    if error_check_ret is not None:
+-        return error_check_ret
+ 
+     if request.logit_bias is not None and len(request.logit_bias) > 0:
+-        # TODO: support logit_bias in vLLM engine.
+         return create_error_response(HTTPStatus.BAD_REQUEST,
+                                      "logit_bias is not currently supported")
+ 
+@@ -438,7 +355,6 @@ async def create_chat_completion(request: ChatCompletionRequest,
+         return await completion_full_generator()
+ 
+ 
+-@app.post("/v1/completions")
+ async def create_completion(request: CompletionRequest, raw_request: Request):
+     """Completion API similar to OpenAI's API.
+ 
+@@ -451,10 +367,6 @@ async def create_completion(request: CompletionRequest, raw_request: Request):
+         - logit_bias (to be supported by vLLM engine)
+     """
+ 
+-    error_check_ret = await check_model(request)
+-    if error_check_ret is not None:
+-        return error_check_ret
+-
+     # OpenAI API supports echoing the prompt when max_tokens is 0.
+     echo_without_generation = request.echo and request.max_tokens == 0
+ 
+@@ -464,7 +376,6 @@ async def create_completion(request: CompletionRequest, raw_request: Request):
+                                      "suffix is not currently supported")
+ 
+     if request.logit_bias is not None and len(request.logit_bias) > 0:
+-        # TODO: support logit_bias in vLLM engine.
+         return create_error_response(HTTPStatus.BAD_REQUEST,
+                                      "logit_bias is not currently supported")
+ 
+@@ -481,7 +392,6 @@ async def create_completion(request: CompletionRequest, raw_request: Request):
+             use_token_ids = True
+             prompt = request.prompt
+         elif isinstance(first_element, (str, list)):
+-            # TODO: handles multiple prompt case in list[list[int]]
+             if len(request.prompt) > 1:
+                 return create_error_response(
+                     HTTPStatus.BAD_REQUEST,
+@@ -713,45 +623,24 @@ async def create_completion(request: CompletionRequest, raw_request: Request):
+     return response
+ 
+ 
+-if __name__ == "__main__":
+-    args = parse_args()
+-
+-    app.add_middleware(
+-        CORSMiddleware,
+-        allow_origins=args.allowed_origins,
+-        allow_credentials=args.allow_credentials,
+-        allow_methods=args.allowed_methods,
+-        allow_headers=args.allowed_headers,
+-    )
+-
++def init_openai_api_server(args, arg_engine):
+     logger.info(f"args: {args}")
+ 
+-    if args.served_model_name is not None:
+-        served_model = args.served_model_name
+-    else:
+-        served_model = args.model
+-
++    global response_role
+     response_role = args.response_role
+ 
+-    engine_args = AsyncEngineArgs.from_cli_args(args)
+-    engine = AsyncLLMEngine.from_engine_args(engine_args)
++    global engine
++    engine = arg_engine
++
+     engine_model_config = asyncio.run(engine.get_model_config())
++
++    global max_model_len
+     max_model_len = engine_model_config.max_model_len
+ 
+     # A separate tokenizer to map token IDs to strings.
++    global tokenizer
+     tokenizer = get_tokenizer(
+         engine_model_config.tokenizer,
+         tokenizer_mode=engine_model_config.tokenizer_mode,
+         trust_remote_code=engine_model_config.trust_remote_code)
+     load_chat_template(args, tokenizer)
+-
+-    # Register labels for metrics
+-    add_global_metrics_labels(model_name=engine_args.model)
+-
+-    uvicorn.run(app,
+-                host=args.host,
+-                port=args.port,
+-                log_level="info",
+-                timeout_keep_alive=TIMEOUT_KEEP_ALIVE,
+-                ssl_keyfile=args.ssl_keyfile,
+-                ssl_certfile=args.ssl_certfile)
 diff --git a/vllm/model_executor/weight_utils.py b/vllm/model_executor/weight_utils.py
-index a9d899a..57f39b5 100644
+index 365c847..eeb9c75 100644
 --- a/vllm/model_executor/weight_utils.py
 +++ b/vllm/model_executor/weight_utils.py
-@@ -3,13 +3,17 @@ import filelock
- import glob
- import json
- import os
+@@ -286,3 +286,181 @@ def initialize_dummy_weights(
+     for param in model.state_dict().values():
+         if torch.is_floating_point(param):
+             param.data.uniform_(low, high)
++
++
 +import time
- from typing import Iterator, List, Optional, Tuple
- 
--from huggingface_hub import snapshot_download
++import boto3
 +from google.cloud import storage
-+from huggingface_hub import hf_hub_download, snapshot_download
- import numpy as np
- import torch
- from tqdm.auto import tqdm
- 
++from huggingface_hub import hf_hub_download
++
 +HF_PREFIX = "hf://"
++MODEL_DIR = "/tmp/vllm_model"
 +
- 
- class Disabledtqdm(tqdm):
- 
-@@ -22,60 +26,90 @@ def hf_model_weights_iterator(
-     cache_dir: Optional[str] = None,
-     use_np_cache: bool = False,
- ) -> Iterator[Tuple[str, torch.Tensor]]:
-+    if use_np_cache:
-+        raise ValueError("Do not support use_np_cache for lazy download.")
 +
-     # Prepare file lock directory to prevent multiple processes from
-     # downloading the same model weights at the same time.
-     lock_dir = cache_dir if cache_dir is not None else "/tmp"
-     lock_file_name = model_name_or_path.replace("/", "-") + ".lock"
-     lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name))
- 
--    # Download model weights from huggingface.
--    is_local = os.path.isdir(model_name_or_path)
--    if not is_local:
--        with lock:
--            hf_folder = snapshot_download(model_name_or_path,
--                                          allow_patterns="*.bin",
--                                          cache_dir=cache_dir,
--                                          tqdm_class=Disabledtqdm)
--    else:
--        hf_folder = model_name_or_path
--
--    hf_bin_files = [
--        x for x in glob.glob(os.path.join(hf_folder, "*.bin"))
--        if not x.endswith("training_args.bin")
--    ]
--
--    if use_np_cache:
--        # Convert the model weights from torch tensors to numpy arrays for
--        # faster loading.
--        np_folder = os.path.join(hf_folder, "np")
--        os.makedirs(np_folder, exist_ok=True)
--        weight_names_file = os.path.join(np_folder, "weight_names.json")
--        with lock:
--            if not os.path.exists(weight_names_file):
--                weight_names = []
--                for bin_file in hf_bin_files:
--                    state = torch.load(bin_file, map_location="cpu")
--                    for name, param in state.items():
--                        param_path = os.path.join(np_folder, name)
--                        with open(param_path, "wb") as f:
--                            np.save(f, param.cpu().detach().numpy())
--                        weight_names.append(name)
--                with open(weight_names_file, "w") as f:
--                    json.dump(weight_names, f)
--
--        with open(weight_names_file, "r") as f:
--            weight_names = json.load(f)
--
--        for name in weight_names:
--            param_path = os.path.join(np_folder, name)
--            with open(param_path, "rb") as f:
--                param = np.load(f)
--            yield name, torch.from_numpy(param)
-+    bin_files = []
++def prepare_hf_model_weights_on_the_fly(
++    model_name_or_path: str,
++    cache_dir: Optional[str] = None,
++    use_safetensors: bool = False,
++    fall_back_to_pt: bool = True,
++    revision: Optional[str] = None,
++) -> Tuple[List[str], bool]:
++    logger.info("Loading weights on the fly.")
++    lock = get_lock(model_name_or_path, cache_dir)
++
++    hf_weights_files = []
++    if use_safetensors:
++        logger.info("Looking for .safetensors files")
++        index_filename = "model.safetensors.index.json"
++        allow_patterns = "*.safetensors"
++    else:
++        logger.info("Looking for .bin files")
++        index_filename = "pytorch_model.bin.index.json"
++        allow_patterns = "*.bin"
 +    if not os.path.isdir(model_name_or_path):
 +        try:
 +            with lock:
 +                index_file = hf_hub_download(repo_id=model_name_or_path,
-+                                             filename="pytorch_model.bin.index.json",
++                                             filename=index_filename,
 +                                             cache_dir=cache_dir)
 +        except:
-+            print("==> The model is in HF hub with 1 bin file, download it directly.", flush=True)
++            logger.info("The model is in HF hub with 1 file, download it directly.")
 +            with lock:
 +                hf_folder = snapshot_download(repo_id=model_name_or_path,
-+                                              allow_patterns="*.bin",
++                                              allow_patterns=allow_patterns,
 +                                              cache_dir=cache_dir,
 +                                              tqdm_class=Disabledtqdm)
-+            bin_files = [x for x in glob.glob(os.path.join(hf_folder, "*.bin"))]
++            hf_weights_files = [x for x in glob.glob(os.path.join(hf_folder, allow_patterns))]
 +        else:
-+            print("==> The model is in HF hub with multiple bin file, do not download it now.", flush=True)
++            logger.info("The model is in HF hub with multiple files, do not download it now.")
 +            with open(index_file, "r") as f:
 +                index = json.loads(f.read())
-+            bin_filenames = set(index["weight_map"].values())
-+            bin_files = [f"{HF_PREFIX}{model_name_or_path}/{bin_filename}" for bin_filename in bin_filenames]
-     else:
--        for bin_file in hf_bin_files:
--            state = torch.load(bin_file, map_location="cpu")
--            for name, param in state.items():
--                yield name, param
-+        print("==> The model is in local disk.", flush=True)
-+        bin_files = [x for x in glob.glob(os.path.join(model_name_or_path, "*.bin"))]
++            weight_filenames = set(index["weight_map"].values())
++            hf_weights_files = [f"{HF_PREFIX}{model_name_or_path}/{weight_filename}" for weight_filename in weight_filenames]
++    else:
++        logger.info("The model is possibly in local disk.")
++        hf_weights_files = [x for x in glob.glob(os.path.join(model_name_or_path, allow_patterns))]
 +
-+    if "training_args.bin" in bin_files:
-+        bin_files.remove("training_args.bin")
-+    bin_files.sort()
-+    print(f"==> Fetched bin files: {bin_files}", flush=True)
++    if not use_safetensors:
++        # Exclude files that are not needed for inference.
++        # https://github.com/huggingface/transformers/blob/v4.34.0/src/transformers/trainer.py#L227-L233
++        blacklist = [
++            "training_args.bin",
++            "optimizer.bin",
++            "optimizer.pt",
++            "scheduler.pt",
++            "scaler.pt",
++        ]
++        hf_weights_files = [
++            f for f in hf_weights_files
++            if not any(f.endswith(x) for x in blacklist)
++        ]
++    hf_weights_files.sort()
 +
-+    model_dir = "/tmp/model"
-+    os.makedirs(model_dir, exist_ok=True)
-+    for bin_file in bin_files:
++    if not hf_weights_files and use_safetensors:
++        return prepare_hf_model_weights_on_the_fly(model_name_or_path,
++                                        cache_dir=cache_dir,
++                                        use_safetensors=False,
++                                        fall_back_to_pt=False,
++                                        revision=revision)
++    if not hf_weights_files:
++        raise RuntimeError(f"No weight files found in {model_name_or_path}")
++    logger.info(f"Fetched weight files: {hf_weights_files}")
++    return hf_weights_files, use_safetensors
++
++
++def hf_model_weights_iterator_download_on_the_fly(
++    model_name_or_path: str,
++    cache_dir: Optional[str] = None,
++    load_format: str = "auto",
++    revision: Optional[str] = None,
++    fall_back_to_pt: Optional[bool] = True,
++) -> Iterator[Tuple[str, torch.Tensor]]:
++    lock = get_lock(model_name_or_path, cache_dir)
++    hf_weights_files, use_safetensors = prepare_hf_model_weights_on_the_fly(
++            model_name_or_path=model_name_or_path,
++            cache_dir=cache_dir,
++            use_safetensors=True,
++            fall_back_to_pt=fall_back_to_pt,
++            revision=revision)
++    os.makedirs(MODEL_DIR, exist_ok=True)
++    for hf_weight_file in hf_weights_files:
 +        delete_download = False
 +
-+        if os.path.exists(bin_file):
-+            if open(bin_file, "rb").read(2) == b"gs":
-+                gcs_path = open(bin_file).read()
-+                bin_filename = gcs_path.split("/")[-1]
-+                local_file = os.path.join(model_dir, bin_filename)
++        if os.path.exists(hf_weight_file):
++            prefix = open(hf_weight_file, "rb").read(2)
++            # Download from GCS.
++            if prefix == b"gs":
++                gcs_path = open(hf_weight_file).read()
++                hf_weight_filename = gcs_path.split("/")[-1]
++                local_file = os.path.join(MODEL_DIR, hf_weight_filename)
 +                with lock:
 +                    if not os.path.exists(local_file):
 +                        client = storage.Client()
 +                        with open(local_file, 'wb') as f:
-+                            print(f"==> Download {gcs_path} to {bin_file}", flush=True)
++                            logger.info(f"Download {gcs_path} to {hf_weight_file}")
 +                            client.download_blob_to_file(gcs_path, f)
-+                bin_file = local_file
++                hf_weight_file = local_file
 +                delete_download = True
++            # Download from S3.
++            elif prefix == b"s3":
++                s3_path = open(hf_weight_file).read()
++                hf_weight_filename = s3_path.split("/")[-1]
++                local_file = os.path.join(MODEL_DIR, hf_weight_filename)
++
++                bucket_name = s3_path.split('/')[2]
++                obj_key = s3_path.split(bucket_name)[1][1:]
++                with lock:
++                    if not os.path.exists(local_file):
++                        access_key_id = os.environ['AWS_ACCESS_KEY_ID']
++                        secret_key = os.environ['AWS_SECRET_ACCESS_KEY']
++                        client = boto3.client(
++                            's3',
++                            aws_access_key_id=access_key_id,
++                            aws_secret_access_key=secret_key,
++)
++                        with open(local_file, 'wb') as f:
++                            logger.info(f"Download {s3_path} to {hf_weight_file}")
++                            client.download_fileobj(bucket_name, obj_key, f)
++                hf_weight_file = local_file
++                delete_download = True
++
 +        else:
-+            assert bin_file.startswith(HF_PREFIX)
-+            bin_filename = os.path.basename(bin_file)
-+            local_file = os.path.join(model_dir, bin_filename)
++            # Download from HF.
++            assert hf_weight_file.startswith(HF_PREFIX)
++            hf_weight_filename = os.path.basename(hf_weight_file)
++            local_file = os.path.join(MODEL_DIR, hf_weight_filename)
 +            with lock:
 +                if not os.path.exists(local_file):
-+                    print(f"==> Download {model_name_or_path}/{bin_filename} to {local_file}", flush=True)
++                    logger.info(f"Download {model_name_or_path}/{hf_weight_filename} to {local_file}")
 +                    hf_hub_download(repo_id=model_name_or_path,
-+                                    filename=bin_filename,
-+                                    local_dir=model_dir,
++                                    filename=hf_weight_filename,
++                                    local_dir=MODEL_DIR,
 +                                    local_dir_use_symlinks=False,
 +                                    force_download=True)
-+            bin_file = local_file
++            hf_weight_file = local_file
 +            delete_download = True
 +
-+        torch.distributed.barrier()
-+        print(f"==> Load {bin_file} to memory.", flush=True)
-+        state = torch.load(bin_file, map_location="cpu")
-+        for name, param in state.items():
-+            yield name, param
-+        torch.distributed.barrier()
++        if use_safetensors:
++            with safe_open(hf_weight_file, framework="pt") as f:
++                for name in f.keys():
++                    param = f.get_tensor(name)
++                    yield name, param
++                torch.distributed.barrier()
++        else:
++            torch.distributed.barrier()
++            logger.info(f"Load {hf_weight_file} to memory.")
++            state = torch.load(hf_weight_file, map_location="cpu")
++            for name, param in state.items():
++                yield name, param
++            del state
++            torch.cuda.empty_cache()
++            torch.distributed.barrier()
 +
 +        if delete_download:
 +            with lock:
-+                if os.path.exists(bin_file):
-+                    print(f"==> Delete {bin_file}", flush=True)
-+                    os.remove(bin_file)
- 
- 
- def load_tensor_parallel_weights(
- 
++                if os.path.exists(hf_weight_file):
++                    logger.info(f"Delete {hf_weight_file}")
++                    os.remove(hf_weight_file)
++
++
++hf_model_weights_iterator = hf_model_weights_iterator_download_on_the_fly

community-content/vertex_model_garden/model_oss/vllm/vllm_startup_prober.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+
+# !/bin/bash
+# The Startup prober is built to check whether the server is ready to
+# serve traffic. The stript returns 0 if succeed. Any other returned
+# value are consider as an error. More detail could be found from
+# [shell script Exit codes](http://shellscript.sh/exitcodes.html).
+
+PORT=7080
+check_model_availability(){
+  curl -s -o /dev/null -w "%{http_code}" "http://0.0.0.0:${PORT}/health" | grep "200" -q
+}
+
+main(){
+  check_model_availability
+  local available=$?
+  if [[ $available -gt 0 ]]
+  then
+    echo "Warning: vLLM server is not yet available."
+    return 1
+  fi
+  return 0
+}
+main
+

community-content/vertex_pipeline_kfpv2/async_prediction.py

@@ -0,0 +1,32 @@
+import numpy as np
+from kfp.v2 import dsl
+
+@dsl.component(base_image='python:3.8',packages_to_install=['google-cloud-aiplatform==1.36.0'])
+def async_predict(
+    endpoint_id: str,
+    instances: dict,
+) -> np.ndarray:
+    import numpy as np
+    from google.cloud import aiplatform
+
+    endpoint = aiplatform.Endpoint(endpoint_id)
+    response = await endpoint.predict_async(instances)
+    predictions = np.asarray(response.predictions)
+    print(predictions.tolist())
+    return predictions
+
+@dsl.pipeline(name='async-prediction')
+def pipeline_prediction():
+    project = "projects/990000000009/locations/us-west1"
+    endpoint_id = project + "/endpoints/2200000000000000002"
+    instances = [{
+        "key1": "value1",
+        "key2": 2
+    }]
+    async_predict(endpoint_id, instances)
+
+if __name__ == "__main__":
+    from kfp.v2 import compiler
+    compiler.Compiler().compile(
+        pipeline_func=pipeline_prediction,
+        package_path='async_prediction.json')

community-content/vertex_pipeline_kfpv2/batch_prediction.py

@@ -0,0 +1,47 @@
+from kfp.v2 import dsl
+
+@dsl.component(base_image='python:3.8',packages_to_install=['google-cloud-aiplatform==1.36.0'])
+def batch_predict(
+    model_id: str,
+    job_name: str,
+    file_format: str,
+    file_sources: List[str],
+    out_sources: str,
+    machine_type: str,
+    min_replica_count: int,
+    max_replica_count: int,
+):
+    import json
+    from google.cloud import aiplatform
+
+    aiplatform.init(
+        project=project_id,
+        location=location,
+    )
+    
+    model = aiplatform.Model(model_id)
+    batch_job = model.batch_predict(
+        job_display_name=job_name,
+        instances_format=file_format,
+        gcs_source=file_sources,
+        gcs_destination_prefix=out_sources,
+        machine_type=machine_typem,
+        starting_replica_count=min_replica_count,
+        max_replica_count=max_replica_count,
+    )
+    batch_job.wait()
+
+
+@dsl.pipeline(name='batch-predict')
+def pipeline_batch_predict():
+    batch_predict('projects/990000000009/locations/us-west1/models/1100000000000000001',
+        'batch-predict-job', 'csv',
+        ['gs://yourbucket/predict/file.csv'], 'gs://yourbucket/results/',
+        'n1-standard-2', 1, 1)
+
+
+if __name__ == "__main__":
+    from kfp.v2 import compiler
+    compiler.Compiler().compile(
+        pipeline_func=pipeline_batch_predict,
+        package_path='batch_predict.json')

community-content/vertex_pipeline_kfpv2/deploy-model.py

@@ -1,37 +1,51 @@
-from kfp.v2 import dsl
-
-@dsl.component(base_image='python:3.8',packages_to_install=['google-cloud-aiplatform==1.36.0'])
-def deploy_model(
-    model_id: str,
-    endpoint_id: str,
-    machine_type: str,
-    min_replica_count: int,
-    max_replica_count: int,
-):
-    import json
-    from google.cloud import aiplatform
-
-    model = aiplatform.Model(model_id)
-    endpoint = aiplatform.Endpoint(endpoint_id)
-
-    endpoint = model.deploy(
-        endpoint=endpoint,
-        machine_type=machine_type,
-        min_replica_count=min_replica_count,
-        max_replica_count=max_replica_count,
-    )
-
-
-@dsl.pipeline(name='deploy-model')
-def pipeline_deploy_model():
-    project = "projects/990000000009/locations/us-west1"
-    model_id = project + "/models/1100000000000000001"
-    endpoint_id = project + "/endpoints/2200000000000000002"
-    deploy_model(model_id, endpoint_id, "n1-standard-2", 1, 1)
-
-
-if __name__ == "__main__":
-    from kfp.v2 import compiler
-    compiler.Compiler().compile(
-        pipeline_func=pipeline_deploy_model,
-        package_path='deploy_model.json')
+from kfp.v2 import dsl
+
+@dsl.component(base_image='python:3.8',packages_to_install=['google-cloud-aiplatform==1.36.0'])
+def deploy_model(
+    model_id: str,
+    endpoint_id: str,
+    machine_type: str,
+    min_replica_count: int,
+    max_replica_count: int,
+):
+    import json
+    from google.cloud import aiplatform
+
+    model = aiplatform.Model(model_id)
+    endpoint = aiplatform.Endpoint(endpoint_id)
+
+    endpoint = model.deploy(
+        endpoint=endpoint,
+        machine_type=machine_type,
+        min_replica_count=min_replica_count,
+        max_replica_count=max_replica_count,
+    )
+
+
+@dsl.component(base_image='python:3.8',packages_to_install=['google-cloud-aiplatform==1.36.0'])
+def delete_endpoint(
+    endpoint_id: str,
+):
+    from google.cloud import aiplatform
+
+    endpoint = aiplatform.Endpoint(endpoint_id)
+    endpoint.undeploy_all()
+    endpoint.delete()
+
+
+@dsl.pipeline(name='deploy-model')
+def pipeline_deploy_model():
+    project = "projects/990000000009/locations/us-west1"
+    model_id = project + "/models/1100000000000000001"
+    endpoint_id = project + "/endpoints/2200000000000000002"
+    deploy_model(model_id, endpoint_id, "n1-standard-2", 1, 1)
+
+    # After serving predictions, recycling computing resources
+    delete_endpoint(endpoint_id)
+
+
+if __name__ == "__main__":
+    from kfp.v2 import compiler
+    compiler.Compiler().compile(
+        pipeline_func=pipeline_deploy_model,
+        package_path='deploy_model.json')

community-content/vertex_pipeline_kfpv2/run-experiment.py

@@ -0,0 +1,33 @@
+from kfp.v2 import dsl
+
+@dsl.component(base_image='python:3.8',packages_to_install=['google-cloud-aiplatform==1.36.0'])
+def run_experiment(
+    project_id: str,
+    location: str,
+    experiment_name: str,
+    run_name: str,
+):
+    import json
+    from google.cloud import aiplatform
+
+    aiplatform.init(
+        project=project_id,
+        location=location,
+    )
+    test_expt = aiplatform.Experiment.create(experiment_name)
+    test_run = aiplatform.ExperimentRun.create(run_name, experiment=test_expt)
+    metric = test_run.get_classification_metrics()[0]
+    print(metric)
+
+
+
+@dsl.pipeline(name='run_experiment')
+def pipeline_run_experiment():
+    run_experiment('990000000009', 'us-west1', 'test-experiment', 'test-run')
+
+
+if __name__ == "__main__":
+    from kfp.v2 import compiler
+    compiler.Compiler().compile(
+        pipeline_func=pipeline_run_experiment,
+        package_path='run_experiment.json')

notebooks/community/CODEOWNERS

@@ -36,15 +36,23 @@
 /notebooks/community/ml_ops/stage3/get_started_with_dataflow_flex_template_component.ipynb @wintwoo
 /notebooks/community/vertex-ai-samples/notebooks/community/model_registry/vertex_ai_model_registry_bqml_custom_model_versioning.ipynb @inardini
 /notebooks/community/vertex-ai-samples/notebooks/community/model_registry/vertex_ai_model_registry_automl_model_versioning.ipynb @inardini
+/notebooks/community/vision/video_warehouse_curl.ipynb @zhangxiaotian @liangyz
 /notebooks/community/vision/image_warehouse_sdk.ipynb @bingwang @zhangxiaotian
 /notebooks/community/vision/video_warehouse_sdk.ipynb @zhangxiaotian
 /notebooks/community/vizier/conversions_vertex_vizier_and_open_source_vizier.ipynb @halio-g
+/notebooks/community/exploratory_data_analysis/eda_with_r_and_biqquery.ipynb @alokpattani
 /notebooks/community/experiments/vertex_ai_model_experimentation.ipynb @inardini @asobran
 /notebooks/community/pipelines/google_cloud_pipeline_components_bqml_pipeline_anomaly_detection.ipynb @inardini
 /notebooks/community/pipelines/google_cloud_pipeline_components_cloud_natural_language_pipeline.ipynb @Narwhalprime
 /notebooks/community/pipelines/google_cloud_pipeline_components_ready_to_go_text_classification_pipeline.ipynb @Narwhalprime
 /notebooks/community/feature_store/get_started_vertex_feature_store.ipynb @junkourata
+/notebooks/community/model_garden/model_garden_codegemma_deployment_on_vertex.ipynb @minwoo33park
+/notebooks/community/model_garden/model_garden_e5.ipynb @pouiscakes
 /notebooks/community/model_garden/model_garden_huggingface_local_inference.ipynb @dstnluong-google
+/notebooks/community/model_garden/model_garden_jax_paligemma_deployment.ipynb @minwoo33park
+/notebooks/community/model_garden/model_garden_jax_paligemma_finetuning.ipynb @minwoo33park
+/notebooks/community/model_garden/model_garden_jax_stable_diffusion_xl.ipynb @weigary
+/notebooks/community/model_garden/model_garden_mammut.ipynb @ivywang9331
 /notebooks/community/model_garden/model_garden_mediapipe_face_stylizer.pynb @schmidt-sebastian
 /notebooks/community/model_garden/model_garden_mediapipe_gesture_recognition.ipynb @schmidt-sebastian
 /notebooks/community/model_garden/model_garden_mediapipe_image_classification.ipynb @schmidt-sebastian
@@ -59,8 +67,13 @@
 /notebooks/community/model_garden/model_garden_tfvision_image_segmentation.ipynb @genquan9
 /notebooks/community/model_garden/model_garden_pytorch_stable_diffusion.ipynb @xiangxu-google
 /notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_2_1.ipynb @bingatgoogle
+/notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_custom.ipynb @weigary
 /notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_inpainting.ipynb @xiangxu-google
 /notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_xl_1_0.ipynb @bingatgoogle
+/notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_xl_lcm.ipynb @weigary
+/notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_xl_lightning.ipynb @xcchen1
+/notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_xl_lora.ipynb @weigary
+/notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_xl_turbo.ipynb @weigary
 /notebooks/community/model_garden/model_garden_pytorch_instructpix2pix.ipynb @xiangxu-google
 /notebooks/community/model_garden/model_garden_pytorch_controlnet.ipynb @xiangxu-google
 /notebooks/community/model_garden/model_garden_pytorch_blip_image_captioning.ipynb @xiangxu-google
@@ -74,6 +87,7 @@
 /notebooks/community/model_garden/model_garden_pytorch_detectron2.ipynb @lavraicse
 /notebooks/community/model_garden/model_garden_pytorch_dolly_v2.ipynb @lavraicse
 /notebooks/community/model_garden/model_garden_pytorch_bart_large_cnn.ipynb @lavraicse
+/notebooks/community/model_garden/model_garden_pytorch_autogluon.ipynb @lavraicse
 /notebooks/community/model_garden/model_garden_pytorch_starcoder.ipynb @xcchen1
 /notebooks/community/model_garden/model_garden_jax_vision_transformer.ipynb @lavraicse
 /notebooks/community/model_garden/model_garden_jax_fvlm.ipynb @lavraicse
@@ -94,7 +108,6 @@
 /notebooks/community/model_garden/model_garden_movinet_clip_classification.ipynb @KCFindstr
 /notebooks/community/model_garden/model_garden_movinet_action_recognition.ipynb @KCFindstr
 /notebooks/community/model_garden/model_garden_pytorch_open_clip.ipynb @lydhr
-/notebooks/community/model_garden/model_garden_pytorch_llama2_peft.ipynb @genquan9
 /notebooks/community/model_garden/model_garden_pytorch_codellama.ipynb @xiangxu-google
 /notebooks/community/model_garden/model_garden_pytorch_nllb.ipynb @weigary
 /notebooks/community/model_garden/model_garden_pytorch_vicuna.ipynb @dstnluong-google
@@ -103,11 +116,39 @@
 /notebooks/community/model_garden/model_garden_pytorch_biomedclip.ipynb @KCFindstr
 /notebooks/community/model_garden/model_garden_pytorch_imagebind.ipynb @kathyyu-google
 /notebooks/community/persistent_resource/00_persistent_resource_getting_started_cli.ipynb @jbrache
+/notebooks/community/persistent_resource/00_persistent_resource_getting_started_sdk.ipynb @jbrache
 /notebooks/community/model_garden/model_garden_pytorch_llama2_deployment.ipynb @genquan9
 /notebooks/community/model_garden/model_garden_pytorch_llama2_peft_finetuning.ipynb @genquan9
 /notebooks/community/model_garden/model_garden_pytorch_llama2_quantization.ipynb @dstnluong-google
+/notebooks/community/model_garden/model_garden_pytorch_llama2_peft_hyperparameter_tuning.ipynb @KCFindstr
 /notebooks/community/model_garden/model_garden_pytorch_llama2_evaluation.ipynb @kathyyu-google
 /notebooks/community/model_garden/model_garden_pytorch_llama2_rlhf_tuning.ipynb @genquan9
+/notebooks/community/model_garden/model_garden_pytorch_llama3_deployment.ipynb @kathyyu-google
+/notebooks/community/model_garden/model_garden_pytorch_llama3_finetuning.ipynb @kathyyu-google
+/notebooks/community/model_garden/model_garden_pytorch_llama3_1_deployment.ipynb @xiangxu-google
+/notebooks/community/model_garden/model_garden_pytorch_llama3_1_finetuning.ipynb @wrzhao-work
 /notebooks/community/model_garden/model_garden_pytorch_wizard_coder.ipynb @KCFindstr
 /notebooks/community/model_registry/get_started_with_vertex_ai_deployer.ipynb angelmontero@ @inardini
 /notebooks/community/model_garden/model_garden_pytorch_wizard_lm.ipynb @KCFindstr
+/notebooks/community/model_garden/model_garden_pytorch_llava.ipynb @py4
+/notebooks/community/model_garden/model_garden_pytorch_lama.ipynb @dstnluong-google
+/notebooks/community/model_garden/model_garden_pytorch_mistral_peft_tuning.ipynb @KCFindstr
+/notebooks/community/model_garden/model_garden_pytorch_mistral_deployment.ipynb @KCFindstr
+/notebooks/community/model_garden/model_garden_pytorch_mixtral_peft_tuning.ipynb @KCFindstr
+/notebooks/community/model_garden/model_garden_pytorch_mixtral_deployment.ipynb @KCFindstr
+/notebooks/community/model_garden/model_garden_gemma_deployment_on_gke.ipynb @vilobhmm
+/notebooks/community/model_garden/model_garden_gemma_deployment_on_vertex.ipynb @kathyyu-google
+/notebooks/community/model_garden/model_garden_gemma2_deployment_on_vertex.ipynb @kathyyu-google
+/notebooks/community/model_garden/model_garden_gemma_finetuning_on_vertex.ipynb @KCFindstr
+/notebooks/community/model_garden/model_garden_pytorch_gemma_peft_finetuning_hf.ipynb @KCFindstr
+/notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_deployment_1_5.ipynb @weigary
+/notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_gradio.ipynb @weigary
+/notebooks/community/model_garden/model_garden_pytorch_sd_2_1_finetuning_dreambooth.ipynb @weigary
+/notebooks/community/model_garden/model_garden_pytorch_sd_xl_finetuning_dreambooth_lora.ipynb @weigary
+/notebooks/community/model_garden/model_garden_pytorch_sd_2_1_local_finetuning_dreambooth.ipynb @weigary
+/notebooks/community/model_garden/model_garden_timesfm_deployment_on_vertex.ipynb @siriuz42
+/notebooks/community/model_garden/model_garden_llama_guard_deployment.ipynb @kathyyu-google
+/notebooks/community/model_garden/model_garden_rag.ipynb @kathyyu-google
+/notebooks/community/model_garden/synthetic_data_generation_using_llama3_1.ipynb @xiangxu-google
+/notebooks/community/model_garden/model_garden_autosxs_evaluation_llama3_1.ipynb @inardini
+/notebooks/community/model_garden/model_garden_openai_api_llama3_1.ipynb @inardini

notebooks/community/exploratory_data_analysis/eda_with_r_and_bigquery.ipynb

@@ -0,0 +1,682 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "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": {
+    "tags": []
+   },
+   "source": [
+    "# Exploratory Data Analysis with R and BigQuery\n",
+    "\n",
+    "**Author**:  [Alok Pattani](https://github.com/alokpattani)\n",
+    "\n",
+    "**Last Updated**:  July 2024\n",
+    "\n",
+    "## Overview\n",
+    "\n",
+    "This notebook illustrates how to perform exploratory data analysis (EDA) using [R](https://www.r-project.org/about.html) on data extracted from [BigQuery](https://cloud.google.com/bigquery). After you analyze and process the data, the transformed data is stored in [Cloud Storage](https://cloud.google.com/storage) for further machine learning (ML) tasks.\n",
+    "\n",
+    "R is one of the most widely used programming languages for statistical modeling. It has a large and active community of data scientists and machine learning (ML) professionals. With more than 20,000 packages in the open-source repository of [CRAN](https://cran.r-project.org/), R has tools for all statistical data analysis applications, ML, and visualization.\n",
+    "\n",
+    "## Dataset\n",
+    "The dataset used in this tutorial is the BigQuery natality dataset. This public dataset includes information about more than 137 million births registered in the United States from 1969 to 2008. The dataset is available [here](https://console.cloud.google.com/bigquery?p=bigquery-public-data&d=samples&t=natality&page=table&_ga=2.99329886.-1705629017.1551465326&_gac=1.109796023.1561476396.CI2rz-z4hOMCFc6RhQods4oEXA).\n",
+    "\n",
+    "In this notebook, we focus on exploratory data analysis and visualization using R and BigQuery, with an eye toward a potential machine learning goal of predicting a baby's weight given a number of factors about the pregnancy and about the baby's mother.\n",
+    "\n",
+    "## Objective\n",
+    "The goal of this tutorial is to:\n",
+    "1. Query and analyze data from BigQuery using the [bigrquery](https://cran.r-project.org/web/packages/bigrquery/index.html) R library.\n",
+    "2. Prepare and store data for ML in Cloud Storage.\n",
+    "\n",
+    "## Costs\n",
+    "This tutorial uses the following billable components of Google Cloud:\n",
+    "1. [BigQuery](https://cloud.google.com/bigquery/pricing)\n",
+    "2. [Cloud Storage](https://cloud.google.com/storage/pricing)\n",
+    "3. [Vertex AI Workbench Instances](https://cloud.google.com/vertex-ai/pricing#notebooks) (if running this notebook there)\n",
+    "\n",
+    "Use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 0. Setup "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "Check the version of R being run."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "version"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Install necessary R packages if not already available in the current session."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# List the necessary packages\n",
+    "needed_packages <- c(\"dplyr\", \"ggplot2\", \"bigrquery\")\n",
+    "\n",
+    "# Check if packages are installed\n",
+    "installed_packages <- .packages(all.available = TRUE)\n",
+    "missing_packages <- needed_packages[!(needed_packages %in% installed_packages)]\n",
+    "\n",
+    "# If any are missing, install them\n",
+    "if (length(missing_packages) > 0) {\n",
+    "  install.packages(missing_packages)\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Load the required packages\n",
+    "lapply(needed_packages, library, character.only = TRUE) "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Use BigQuery out-of-band authentication"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "bq_auth(use_oob = TRUE)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Set a variable to the name of the project that you want to use for this tutorial."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Set the project ID\n",
+    "PROJECT_ID <- \"[YOUR-PROJECT-ID]\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Set a variable to the name of the Cloud Storage bucket that you want to use later to store the output data. The name must be globally unique."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Set your Cloud Storage bucket name\n",
+    "BUCKET_NAME <- \"[YOUR-BUCKET-NAME]\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Set default height/width for plots generated\n",
+    "options(repr.plot.height = 9, repr.plot.width = 16)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 1. Querying Data from BigQuery "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 1.1. Prepare the BigQuery query"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "sql_query_template <- \"\n",
+    "    SELECT\n",
+    "      TIMESTAMP_DIFF(dropoff_datetime, pickup_datetime, MINUTE) AS trip_time_minutes,  \n",
+    "\n",
+    "      passenger_count,\n",
+    "\n",
+    "      ROUND(trip_distance, 1) AS trip_distance_miles,\n",
+    "\n",
+    "      rate_code,\n",
+    "      /* Mapping from rate code to type from description column in BQ table schema */\n",
+    "      (CASE \n",
+    "        WHEN rate_code = '1.0'\n",
+    "          THEN 'Standard rate'\n",
+    "        WHEN rate_code = '2.0'\n",
+    "          THEN 'JFK'\n",
+    "        WHEN rate_code = '3.0'\n",
+    "          THEN 'Newark'\n",
+    "        WHEN rate_code = '4.0'\n",
+    "          THEN 'Nassau or Westchester'\n",
+    "        WHEN rate_code = '5.0'\n",
+    "          THEN 'Negotiated fare'\n",
+    "        WHEN rate_code = '6.0'\n",
+    "          THEN 'Group ride'\n",
+    "        /* Several NULL AND some '99.0' values go here */\n",
+    "        ELSE 'Unknown'\n",
+    "        END)\n",
+    "        AS rate_type,\n",
+    "\n",
+    "      fare_amount,\n",
+    "\n",
+    "      CAST(ABS(FARM_FINGERPRINT(\n",
+    "        CONCAT(\n",
+    "          CAST(trip_distance AS STRING), \n",
+    "          CAST(fare_amount AS STRING)\n",
+    "          )\n",
+    "        ))\n",
+    "        AS STRING)\n",
+    "        AS key\n",
+    "\n",
+    "    FROM\n",
+    "      `bigquery-public-data.new_york_taxi_trips.tlc_yellow_trips_2022`\n",
+    "\n",
+    "    /* Filter out some outlier or hard to understand values */\n",
+    "    WHERE\n",
+    "      (TIMESTAMP_DIFF(dropoff_datetime, pickup_datetime, MINUTE)\n",
+    "        BETWEEN 0.01 AND 120)\n",
+    "      AND\n",
+    "      (passenger_count BETWEEN 1 AND 10)\n",
+    "      AND\n",
+    "      (trip_distance BETWEEN 0.01 AND 100)\n",
+    "      AND\n",
+    "      (fare_amount BETWEEN 0.01 AND 250)\n",
+    "\n",
+    "    LIMIT %s\n",
+    "\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 1.2. Execute the query \n",
+    "The data will be retreived from BigQuery, and the results will be stored in an in-memory [tibble](https://tibble.tidyverse.org/) (like a data frame)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "sample_size <- 10000\n",
+    "\n",
+    "sql_query <- sprintf(sql_query_template, sample_size)\n",
+    "\n",
+    "taxi_trip_data <- bq_table_download(\n",
+    "    bq_project_query(\n",
+    "        PROJECT_ID, \n",
+    "        query = sql_query\n",
+    "    )\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### 1.3. View the query results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# View the query result\n",
+    "head(taxi_trip_data)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Show # of rows and data types of each column\n",
+    "str(taxi_trip_data)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# View the results summary\n",
+    "summary(taxi_trip_data)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 2. Visualizing retrieved data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Display the distribution of fare amounts using a histogram\n",
+    "ggplot(\n",
+    "    data = taxi_trip_data, \n",
+    "    aes(x = fare_amount)\n",
+    "    ) + \n",
+    "geom_histogram(bins = 100)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Display the relationship between trip distance and fare amount\n",
+    "ggplot(\n",
+    "    data = taxi_trip_data, \n",
+    "    aes(x = trip_distance_miles, y = fare_amount)\n",
+    "    ) + \n",
+    "geom_point() + \n",
+    "geom_smooth(method = \"lm\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Performing the processing in BigQuery\n",
+    "Create a function that finds the number of trips and the average fare amount for each value of the chosen column."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "get_distinct_value_aggregates <- function(column) {\n",
+    "    query <- paste0(\n",
+    "        'SELECT ', \n",
+    "          column, \n",
+    "          ', \n",
+    "          COUNT(1) AS num_trips,\n",
+    "          AVG(fare_amount) AS avg_fare_amount\n",
+    "        \n",
+    "        FROM\n",
+    "          `bigquery-public-data.new_york_taxi_trips.tlc_yellow_trips_2022`\n",
+    "        \n",
+    "        WHERE\n",
+    "          (TIMESTAMP_DIFF(dropoff_datetime, pickup_datetime, MINUTE) \n",
+    "            BETWEEN 0.01 AND 120)\n",
+    "          AND\n",
+    "          (passenger_count BETWEEN 1 AND 10)\n",
+    "          AND\n",
+    "          (trip_distance BETWEEN 0.01 AND 100)\n",
+    "          AND\n",
+    "          (fare_amount BETWEEN 0.01 AND 250)\n",
+    "        \n",
+    "        GROUP BY 1\n",
+    "        '\n",
+    "        )\n",
+    "    \n",
+    "    bq_table_download(\n",
+    "        bq_project_query(\n",
+    "            PROJECT_ID, \n",
+    "            query = query\n",
+    "        )\n",
+    "    )\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Apply the function to get distinct values for various columns and plot them to study patterns."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "df <- get_distinct_value_aggregates(\n",
+    "  'TIMESTAMP_DIFF(dropoff_datetime, pickup_datetime, MINUTE) AS trip_time_minutes')\n",
+    "\n",
+    "ggplot(\n",
+    "    data = df, \n",
+    "    aes(x = trip_time_minutes, y = num_trips)\n",
+    "    ) + \n",
+    "geom_line()\n",
+    "\n",
+    "ggplot(\n",
+    "    data = df,\n",
+    "    aes(x = trip_time_minutes, y = avg_fare_amount)\n",
+    "    ) + \n",
+    "geom_line()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "df <- get_distinct_value_aggregates('passenger_count')\n",
+    "\n",
+    "ggplot(\n",
+    "    data = df, \n",
+    "    aes(x = passenger_count, y = num_trips)\n",
+    "    ) + \n",
+    "geom_col() +\n",
+    "scale_x_continuous(breaks = 1:10)\n",
+    "\n",
+    "ggplot(\n",
+    "    data = df, \n",
+    "    aes(x = passenger_count, y = avg_fare_amount)\n",
+    "    ) + \n",
+    "geom_col() +\n",
+    "scale_x_continuous(breaks = 1:10)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "df <- get_distinct_value_aggregates('ROUND(trip_distance, 0) AS trip_distance_miles')\n",
+    "\n",
+    "ggplot(\n",
+    "    data = df, \n",
+    "    aes(x = trip_distance_miles, y = num_trips)\n",
+    "    ) + \n",
+    "geom_line()\n",
+    "\n",
+    "ggplot(\n",
+    "    data = df,\n",
+    "    aes(x = trip_distance_miles, y = avg_fare_amount)\n",
+    "    ) + \n",
+    "geom_line()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "df <- get_distinct_value_aggregates(\"\n",
+    "  (CASE \n",
+    "    WHEN rate_code = '1.0'\n",
+    "      THEN 'Standard rate'\n",
+    "    WHEN rate_code = '2.0'\n",
+    "      THEN 'JFK'\n",
+    "    WHEN rate_code = '3.0'\n",
+    "      THEN 'Newark'\n",
+    "    WHEN rate_code = '4.0'\n",
+    "      THEN 'Nassau or Westchester'\n",
+    "    WHEN rate_code = '5.0'\n",
+    "      THEN 'Negotiated fare'\n",
+    "    WHEN rate_code = '6.0'\n",
+    "      THEN 'Group ride'\n",
+    "    /* Several NULL AND some '99.0' values go here */\n",
+    "    ELSE 'Unknown'\n",
+    "    END)\n",
+    "    AS rate_type\n",
+    "  \")\n",
+    "\n",
+    "ggplot(\n",
+    "    data = df,\n",
+    "    aes(x = rate_type, y = num_trips)\n",
+    "    ) + \n",
+    "geom_col()\n",
+    "\n",
+    "ggplot(\n",
+    "    data = df,\n",
+    "    aes(x = rate_type, y = avg_fare_amount)\n",
+    "    ) + \n",
+    "geom_col()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 3. Saving the data as CSVs to Cloud Storage"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Prepare training and evaluation data from BigQuery\n",
+    "sample_size <- 10000\n",
+    "\n",
+    "sql_query <- sprintf(sql_query_template, sample_size)\n",
+    "\n",
+    "# Split data into 75% training, 25% evaluation\n",
+    "train_query <- paste('SELECT * FROM (', sql_query, \n",
+    "  ') WHERE MOD(CAST(key AS INT64), 100) <= 75')\n",
+    "eval_query <- paste('SELECT * FROM (', sql_query,\n",
+    "  ') WHERE MOD(CAST(key AS INT64), 100) > 75')\n",
+    "\n",
+    "# Load training data to data frame\n",
+    "train_data <- bq_table_download(\n",
+    "    bq_project_query(\n",
+    "        PROJECT_ID, \n",
+    "        query = train_query\n",
+    "    )\n",
+    ")\n",
+    "\n",
+    "# Load evaluation data to data frame\n",
+    "eval_data <- bq_table_download(\n",
+    "    bq_project_query(\n",
+    "        PROJECT_ID, \n",
+    "        query = eval_query\n",
+    "    )\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "print(paste0(\"Training instances count: \", nrow(train_data)))\n",
+    "\n",
+    "print(paste0(\"Evaluation instances count: \", nrow(eval_data)))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Write data frames to local CSV files, with headers\n",
+    "dir.create(file.path('data'), showWarnings = FALSE)\n",
+    "\n",
+    "write.table(train_data, \"data/train_data.csv\", \n",
+    "   row.names = FALSE, col.names = TRUE, sep = \",\")\n",
+    "\n",
+    "write.table(eval_data, \"data/eval_data.csv\", \n",
+    "   row.names = FALSE, col.names = TRUE, sep = \",\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "# Upload CSV data to Cloud Storage by passing gsutil commands to system\n",
+    "gcs_url <- paste0(\"gs://\", BUCKET_NAME, \"/\")\n",
+    "\n",
+    "command <- paste(\"gsutil mb\", gcs_url)\n",
+    "\n",
+    "system(command)\n",
+    "\n",
+    "gcs_data_dir <- paste0(\"gs://\", BUCKET_NAME, \"/data\")\n",
+    "\n",
+    "command <- paste(\"gsutil cp data/*_data.csv\", gcs_data_dir)\n",
+    "\n",
+    "system(command)\n",
+    "\n",
+    "command <- paste(\"gsutil ls -l\", gcs_data_dir)\n",
+    "\n",
+    "system(command, intern = TRUE)"
+   ]
+  }
+ ],
+ "metadata": {
+  "environment": {
+   "kernel": "conda-env-r-r",
+   "name": "workbench-notebooks.m123",
+   "type": "gcloud",
+   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m123"
+  },
+  "kernelspec": {
+   "display_name": "R (Local)",
+   "language": "R",
+   "name": "conda-env-r-r"
+  },
+  "language_info": {
+   "codemirror_mode": "r",
+   "file_extension": ".r",
+   "mimetype": "text/x-r-source",
+   "name": "R",
+   "pygments_lexer": "r",
+   "version": "4.3.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

notebooks/community/model_garden/model_garden_autosxs_evaluation_llama3_1.ipynb

@@ -0,0 +1,741 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ur8xi4C7S06n"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2024 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "JAPoU8Sm5E6e"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Evaluate Llama 3.1 models using Vertex AI AutoSxS\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_autosxs_evaluation_llama3_1.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
+        "    </a>\n",
+        "  </td>\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_autosxs_evaluation_llama3_1.ipynb\"\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_autosxs_evaluation_llama3_1.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\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_autosxs_evaluation_llama3_1.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demostrates how to use the Vertex AI automatic side-by-side (AutoSxS) tool to evaluate Llama 3.1 models for a question-answering task.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Choose the Llama 3.1 models you want to compare.\n",
+        "\n",
+        "- Create an evaluation dataset with question-answer data.\n",
+        "\n",
+        "- Create and run a Vertex AI AutoSxS pipeline that generates judgments and a set of AutoSxS metrics using the generated judgments.\n",
+        "\n",
+        "- Print the judgments and AutoSxS metrics.\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "61RBz8LLbxCR"
+      },
+      "source": [
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "No17Cw5hgx12"
+      },
+      "source": [
+        "### Install Vertex AI SDK for Python and other required packages\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tFy3H3aPgx12"
+      },
+      "outputs": [],
+      "source": [
+        "! pip3 install --upgrade --user --quiet google-cloud-aiplatform google-cloud-pipeline-components\n",
+        "! pip3 install --upgrade --user --quiet openai gcsfs"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "R5Xep4W9lq-Z"
+      },
+      "source": [
+        "### Restart runtime (Colab only)\n",
+        "\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "XRvKdaPDTznN"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "SbmM4z7FOBpM"
+      },
+      "source": [
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dmWOrTJ3gx13"
+      },
+      "source": [
+        "### Authenticate your notebook environment (Colab only)\n",
+        "\n",
+        "Authenticate your environment on Google Colab.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NyKGtVQjgx13"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "DF4l8DTdWgPY"
+      },
+      "source": [
+        "### Set Google Cloud project information\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Nqwi-5ufWp_B"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"<your-project-id>\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Set the region of the instance\n",
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Create a Cloud Storage bucket\n",
+        "\n",
+        "Create a storage bucket to store tutorial artifacts."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MzGDU7TWdts_"
+      },
+      "outputs": [],
+      "source": [
+        "BUCKET_NAME = \"<your-bucket-name>\"  # @param {type:\"string\"}\n",
+        "\n",
+        "BUCKET_URI = f\"gs://{BUCKET_NAME}\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-EcIXiGsCePi"
+      },
+      "source": [
+        "**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": "NIq7R4HZCfIc"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0Wn8ZkcV86KR"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "B8DawN9D9NLU"
+      },
+      "outputs": [],
+      "source": [
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "jVYoyDl165EE"
+      },
+      "source": [
+        "### Import libraries\n",
+        "\n",
+        "Import libraries to use in this tutorial."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c1tEW-U968h8"
+      },
+      "outputs": [],
+      "source": [
+        "import random\n",
+        "import string\n",
+        "\n",
+        "import openai\n",
+        "import pandas as pd\n",
+        "from google.auth import default, transport\n",
+        "from google.cloud import aiplatform\n",
+        "from google_cloud_pipeline_components.v1 import model_evaluation\n",
+        "from kfp import compiler"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ZXnx1_CtEV5L"
+      },
+      "source": [
+        "### Set variables\n",
+        "\n",
+        "Before starting, you must decide how to access Llama 3.1 models. You can access Llama 3.1 models in just a few clicks using Model-as-a-Service (MaaS) without any setup or infrastructure hassles. You can also access Llama models for self-service in Vertex AI Model Garden, allowing you to choose your preferred infrastructure.\n",
+        "\n",
+        "This tutorial assumes that you deploy a self-managed instance of the Llama 3.1 model and compare it with Llama 3 405b using Model-as-a-Service (MaaS). Notice, only `us-central1` is supported region for Llama 3.1 models using Model-as-a-Service (MaaS).\n",
+        "\n",
+        "[Check out Llama 3 model card](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama3?_ga=2.31261500.2048242469.1721714335-1107467625.1721655511) to learn how to deploy a Llama 3.1 models on Vertex AI."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "XASp0SPNEX10"
+      },
+      "outputs": [],
+      "source": [
+        "SELF_DEPLOYED_ENDPOINT_REGION = \"<your-endpoint-region>\"  # @param {type:\"string\"}\n",
+        "SELF_DEPLOYED_ENDPOINT_ID = \"<your-endpoint-id>\"  # @param {type:\"string\"}\n",
+        "MODEL_LOCATION = \"us-central1\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "_sdmrDed2aHd"
+      },
+      "source": [
+        "### Helpers"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "pACHdEUf2bfq"
+      },
+      "outputs": [],
+      "source": [
+        "def self_model_generate(\n",
+        "    question,\n",
+        "    context,\n",
+        "    endpoint_id=SELF_DEPLOYED_ENDPOINT_ID,\n",
+        "    endpoint_location=SELF_DEPLOYED_ENDPOINT_REGION,\n",
+        "    **model_kwargs,\n",
+        "):\n",
+        "    \"\"\"Generate a response from a self-managed Llama 3.1 model.\"\"\"\n",
+        "\n",
+        "    aiplatform.init(project=PROJECT_ID, location=endpoint_location)\n",
+        "\n",
+        "    prompt = \"\"\"You are an AI assistant. Your goal is to answer questions using the pieces of context. \"\"\"\n",
+        "    prompt += f\"\"\"Question: {question}.\"\"\"\n",
+        "    prompt += f\"\"\"Context: {context}.\"\"\"\n",
+        "    prompt += \"\"\"Answer:\"\"\"\n",
+        "\n",
+        "    instance = {\"prompt\": prompt}\n",
+        "    instance.update(model_kwargs)\n",
+        "    instances = [instance]\n",
+        "\n",
+        "    endpoint = aiplatform.Endpoint(endpoint_id)\n",
+        "    response = endpoint.predict(instances=instances)\n",
+        "    return response.predictions[0][len(prompt) + 1 :]\n",
+        "\n",
+        "\n",
+        "def maas_generate(\n",
+        "    question,\n",
+        "    context,\n",
+        "    model=\"meta/llama3-405b-instruct-maas\",\n",
+        "    model_location=MODEL_LOCATION,\n",
+        "    **model_kwargs,\n",
+        "):\n",
+        "    \"\"\"Generate a response from a MaaS Llama 3.1 model.\"\"\"\n",
+        "\n",
+        "    creds, _ = default()\n",
+        "    auth_req = transport.requests.Request()\n",
+        "    creds.refresh(auth_req)\n",
+        "    if model_kwargs is None:\n",
+        "        model_kwargs = {}\n",
+        "\n",
+        "    client = openai.OpenAI(\n",
+        "        base_url=f\"https://{model_location}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{model_location}/endpoints/openapi/chat/completions?\",\n",
+        "        api_key=creds.token,\n",
+        "    )\n",
+        "\n",
+        "    response = client.chat.completions.create(\n",
+        "        model=model,\n",
+        "        messages=[\n",
+        "            {\n",
+        "                \"role\": \"system\",\n",
+        "                \"content\": \"\"\"You are an AI assistant. Your goal is to answer questions using the pieces of context. If you don't know the answer, say that you don't know.\"\"\",\n",
+        "            },\n",
+        "            {\"role\": \"user\", \"content\": question},\n",
+        "            {\"role\": \"assistant\", \"content\": context},\n",
+        "        ],\n",
+        "        **model_kwargs,\n",
+        "    )\n",
+        "\n",
+        "    return response.choices[0].message.content\n",
+        "\n",
+        "\n",
+        "def generate_uuid(length: int = 8) -> str:\n",
+        "    \"\"\"Generate a uuid of a specified length (default=8).\"\"\"\n",
+        "    return \"\".join(random.choices(string.ascii_lowercase + string.digits, k=length))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "eha2l9nkNxZs"
+      },
+      "source": [
+        "### Generate evaluation dataset for AutoSxS\n",
+        "\n",
+        "Below you create your evaluation dataset, you specify a set of prompts to evaluate on.\n",
+        "\n",
+        "In this notebook, you:\n",
+        "\n",
+        "- Use 10 examples from the original dataset to create an evaluation dataset for AutoSxS.\n",
+        "  - Data in the `contexts` column will be treated as model context.\n",
+        "  - Data in the `questions` column will be treated as model instruction.\n",
+        "  - Data in the `response_a` column will be treated as responses for model A.\n",
+        "  - Data in the `response_b` will be treated as responses for model B.\n",
+        "\n",
+        "- Store the data in a JSON file in Google sCloud Storage.\n",
+        "\n",
+        "#### **Note: For the best results we recommend using at least 100 examples. There are diminishing returns when using more than 400 examples.**"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kuVd8Y7GHbp8"
+      },
+      "source": [
+        "#### Provide context and question"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "j_OuH0yh_PMe"
+      },
+      "outputs": [],
+      "source": [
+        "contexts = [\n",
+        "    \"Beginning in the late 1910s and early 1920s, Whitehead gradually turned his attention from mathematics to philosophy of science, and finally to metaphysics. He developed a comprehensive metaphysical system which radically departed from most of western philosophy. Whitehead argued that reality consists of processes rather than material objects, and that processes are best defined by their relations with other processes, thus rejecting the theory that reality is fundamentally constructed by bits of matter that exist independently of one another. Today Whitehead's philosophical works – particularly Process and Reality – are regarded as the foundational texts of process philosophy.\",\n",
+        "    \"The gills have an adnate attachment to the cap, are narrow to moderately broad, closely spaced, and eventually separate from the stem. Young gills are cinnamon-brown in color, with lighter edges, but darken in maturity because they become covered with the dark spores. The stem is 6 to 8 cm (2+3⁄8 to 3+1⁄8 in) long by 1.5 to 2 mm (1⁄16 to 3⁄32 in) thick, and roughly equal in width throughout except for a slightly enlarged base. The lower region of the stem is brownish in color and has silky 'hairs' pressed against the stem; the upper region is grayish and pruinose (lightly dusted with powdery white granules). The flesh turns slightly bluish or greenish where it has been injured. The application of a drop of dilute potassium hydroxide solution on the cap or flesh will cause a color change to pale to dark yellowish to reddish brown; a drop on the stem produces a less intense or no color change.\",\n",
+        "    \"Go to Device Support. Choose your device. Scroll to Getting started and select Hardware & phone details. Choose Insert or remove SIM card and follow the steps. Review the Account Summary page for details. Image 13 Activate online Go to att.com/activateprepaid ((att.com/activarprepaid for Spanish)) and follow the prompts. Activate over the phone Call us at 877.426.0525 for automated instructions. You will need to know your SIM/eSIM ICCID & IMEI number for activation. Note: Look for your SIM (( ICCID )) number on your box or SIM card Now youre ready to activate your phone 1. Start with your new device powered off. 2. To activate a new line of service or a replacement device, please go to the AT&T Activation site or call 866.895.1099. You download the eSIM to your device over Wi-Fi®. The eSIM connects your device to our wireless network. How do I activate my phone with an eSIM? Turn your phone on, connect to Wi-Fi, and follow the prompts. Swap active SIM cards AT&T Wireless SM SIM Card Turn your device off. Remove the old SIM card. Insert the new one. Turn on your device.\",\n",
+        "    \"According to chief astronaut Deke Slayton's autobiography, he chose Bassett for Gemini 9 because he was 'strong enough to carry' both himself and See. Slayton had also assigned Bassett as command module pilot for the second backup Apollo crew, alongside Frank Borman and William Anders.\",\n",
+        "    \"Adaptation of the endosymbiont to the host's lifestyle leads to many changes in the endosymbiont–the foremost being drastic reduction in its genome size. This is due to many genes being lost during the process of metabolism, and DNA repair and recombination. While important genes participating in the DNA to RNA transcription, protein translation and DNA/RNA replication are retained. That is, a decrease in genome size is due to loss of protein coding genes and not due to lessening of inter-genic regions or open reading frame (ORF) size. Thus, species that are naturally evolving and contain reduced sizes of genes can be accounted for an increased number of noticeable differences between them, thereby leading to changes in their evolutionary rates. As the endosymbiotic bacteria related with these insects are passed on to the offspring strictly via vertical genetic transmission, intracellular bacteria goes through many hurdles during the process, resulting in the decrease in effective population sizes when compared to the free living bacteria. This incapability of the endosymbiotic bacteria to reinstate its wild type phenotype via a recombination process is called as Muller's ratchet phenomenon. Muller's ratchet phenomenon together with less effective population sizes has led to an accretion of deleterious mutations in the non-essential genes of the intracellular bacteria. This could have been due to lack of selection mechanisms prevailing in the rich environment of the host.\",\n",
+        "    \"The National Archives Building in downtown Washington holds record collections such as all existing federal census records, ships' passenger lists, military unit records from the American Revolution to the Philippine–American War, records of the Confederate government, the Freedmen's Bureau records, and pension and land records.\",\n",
+        "    \"Standard 35mm photographic film used for cinema projection has a much higher image resolution than HDTV systems, and is exposed and projected at a rate of 24 frames per second (frame/s). To be shown on standard television, in PAL-system countries, cinema film is scanned at the TV rate of 25 frame/s, causing a speedup of 4.1 percent, which is generally considered acceptable. In NTSC-system countries, the TV scan rate of 30 frame/s would cause a perceptible speedup if the same were attempted, and the necessary correction is performed by a technique called 3:2 Pulldown: Over each successive pair of film frames, one is held for three video fields (1/20 of a second) and the next is held for two video fields (1/30 of a second), giving a total time for the two frames of 1/12 of a second and thus achieving the correct average film frame rate.\",\n",
+        "    \"Maria Deraismes was initiated into Freemasonry in 1882, then resigned to allow her lodge to rejoin their Grand Lodge. Having failed to achieve acceptance from any masonic governing body, she and Georges Martin started a mixed masonic lodge that actually worked masonic ritual. Annie Besant spread the phenomenon to the English speaking world. Disagreements over ritual led to the formation of exclusively female bodies of Freemasons in England, which spread to other countries. Meanwhile, the French had re-invented Adoption as an all-female lodge in 1901, only to cast it aside again in 1935. The lodges, however, continued to meet, which gave rise, in 1959, to a body of women practising continental Freemasonry.\",\n",
+        "    \"Excavation of the foundations began in November 1906, with an average of 275 workers during the day shift and 100 workers during the night shift. The excavation was required to be completed in 120 days. To remove the spoils from the foundation, three temporary wooden platforms were constructed to street level. Hoisting engines were installed to place the beams for the foundation, while the piers were sunk into the ground under their own weight. Because of the lack of space in the area, the contractors' offices were housed beneath the temporary platforms. During the process of excavation, the Gilsey Building's foundations were underpinned or shored up, because that building had relatively shallow foundations descending only 18 feet (5.5 m) below Broadway.\",\n",
+        "    \"Dopamine consumed in food cannot act on the brain, because it cannot cross the blood–brain barrier. However, there are also a variety of plants that contain L-DOPA, the metabolic precursor of dopamine. The highest concentrations are found in the leaves and bean pods of plants of the genus Mucuna, especially in Mucuna pruriens (velvet beans), which have been used as a source for L-DOPA as a drug. Another plant containing substantial amounts of L-DOPA is Vicia faba, the plant that produces fava beans (also known as 'broad beans'). The level of L-DOPA in the beans, however, is much lower than in the pod shells and other parts of the plant. The seeds of Cassia and Bauhinia trees also contain substantial amounts of L-DOPA.\",\n",
+        "]\n",
+        "\n",
+        "questions = [\n",
+        "    \"What was the predominant theory of reality that Whitehead opposed?\",\n",
+        "    \"Why do the gills on the Psilocybe pelliculosa mushroom darken as they mature?\",\n",
+        "    \"user: How do I provision my AT&T SIM card?\",\n",
+        "    \"Why did chief astronaut Deke Slayton choose Charles Bassett for Gemini 9, according to Slayton's autobiography?\",\n",
+        "    \"What is the main alteration in an endosymbiont when it adapts to a host?\",\n",
+        "    \"What's the earliest war The National Archives Building has military unit records for\",\n",
+        "    \"To be shown on SDTV in PAL-system countries, at what rate is cinema film scanned?\",\n",
+        "    \"What year was the all-female masonic lodge cast aside?\",\n",
+        "    \"Why did the Gilsey Building have underpinned and shored up foundations?\",\n",
+        "    \"Why can dopamine consumed in food not act on the brain?\",\n",
+        "]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Oiwr677h_cSk"
+      },
+      "outputs": [],
+      "source": [
+        "examples = pd.DataFrame(\n",
+        "    {\n",
+        "        \"questions\": questions,\n",
+        "        \"context\": contexts,\n",
+        "    }\n",
+        ")\n",
+        "examples.head()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2je-Rs8e_65p"
+      },
+      "outputs": [],
+      "source": [
+        "examples[\"response_a\"] = examples.apply(\n",
+        "    lambda x: self_model_generate(\n",
+        "        x[\"questions\"], x[\"context\"], max_tokens=2500, temperature=0.5\n",
+        "    ),\n",
+        "    axis=1,\n",
+        ")\n",
+        "examples.head()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "VGIpgBk9Br_G"
+      },
+      "outputs": [],
+      "source": [
+        "examples[\"response_b\"] = examples.apply(\n",
+        "    lambda x: maas_generate(\n",
+        "        x[\"questions\"], x[\"context\"], max_tokens=2500, temperature=0.5\n",
+        "    ),\n",
+        "    axis=1,\n",
+        ")\n",
+        "examples.head()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "NY1Jsj4aOCe1"
+      },
+      "source": [
+        "#### Upload your dataset to Cloud Storage\n",
+        "\n",
+        "Finally, we upload our evaluation dataset to Cloud Storage to be used as input for AutoSxS."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "vykmkhp-ODKg"
+      },
+      "outputs": [],
+      "source": [
+        "examples.to_json(f\"{BUCKET_URI}/evaluation_dataset.json\", orient=\"records\", lines=True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Qgdk_qNIOFik"
+      },
+      "source": [
+        "### Create and run AutoSxS job\n",
+        "\n",
+        "In order to run AutoSxS, we need to define a `autosxs_pipeline` job with the following parameters.\n",
+        "\n",
+        "More details of the AutoSxS pipeline configuration can be found [here](https://google-cloud-pipeline-components.readthedocs.io/en/google-cloud-pipeline-components-2.9.0/api/preview/model_evaluation.html#preview.model_evaluation.autosxs_pipeline)."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "veq26QZ7OMoC"
+      },
+      "source": [
+        "First, compile the AutoSxS pipeline locally."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "C2NGZzOMOJPV"
+      },
+      "outputs": [],
+      "source": [
+        "template_uri = \"pipeline.yaml\"\n",
+        "compiler.Compiler().compile(\n",
+        "    pipeline_func=model_evaluation.autosxs_pipeline,\n",
+        "    package_path=template_uri,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "I0aMBhoqOTXF"
+      },
+      "source": [
+        "The following code starts a Vertex Pipeline job, viewable from the Vertex UI. This pipeline job will take ~15 mins. This pipeline is made for batch prediction at a much larger scale than this example, so the time won't scale up linearly if there were thousands of Q&A pairs."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tRdA3ovUOV6j"
+      },
+      "outputs": [],
+      "source": [
+        "display_name = f\"autosxs-qa-{generate_uuid()}\"\n",
+        "context_column = \"context\"\n",
+        "question_column = \"questions\"\n",
+        "response_column_a = \"response_a\"\n",
+        "response_column_b = \"response_b\"\n",
+        "\n",
+        "parameters = {\n",
+        "    \"evaluation_dataset\": BUCKET_URI + \"/evaluation_dataset.json\",\n",
+        "    \"id_columns\": [question_column],\n",
+        "    \"autorater_prompt_parameters\": {\n",
+        "        \"inference_context\": {\"column\": context_column},\n",
+        "        \"inference_instruction\": {\"column\": question_column},\n",
+        "    },\n",
+        "    \"task\": \"question_answering\",\n",
+        "    \"response_column_a\": response_column_a,\n",
+        "    \"response_column_b\": response_column_b,\n",
+        "}\n",
+        "\n",
+        "job = aiplatform.PipelineJob(\n",
+        "    job_id=display_name,\n",
+        "    display_name=display_name,\n",
+        "    pipeline_root=BUCKET_URI + \"/pipeline\",\n",
+        "    template_path=template_uri,\n",
+        "    parameter_values=parameters,\n",
+        "    enable_caching=False,\n",
+        "    project=PROJECT_ID,\n",
+        "    location=LOCATION,\n",
+        ")\n",
+        "job.run()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "EinPbr3XOYPQ"
+      },
+      "source": [
+        "### Get the judgments and AutoSxS metrics\n",
+        "Next, you can review judgments from the completed AutoSxS job."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "V_9yMfhrOZDk"
+      },
+      "outputs": [],
+      "source": [
+        "for details in job.task_details:\n",
+        "    if details.task_name == \"online-evaluation-pairwise\":\n",
+        "        break\n",
+        "\n",
+        "judgments_uri = details.outputs[\"judgments\"].artifacts[0].uri\n",
+        "judgments_df = pd.read_json(judgments_uri, lines=True)\n",
+        "judgments_df.head()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BlKXu5Ze4tD3"
+      },
+      "source": [
+        "You can also review AutoSxS metrics computed from the judgments.\n",
+        "\n",
+        "You can find more details of AutoSxS metrics [here](https://cloud.google.com/vertex-ai/generative-ai/docs/models/side-by-side-eval#aggregate-metrics)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "G7meI2Eq4muT"
+      },
+      "outputs": [],
+      "source": [
+        "for details in job.task_details:\n",
+        "    if details.task_name == \"model-evaluation-text-generation-pairwise\":\n",
+        "        break\n",
+        "pd.DataFrame([details.outputs[\"autosxs_metrics\"].artifacts[0].metadata])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "TpV-iwP9qw9c"
+      },
+      "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:\n",
+        "\n",
+        "Set `delete_bucket` to **True** to delete the Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "sx_vKniMq9ZX"
+      },
+      "outputs": [],
+      "source": [
+        "delete_pipeline_job = False  # @param {type:\"boolean\"}\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "if delete_pipeline_job:\n",
+        "    job.delete()\n",
+        "\n",
+        "if delete_bucket:\n",
+        "    ! gsutil rm -r gs://{BUCKET_NAME}"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_autosxs_evaluation_llama3_1.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_camp_zipnerf.ipynb

@@ -0,0 +1,691 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "-e1HpvsDh34Q"
+      },
+      "outputs": [],
+      "source": [
+        "# 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": "L5o1Ggr5h34U"
+      },
+      "source": [
+        "# Vertex AI Model Garden - CamP ZipNeRF (Jax) Notebook\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_camp_zipnerf.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_camp_zipnerf.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": "U-SERmqUh34V"
+      },
+      "source": [
+        "**_NOTE_**: This notebook has been tested in the following environment:\n",
+        "\n",
+        "* Python version = 3.9"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "V6QmW0Doh34W"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates a [jax implementation](https://github.com/jonbarron/camp_zipnerf) of [CamP: Camera Preconditioning\n",
+        "for Neural Radiance Fields](https://camp-nerf.github.io/) for training and rendering Neural Radiance Fields (NeRFs) more efficiently. It is primarily aimed at addressing some of the limitations of traditional NeRF techniques, which, while powerful for creating detailed 3D models from 2D images, can be computationally intensive and slow."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "vkSMThcKh34W"
+      },
+      "source": [
+        "## Objective\n",
+        "\n",
+        "In this tutorial, you learn how to:\n",
+        "\n",
+        "- Use [COLMAP](https://colmap.github.io/) to perform Structure from Motion (SfM), a technique that estimates the three-dimensional structure of a scene from a series of two-dimensional images.\n",
+        "- Calibrate, train and render NERF scenes using [Vertex AI custom jobs](https://cloud.google.com/vertex-ai/docs/samples/aiplatform-create-custom-job-sample).\n",
+        "- Render a video along a custom camera path using a series of keyframe photos.\n",
+        "\n",
+        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "\n",
+        "- Vertex AI Training\n",
+        "- Vertex AI Custom Job"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "myi4N60Xh34W"
+      },
+      "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 pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "vofRExleAA8k"
+      },
+      "source": [
+        "## Setup"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "qayv5ifRh34Y"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. 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\"). If not set, a unique GCS bucket will be created instead.\n",
+        "\n",
+        "import os\n",
+        "from datetime import datetime\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Cloud Storage bucket for storing the experiment artifacts.\n",
+        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
+        "# prefer using your own GCS bucket, please change the value yourself below.\n",
+        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
+        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "\n",
+        "# Create a unique GCS bucket for this notebook, if not specified by the user.\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
+        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
+        "else:\n",
+        "    shell_output = ! gsutil ls -Lb {BUCKET_URI} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
+        "    bucket_region = shell_output[0].strip().lower()\n",
+        "    if bucket_region != REGION:\n",
+        "        raise ValueError(\n",
+        "            \"Bucket region %s is different from notebook region %s\"\n",
+        "            % (bucket_region, REGION)\n",
+        "        )\n",
+        "\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
+        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "staging_bucket = os.path.join(BUCKET_URI, \"zipnerf_staging\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=staging_bucket)\n",
+        "\n",
+        "# The pre-built calibration docker image.\n",
+        "CALIBRATION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-cloudnerf-calibrate:latest\"\n",
+        "# The pre-built training docker image.\n",
+        "TRAINING_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-cloudnerf-train:latest\"\n",
+        "# The pre-built rendering docker image.\n",
+        "RENDERING_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-cloudnerf-render:latest\"\n",
+        "\n",
+        "import subprocess\n",
+        "from datetime import datetime\n",
+        "from typing import Any, List\n",
+        "\n",
+        "IMAGE_EXTENSIONS = (\".png\", \".jpg\", \".jpeg\", \".gif\", \".bmp\")\n",
+        "GCS_API_ENDPOINT = \"https://storage.cloud.google.com/\"\n",
+        "\n",
+        "\n",
+        "def get_job_name_with_datetime(prefix: str) -> str:\n",
+        "    \"\"\"Gets the job 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 get_mp4_video_link(mp4_rendering_path: str) -> str:\n",
+        "    # Define the gsutil command.\n",
+        "    command = f\"gsutil ls {mp4_rendering_path}\"\n",
+        "\n",
+        "    # Run the command and capture the output.\n",
+        "    try:\n",
+        "        result = subprocess.check_output(command, shell=True, text=True)\n",
+        "        # Split the result by newlines to get a list of files.\n",
+        "        file_list = result.strip().split(\"\\n\")\n",
+        "    except subprocess.CalledProcessError as e:\n",
+        "        print(f\"An error occurred: {e}\")\n",
+        "        file_list = []\n",
+        "    mp4_video_link = file_list[0].replace(\"gs://\", GCS_API_ENDPOINT)\n",
+        "    return mp4_video_link\n",
+        "\n",
+        "\n",
+        "def write_keyframe_list_to_gcs(\n",
+        "    bucket_path: str, output_gcs_file: str, max_files: int = 10\n",
+        ") -> List[Any]:\n",
+        "    # Get the list of files in the GCS bucket.\n",
+        "    cmd = f\"gsutil ls {bucket_path}\"\n",
+        "    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)\n",
+        "\n",
+        "    if result.returncode != 0:\n",
+        "        print(\"Error listing GCS bucket:\", result.stderr)\n",
+        "        return []\n",
+        "\n",
+        "    # Filter for image files and extract file names.\n",
+        "    files = result.stdout.splitlines()\n",
+        "    image_files = [\n",
+        "        os.path.basename(f) for f in files if f.lower().endswith(IMAGE_EXTENSIONS)\n",
+        "    ]\n",
+        "\n",
+        "    output_file = \"out.txt\"\n",
+        "    with open(output_file, \"w\") as file:\n",
+        "        for name in image_files[:max_files]:\n",
+        "            file.write(name + \"\\n\")\n",
+        "\n",
+        "    cmd = f\"gsutil cp {output_file} {output_gcs_file}\"\n",
+        "    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)\n",
+        "\n",
+        "    if result.returncode != 0:\n",
+        "        print(\"Error listing GCS bucket:\", result.stderr)\n",
+        "        return []"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "OBhvKerXh34a"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Prepare dataset\n",
+        "# @markdown Mip-NeRF 360 dataset contains the following 9 scenes:\n",
+        "# @markdown - `bicycle`\n",
+        "# @markdown - `bonsai`\n",
+        "# @markdown - `counter`\n",
+        "# @markdown - `flowers`\n",
+        "# @markdown - `garden`\n",
+        "# @markdown - `kitchen`\n",
+        "# @markdown - `room`\n",
+        "# @markdown - `stump`\n",
+        "# @markdown - `treehill`\n",
+        "\n",
+        "# @markdown Please note that `flowers` and `treehill` require author's permission. Each scene comes preprocessed with COLMAP information so the calibration step in the following section is optional.\n",
+        "# @markdown If you need to prepare your dataset and store it on Cloud Storage, then the following example shows how to do this for the [mipnerf360 dataset](https://jonbarron.info/mipnerf360/).\n",
+        "\n",
+        "\n",
+        "mipnerf_dataset_directory = \"mipnerf360_dataset\"  # @param {type:\"string\"}\n",
+        "MIPNERF_DATA_GCS_PATH = os.path.join(BUCKET_URI, mipnerf_dataset_directory)\n",
+        "\n",
+        "# Download the bicycle scene data to a local directory.\n",
+        "! rm -rf $mipnerf_dataset_directory\n",
+        "! mkdir -p $mipnerf_dataset_directory\n",
+        "! wget -P $mipnerf_dataset_directory http://storage.googleapis.com/gresearch/refraw360/garden.zip\n",
+        "\n",
+        "# Unzip the mipnerf360 garden dataset.\n",
+        "! unzip $mipnerf_dataset_directory/garden.zip -d $mipnerf_dataset_directory\n",
+        "\n",
+        "# Move mipnerf360 data from local directory to Cloud Storage.\n",
+        "# This step takes a few minutes to finish.\n",
+        "! gsutil -m cp -R $mipnerf_dataset_directory/* $MIPNERF_DATA_GCS_PATH"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WdWHe-RmSEn6"
+      },
+      "source": [
+        "## NERF pipeline"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "j1evctm2h34g"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Run Camera Pose Estimation Custom Job\n",
+        "\n",
+        "# @markdown Once data and experiment paths have been configured, run the custom job below.\n",
+        "\n",
+        "# @markdown The following parameters are required:\n",
+        "\n",
+        "# @markdown * `use_gpu`: Whether to use GPU or not.\n",
+        "# @markdown * `gcs_dataset_path`: Path to image folder in GCS dataset.\n",
+        "# @markdown * `gcs_experiment_path`: GCS path for storing experiment outputs.\n",
+        "# @markdown * `camera`: Type of camera used. `OPENCV` for perspective, `OPENCV_FISHEYE` for fisheye.\n",
+        "\n",
+        "# @markdown The custom job will run on the images in the `gcs_dataset_path` folder and store the colmap outputs in the `gcs_experiment_path/data` folder.\n",
+        "\n",
+        "# @markdown On the scenes in this current dataset, this step takes about 30 minutes.\n",
+        "\n",
+        "# Folder containing all the images of the garden scene.\n",
+        "# e.g. f\"{BUCKET_URI}/{mipnerf_dataset_directory}/garden/images\"\n",
+        "INPUT_IMAGES_FOLDER = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Folder for storing experiment outputs for calibration, training and rendering.\n",
+        "# e.g. f\"{BUCKET_URI}/{mipnerf_dataset_directory}/exp/garden\"\n",
+        "OUTPUT_FOLDER = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# This job will run colmap camera pose estimation.\n",
+        "data_calibration_job_name = get_job_name_with_datetime(\"colmap\")\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "machine_type = \"n1-highmem-64\"\n",
+        "num_nodes = 1\n",
+        "gpu_type = \"NVIDIA_TESLA_V100\"\n",
+        "num_gpus = 8\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": machine_type,\n",
+        "            \"accelerator_type\": gpu_type,\n",
+        "            \"accelerator_count\": num_gpus,\n",
+        "        },\n",
+        "        \"replica_count\": num_nodes,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": CALIBRATION_DOCKER_URI,\n",
+        "            \"args\": [\n",
+        "                \"-use_gpu\",\n",
+        "                \"1\",\n",
+        "                \"-gcs_dataset_path\",\n",
+        "                INPUT_IMAGES_FOLDER,\n",
+        "                \"-gcs_experiment_path\",\n",
+        "                OUTPUT_FOLDER,\n",
+        "                \"-camera\",\n",
+        "                \"OPENCV\",\n",
+        "            ],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "data_calibration_custom_job = aiplatform.CustomJob(\n",
+        "    display_name=data_calibration_job_name,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=staging_bucket,\n",
+        ")\n",
+        "\n",
+        "data_calibration_custom_job.run()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "q7ZWhSpjh34g"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Training the ZipNeRF model\n",
+        "\n",
+        "# @markdown Once the Colmap pose calibration is completed, we can run training.\n",
+        "\n",
+        "# @markdown The following parameters are required:\n",
+        "\n",
+        "# @markdown * `gcs_experiment_path`: GCS path for loading processed dataset and storing experiment outputs.\n",
+        "# @markdown * `factor`: A factor of the downsampled images in the preprocessing step that affects the resolution or detail level of the training pixel ground truth and rendered images. A factor of 2 is recommended for indoor scenes and a factor of 4 for outdoor scenes.\n",
+        "\n",
+        "# @markdown The custom job will run on the images in the `gcs_experiment_path/data` colmap dataset and outputs in the checkpoints in `gcs_experiment_path/checkpoints` folder.\n",
+        "\n",
+        "# @markdown Depending on the configuration, this step could take up to 3 hours.\n",
+        "\n",
+        "# This job will run zipnerf training.\n",
+        "\n",
+        "# This is the nerf training job name. You will use it to load the checkpoints\n",
+        "# in the rendering job for the current run.\n",
+        "nerf_training_job_name = get_job_name_with_datetime(\"nerf_training\")\n",
+        "\n",
+        "FACTOR = 0  # @param [0, 2, 4, 8]\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "machine_type = \"n1-highmem-64\"\n",
+        "num_nodes = 1\n",
+        "gpu_type = \"NVIDIA_TESLA_V100\"\n",
+        "num_gpus = 8\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": machine_type,\n",
+        "            \"accelerator_type\": gpu_type,\n",
+        "            \"accelerator_count\": num_gpus,\n",
+        "        },\n",
+        "        \"replica_count\": num_nodes,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": TRAINING_DOCKER_URI,\n",
+        "            \"args\": [\n",
+        "                \"-training_job_name\",\n",
+        "                nerf_training_job_name,\n",
+        "                \"-gcs_experiment_path\",\n",
+        "                OUTPUT_FOLDER,\n",
+        "                \"-factor\",\n",
+        "                str(FACTOR),\n",
+        "            ],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "nerf_training_custom_job = aiplatform.CustomJob(\n",
+        "    display_name=nerf_training_job_name,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=staging_bucket,\n",
+        ")\n",
+        "\n",
+        "nerf_training_custom_job.run(enable_web_access=True)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "7lufsDWzh34g"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Rendering the ZipNeRF model (360)\n",
+        "\n",
+        "# @markdown Once the training is completed, we can run rendering.\n",
+        "\n",
+        "# @markdown The following parameters are required:\n",
+        "\n",
+        "# @markdown * `gcs_experiment_path`: GCS path for loading processed dataset and storing experiment outputs.\n",
+        "# @markdown * `render_video_fps`: Frame rate of rendered video.\n",
+        "# @markdown * `render_path_frames`: Number of frames to render for a path.\n",
+        "# @markdown * `render_resolution`: Standard display resolutions, for example: (VIDEO_WIDTH, VIDEO_HEIGHT).\n",
+        "\n",
+        "# @markdown The custom job will run on the images in the `gcs_experiment_path/data` colmap dataset and outputs in the checkpoints in `gcs_experiment_path/checkpoints` folder.\n",
+        "\n",
+        "# This job will run zipnerf rendering.\n",
+        "nerf_rendering_job_name = get_job_name_with_datetime(\"nerf_rendering\")\n",
+        "VIDEO_WIDTH = 1280  # @param {type:\"integer\"}\n",
+        "VIDEO_HEIGHT = 720  # @param {type:\"integer\"}\n",
+        "RENDER_PATH_FRAMES = 150  # @param {type:\"integer\"}\n",
+        "RENDER_VIDEO_FPS = 30  # @param {type:\"integer\"}\n",
+        "VIDEO_RESOLUTION = f\"({VIDEO_WIDTH}, {VIDEO_HEIGHT})\"\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "machine_type = \"n1-highmem-64\"\n",
+        "num_nodes = 1\n",
+        "gpu_type = \"NVIDIA_TESLA_V100\"\n",
+        "num_gpus = 8\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": machine_type,\n",
+        "            \"accelerator_type\": gpu_type,\n",
+        "            \"accelerator_count\": num_gpus,\n",
+        "        },\n",
+        "        \"replica_count\": num_nodes,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": RENDERING_DOCKER_URI,\n",
+        "            \"args\": [\n",
+        "                \"-rendering_job_name\",\n",
+        "                nerf_rendering_job_name,\n",
+        "                \"-training_job_name\",\n",
+        "                nerf_training_job_name,\n",
+        "                \"-gcs_experiment_path\",\n",
+        "                OUTPUT_FOLDER,\n",
+        "                \"-render_video_fps\",\n",
+        "                str(RENDER_VIDEO_FPS),\n",
+        "                \"-render_path_frames\",\n",
+        "                str(RENDER_PATH_FRAMES),\n",
+        "                \"-render_resolution\",\n",
+        "                VIDEO_RESOLUTION,\n",
+        "            ],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "nerf_rendering_custom_job = aiplatform.CustomJob(\n",
+        "    display_name=nerf_rendering_job_name,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=staging_bucket,\n",
+        ")\n",
+        "\n",
+        "nerf_rendering_custom_job.run(enable_web_access=True)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "ifhDb9xeh34g"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Show rendered video from GCS\n",
+        "\n",
+        "from IPython.display import Video\n",
+        "\n",
+        "MP4_RENDERING_PATH = (\n",
+        "    f\"{OUTPUT_FOLDER}/render/{nerf_rendering_job_name}/path_videos/videos/*color.mp4\"\n",
+        ")\n",
+        "mp4_video_link = get_mp4_video_link(MP4_RENDERING_PATH)\n",
+        "Video(mp4_video_link)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "Gai5cc-bh34g"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Rendering the ZipNeRF model (custom camera trajectory)\n",
+        "\n",
+        "# @markdown Create keyframe file list for rendering custom camera trajectories.\n",
+        "\n",
+        "# @markdown To create a custom camera trajectory in a Neural Radiance Field (NeRF) model using images from the same dataset used for training, you can generate a keyframe file list where each keyframe corresponds to the name of an image file stored in a Google Cloud Storage (GCS) bucket. This section will guide you through creating this keyframe file list.\n",
+        "\n",
+        "# @markdown Step 1: Identifying keyframe images\n",
+        "# @markdown First, identify the images within your dataset that you want to use as keyframes. These images should ideally represent the significant views or angles that you want your camera trajectory to include.\n",
+        "\n",
+        "# @markdown Step 2: Creating a list of image file names\n",
+        "# @markdown Access Your GCS Bucket: Navigate to your GCS bucket where the dataset is stored.\n",
+        "\n",
+        "# @markdown Select Image Files: Choose the specific image files that you want to use as keyframes. Remember, these should be files used in training the NeRF model, as they will have corresponding camera parameters already defined.\n",
+        "\n",
+        "# @markdown Compile File Names: Create a list of the file names (not the paths) of these selected images. Ensure that each file name is on a separate line. For example:\n",
+        "\n",
+        "\n",
+        "# This job will run zipnerf rendering.\n",
+        "nerf_custom_rendering_job_name = get_job_name_with_datetime(\"nerf_custom_rendering\")\n",
+        "\n",
+        "# Example usage.\n",
+        "KEYFRAME_IMAGE_FILELIST = (\n",
+        "    f\"{OUTPUT_FOLDER}/keyframe_list_{nerf_custom_rendering_job_name}.txt\"\n",
+        ")\n",
+        "max_files = 30  # Set this to the number of files you want\n",
+        "write_keyframe_list_to_gcs(\n",
+        "    INPUT_IMAGES_FOLDER, KEYFRAME_IMAGE_FILELIST, max_files=max_files\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "Hi0oIbyZh34h"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Run rendering for custom path\n",
+        "\n",
+        "# @markdown Once the training is completed, we can run rendering.\n",
+        "\n",
+        "# @markdown The following parameters are required:\n",
+        "\n",
+        "# @markdown * `gcs_experiment_path`: GCS path for loading processed dataset and storing experiment outputs.\n",
+        "# @markdown * `render_video_fps`: Frame rate of rendered video.\n",
+        "# @markdown * `render_resolution`: Standard display resolutions, for example: (VIDEO_WIDTH, VIDEO_HEIGHT).\n",
+        "# @markdown * `keyframe_image_list`: List of image filename, one per line, for rendering custom camera path.\n",
+        "\n",
+        "# @markdown With keyframes, an interpolated path is generated. This path represents a smoothly contoured spline that interconnects the specified keyframe camera poses. The process utilizes a configuration variable, `render_spline_n_interp`, which is preset to a default value of 30. As a result, the finalized interpolated path comprises a total of `render_spline_n_interp` * (n - 1) poses. In the specific scenario under discussion, the config.render_spline_n_interp is configured to 30. **With an input of 30 keyframes, the calculation yields a total of 30 * 29, amounting to 870 poses**.\n",
+        "\n",
+        "# This job will run zipnerf rendering.\n",
+        "# Worker pool spec.\n",
+        "machine_type = \"n1-highmem-64\"\n",
+        "num_nodes = 1\n",
+        "gpu_type = \"NVIDIA_TESLA_V100\"\n",
+        "num_gpus = 8\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": machine_type,\n",
+        "            \"accelerator_type\": gpu_type,\n",
+        "            \"accelerator_count\": num_gpus,\n",
+        "        },\n",
+        "        \"replica_count\": num_nodes,\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": RENDERING_DOCKER_URI,\n",
+        "            \"args\": [\n",
+        "                \"-rendering_job_name\",\n",
+        "                nerf_custom_rendering_job_name,\n",
+        "                \"-training_job_name\",\n",
+        "                nerf_training_job_name,\n",
+        "                \"-gcs_experiment_path\",\n",
+        "                OUTPUT_FOLDER,\n",
+        "                \"-render_resolution\",\n",
+        "                VIDEO_RESOLUTION,\n",
+        "                \"-gcs_keyframes_file\",\n",
+        "                KEYFRAME_IMAGE_FILELIST,\n",
+        "            ],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "nerf_custom_rendering_custom_job = aiplatform.CustomJob(\n",
+        "    display_name=nerf_custom_rendering_job_name,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=staging_bucket,\n",
+        ")\n",
+        "\n",
+        "nerf_custom_rendering_custom_job.run(enable_web_access=True)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "xeup-oLAh34h"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Show rendered video from GCS\n",
+        "\n",
+        "from IPython.display import Video\n",
+        "\n",
+        "MP4_RENDERING_PATH = (\n",
+        "    f\"{OUTPUT_FOLDER}/render/{nerf_rendering_job_name}/path_videos/videos/*color.mp4\"\n",
+        ")\n",
+        "mp4_video_link = get_mp4_video_link(MP4_RENDERING_PATH)\n",
+        "Video(mp4_video_link)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "K9K-sK6INmDP"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Clean up resources\n",
+        "# @markdown Delete the experiment finished jobs and bucket to avoid\n",
+        "# @markdown unnecessary continouous charges that may incur.\n",
+        "\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_URI\n",
+        "\n",
+        "# Delete pose estimation, training and rendering custom jobs.\n",
+        "if data_calibration_custom_job.list(\n",
+        "    filter=f'display_name=\"{data_calibration_job_name}\"'\n",
+        "):\n",
+        "    data_calibration_custom_job.delete()\n",
+        "if nerf_training_custom_job.list(filter=f'display_name=\"{nerf_training_job_name}\"'):\n",
+        "    nerf_training_custom_job.delete()\n",
+        "if nerf_rendering_custom_job.list(filter=f'display_name=\"{nerf_rendering_job_name}\"'):\n",
+        "    nerf_rendering_custom_job.delete()\n",
+        "if nerf_custom_rendering_custom_job.list(\n",
+        "    filter=f'display_name=\"{nerf_custom_rendering_job_name}\"'\n",
+        "):\n",
+        "    nerf_custom_rendering_custom_job.delete()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_camp_zipnerf.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_codegemma_deployment_on_vertex.ipynb

@@ -0,0 +1,636 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "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 - CodeGemma Model (Deployment)\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_codegemma_deployment_on_vertex.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_codegemma_deployment_on_vertex.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 deploying CodeGemma models\n",
+        " * on TPU using **Hex-LLM**, a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel serving solution built with **XLA** that is being developed by Google Cloud, and\n",
+        " * on GPU using [vLLM](https://github.com/vllm-project/vllm), the state-of-the-art open source LLM serving solution on GPU.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Deploy CodeGemma with Hex-LLM on TPU\n",
+        "- Deploy CodeGemma with [vLLM](https://github.com/vllm-project/vllm) on GPU\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "264c07757582"
+      },
+      "source": [
+        "## Run the notebook"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "494697c28ee7"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Request for TPU quota\n",
+        "\n",
+        "# @markdown By default, the quota for TPU deployment `Custom model serving TPU v5e cores per region` is 4. TPU quota is only available in `us-west1`. You can request for higher TPU quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "nLsuvskfhOv4"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. 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\"). If not set, a unique GCS bucket will be created instead.\n",
+        "\n",
+        "import os\n",
+        "from datetime import datetime\n",
+        "from typing import Tuple\n",
+        "\n",
+        "from google.cloud import aiplatform  # Get the default cloud project id.\n",
+        "\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\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",
+        "# Cloud Storage bucket for storing the experiment artifacts.\n",
+        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
+        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
+        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
+        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "# Gets the default BUCKET_URI and SERVICE_ACCOUNT if they were not specified by the user.\n",
+        "SERVICE_ACCOUNT = None\n",
+        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
+        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
+        "\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
+        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "45c8c5438737"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Access CodeGemma Models\n",
+        "\n",
+        "# @markdown If you already obtained access to CodeGemma models on [Hugging Face](https://huggingface.co/), you can load models from there.\n",
+        "# @markdown Alternatively, you can also load the original CodeGemma models for serving from Vertex AI after accepting the agreement.\n",
+        "# @markdown **Please only select and fill one of the two following sections.**\n",
+        "LOAD_MODEL_FROM = \"Google Cloud\"  # @param [\"Hugging Face\", \"Google Cloud\"] {isTemplate:true}\n",
+        "\n",
+        "# @markdown #### Access CodeGemma models on Vertex AI\n",
+        "# @markdown Accept the model agreement to access the models:\n",
+        "# @markdown 1. Open the [CodeGemma model card](https://console.cloud.google.com/vertex-ai/publishers/google/model-garden/364) from [Vertex AI Model Garden](https://cloud.google.com/model-garden).\n",
+        "# @markdown 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",
+        "# @markdown 3. After accepting the agreement of CodeGemma, a `https://` link containing CodeGemma pretrained and finetuned models will be shared.\n",
+        "# @markdown 4. Paste the link in the `VERTEX_MODEL_GARDEN_CODEGEMMA` field below.\n",
+        "# @markdown **Note:** This will unzip and copy the CodeGemma model artifacts to your Cloud Storage bucket, which will take around 30 minutes.\n",
+        "\n",
+        "VERTEX_MODEL_GARDEN_CODEGEMMA = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "VERTEX_MODEL_GARDEN_CODEGEMMA = VERTEX_MODEL_GARDEN_CODEGEMMA.replace(\n",
+        "    \"gs://\", \"https://storage.googleapis.com/\", 1\n",
+        ")\n",
+        "\n",
+        "# @markdown *--- Or ---*\n",
+        "\n",
+        "# @markdown #### Access CodeGemma models on HuggingFace\n",
+        "# @markdown You must provide a Hugging Face User Access Token (read) to access the CodeGemma models. You can follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create a **read** access token and put it in the `HF_TOKEN` field below.\n",
+        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "if LOAD_MODEL_FROM == \"Hugging Face\":\n",
+        "    assert (\n",
+        "        HF_TOKEN\n",
+        "    ), \"Please provide a read HF_TOKEN to load models from Hugging Face, or select a different model source.\"\n",
+        "\n",
+        "if LOAD_MODEL_FROM == \"Google Cloud\":\n",
+        "    assert (\n",
+        "        VERTEX_MODEL_GARDEN_CODEGEMMA\n",
+        "    ), \"Please click the agreement of CodeGemma in Vertex AI Model Garden, and get the URL to CodeGemma model artifacts.\"\n",
+        "\n",
+        "    # Only use the last part in case a full command is pasted.\n",
+        "    signed_url = VERTEX_MODEL_GARDEN_CODEGEMMA.split(\" \")[-1].strip('\"')\n",
+        "\n",
+        "    ! mkdir -p ./codegemma\n",
+        "    ! curl -X GET \"{signed_url}\" | tar -xzvf - -C ./codegemma/\n",
+        "    ! gsutil -m cp -R ./codegemma/* {BUCKET_URI}\n",
+        "\n",
+        "    model_path_prefix = BUCKET_URI.strip(\"/\") + \"/codegemma\"\n",
+        "else:\n",
+        "    model_path_prefix = \"google/\"\n",
+        "\n",
+        "# The pre-built serving docker images with Hex-LLM and vLLM\n",
+        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:deploy\"\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240220_0936_RC01\"\n",
+        "\n",
+        "\n",
+        "def get_job_name_with_datetime(prefix: str) -> str:\n",
+        "    \"\"\"Gets the job name with date time when triggering deployment jobs.\"\"\"\n",
+        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
+        "\n",
+        "\n",
+        "def deploy_model_hexllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    service_account: str,\n",
+        "    machine_type: str = \"ct5lp-hightpu-1t\",\n",
+        "    tensor_parallel_size: int = 1,\n",
+        "    hbm_utilization_factor: float = 0.6,\n",
+        "    max_running_seqs: int = 256,\n",
+        "    endpoint_id: str = \"\",\n",
+        "    min_replica_count: int = 1,\n",
+        "    max_replica_count: int = 1,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys models with Hex-LLM on TPU in Vertex AI.\"\"\"\n",
+        "    if endpoint_id:\n",
+        "        aip_endpoint_name = (\n",
+        "            f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
+        "        )\n",
+        "        endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "    else:\n",
+        "        endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "\n",
+        "    hexllm_args = [\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=7080\",\n",
+        "        \"--log_level=INFO\",\n",
+        "        \"--enable_jit\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        \"--load_format=auto\",\n",
+        "        f\"--tensor_parallel_size={tensor_parallel_size}\",\n",
+        "        f\"--hbm_utilization_factor={hbm_utilization_factor}\",\n",
+        "        f\"--max_running_seqs={max_running_seqs}\",\n",
+        "    ]\n",
+        "    hexllm_envs = {\n",
+        "        \"PJRT_DEVICE\": \"TPU\",\n",
+        "        \"RAY_DEDUP_LOGS\": \"0\",\n",
+        "        \"RAY_USAGE_STATS_ENABLED\": \"0\",\n",
+        "        \"MODEL_ID\": model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "    if HF_TOKEN:\n",
+        "        hexllm_envs.update({\"HF_TOKEN\": HF_TOKEN})\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=HEXLLM_DOCKER_URI,\n",
+        "        serving_container_command=[\"python\", \"-m\", \"hex_llm.server.api_server\"],\n",
+        "        serving_container_args=hexllm_args,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=hexllm_envs,\n",
+        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
+        "        serving_container_deployment_timeout=7200,\n",
+        "    )\n",
+        "\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "        min_replica_count=min_replica_count,\n",
+        "        max_replica_count=max_replica_count,\n",
+        "    )\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "def deploy_model_vllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    service_account: str,\n",
+        "    machine_type: str = \"g2-standard-12\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    max_model_len: int = 8192,\n",
+        "    gpu_memory_utilization=0.9,\n",
+        "    dtype: str = \"bfloat16\",\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys models with vLLM on GPU in Vertex AI.\"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "\n",
+        "    vllm_args = [\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\"--dtype={dtype}\",\n",
+        "        \"--disable-log-stats\",\n",
+        "    ]\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"MODEL_ID\": model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "    if HF_TOKEN:\n",
+        "        env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
+        "        serving_container_command=[\"python\", \"-m\", \"vllm.entrypoints.api_server\"],\n",
+        "        serving_container_args=vllm_args,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        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",
+        "\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "    )\n",
+        "    return model, endpoint"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8neJc8CnDDpu"
+      },
+      "source": [
+        "## Deploy CodeGemma models with Hex-LLM on TPU\n",
+        "\n",
+        "**Hex-LLM** is a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel (LLM) TPU serving solution built with **XLA**, which is being developed by Google Cloud.\n",
+        "\n",
+        "Refer to the \"Request for TPU quota\" section for TPU quota."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "E8OiHHNNE_wj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "\n",
+        "# @markdown Set the model to deploy.\n",
+        "\n",
+        "MODEL_ID = \"codegemma-7b-it\"  # @param [\"codegemma-2b\", \"codegemma-7b\", \"codegemma-7b-it\"]\n",
+        "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
+        "\n",
+        "# @markdown Find Vertex AI prediction TPUv5e machine types in\n",
+        "# @markdown https://cloud.google.com/vertex-ai/docs/predictions/use-tpu#deploy_a_model.\n",
+        "if \"2b\" in model_id:\n",
+        "    # Sets ct5lp-hightpu-1t (1 TPU chip) to deploy CodeGemma 2B models.\n",
+        "    machine_type = \"ct5lp-hightpu-1t\"\n",
+        "    accelerator_type = \"TPU_V5e\"\n",
+        "    # Note: 1 TPU V5 chip has only one core.\n",
+        "    accelerator_count = 1\n",
+        "else:\n",
+        "    # Sets ct5lp-hightpu-4t (4 TPU chips) to deploy CodeGemma 7B models.\n",
+        "    machine_type = \"ct5lp-hightpu-4t\"\n",
+        "    accelerator_type = \"TPU_V5e\"\n",
+        "    # Note: 1 TPU V5 chip has only one core.\n",
+        "    accelerator_count = 4\n",
+        "\n",
+        "# Server parameters.\n",
+        "tensor_parallel_size = accelerator_count\n",
+        "hbm_utilization_factor = 0.6  # Fraction of HBM memory allocated for KV cache after model loading. A larger value improves throughput but gives higher risk of TPU out-of-memory errors with long prompts.\n",
+        "max_running_seqs = 256  # Maximum number of running sequences in a continuous batch.\n",
+        "\n",
+        "# Endpoint configurations.\n",
+        "min_replica_count = 1\n",
+        "max_replica_count = 1\n",
+        "\n",
+        "model_hexllm, endpoint_hexllm = deploy_model_hexllm(\n",
+        "    model_name=get_job_name_with_datetime(prefix=MODEL_ID),\n",
+        "    model_id=model_id,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    machine_type=machine_type,\n",
+        "    tensor_parallel_size=tensor_parallel_size,\n",
+        "    hbm_utilization_factor=hbm_utilization_factor,\n",
+        "    max_running_seqs=max_running_seqs,\n",
+        "    min_replica_count=min_replica_count,\n",
+        "    max_replica_count=max_replica_count,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "nkUaMxIus6Pv"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. The first few requests may have high latency. This is because the server needs to warm up with the initial requests. The following requests should not have the same delay.\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown > What is a car?\n",
+        "# @markdown > A car is a four-wheeled vehicle designed for the transportation of passengers and their belongings.\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown 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.\n",
+        "\n",
+        "# Loads an existing endpoint instance using the endpoint name:\n",
+        "# - Using `endpoint_name = endpoint_hexllm.name` allows us to get the endpoint\n",
+        "#   name of the endpoint `endpoint_hexllm` created in the cell 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:\n",
+        "# endpoint_name = endpoint_hexllm.name\n",
+        "# # endpoint_name = \"\"  # @param {type:\"string\"}\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint_hexllm = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "\n",
+        "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",
+        "instances = [\n",
+        "    {\n",
+        "        \"prompt\": prompt,\n",
+        "        \"max_tokens\": max_tokens,\n",
+        "        \"temperature\": temperature,\n",
+        "        \"top_p\": top_p,\n",
+        "        \"top_k\": top_k,\n",
+        "    },\n",
+        "]\n",
+        "response = endpoint_hexllm.predict(instances=instances)\n",
+        "\n",
+        "prediction = response.predictions[0]\n",
+        "print(prediction)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "YKZ4CBJ2kYaW"
+      },
+      "source": [
+        "## Deploy CodeGemma models with vLLM on GPU\n",
+        "\n",
+        "[vLLM](https://github.com/vllm-project/vllm) is a high-throughput GPU Large Language Model (LLM) serving library which implements a number of optimizations including paged attention and continuous batching.\n",
+        "\n",
+        "Note that V100 GPUs generally offer better throughput and latency performance than L4 GPUs, while L4 GPUs are generally more cost efficient than V100 GPUs. The serving efficiency of L4, V100 and T4 GPUs is inferior to that of A100 GPUs, but L4, V100 and T4 GPUs are nevertheless good serving solutions if you do not have A100 quota.\n",
+        "\n",
+        "CodeGemma model weights are stored in bfloat16 precision. L4 and A100 GPUs are needed for vLLM serving at bfloat16 precision. V100 and T4 GPUs can support vLLM serving at float32 and float16 precision, and they are also meaningful deployment configurations."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "tQIEisUajS6t"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "\n",
+        "# @markdown This section uploads prebuilt CodeGemma model to Model Registry and deploys with [vLLM](https://github.com/vllm-project/vllm) to a Vertex AI Endpoint. It takes 15 to 30 minutes to finish depending on the model and the accelerator.\n",
+        "\n",
+        "# @markdown Set the model to deploy.\n",
+        "\n",
+        "MODEL_ID = \"codegemma-7b-it\"  # @param [\"codegemma-2b\", \"codegemma-7b\", \"codegemma-7b-it\"]\n",
+        "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
+        "\n",
+        "# Find Vertex AI prediction supported accelerators and regions in\n",
+        "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
+        "\n",
+        "# @markdown L4 GPUs are good serving solutions and are more cost effective than A100s.\n",
+        "\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_A100\"]\n",
+        "\n",
+        "if accelerator_type == \"NVIDIA_L4\":\n",
+        "    machine_type = \"g2-standard-12\"\n",
+        "    accelerator_count = 1\n",
+        "elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
+        "    machine_type = \"a2-highgpu-1g\"\n",
+        "    accelerator_count = 1\n",
+        "\n",
+        "# Larger setting of `max-model-len` can lead to higher requirements on\n",
+        "# `gpu-memory-utilization` and GPU configuration. Larger setting of\n",
+        "# `gpu-memory-utilization` increases the risk of running out of GPU memory with\n",
+        "# long prompts.\n",
+        "max_model_len = 2048\n",
+        "gpu_memory_utilization = 0.9\n",
+        "\n",
+        "model_vllm, endpoint_vllm = deploy_model_vllm(\n",
+        "    model_name=get_job_name_with_datetime(prefix=\"codegemma-serve-vllm\"),\n",
+        "    model_id=model_id,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    max_model_len=max_model_len,\n",
+        "    gpu_memory_utilization=gpu_memory_utilization,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "RRR11SWykYaX"
+      },
+      "source": [
+        "Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://github.com/vllm-project/vllm/blob/2e8e49fce3775e7704d413b2f02da6d7c99525c9/vllm/sampling_params.py#L23-L64). Setting `raw_response` to `True` allows you to obtain raw outputs."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "3f5a1e1de60d"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://github.com/vllm-project/vllm/blob/2e8e49fce3775e7704d413b2f02da6d7c99525c9/vllm/sampling_params.py#L23-L64).\n",
+        "\n",
+        "# Loads an existing endpoint instance using the endpoint name:\n",
+        "# - Using `endpoint_name = endpoint_vllm.name` allows us to get the endpoint name of\n",
+        "#   the endpoint `endpoint_vllm` created in the cell 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.\n",
+        "\n",
+        "# endpoint_name = endpoint_vllm.name\n",
+        "# # endpoint_name = \"\"  # @param {type:\"string\"}\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint_vllm = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "\n",
+        "prompt = (\n",
+        "    \"Write a function to list n Fibonacci numbers in Python.\"  # @param {type: \"string\"}\n",
+        ")\n",
+        "max_tokens = 500  # @param {type:\"integer\"}\n",
+        "temperature = 1.0  # @param {type:\"number\"}\n",
+        "top_p = 1.0  # @param {type:\"number\"}\n",
+        "top_k = 10  # @param {type:\"integer\"}\n",
+        "\n",
+        "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\": True,\n",
+        "    },\n",
+        "]\n",
+        "response = endpoint_vllm.predict(instances=instances)\n",
+        "\n",
+        "# \"<|file_separator|>\" is the end of the file token.\n",
+        "for prediction in response.predictions:\n",
+        "    print(prediction.split(\"<|file_separator|>\")[0])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Clean up resources\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
+        "\n",
+        "# Undeploy models and delete endpoints.\n",
+        "endpoint_hexllm.delete(force=True)\n",
+        "endpoint_vllm.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "model_hexllm.delete()\n",
+        "model_vllm.delete()\n",
+        "\n",
+        "# Delete Cloud Storage objects.\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_URI"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_codegemma_deployment_on_vertex.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_e5.ipynb

@@ -0,0 +1,488 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "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 - E5 Text Embedding Models\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_e5.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_e5.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 deploying E5 text embedding models in Vertex AI.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Deploy prebuilt E5 models with Hugging Face [Text Embeddings Inference](https://github.com/huggingface/text-embeddings-inference) (TEI) docker image on a Vertex AI Endpoint\n",
+        "    - [intfloat/multilingual-e5-large-instruct](https://huggingface.co/intfloat/multilingual-e5-large-instruct): 560M params, instruction-tuned\n",
+        "    - [intfloat/multilingual-e5-large](https://huggingface.co/intfloat/multilingual-e5-large): 560M params\n",
+        "    - [intfloat/e5-large-v2](https://huggingface.co/intfloat/e5-large-v2): 335M params\n",
+        "    - [intfloat/multilingual-e5-small](https://huggingface.co/intfloat/multilingual-e5-small): 118M params\n",
+        "    - [intfloat/e5-base-v2](https://huggingface.co/intfloat/e5-base-v2): 109M params\n",
+        "    - [intfloat/e5-small-v2](https://huggingface.co/intfloat/e5-small-v2): 33M params\n",
+        "- Run inference on the deployed Vertex AI Endpoint\n",
+        "\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), [Cloud NL API pricing](https://cloud.google.com/natural-language/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "264c07757582"
+      },
+      "source": [
+        "## Run the notebook"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "LyEVDkHhAUHF"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. 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\"). If not set, a unique GCS bucket will be created instead.\n",
+        "\n",
+        "import os\n",
+        "import uuid\n",
+        "from datetime import datetime\n",
+        "from typing import Tuple\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Cloud Storage bucket for storing the experiment artifacts.\n",
+        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
+        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "BUCKET_URI = \"gs://\"  # @param {type: \"string\"}\n",
+        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "\n",
+        "# @markdown Click \"Show code\" to see more details.\n",
+        "\n",
+        "# Create a unique GCS bucket for this notebook, if not specified by the user.\n",
+        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
+        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "else:\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
+        "    bucket_region = shell_output[0].strip().lower()\n",
+        "    if bucket_region != REGION:\n",
+        "        raise ValueError(\n",
+        "            \"Bucket region %s is different from notebook region %s\"\n",
+        "            % (bucket_region, REGION)\n",
+        "        )\n",
+        "\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "! gcloud services enable language.googleapis.com\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "\n",
+        "# Gets the default BUCKET_URI and SERVICE_ACCOUNT if they were not specified by the user.\n",
+        "\n",
+        "SERVICE_ACCOUNT = None\n",
+        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
+        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "\n",
+        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
+        "\n",
+        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "\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_tei(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    service_account: str,\n",
+        "    docker_uri: str,\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    max_model_len: int = 512,\n",
+        "    gpu_memory_utilization: float = 0.9,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys E5 models with TEI on Vertex AI.\n",
+        "\n",
+        "    Args:\n",
+        "        model_name: Display name of the model.\n",
+        "        model_id: Model ID or path to model weights.\n",
+        "        service_account: Service account for model uploading and deployment.\n",
+        "        machine_type: Deployment machine type.\n",
+        "        accelerator_type: Deployment accelerator type.\n",
+        "        accelerator_count: Number of accelerators to use.\n",
+        "        max_model_len: Maximum model length.\n",
+        "        gpu_memory_utilization: Fraction of GPU memory to be used for the model\n",
+        "            executor.\n",
+        "\n",
+        "    Returns:\n",
+        "        Model instance and endpoint instance.\n",
+        "    \"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "\n",
+        "    tei_args = [\n",
+        "        f\"--model-id={model_id}\",\n",
+        "    ]\n",
+        "    serving_env = {\n",
+        "        \"MODEL_ID\": model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=docker_uri,\n",
+        "        serving_container_args=tei_args,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_environment_variables=serving_env,\n",
+        "    )\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=1800,\n",
+        "        service_account=service_account,\n",
+        "    )\n",
+        "    return model, endpoint"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "kg5MwMIfB9Uj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "# @markdown This section uploads a prebuilt model to Model Registry and deploys it on the Endpoint. The model deployment step will take ~15 minutes to complete.\n",
+        "\n",
+        "prebuilt_model_id = \"intfloat/e5-small-v2\"  # @param [\"intfloat/multilingual-e5-large-instruct\", \"intfloat/multilingual-e5-large\", \"intfloat/e5-large-v2\", \"intfloat/multilingual-e5-small\", \"intfloat/e5-base-v2\", \"intfloat/e5-small-v2\"]\n",
+        "\n",
+        "# @markdown Specify a processor for the TEI docker image. E5 models can be run on either GPU or CPU.\n",
+        "processor = \"NVIDIA_L4\"  # @param[\"NVIDIA_TESLA_V100\", \"NVIDIA_L4\", \"NVIDIA_TESLA_A100\", \"CPU\"]\n",
+        "\n",
+        "if processor == \"NVIDIA_TESLA_V100\":\n",
+        "    accelerator_type = \"NVIDIA_TESLA_V100\"\n",
+        "    machine_type = \"n1-highmem-16\"\n",
+        "    accelerator_count = 2\n",
+        "elif processor == \"NVIDIA_L4\":\n",
+        "    accelerator_type = \"NVIDIA_L4\"\n",
+        "    machine_type = \"g2-standard-8\"\n",
+        "    accelerator_count = 1\n",
+        "elif processor == \"NVIDIA_TESLA_A100\":\n",
+        "    accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "    machine_type = \"a2-highgpu-1g\"\n",
+        "    accelerator_count = 1\n",
+        "elif processor == \"CPU\":\n",
+        "    accelerator_type = None\n",
+        "    machine_type = None\n",
+        "    accelerator_count = None\n",
+        "else:\n",
+        "    raise ValueError(f\"Unsupported processor: {processor}\")\n",
+        "\n",
+        "\n",
+        "# The pre-built serving docker images with TEI.\n",
+        "if processor == \"CPU\":\n",
+        "    TEI_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-embeddings-inference-cpu.1-2\"\n",
+        "else:\n",
+        "    TEI_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-embeddings-inference-cu122.1-2.ubuntu2204\"\n",
+        "\n",
+        "# @markdown Click \"Show code\" to see more details.\n",
+        "\n",
+        "# Finds Vertex AI prediction supported accelerators and regions in\n",
+        "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
+        "\n",
+        "model, endpoint = deploy_model_tei(\n",
+        "    model_name=create_name_with_datetime(prefix=\"e5-serve-tei\"),\n",
+        "    model_id=prebuilt_model_id,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    docker_uri=TEI_DOCKER_URI,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        ")\n",
+        "\n",
+        "print(\"endpoint_name:\", endpoint.name)\n",
+        "print(\"model_name:\", model.display_name)\n",
+        "print(\"model_id:\", model.resource_name)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cUGriDwUQyx4"
+      },
+      "source": [
+        "### Predict\n",
+        "\n",
+        "Once deployment succeeds, you can send requests to the endpoint with text prompts."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "UKfBeMJCN2wl"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Run sample prompt\n",
+        "# @markdown Below is an example to encode queries and passages from the [MS-MARCO passage ranking](https://github.com/microsoft/MSMARCO-Passage-Ranking) dataset.\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown query: how much protein should a female eat\n",
+        "# @markdown query: summit define\n",
+        "# @markdown passage: As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.\n",
+        "# @markdown passage: Definition of summit for English Language Learners. : 1  the highest point of a mountain : the top of a mountain. : 2  the highest level. : 3  a meeting or series of meetings between the leaders of two or more governments.\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown NOTE: Inputs are not limited to 2 queries and 2 passages. To add more inputs, you may modify the code directly.\n",
+        "\n",
+        "query1 = \"how much protein should a female eat?\"  # @param {type: \"string\"}\n",
+        "query2 = \"summit define\"  # @param {type: \"string\"}\n",
+        "passage1 = \"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.\"  # @param {type: \"string\"}\n",
+        "passage2 = \"Definition of summit for English Language Learners. : 1  the highest point of a mountain : the top of a mountain. : 2  the highest level. : 3  a meeting or series of meetings between the leaders of two or more governments.\"  # @param {type: \"string\"}\n",
+        "\n",
+        "# @markdown Click \"Show code\" to see more details.\n",
+        "\n",
+        "# 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.\n",
+        "\n",
+        "# 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",
+        "\n",
+        "from torch import Tensor\n",
+        "\n",
+        "# Each input text should start with \"query: \" or \"passage: \".\n",
+        "# For tasks other than retrieval, you can simply use the \"query: \" prefix.\n",
+        "instances = [\n",
+        "    {\n",
+        "        \"inputs\": [\n",
+        "            f\"query: {query1}\",\n",
+        "            f\"query: {query2}\",\n",
+        "            f\"passage: {passage1}\",\n",
+        "            f\"passage: {passage2}\",\n",
+        "        ],\n",
+        "    },\n",
+        "]\n",
+        "response = endpoint.predict(instances=instances)\n",
+        "\n",
+        "for prediction in response.predictions:\n",
+        "    embeddings = Tensor(prediction)\n",
+        "    scores = (embeddings[:2] @ embeddings[2:].T) * 100\n",
+        "    print(scores.tolist())"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "3hHR7xxyOklz"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Run sample prompt for instruction-tuned models\n",
+        "# @markdown For instruction-tuned models (e.g. intfloat/multilingual-e5-large-instruct), the task definition should be a one-sentence instruction that describes the task. This is a way to customize text embeddings for different scenarios through natural language instructions.\n",
+        "\n",
+        "# @markdown Below is an example to encode queries and passages from the [MS-MARCO passage ranking](https://github.com/microsoft/MSMARCO-Passage-Ranking) dataset.\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown Instruct: Given a web search query, retrieve relevant passages that answer the query\n",
+        "# @markdown Query: how much protein should a female eat\n",
+        "# @markdown Instruct: Given a web search query, retrieve relevant passages that answer the query\n",
+        "# @markdown Query: 南瓜的家常做法\n",
+        "# @markdown As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.\n",
+        "# @markdown 1.清炒南瓜丝 原料:嫩南瓜半个 调料:葱、盐、白糖、鸡精 做法: 1、南瓜用刀薄薄的削去表面一层皮,用勺子刮去瓤 2、擦成细丝(没有擦菜板就用刀慢慢切成细丝) 3、锅烧热放油,入葱花煸出香味 4、入南瓜丝快速翻炒一分钟左右,放盐、一点白糖和鸡精调味出锅 2.香葱炒南瓜 原料:南瓜1只 调料:香葱、蒜末、橄榄油、盐 做法: 1、将南瓜去皮,切成片 2、油锅8成热后,将蒜末放入爆香 3、爆香后,将南瓜片放入,翻炒 4、在翻炒的同时,可以不时地往锅里加水,但不要太多 5、放入盐,炒匀 6、南瓜差不多软和绵了之后,就可以关火 7、撒入香葱,即可出锅\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown NOTE: Inputs are not limited to 1 instruction, 2 queries, and 2 documents. To add more inputs, you may modify the code directly.\n",
+        "\n",
+        "instruction = \"Given a web search query, retrieve relevant passages that answer the query\"  # @param {type: \"string\"}\n",
+        "query1 = \"how much protein should a female eat\"  # @param {type: \"string\"}\n",
+        "query2 = \"南瓜的家常做法\"  # @param {type: \"string\"}\n",
+        "document1 = \"As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.\"  # @param {type: \"string\"}\n",
+        "document2 = \"1.清炒南瓜丝 原料:嫩南瓜半个 调料:葱、盐、白糖、鸡精 做法: 1、南瓜用刀薄薄的削去表面一层皮,用勺子刮去瓤 2、擦成细丝(没有擦菜板就用刀慢慢切成细丝) 3、锅烧热放油,入葱花煸出香味 4、入南瓜丝快速翻炒一分钟左右,放盐、一点白糖和鸡精调味出锅 2.香葱炒南瓜 原料:南瓜1只 调料:香葱、蒜末、橄榄油、盐 做法: 1、将南瓜去皮,切成片 2、油锅8成热后,将蒜末放入爆香 3、爆香后,将南瓜片放入,翻炒 4、在翻炒的同时,可以不时地往锅里加水,但不要太多 5、放入盐,炒匀 6、南瓜差不多软和绵了之后,就可以关火 7、撒入香葱,即可出锅\"  # @param {type: \"string\"}\n",
+        "\n",
+        "# @markdown Click \"Show code\" to see more details.\n",
+        "\n",
+        "# 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.\n",
+        "\n",
+        "# 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",
+        "\n",
+        "from torch import Tensor\n",
+        "\n",
+        "\n",
+        "def get_detailed_instruct(task_description: str, query: str) -> str:\n",
+        "    return f\"Instruct: {task_description}\\nQuery: {query}\"\n",
+        "\n",
+        "\n",
+        "# Each query must come with a one-sentence instruction that describes the task\n",
+        "queries = [\n",
+        "    get_detailed_instruct(instruction, query1),\n",
+        "    get_detailed_instruct(instruction, query2),\n",
+        "]\n",
+        "# No need to add instruction for retrieval documents\n",
+        "documents = [\n",
+        "    document1,\n",
+        "    document2,\n",
+        "]\n",
+        "\n",
+        "instances = [{\"inputs\": queries + documents}]\n",
+        "\n",
+        "response = endpoint.predict(instances=instances)\n",
+        "\n",
+        "for prediction in response.predictions:\n",
+        "    embeddings = Tensor(prediction)\n",
+        "    scores = (embeddings[:2] @ embeddings[2:].T) * 100\n",
+        "    print(scores.tolist())"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Nun3w71JYbss"
+      },
+      "source": [
+        "### End"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "EOL0Qt_0YT5D"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Clean up resources\n",
+        "\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
+        "\n",
+        "endpoint.delete(force=True)\n",
+        "model.delete()\n",
+        "\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_URI"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_e5.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_gemma2_deployment_on_vertex.ipynb

@@ -0,0 +1,725 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "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 - Gemma 2 (Deployment)\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_gemma2_deployment_on_vertex.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_gemma2_deployment_on_vertex.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 deploying Gemma 2 models\n",
+        " * on TPU using **Hex-LLM**, a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel serving solution built with **XLA** that is being developed by Google Cloud, and\n",
+        " * on GPU using **TGI** ([text-generation-inference](https://github.com/huggingface/text-generation-inference)), the state-of-the-art open source LLM serving solution on GPU.\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Deploy Gemma 2 with Hex-LLM on TPU\n",
+        "- Deploy Gemma with [TGI](https://github.com/huggingface/text-generation-inference) on GPU\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "264c07757582"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "C_wC61dhpWXj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Request for TPU quota\n",
+        "\n",
+        "# @markdown By default, the quota for TPU deployment `Custom model serving TPU v5e cores per region` is 4. TPU quota is only available in `us-west1`. You can request for higher TPU quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "2707b02ef5df"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown **[Optional]** Set the GCS BUCKET_URI to store the experiment artifacts, if you want to use your own bucket. **If not set, a unique GCS bucket will be created automatically on your behalf**.\n",
+        "\n",
+        "import json\n",
+        "import os\n",
+        "from datetime import datetime\n",
+        "from typing import Tuple\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Cloud Storage bucket for storing the experiment artifacts.\n",
+        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
+        "# prefer using your own GCS bucket, please change the value yourself below.\n",
+        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    # Create a unique GCS bucket for this notebook if not specified\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
+        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
+        "else:\n",
+        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
+        "    bucket_region = shell_output[0].strip().lower()\n",
+        "    if bucket_region != REGION:\n",
+        "        raise ValueError(\n",
+        "            \"Bucket region %s is different from notebook region %s\"\n",
+        "            % (bucket_region, REGION)\n",
+        "        )\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"gemma2\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "# Gets the default SERVICE_ACCOUNT.\n",
+        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
+        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
+        "\n",
+        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
+        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "\n",
+        "# Enable Vertex AI and Cloud Compute APIs.\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
+        "\n",
+        "# @markdown ## Access Gemma 2 Models\n",
+        "\n",
+        "# @markdown You must provide a Hugging Face User Access Token (read) to access the Gemma 2 models. You can follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create a **read** access token and put it in the `HF_TOKEN` field below.\n",
+        "\n",
+        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "assert (\n",
+        "    HF_TOKEN\n",
+        "), \"Please provide a read HF_TOKEN to load models from Hugging Face, or select a different model source.\"\n",
+        "\n",
+        "model_path_prefix = \"google/\"\n",
+        "\n",
+        "# The pre-built serving docker images.\n",
+        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:gemma2\"\n",
+        "TGI_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu121.2-1.ubuntu2204.py310\"\n",
+        "\n",
+        "SERVICE_ENDPOINT = \"aiplatform.googleapis.com\"\n",
+        "\n",
+        "\n",
+        "def get_job_name_with_datetime(prefix: str) -> str:\n",
+        "    \"\"\"Gets the job name with date time when triggering deployment jobs.\"\"\"\n",
+        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
+        "\n",
+        "\n",
+        "def deploy_model_hexllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    service_account: str,\n",
+        "    machine_type: str = \"ct5lp-hightpu-1t\",\n",
+        "    tensor_parallel_size: int = 1,\n",
+        "    hbm_utilization_factor: float = 0.6,\n",
+        "    max_running_seqs: int = 256,\n",
+        "    endpoint_id: str = \"\",\n",
+        "    min_replica_count: int = 1,\n",
+        "    max_replica_count: int = 1,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys models with Hex-LLM on TPU in Vertex AI.\"\"\"\n",
+        "    if endpoint_id:\n",
+        "        aip_endpoint_name = (\n",
+        "            f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
+        "        )\n",
+        "        endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "    else:\n",
+        "        endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "\n",
+        "    hexllm_args = [\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=7080\",\n",
+        "        \"--log_level=INFO\",\n",
+        "        \"--enable_jit\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        \"--load_format=auto\",\n",
+        "        f\"--tensor_parallel_size={tensor_parallel_size}\",\n",
+        "        f\"--hbm_utilization_factor={hbm_utilization_factor}\",\n",
+        "        f\"--max_running_seqs={max_running_seqs}\",\n",
+        "    ]\n",
+        "    hexllm_envs = {\n",
+        "        \"PJRT_DEVICE\": \"TPU\",\n",
+        "        \"RAY_DEDUP_LOGS\": \"0\",\n",
+        "        \"RAY_USAGE_STATS_ENABLED\": \"0\",\n",
+        "        \"MODEL_ID\": model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "    if HF_TOKEN:\n",
+        "        hexllm_envs.update({\"HF_TOKEN\": HF_TOKEN})\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=HEXLLM_DOCKER_URI,\n",
+        "        serving_container_command=[\"python\", \"-m\", \"hex_llm.server.api_server\"],\n",
+        "        serving_container_args=hexllm_args,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=hexllm_envs,\n",
+        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
+        "        serving_container_deployment_timeout=7200,\n",
+        "    )\n",
+        "\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "        min_replica_count=min_replica_count,\n",
+        "        max_replica_count=max_replica_count,\n",
+        "    )\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "def deploy_model_tgi(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    service_account: str,\n",
+        "    machine_type: str = \"g2-standard-24\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 2,\n",
+        "    max_input_length: int = 1562,\n",
+        "    max_total_tokens: int = 2048,\n",
+        "    max_batch_prefill_tokens: int = 2048,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys models with TGI on GPU in Vertex AI.\"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"AIP_HTTP_PORT\": 7080,\n",
+        "        \"MODEL_ID\": model_id,\n",
+        "        \"NUM_SHARD\": f\"{accelerator_count}\",\n",
+        "        \"MAX_INPUT_LENGTH\": f\"{max_input_length}\",\n",
+        "        \"MAX_TOTAL_TOKENS\": f\"{max_total_tokens}\",\n",
+        "        \"MAX_BATCH_PREFILL_TOKENS\": f\"{max_batch_prefill_tokens}\",\n",
+        "        \"CUDA_MEMORY_FRACTION\": 0.93,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    if HF_TOKEN:\n",
+        "        env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=TGI_DOCKER_URI,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_environment_variables=env_vars,\n",
+        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
+        "    )\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=1800,\n",
+        "        service_account=service_account,\n",
+        "    )\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "def get_quota(project_id: str, region: str, resource_id: str) -> int:\n",
+        "    \"\"\"Returns the quota for a resource in a region. Returns -1 if can not figure out the quota.\"\"\"\n",
+        "    quota_list_output = !gcloud alpha services quota list --service=$SERVICE_ENDPOINT  --consumer=projects/$project_id --filter=\"$SERVICE_ENDPOINT/$resource_id\" --format=json\n",
+        "    # Use '.s' on the command output because it is an SList type.\n",
+        "    quota_data = json.loads(quota_list_output.s)\n",
+        "    if len(quota_data) == 0 or \"consumerQuotaLimits\" not in quota_data[0]:\n",
+        "        return -1\n",
+        "    if (\n",
+        "        len(quota_data[0][\"consumerQuotaLimits\"]) == 0\n",
+        "        or \"quotaBuckets\" not in quota_data[0][\"consumerQuotaLimits\"][0]\n",
+        "    ):\n",
+        "        return -1\n",
+        "    all_regions_data = quota_data[0][\"consumerQuotaLimits\"][0][\"quotaBuckets\"]\n",
+        "    for region_data in all_regions_data:\n",
+        "        if (\n",
+        "            region_data.get(\"dimensions\")\n",
+        "            and region_data[\"dimensions\"][\"region\"] == region\n",
+        "        ):\n",
+        "            if \"effectiveLimit\" in region_data:\n",
+        "                return int(region_data[\"effectiveLimit\"])\n",
+        "            else:\n",
+        "                return 0\n",
+        "    return -1\n",
+        "\n",
+        "\n",
+        "def get_resource_id(accelerator_type: str, is_for_training: bool) -> str:\n",
+        "    \"\"\"Returns the resource id for a given accelerator type and the use case.\n",
+        "    Args:\n",
+        "      accelerator_type: The accelerator type.\n",
+        "      is_for_training: Whether the resource is used for training. Set false\n",
+        "      for serving use case.\n",
+        "    Returns:\n",
+        "      The resource id.\n",
+        "    \"\"\"\n",
+        "    training_accelerator_map = {\n",
+        "        \"NVIDIA_TESLA_V100\": \"custom_model_training_nvidia_v100_gpus\",\n",
+        "        \"NVIDIA_L4\": \"custom_model_training_nvidia_l4_gpus\",\n",
+        "        \"NVIDIA_TESLA_A100\": \"custom_model_training_nvidia_a100_gpus\",\n",
+        "        \"NVIDIA_TESLA_T4\": \"custom_model_training_nvidia_t4_gpus\",\n",
+        "        \"TPU_V5e\": \"custom_model_training_tpu_v5e\",\n",
+        "        \"TPU_V3\": \"custom_model_training_tpu_v3\",\n",
+        "    }\n",
+        "    serving_accelerator_map = {\n",
+        "        \"NVIDIA_TESLA_V100\": \"custom_model_serving_nvidia_v100_gpus\",\n",
+        "        \"NVIDIA_L4\": \"custom_model_serving_nvidia_l4_gpus\",\n",
+        "        \"NVIDIA_TESLA_A100\": \"custom_model_serving_nvidia_a100_gpus\",\n",
+        "        \"NVIDIA_TESLA_T4\": \"custom_model_serving_nvidia_t4_gpus\",\n",
+        "        \"TPU_V5e\": \"custom_model_serving_tpu_v5e\",\n",
+        "    }\n",
+        "    if is_for_training:\n",
+        "        if accelerator_type in training_accelerator_map:\n",
+        "            return training_accelerator_map[accelerator_type]\n",
+        "        else:\n",
+        "            raise ValueError(\n",
+        "                f\"Could not find accelerator type: {accelerator_type} for training.\"\n",
+        "            )\n",
+        "    else:\n",
+        "        if accelerator_type in serving_accelerator_map:\n",
+        "            return serving_accelerator_map[accelerator_type]\n",
+        "        else:\n",
+        "            raise ValueError(\n",
+        "                f\"Could not find accelerator type: {accelerator_type} for serving.\"\n",
+        "            )\n",
+        "\n",
+        "\n",
+        "def check_quota(\n",
+        "    project_id: str,\n",
+        "    region: str,\n",
+        "    accelerator_type: str,\n",
+        "    accelerator_count: int,\n",
+        "    is_for_training: bool,\n",
+        "):\n",
+        "    \"\"\"Checks if the project and the region has the required quota.\"\"\"\n",
+        "    resource_id = get_resource_id(accelerator_type, is_for_training)\n",
+        "    quota = get_quota(project_id, region, resource_id)\n",
+        "    quota_request_instruction = (\n",
+        "        \"Either use \"\n",
+        "        \"a different region or request additional quota. Follow \"\n",
+        "        \"instructions here \"\n",
+        "        \"https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota\"\n",
+        "        \" to check quota in a region or request additional quota for \"\n",
+        "        \"your project.\"\n",
+        "    )\n",
+        "    if quota == -1:\n",
+        "        raise ValueError(\n",
+        "            f\"\"\"Quota not found for: {resource_id} in {region}.\n",
+        "            {quota_request_instruction}\"\"\"\n",
+        "        )\n",
+        "    if quota < accelerator_count:\n",
+        "        raise ValueError(\n",
+        "            f\"\"\"Quota not enough for {resource_id} in {region}:\n",
+        "            {quota} < {accelerator_count}.\n",
+        "            {quota_request_instruction}\"\"\"\n",
+        "        )"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8neJc8CnDDpu"
+      },
+      "source": [
+        "## Deploy Gemma 2 models with Hex-LLM on TPU\n",
+        "\n",
+        "**Hex-LLM** is a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel (LLM) TPU serving solution built with **XLA**, which is being developed by Google Cloud.\n",
+        "\n",
+        "Refer to the \"Request for TPU quota\" section for TPU quota."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "E8OiHHNNE_wj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "# @markdown Set the model ID. Model weights can be loaded from HuggingFace or from a GCS bucket.\n",
+        "\n",
+        "# @markdown Select one of the four model variations.\n",
+        "MODEL_ID = \"gemma-2-9b\"  # @param [\"gemma-2-9b\", \"gemma-2-9b-it\", \"gemma-2-27b\", \"gemma-2-27b-it\"] {allow-input: true, isTemplate: true}\n",
+        "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
+        "\n",
+        "# @markdown Find Vertex AI prediction TPUv5e machine types in\n",
+        "# @markdown https://cloud.google.com/vertex-ai/docs/predictions/use-tpu#deploy_a_model.\n",
+        "if \"9b\" in model_id:\n",
+        "    # Sets ct5lp-hightpu-4t (4 TPU chips) to deploy Gemma 2 9B models.\n",
+        "    machine_type = \"ct5lp-hightpu-4t\"\n",
+        "    accelerator_type = \"TPU_V5e\"\n",
+        "    # Note: 1 TPU V5 chip has only one core.\n",
+        "    accelerator_count = 4\n",
+        "else:\n",
+        "    # Sets ct5lp-hightpu-8t (8 TPU chips) to deploy Gemma 2 27B models.\n",
+        "    machine_type = \"ct5lp-hightpu-8t\"\n",
+        "    accelerator_type = \"TPU_V5e\"\n",
+        "    # Note: 1 TPU V5 chip has only one core.\n",
+        "    accelerator_count = 8\n",
+        "\n",
+        "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",
+        "\n",
+        "# Server parameters.\n",
+        "tensor_parallel_size = accelerator_count\n",
+        "hbm_utilization_factor = 0.6  # Fraction of HBM memory allocated for KV cache after model loading. A larger value improves throughput but gives higher risk of TPU out-of-memory errors with long prompts.\n",
+        "max_running_seqs = 256  # Maximum number of running sequences in a continuous batch.\n",
+        "\n",
+        "# Endpoint configurations.\n",
+        "min_replica_count = 1\n",
+        "max_replica_count = 1\n",
+        "\n",
+        "model_hexllm, endpoint_hexllm = deploy_model_hexllm(\n",
+        "    model_name=get_job_name_with_datetime(prefix=MODEL_ID),\n",
+        "    model_id=model_id,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    machine_type=machine_type,\n",
+        "    tensor_parallel_size=tensor_parallel_size,\n",
+        "    hbm_utilization_factor=hbm_utilization_factor,\n",
+        "    max_running_seqs=max_running_seqs,\n",
+        "    min_replica_count=min_replica_count,\n",
+        "    max_replica_count=max_replica_count,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "nkUaMxIus6Pv"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. The first few requests may have high latency. This is because the server needs to warm up with the initial requests. The following requests should not have the same delay.\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown > What is a car?\n",
+        "# @markdown > A car is a four-wheeled vehicle designed for the transportation of passengers and their belongings.\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown 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.\n",
+        "\n",
+        "# Loads an existing endpoint instance using the endpoint name:\n",
+        "# - Using `endpoint_name = endpoint_hexllm.name` allows us to get the endpoint\n",
+        "#   name of the endpoint `endpoint_hexllm` created in the cell 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:\n",
+        "# endpoint_name = endpoint_without_peft.name\n",
+        "# # endpoint_name = \"\"  # @param {type:\"string\"}\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint_hexllm = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "\n",
+        "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",
+        "instances = [\n",
+        "    {\n",
+        "        \"prompt\": prompt,\n",
+        "        \"max_tokens\": max_tokens,\n",
+        "        \"temperature\": temperature,\n",
+        "        \"top_p\": top_p,\n",
+        "        \"top_k\": top_k,\n",
+        "    },\n",
+        "]\n",
+        "response = endpoint_hexllm.predict(instances=instances)\n",
+        "\n",
+        "prediction = response.predictions[0]\n",
+        "print(prediction)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "GFfqpQm8BNwZ"
+      },
+      "source": [
+        "## Deploy Gemma models with TGI on GPU\n",
+        "\n",
+        "[TGI](https://github.com/huggingface/text-generation-inference) stands for Text Generation Inference. It's a powerful library designed specifically for running large language models on GPUs efficiently. TGI utilizes techniques like \"paged attention\" and \"continuous batching\" to improve the speed and throughput of LLMs.\n",
+        "\n",
+        "Currently, only L4 GPUs are demonstrated in this notebook. Functionality on other GPU types will be confirmed and added in the future.\n",
+        "\n",
+        "Gemma2 9B models require at least 2 L4 GPUs for deployment. Gemma2 27B models require at least 4 L4 GPUs for deployment."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "TBNJYZMlBNwZ"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "MODEL_ID = \"gemma-2-9b\"  # @param [\"gemma-2-9b\", \"gemma-2-9b-it\", \"gemma-2-27b\", \"gemma-2-27b-it\"] {allow-input: true, isTemplate: true}\n",
+        "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
+        "\n",
+        "# @markdown Finds Vertex AI prediction supported accelerators and regions in\n",
+        "# @markdown https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
+        "\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\"] {isTemplate: true}\n",
+        "\n",
+        "if \"9b\" in MODEL_ID:\n",
+        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "        # Sets 2 L4 (24G) to deploy Gemma 9B models.\n",
+        "        machine_type = \"g2-standard-24\"\n",
+        "        accelerator_count = 2\n",
+        "    else:\n",
+        "        raise ValueError(\n",
+        "            \"Recommended machine settings not found for accelerator type: %s\"\n",
+        "            % accelerator_type\n",
+        "        )\n",
+        "elif \"27b\" in MODEL_ID:\n",
+        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "        # Sets 4 L4 (24G) to deploy Gemma 27B models.\n",
+        "        machine_type = \"g2-standard-48\"\n",
+        "        accelerator_count = 4\n",
+        "    else:\n",
+        "        raise ValueError(\n",
+        "            \"Recommended machine settings not found for accelerator type: %s\"\n",
+        "            % accelerator_type\n",
+        "        )\n",
+        "else:\n",
+        "    raise ValueError(\"Recommended machine settings not found for model: %s\" % MODEL_ID)\n",
+        "\n",
+        "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",
+        "\n",
+        "# Note that larger token counts will require more GPU memory. For example, if you'd\n",
+        "# like to increase the `max_total_tokens` and `max_batch_prefill_tokens` to 8192,\n",
+        "# you may need 4 L4s for the 9b model, and 8 L4s for the 27b model.\n",
+        "max_input_length = 1562\n",
+        "max_total_tokens = 2048\n",
+        "max_batch_prefill_tokens = 2048\n",
+        "\n",
+        "model_tgi, endpoint_tgi = deploy_model_tgi(\n",
+        "    model_name=get_job_name_with_datetime(prefix=MODEL_ID),\n",
+        "    model_id=model_id,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    max_input_length=max_input_length,\n",
+        "    max_total_tokens=max_total_tokens,\n",
+        "    max_batch_prefill_tokens=max_batch_prefill_tokens,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "ZCwACkjuBNwZ"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts.\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown > What is a car?\n",
+        "# @markdown > A car is a four-wheeled vehicle designed for the transportation of passengers and their belongings.\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown 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.\n",
+        "\n",
+        "\n",
+        "# @markdown Please click \"Show Code\" to see more details.\n",
+        "\n",
+        "# Loads an existing endpoint instance using the endpoint name:\n",
+        "# - Using `endpoint_name = endpoint_tgi.name` allows us to get the\n",
+        "#   endpoint name of the endpoint `endpoint_tgi` 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.\n",
+        "\n",
+        "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint_tgi = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "\n",
+        "prompt = \"What is a car?\"  # @param {type: \"string\"}\n",
+        "max_new_tokens = 128  # @param {type:\"integer\"}\n",
+        "temperature = 1.0  # @param {type:\"number\"}\n",
+        "top_p = 0.9  # @param {type:\"number\"}\n",
+        "top_k = 1  # @param {type:\"integer\"}\n",
+        "\n",
+        "# Overides max_new_tokens and top_k parameters during inferences.\n",
+        "# If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`,\n",
+        "# you can reduce the max length, such as set max_new_tokens as 20.\n",
+        "instances = [\n",
+        "    {\n",
+        "        \"inputs\": f\"### Human: {prompt}### Assistant: \",\n",
+        "        \"parameters\": {\n",
+        "            \"max_new_tokens\": max_new_tokens,\n",
+        "            \"temperature\": temperature,\n",
+        "            \"top_p\": top_p,\n",
+        "            \"top_k\": top_k,\n",
+        "        },\n",
+        "    },\n",
+        "]\n",
+        "\n",
+        "response = endpoint_tgi.predict(instances=instances)\n",
+        "\n",
+        "for prediction in response.predictions:\n",
+        "    print(prediction)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "af21a3cff1e0"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Delete the models and endpoints\n",
+        "\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
+        "# Undeploy models and delete endpoints.\n",
+        "endpoint_hexllm.delete(force=True)\n",
+        "endpoint_tgi.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "model_hexllm.delete()\n",
+        "model_tgi.delete()\n",
+        "\n",
+        "# Delete Cloud Storage objects.\n",
+        "delete_bucket = False  # @param {type:\"boolean\", isTemplate: true}\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_URI"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_gemma2_deployment_on_vertex.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_gemma_deployment_on_gke.ipynb

@@ -0,0 +1,395 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "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": [
+        "# Gemma deployment to GKE using TGI on GPU\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_gemma_deployment_on_gke.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_gemma_deployment_on_gke.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 downloading and deploying Gemma, open models from Google DeepMind using Text Generation Inference [TGI](https://github.com/), an efficient serving option to improve serving throughput. In this notebook we will deploy and serve TGI on GPUs. In this guide we specifically use L4 GPUs but this guide should also work for A100(40 GB), A100(80 GB), H100(80 GB) GPUs.\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "Deploy and run inference for serving Gemma with TGI on GPUs.\n",
+        "\n",
+        "### GPUs\n",
+        "\n",
+        "GPUs let you accelerate specific workloads running on your nodes such as machine learning and data processing. GKE provides a range of machine type options for node configuration, including machine types with NVIDIA H100, L4, and A100 GPUs.\n",
+        "\n",
+        "Before you use GPUs in GKE, we recommend that you complete the following learning path:\n",
+        "\n",
+        "Learn about [current GPU version availability](https://cloud.google.com/compute/docs/gpus)\n",
+        "\n",
+        "Learn about [GPUs in GKE](https://cloud.google.com/kubernetes-engine/docs/concepts/gpus)\n",
+        "\n",
+        "\n",
+        "### TGI\n",
+        "\n",
+        "TGI is a highly optimized open-source LLM serving framework that can increase serving throughput on GPUs. TGI includes features such as:\n",
+        "\n",
+        "Optimized transformer implementation with PagedAttention\n",
+        "Continuous batching to improve the overall serving throughput\n",
+        "Tensor parallelism and distributed serving on multiple GPUs\n",
+        "\n",
+        "To learn more, refer to the [TGI documentation](https://github.com/huggingface/text-generation-inference/blob/main/README.md)\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "264c07757582"
+      },
+      "source": [
+        "## Run the notebook"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "855d6b96f291"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. Set Hugging Face access token in `HF_TOKEN` field. If you don't already have a \"read\" access token, follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create an access token with \"read\" permission. You can find your existing access tokens in the Hugging Face [Access Token](https://huggingface.co/settings/tokens) page.\n",
+        "\n",
+        "# @markdown 3. **[Optional]** Set `CLUSTER_NAME` if you want to use your own GKE cluster. If not set, this example will create a standard cluster with 2 NVIDIA L4 GPU accelerators.\n",
+        "\n",
+        "import os\n",
+        "from datetime import datetime\n",
+        "\n",
+        "# The HuggingFace token used to download models.\n",
+        "HF_TOKEN = \"\"  # @param {type:\"string\"}\n",
+        "assert HF_TOKEN, \"Please set Hugging Face access token in `HF_TOKEN`.\"\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Set up gcloud.\n",
+        "! gcloud config set project \"$PROJECT_ID\"\n",
+        "! gcloud services enable container.googleapis.com\n",
+        "\n",
+        "# Add kubectl to the set of available tools.\n",
+        "! mkdir -p /tools/google-cloud-sdk/.install\n",
+        "! gcloud components install kubectl --quiet\n",
+        "\n",
+        "# The cluster name to create\n",
+        "CLUSTER_NAME = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Use existing GKE cluster or create a new cluster.\n",
+        "if CLUSTER_NAME:\n",
+        "    ! gcloud container clusters get-credentials {CLUSTER_NAME} --location {REGION}\n",
+        "else:\n",
+        "    now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "    CLUSTER_NAME=f\"gke-gemma-cluster-{now}\"\n",
+        "    ! gcloud container clusters create {CLUSTER_NAME} \\\n",
+        "        --project={PROJECT_ID} \\\n",
+        "        --region={REGION} \\\n",
+        "        --workload-pool={PROJECT_ID}.svc.id.goog \\\n",
+        "        --release-channel=rapid \\\n",
+        "        --num-nodes=4\n",
+        "    ! gcloud container node-pools create gpupool \\\n",
+        "        --accelerator=type=nvidia-l4,count=2,gpu-driver-version=latest \\\n",
+        "        --project={PROJECT_ID} \\\n",
+        "        --location={REGION} \\\n",
+        "        --node-locations={REGION}-a \\\n",
+        "        --cluster={CLUSTER_NAME} \\\n",
+        "        --machine-type=g2-standard-24 \\\n",
+        "        --num-nodes=1\n",
+        "\n",
+        "# Create Kubernetes secret for Hugging Face credentials\n",
+        "! kubectl create secret generic hf-secret \\\n",
+        "    --from-literal=hf_api_token={HF_TOKEN} \\\n",
+        "    --dry-run=client -o yaml > hf-secret.yaml\n",
+        "\n",
+        "! kubectl apply -f hf-secret.yaml"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "6psJZY_zUDgj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy TGI\n",
+        "\n",
+        "# @markdown This section deploys Gemma on TGI.\n",
+        "\n",
+        "# @markdown Select one of the following model version and size options:\n",
+        "\n",
+        "# The size of the model to launch\n",
+        "MODEL_VERSION = \"1.1\"  # @param [\"1.0\", \"1.1\"]\n",
+        "if MODEL_VERSION == \"1.1\":\n",
+        "    version_string = \"-1.1\"\n",
+        "else:\n",
+        "    version_string = \"\"\n",
+        "# The size of the model to launch\n",
+        "MODEL_SIZE = \"2b\"  # @param [\"2b\", \"7b\"]\n",
+        "\n",
+        "# @markdown After the container is up, there will be another ~5 minutes to download the needed artifacts, the time depends on what runtime you are using to run your colab environment.\n",
+        "\n",
+        "# The number of GPUs to run: 1 for 2b, 2 for 7b\n",
+        "GPU_COUNT = 1\n",
+        "if MODEL_SIZE == \"7b\":\n",
+        "    GPU_COUNT = 2\n",
+        "\n",
+        "# Ephemeral storage\n",
+        "EPHEMERAL_STORAGE_SIZE = \"20Gi\"\n",
+        "if MODEL_SIZE == \"7b\":\n",
+        "    EPHEMERAL_STORAGE_SIZE = \"40Gi\"\n",
+        "\n",
+        "# Memory size\n",
+        "MEMORY_SIZE = \"7Gi\"\n",
+        "if MODEL_SIZE == \"7b\":\n",
+        "    MEMORY_SIZE = \"25Gi\"\n",
+        "\n",
+        "GPU_SHARD = 1\n",
+        "if MODEL_SIZE == \"7b\":\n",
+        "    GPU_SHARD = 2\n",
+        "\n",
+        "CPU_LIMITS = 2\n",
+        "if MODEL_SIZE == \"7b\":\n",
+        "    CPU_LIMITS = 10\n",
+        "\n",
+        "K8S_YAML = f\"\"\"\n",
+        "apiVersion: apps/v1\n",
+        "kind: Deployment\n",
+        "metadata:\n",
+        "  name: tgi-gemma-deployment\n",
+        "spec:\n",
+        "  replicas: 1\n",
+        "  selector:\n",
+        "    matchLabels:\n",
+        "      app: gemma-server\n",
+        "  template:\n",
+        "    metadata:\n",
+        "      labels:\n",
+        "        app: gemma-server\n",
+        "        ai.gke.io/model: gemma{version_string}-{MODEL_SIZE}\n",
+        "        ai.gke.io/inference-server: text-generation-inference\n",
+        "        examples.ai.gke.io/source: user-guide\n",
+        "    spec:\n",
+        "      containers:\n",
+        "      - name: inference-server\n",
+        "        image: us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-hf-tgi-serve:20240220_0936_RC01\n",
+        "        resources:\n",
+        "          requests:\n",
+        "            cpu: \"2\"\n",
+        "            memory: {MEMORY_SIZE}\n",
+        "            ephemeral-storage: {EPHEMERAL_STORAGE_SIZE}\n",
+        "            nvidia.com/gpu: {GPU_COUNT}\n",
+        "          limits:\n",
+        "            cpu: {CPU_LIMITS}\n",
+        "            memory: {MEMORY_SIZE}\n",
+        "            ephemeral-storage: {EPHEMERAL_STORAGE_SIZE}\n",
+        "            nvidia.com/gpu: {GPU_COUNT}\n",
+        "        args:\n",
+        "        - --model-id=$(MODEL_ID)\n",
+        "        - --num-shard={GPU_SHARD}\n",
+        "        env:\n",
+        "        - name: MODEL_ID\n",
+        "          value: google/gemma{version_string}-{MODEL_SIZE}-it\n",
+        "        - name: PORT\n",
+        "          value: \"8000\"\n",
+        "        - name: HUGGING_FACE_HUB_TOKEN\n",
+        "          valueFrom:\n",
+        "            secretKeyRef:\n",
+        "              name: hf-secret\n",
+        "              key: hf_api_token\n",
+        "        volumeMounts:\n",
+        "        - mountPath: /dev/shm\n",
+        "          name: dshm\n",
+        "      volumes:\n",
+        "      - name: dshm\n",
+        "        emptyDir:\n",
+        "          medium: Memory\n",
+        "      nodeSelector:\n",
+        "        cloud.google.com/gke-accelerator: nvidia-l4\n",
+        "---\n",
+        "apiVersion: v1\n",
+        "kind: Service\n",
+        "metadata:\n",
+        "  name: llm-service\n",
+        "spec:\n",
+        "  selector:\n",
+        "    app: gemma-server\n",
+        "  type: ClusterIP\n",
+        "  ports:\n",
+        "  - protocol: TCP\n",
+        "    port: 8000\n",
+        "    targetPort: 8000\n",
+        "\"\"\"\n",
+        "\n",
+        "with open(\"tgi.yaml\", \"w\") as f:\n",
+        "    f.write(K8S_YAML)\n",
+        "\n",
+        "! kubectl apply -f tgi.yaml\n",
+        "\n",
+        "# Wait for container to be created.\n",
+        "import time\n",
+        "\n",
+        "print(\"Waiting for container to be created...\\n\")\n",
+        "while True:\n",
+        "    shell_output = ! kubectl get pod\n",
+        "    container_status = \"\\n\".join(shell_output)\n",
+        "    if \"1/1\" in container_status:\n",
+        "        break\n",
+        "    time.sleep(5)\n",
+        "\n",
+        "print(container_status)\n",
+        "\n",
+        "# Wait for downloading artifacts.\n",
+        "print(\"\\nDownloading artifacts...\")\n",
+        "while True:\n",
+        "    shell_output = ! kubectl logs -l app=gemma-server\n",
+        "    logs = \"\\n\".join(shell_output)\n",
+        "    if \"Connected\" in logs:\n",
+        "        break\n",
+        "    time.sleep(5)\n",
+        "\n",
+        "print(\"Server is up and running.\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "QDTasPgGW7EG"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Prediction\n",
+        "\n",
+        "# @markdown Once the server is up and running, you may send prompts to local server for prediction.\n",
+        "\n",
+        "import json\n",
+        "\n",
+        "prompt = \"What are the top 5 most popular programming languages? Please be brief.\"  # @param {type: \"string\"}\n",
+        "temperature = 0.40  # @param {type: \"number\"}\n",
+        "top_p = 0.1  # @param {type: \"number\"}\n",
+        "max_tokens = 250  # @param {type: \"number\"}\n",
+        "\n",
+        "request = {\n",
+        "    \"inputs\": prompt,\n",
+        "    \"temperature\": temperature,\n",
+        "    \"top_p\": top_p,\n",
+        "    \"max_tokens\": max_tokens,\n",
+        "}\n",
+        "\n",
+        "command = f\"\"\"kubectl exec -t $( kubectl get pod -l app=gemma-server -o jsonpath=\"{{.items[0].metadata.name}}\" ) -c inference-server -- curl -X POST http://localhost:8000/generate \\\n",
+        "   -H \"Content-Type: application/json\" \\\n",
+        "   -d '{json.dumps(request)}' \\\n",
+        "   2> /dev/null\"\"\"\n",
+        "\n",
+        "output = !{command}\n",
+        "print(\"Output:\")\n",
+        "print(json.loads(output[0])[\"generated_text\"])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "wbRmgoOZF6es"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
+        "\n",
+        "! kubectl delete deployments tgi-gemma-deployment\n",
+        "! kubectl delete services llm-service\n",
+        "! kubectl delete secrets hf-secret\n",
+        "\n",
+        "DELETE_CLUSTER = False # @param {type: \"boolean\"}\n",
+        "\n",
+        "if DELETE_CLUSTER:\n",
+        "  ! gcloud container clusters delete {CLUSTER_NAME} \\\n",
+        "    --region={REGION} \\\n",
+        "    --quiet"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_gemma_deployment_on_gke.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_gemma_deployment_on_vertex.ipynb

@@ -0,0 +1,811 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "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 - Gemma (Deployment)\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_gemma_deployment_on_vertex.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_gemma_deployment_on_vertex.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 deploying Gemma models\n",
+        " * on TPU using **Hex-LLM**, a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel serving solution built with **XLA** that is being developed by Google Cloud, and\n",
+        " * on GPU using [vLLM](https://github.com/vllm-project/vllm), the state-of-the-art open source LLM serving solution on GPU.\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Deploy Gemma with Hex-LLM on TPU\n",
+        "- Deploy Gemma with [vLLM](https://github.com/vllm-project/vllm) on GPU\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "264c07757582"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "2707b02ef5df"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. 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\"). If not set, a unique GCS bucket will be created instead.\n",
+        "\n",
+        "# @markdown 3. By default, the quota for TPU deployment `Custom model serving TPU v5e cores per region` is 4. TPU quota is only available in `us-west1`. You can request for higher TPU quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
+        "\n",
+        "# Import the necessary packages\n",
+        "\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "import importlib\n",
+        "import os\n",
+        "from datetime import datetime\n",
+        "from typing import Tuple\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\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",
+        "# Cloud Storage bucket for storing the experiment artifacts.\n",
+        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
+        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
+        "\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
+        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
+        "else:\n",
+        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
+        "    bucket_region = shell_output[0].strip().lower()\n",
+        "    if bucket_region != REGION:\n",
+        "        raise ValueError(\n",
+        "            \"Bucket region %s is different from notebook region %s\"\n",
+        "            % (bucket_region, REGION)\n",
+        "        )\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"gemma\")\n",
+        "\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "# Gets the default SERVICE_ACCOUNT.\n",
+        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
+        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
+        "\n",
+        "\n",
+        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
+        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "# @markdown ## Access Gemma Models\n",
+        "# @markdown Choose between accessing Gemma models on [Hugging Face](https://huggingface.co/)\n",
+        "# @markdown or Vertex AI as described below.\n",
+        "\n",
+        "# @markdown If you already obtained access to Gemma models on [Hugging Face](https://huggingface.co/), you can load models from there.\n",
+        "# @markdown Alternatively, you can also load the original Gemma models for serving from Vertex AI after accepting the agreement.\n",
+        "\n",
+        "# @markdown **Please only select and fill one of the two following sections.**\n",
+        "LOAD_MODEL_FROM = (\n",
+        "    \"Hugging Face\"  # @param [\"Hugging Face\", \"Google Cloud\"] {isTemplate:true}\n",
+        ")\n",
+        "\n",
+        "# @markdown ---\n",
+        "\n",
+        "# @markdown ### Access Gemma models on Hugging Face\n",
+        "# @markdown You must provide a Hugging Face User Access Token (read) to access the Gemma models. You can follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create a **read** access token and put it in the `HF_TOKEN` field below.\n",
+        "\n",
+        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "if LOAD_MODEL_FROM == \"Hugging Face\":\n",
+        "    assert (\n",
+        "        HF_TOKEN\n",
+        "    ), \"Please provide a read HF_TOKEN to load models from Hugging Face, or select a different model source.\"\n",
+        "\n",
+        "# @markdown *--- Or ---*\n",
+        "# @markdown ### Access Gemma models on Vertex AI\n",
+        "# @markdown Accept the model agreement to access the models:\n",
+        "# @markdown 1. Open the [Gemma model card](https://console.cloud.google.com/vertex-ai/publishers/google/model-garden/335) from [Vertex AI Model Garden](https://cloud.google.com/model-garden).\n",
+        "# @markdown 1. 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",
+        "# @markdown 1. After accepting the agreement of Gemma, a `https://` link containing Gemma pretrained and finetuned models will be shared.\n",
+        "# @markdown 1. Paste the link in the `VERTEX_AI_MODEL_GARDEN_GEMMA` field below.\n",
+        "# @markdown **Note:** This will unzip and copy the Gemma model artifacts to your Cloud Storage bucket, which will take around 1 hour.\n",
+        "\n",
+        "VERTEX_AI_MODEL_GARDEN_GEMMA = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "\n",
+        "\n",
+        "if LOAD_MODEL_FROM == \"Google Cloud\":\n",
+        "    assert (\n",
+        "        VERTEX_AI_MODEL_GARDEN_GEMMA\n",
+        "    ), \"Please accept the agreement of Gemma in Vertex AI Model Garden and get the URL to Gemma model artifacts, or select a different model source.\"\n",
+        "\n",
+        "    # Only use the last part in case a full command is pasted.\n",
+        "    signed_url = VERTEX_AI_MODEL_GARDEN_GEMMA.split(\" \")[-1].strip('\"')\n",
+        "\n",
+        "    ! mkdir -p ./gemma\n",
+        "    ! curl -X GET \"{signed_url}\" | tar -xzvf - -C ./gemma/\n",
+        "    ! gsutil -m cp -R ./gemma/* {MODEL_BUCKET}\n",
+        "\n",
+        "    model_path_prefix = MODEL_BUCKET\n",
+        "    HF_TOKEN = \"\"\n",
+        "else:\n",
+        "    model_path_prefix = \"google/\"\n",
+        "\n",
+        "# @markdown ---\n",
+        "\n",
+        "# The pre-built serving docker images.\n",
+        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:deploy\"\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240508_0916_RC02\"\n",
+        "\n",
+        "\n",
+        "def deploy_model_hexllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    service_account: str,\n",
+        "    base_model_id: str = None,\n",
+        "    tensor_parallel_size: int = 1,\n",
+        "    machine_type: str = \"ct5lp-hightpu-1t\",\n",
+        "    hbm_utilization_factor: float = 0.6,\n",
+        "    max_running_seqs: int = 256,\n",
+        "    endpoint_id: str = \"\",\n",
+        "    min_replica_count: int = 1,\n",
+        "    max_replica_count: int = 1,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys models with Hex-LLM on TPU in Vertex AI.\"\"\"\n",
+        "    if endpoint_id:\n",
+        "        aip_endpoint_name = (\n",
+        "            f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
+        "        )\n",
+        "        endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "    else:\n",
+        "        endpoint = aiplatform.Endpoint.create(\n",
+        "            display_name=f\"{model_name}-endpoint\",\n",
+        "            location=TPU_DEPLOYMENT_REGION,\n",
+        "        )\n",
+        "\n",
+        "    if not base_model_id:\n",
+        "        base_model_id = model_id\n",
+        "\n",
+        "    if not tensor_parallel_size:\n",
+        "        tensor_parallel_size = int(machine_type[-2])\n",
+        "\n",
+        "    hexllm_args = [\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=7080\",\n",
+        "        \"--log_level=INFO\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        f\"--tensor_parallel_size={tensor_parallel_size}\",\n",
+        "        \"--enable_jit\",\n",
+        "        \"--load_format=auto\",\n",
+        "        f\"--hbm_utilization_factor={hbm_utilization_factor}\",\n",
+        "        f\"--max_running_seqs={max_running_seqs}\",\n",
+        "    ]\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"MODEL_ID\": base_model_id,\n",
+        "        \"PJRT_DEVICE\": \"TPU\",\n",
+        "        \"RAY_DEDUP_LOGS\": \"0\",\n",
+        "        \"RAY_USAGE_STATS_ENABLED\": \"0\",\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    try:\n",
+        "        if HF_TOKEN:\n",
+        "            env_vars.update({\"HF_TOKEN\": HF_TOKEN})\n",
+        "    except:\n",
+        "        pass\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=HEXLLM_DOCKER_URI,\n",
+        "        serving_container_command=[\"python\", \"-m\", \"hex_llm.server.api_server\"],\n",
+        "        serving_container_args=hexllm_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",
+        "        location=TPU_DEPLOYMENT_REGION,\n",
+        "    )\n",
+        "\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "        min_replica_count=min_replica_count,\n",
+        "        max_replica_count=max_replica_count,\n",
+        "    )\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "def deploy_model_vllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    service_account: str,\n",
+        "    base_model_id: str = None,\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 = 4096,\n",
+        "    dtype: str = \"auto\",\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",
+        "    if not base_model_id:\n",
+        "        base_model_id = model_id\n",
+        "\n",
+        "    vllm_args = [\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\"--dtype={dtype}\",\n",
+        "        \"--disable-log-stats\",\n",
+        "    ]\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"MODEL_ID\": base_model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    try:\n",
+        "        if HF_TOKEN:\n",
+        "            env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
+        "    except:\n",
+        "        pass\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
+        "        serving_container_command=[\"python\", \"-m\", \"vllm.entrypoints.api_server\"],\n",
+        "        serving_container_args=vllm_args,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        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=1800,\n",
+        "        service_account=service_account,\n",
+        "    )\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
+        "    return model, endpoint"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8neJc8CnDDpu"
+      },
+      "source": [
+        "## Deploy Gemma models with Hex-LLM on TPU\n",
+        "\n",
+        "**Hex-LLM** is a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel (LLM) TPU serving solution built with **XLA**, which is being developed by Google Cloud.\n",
+        "\n",
+        "Refer to the \"Request for TPU quota\" section for TPU quota."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "E8OiHHNNE_wj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "# @markdown Set the model ID. Model weights can be loaded from HuggingFace or from a GCS bucket.\n",
+        "\n",
+        "# @markdown Select one of the six model variations.\n",
+        "MODEL_ID = \"gemma-1.1-2b-it\"  # @param [\"gemma-2b\", \"gemma-2b-it\", \"gemma-7b\", \"gemma-7b-it\", \"gemma-1.1-2b-it\", \"gemma-1.1-7b-it\"] {allow-input: true, isTemplate: true}\n",
+        "TPU_DEPLOYMENT_REGION = \"us-west1\"  # @param [\"us-west1\"] {isTemplate:true}\n",
+        "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
+        "\n",
+        "# @markdown Find Vertex AI prediction TPUv5e machine types in\n",
+        "# @markdown https://cloud.google.com/vertex-ai/docs/predictions/use-tpu#deploy_a_model.\n",
+        "if \"2b\" in model_id:\n",
+        "    # Sets ct5lp-hightpu-1t (1 TPU chip) to deploy Gemma 2B models.\n",
+        "    machine_type = \"ct5lp-hightpu-1t\"\n",
+        "    accelerator_type = \"TPU_V5e\"\n",
+        "    # Note: 1 TPU V5 chip has only one core.\n",
+        "    accelerator_count = 1\n",
+        "else:\n",
+        "    # Sets ct5lp-hightpu-4t (4 TPU chips) to deploy Gemma 7B models.\n",
+        "    machine_type = \"ct5lp-hightpu-4t\"\n",
+        "    accelerator_type = \"TPU_V5e\"\n",
+        "    # Note: 1 TPU V5 chip has only one core.\n",
+        "    accelerator_count = 4\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",
+        "\n",
+        "# Server parameters.\n",
+        "hbm_utilization_factor = 0.6  # A larger value improves throughput but gives higher risk of TPU out-of-memory errors with long prompts.\n",
+        "max_running_seqs = 256\n",
+        "\n",
+        "# Endpoint configurations.\n",
+        "min_replica_count = 1\n",
+        "max_replica_count = 1\n",
+        "\n",
+        "models[\"hexllm_tpu\"], endpoints[\"hexllm_tpu\"] = deploy_model_hexllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=MODEL_ID),\n",
+        "    model_id=model_id,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    machine_type=machine_type,\n",
+        "    hbm_utilization_factor=hbm_utilization_factor,\n",
+        "    max_running_seqs=max_running_seqs,\n",
+        "    min_replica_count=min_replica_count,\n",
+        "    max_replica_count=max_replica_count,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "nkUaMxIus6Pv"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts based on your `template`. Note that the first few prompts will take longer to execute.\n",
+        "\n",
+        "# @markdown 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.\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown > What is a car?\n",
+        "# @markdown > A car is a four-wheeled vehicle designed for the transportation of passengers and their belongings.\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown 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.\n",
+        "\n",
+        "# Loads an existing endpoint instance using the endpoint name:\n",
+        "# - Using `endpoint_name = endpoint.name` allows us to get the endpoint\n",
+        "#   name of the endpoint `endpoint` created in the cell 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:\n",
+        "# endpoint_name = endpoint_without_peft.name\n",
+        "# # 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",
+        "\n",
+        "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",
+        "instances = [\n",
+        "    {\n",
+        "        \"prompt\": prompt,\n",
+        "        \"max_tokens\": max_tokens,\n",
+        "        \"temperature\": temperature,\n",
+        "        \"top_p\": top_p,\n",
+        "        \"top_k\": top_k,\n",
+        "    },\n",
+        "]\n",
+        "response = endpoints[\"hexllm_tpu\"].predict(instances=instances)\n",
+        "\n",
+        "for prediction in response.predictions:\n",
+        "    print(prediction)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f615c03d6638"
+      },
+      "source": [
+        "### Build chat applications with Gemma\n",
+        "\n",
+        "You can build chat applications with the instruction finetuned Gemma models.\n",
+        "\n",
+        "The instruction tuned Gemma models were trained with a specific formatter that annotates instruction tuning examples with extra information, both during training and inference. The annotations (1) indicate roles in a conversation, and (2) delineate tunes in a conversation. Below we show a sample code snippet for formatting the model prompt using the user and model chat templates for a multi-turn conversation. The relevant tokens are:\n",
+        "- `user`: user turn\n",
+        "- `model`: model turn\n",
+        "- `<start_of_turn>`: beginning of dialogue turn\n",
+        "- `<end_of_turn>`: end of dialogue turn\n",
+        "\n",
+        "An example set of dialogues is:\n",
+        "```\n",
+        "<start_of_turn>user\n",
+        "knock knock<end_of_turn>\n",
+        "<start_of_turn>model\n",
+        "who is there<end_of_turn>\n",
+        "<start_of_turn>user\n",
+        "LaMDA<end_of_turn>\n",
+        "<start_of_turn>model\n",
+        "LaMDA who?<end_of_turn>\n",
+        "```\n",
+        "where `<end_of_turn>\\n` is the turn separator and `<start_of_turn>model\\n` is the prompt prefix. This means if we would like to prompt the model with a question like, `What is Cramer's Rule?`, we should use:\n",
+        "```\n",
+        "<start_of_turn>user\n",
+        "What is Cramer's Rule?<end_of_turn>\n",
+        "<start_of_turn>model\n",
+        "```"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "e59377392346"
+      },
+      "outputs": [],
+      "source": [
+        "# Chat templates.\n",
+        "USER_CHAT_TEMPLATE = \"<start_of_turn>user\\n{prompt}<end_of_turn>\\n\"\n",
+        "MODEL_CHAT_TEMPLATE = \"<start_of_turn>model\\n{prompt}<end_of_turn>\\n\"\n",
+        "\n",
+        "# Sample formatted prompt.\n",
+        "prompt = (\n",
+        "    USER_CHAT_TEMPLATE.format(prompt=\"What is a good place for travel in the US?\")\n",
+        "    + MODEL_CHAT_TEMPLATE.format(prompt=\"California.\")\n",
+        "    + USER_CHAT_TEMPLATE.format(prompt=\"What can I do in California?\")\n",
+        "    + \"<start_of_turn>model\\n\"\n",
+        ")\n",
+        "print(\"Chat prompt:\\n\", prompt)\n",
+        "\n",
+        "instances = [\n",
+        "    {\n",
+        "        \"prompt\": prompt,\n",
+        "        \"max_tokens\": 50,\n",
+        "        \"temperature\": 1.0,\n",
+        "        \"top_p\": 1.0,\n",
+        "        \"top_k\": 1,\n",
+        "    },\n",
+        "]\n",
+        "response = endpoints[\"hexllm_tpu\"].predict(instances=instances)\n",
+        "\n",
+        "prediction = response.predictions[0]\n",
+        "print(prediction)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "YKZ4CBJ2kYaW"
+      },
+      "source": [
+        "## Deploy Gemma models with vLLM on GPU\n",
+        "\n",
+        "[vLLM](https://github.com/vllm-project/vllm) is a high-throughput GPU Large Language Model (LLM) serving library which implements a number of optimizations including paged attention and continuous batching.\n",
+        "\n",
+        "Note that V100 GPUs generally offer better throughput and latency performance than L4 GPUs, while L4 GPUs are generally more cost efficient than V100 GPUs. The serving efficiency of L4, V100 and T4 GPUs is inferior to that of A100 GPUs, but L4, V100 and T4 GPUs are nevertheless good serving solutions if you do not have A100 quota.\n",
+        "\n",
+        "Gemma model weights are stored in bfloat16 precision. L4 and A100 GPUs are needed for vLLM serving at bfloat16 precision. V100 and T4 GPUs can support vLLM serving at float32 and float16 precision, and they are also meaningful deployment configurations."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "03d504bcd60b"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "MODEL_ID = \"gemma-1.1-2b-it\"  # @param [\"gemma-2b\", \"gemma-2b-it\", \"gemma-7b\", \"gemma-7b-it\", \"gemma-1.1-2b-it\", \"gemma-1.1-7b-it\"] {isTemplate: true}\n",
+        "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
+        "\n",
+        "# @markdown Finds Vertex AI prediction supported accelerators and regions in\n",
+        "# @markdown https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
+        "\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_V100\", \"NVIDIA_TESLA_T4\", \"NVIDIA_TESLA_A100\"] {isTemplate: true}\n",
+        "\n",
+        "if \"2b\" in MODEL_ID:\n",
+        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "        # Sets 1 L4 (24G) to deploy Gemma 2B models.\n",
+        "        machine_type = \"g2-standard-8\"\n",
+        "        accelerator_count = 1\n",
+        "        vllm_dtype = \"bfloat16\"\n",
+        "    elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
+        "        # Sets 1 V100 (16G) to deploy Gemma 2B models.\n",
+        "        machine_type = \"n1-standard-8\"\n",
+        "        accelerator_count = 1\n",
+        "        vllm_dtype = \"float32\"\n",
+        "    elif accelerator_type == \"NVIDIA_TESLA_T4\":\n",
+        "        # Sets 1 T4 (16G) to deploy Gemma 2B models.\n",
+        "        machine_type = \"n1-standard-8\"\n",
+        "        accelerator_count = 1\n",
+        "        vllm_dtype = \"float32\"\n",
+        "    elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
+        "        # Sets 1 A100 (40G) to deploy Gemma 2B models.\n",
+        "        machine_type = \"a2-highgpu-1g\"\n",
+        "        accelerator_count = 1\n",
+        "        vllm_dtype = \"bfloat16\"\n",
+        "    else:\n",
+        "        raise ValueError(\n",
+        "            \"Recommended machine settings not found for accelerator type: %s\"\n",
+        "            % accelerator_type\n",
+        "        )\n",
+        "elif \"7b\" in MODEL_ID:\n",
+        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "        # Sets 1 L4 (24G) to deploy Gemma 7B models.\n",
+        "        machine_type = \"g2-standard-12\"\n",
+        "        accelerator_count = 1\n",
+        "        vllm_dtype = \"bfloat16\"\n",
+        "    elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
+        "        # Sets 1 A100 (40G) to deploy Gemma 7B models.\n",
+        "        machine_type = \"a2-highgpu-1g\"\n",
+        "        accelerator_count = 1\n",
+        "        vllm_dtype = \"bfloat16\"\n",
+        "    else:\n",
+        "        raise ValueError(\n",
+        "            \"Recommended machine settings not found for accelerator type: %s\"\n",
+        "            % accelerator_type\n",
+        "        )\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        \"Recommended machine settings not found for accelerator type: %s\"\n",
+        "        % accelerator_type\n",
+        "    )\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",
+        "\n",
+        "# Note that a larger max_model_len will require more GPU memory.\n",
+        "max_model_len = 2048\n",
+        "\n",
+        "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"gemma-serve-vllm\"),\n",
+        "    model_id=model_id,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    max_model_len=max_model_len,\n",
+        "    dtype=vllm_dtype,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "RRR11SWykYaX"
+      },
+      "source": [
+        "Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://github.com/vllm-project/vllm/blob/2e8e49fce3775e7704d413b2f02da6d7c99525c9/vllm/sampling_params.py#L23-L64). Setting `raw_response` to `True` allows you to obtain raw outputs."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "3f5a1e1de60d"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "\n",
+        "# @markdown 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",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown Human: What is a car?\n",
+        "# @markdown 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",
+        "# @markdown ```\n",
+        "# @markdown 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.\n",
+        "\n",
+        "# 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.\n",
+        "\n",
+        "# 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",
+        "\n",
+        "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\"}\n",
+        "\n",
+        "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)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "104fe2c03812"
+      },
+      "source": [
+        "### Apply chat templates\n",
+        "\n",
+        "Chat templates can be applied to model predictions generated by the vLLM endpoint as well. You may use the same code snippets as for the Hex-LLM endpoint. They are not repeated here for brevity."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "af21a3cff1e0"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
+        "\n",
+        "# Undeploy model and delete endpoint.\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "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_URI"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_gemma_deployment_on_vertex.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_gemma_evaluation.ipynb

@@ -0,0 +1,380 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "language": "python",
+      "metadata": {
+        "id": "B8S-yo8qTIcO"
+      },
+      "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": "MTRywGxLTZfU"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Gemma Evaluation\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_gemma_evaluation.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_gemma_evaluation.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": "2CXS0vZfT8_7"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates evaluating pre-trained and instruction-tuned Gemma models in Vertex AI.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Evaluate pre-trained and instruction-tuned Gemma model on any of the benchmark datasets\n",
+        "- Clean up the resources\n",
+        "\n",
+        "| Models |\n",
+        "| :- |\n",
+        "| [google/gemma-2b](https://huggingface.co/google/gemma-2b)\n",
+        "| [google/gemma-2b-it](https://huggingface.co/google/gemma-2b-it)\n",
+        "| [google/gemma-7b](https://huggingface.co/google/gemma-7b)\n",
+        "| [google/gemma-7b-it](https://huggingface.co/google/gemma-7b-it)\n",
+        "| [google/gemma-1.1-2b-it](https://huggingface.co/google/gemma-1.1-2b-it)\n",
+        "| [google/gemma-1.1-7b-it](https://huggingface.co/google/gemma-1.1-7b-it)\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "HCY8PGrFUbT1"
+      },
+      "source": [
+        "## Run the notebook"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "language": "python",
+      "metadata": {
+        "cellView": "form",
+        "id": "81CC3tL1T_TL"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. 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\"). If not set, a unique GCS bucket will be created instead.\n",
+        "\n",
+        "# Import the necessary packages\n",
+        "\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "import importlib\n",
+        "import os\n",
+        "from datetime import datetime\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\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",
+        "# Cloud Storage bucket for storing the experiment artifacts.\n",
+        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
+        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
+        "\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
+        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
+        "else:\n",
+        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
+        "    bucket_region = shell_output[0].strip().lower()\n",
+        "    if bucket_region != REGION:\n",
+        "        raise ValueError(\n",
+        "            \"Bucket region %s is different from notebook region %s\"\n",
+        "            % (bucket_region, REGION)\n",
+        "        )\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"gemma\")\n",
+        "\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "# Gets the default SERVICE_ACCOUNT.\n",
+        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
+        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
+        "\n",
+        "\n",
+        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
+        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "# The evaluation docker image.\n",
+        "EVAL_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-lm-evaluation-harness:20240320_0655_RC00\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "language": "python",
+      "metadata": {
+        "cellView": "form",
+        "id": "pNHMbjr0UjrK"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Evaluate Gemma models\n",
+        "\n",
+        "# @markdown This section demonstrates how to evaluate the Gemma models with and without finetuned LoRA adapters using EleutherAI's [Language Model Evaluation Harness (lm-evaluation-harness)](https://github.com/EleutherAI/lm-evaluation-harness) with Vertex CustomJob. Please reference the peak GPU memory usage for serving and adjust the machine type, accelerator type and accelerator count accordingly.\n",
+        "\n",
+        "# @markdown You must provide a Hugging Face User Access Token (read) to access the Gemma models. You can follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create a **read** access token and put it in the `HF_TOKEN` field below.\n",
+        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "\n",
+        "# @markdown This example uses the dataset [HellaSwag](https://arxiv.org/abs/1905.07830). All supported tasks are listed in [this task table](https://github.com/EleutherAI/lm-evaluation-harness/blob/master/docs/task_table.md).\n",
+        "# @markdown Set evaluation dataset.\n",
+        "eval_dataset = \"hellaswag\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "# Find Vertex AI supported accelerators and regions in:\n",
+        "#  https://cloud.google.com/vertex-ai/docs/training/configure-compute\n",
+        "\n",
+        "\n",
+        "# Setup evaluation job.\n",
+        "# @markdown Set the base model id.\n",
+        "base_model_id = \"google/gemma-1.1-2b-it\"  # @param[\"google/gemma-2b\", \"google/gemma-2b-it\", \"google/gemma-7b\", \"google/gemma-7b-it\", \"google/gemma-1.1-2b-it\", \"google/gemma-1.1-7b-it\"] {isTemplate:true}\n",
+        "job_name = common_util.get_job_name_with_datetime(prefix=\"gemma-eval\")\n",
+        "eval_output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
+        "eval_output_dir_gcsfuse = eval_output_dir.replace(\"gs://\", \"/gcs/\")\n",
+        "\n",
+        "# @markdown Set the accelerator type.\n",
+        "accelerator_type = (\n",
+        "    \"NVIDIA_L4\"  # @param[\"NVIDIA_TESLA_V100\", \"NVIDIA_L4\", \"NVIDIA_TESLA_A100\"]\n",
+        ")\n",
+        "\n",
+        "# @markdown To evaluate a PEFT-finetuned model, enter the PEFT output directory to the LoRA adapter below.\n",
+        "# @markdown Otherwise, leave it empty.\n",
+        "# @markdown See the [finetuning notebook](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_gemma_finetuning_on_vertex.ipynb) for more details.\n",
+        "# @markdown Set the PEFT output directory.\n",
+        "peft_output_dir = \"\"  # @param {type:\"string\"}\n",
+        "peft_output_dir_gcsfuse = peft_output_dir.replace(\"gs://\", \"/gcs/\")\n",
+        "\n",
+        "if accelerator_type == \"NVIDIA_TESLA_A100\":\n",
+        "    machine_type = \"a2-highgpu-1g\"\n",
+        "    accelerator_count = 1\n",
+        "elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
+        "    machine_type = \"n1-standard-8\"\n",
+        "    accelerator_count = 2\n",
+        "elif accelerator_type == \"NVIDIA_L4\":\n",
+        "    machine_type = \"g2-standard-8\"\n",
+        "    accelerator_count = 1\n",
+        "else:\n",
+        "    print(f\"Unsupported accelerator type: {accelerator_type}\")\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",
+        "# Prepare evaluation command that runs the evaluation harness.\n",
+        "# Set `trust_remote_code = True` because evaluating the model requires\n",
+        "# executing code from the model repository.\n",
+        "# Set `use_accelerate = True` to enable evaluation across multiple GPUs.\n",
+        "eval_command = [\n",
+        "    \"lm_eval\",\n",
+        "    \"--model\",\n",
+        "    \"hf\",\n",
+        "    \"--tasks\",\n",
+        "    f\"{eval_dataset}\",\n",
+        "    \"--output_path\",\n",
+        "    f\"{eval_output_dir_gcsfuse}\",\n",
+        "]\n",
+        "\n",
+        "if peft_output_dir_gcsfuse:\n",
+        "    eval_command += [\n",
+        "        \"--model_args\",\n",
+        "        f\"pretrained={base_model_id},peft={peft_output_dir_gcsfuse},trust_remote_code=True,parallelize=True,device_map_option=auto\",\n",
+        "    ]\n",
+        "else:\n",
+        "    eval_command += [\n",
+        "        \"--model_args\",\n",
+        "        f\"pretrained={base_model_id},trust_remote_code=True,parallelize=True,device_map_option=auto\",\n",
+        "    ]\n",
+        "\n",
+        "# Pass evaluation arguments and launch job.\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": machine_type,\n",
+        "            \"accelerator_type\": accelerator_type,\n",
+        "            \"accelerator_count\": accelerator_count,\n",
+        "        },\n",
+        "        \"replica_count\": replica_count,\n",
+        "        \"disk_spec\": {\n",
+        "            \"boot_disk_size_gb\": 500,\n",
+        "        },\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": EVAL_DOCKER_URI,\n",
+        "            \"env\": [\n",
+        "                {\n",
+        "                    \"name\": \"HF_TOKEN\",\n",
+        "                    \"value\": HF_TOKEN,\n",
+        "                }\n",
+        "            ],\n",
+        "            \"command\": eval_command,\n",
+        "            \"args\": [],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "eval_job = aiplatform.CustomJob(\n",
+        "    display_name=job_name,\n",
+        "    worker_pool_specs=worker_pool_specs,\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": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "CVBxGpwWU3kY"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Fetch and print evaluation results\n",
+        "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 :] + \"/results.json\"\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}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "execution_count": null,
+      "metadata": {
+        "id": "unjukbcjEBOd"
+      },
+      "outputs": [],
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "qWN3cl_VU7pa"
+      },
+      "outputs": [],
+      "source": [
+        "# Delete evaluation job.\n",
+        "\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_URI\n",
+        "    # Uncomment below to delete all artifacts\n",
+        "    # !gsutil -m rm -r $STAGING_BUCKET $MODEL_BUCKET $EXPERIMENT_BUCKET\n",
+        "\n",
+        "eval_job.delete()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_gemma_evaluation.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_huggingface_local_inference.ipynb

@@ -24,7 +24,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "4dc4391f6be7"
@@ -55,7 +54,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "4e8a0fdd6f44"
@@ -79,7 +77,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "69453bf7230e"
@@ -89,7 +86,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "68990d91bc5f"
@@ -140,7 +136,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "05e23144b125"
@@ -158,7 +153,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "ad1a690839d5"
@@ -168,7 +162,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "0a4008240483"
@@ -200,7 +193,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "ae94b9b23a52"
@@ -245,7 +237,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "e76b3fe8d10c"
@@ -292,7 +283,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "1ade95a9b20e"

notebooks/community/model_garden/model_garden_jax_fvlm.ipynb

@@ -1318,12 +1318,18 @@
       },
       "outputs": [],
       "source": [
+        "serving_env = {\n",
+        "    \"MODEL_ID\": \"F-VLM-JAX-\",\n",
+        "    \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "}\n",
+        "\n",
         "jax_fvlm_model = aiplatform.Model.upload(\n",
         "    display_name=\"jax_fvlm\",\n",
         "    artifact_uri=GCS_CONVERTED_SAVED_MODEL_DIR,\n",
         "    serving_container_image_uri=OPTIMIZED_TF_RUNTIME_IMAGE_URI,\n",
         "    serving_container_args=[],\n",
         "    location=REGION,\n",
+        "    serving_container_environment_variables=serving_env,\n",
         ")\n",
         "\n",
         "jax_fvlm_endpoint = jax_fvlm_model.deploy(\n",

notebooks/community/model_garden/model_garden_jax_owl_vit_v2.ipynb

@@ -634,12 +634,18 @@
       },
       "outputs": [],
       "source": [
+        "serving_env = {\n",
+        "    \"MODEL_ID\": \"jax-owl-vit-v2\",\n",
+        "    \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "}\n",
+        "\n",
         "jax_owl2_model = aiplatform.Model.upload(\n",
         "    display_name=\"jax_owl2\",\n",
         "    artifact_uri=GCS_CONVERTED_SAVED_MODEL_DIR,\n",
         "    serving_container_image_uri=OPTIMIZED_TF_RUNTIME_IMAGE_URI,\n",
         "    serving_container_args=[],\n",
         "    location=REGION,\n",
+        "    serving_container_environment_variables=serving_env,\n",
         ")\n",
         "\n",
         "jax_owl2_endpoint = jax_owl2_model.deploy(\n",

notebooks/community/model_garden/model_garden_jax_paligemma_finetuning.ipynb

@@ -0,0 +1,884 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "OdZIyZwjgsQcOXnmE8X0xy40",
+      "metadata": {
+        "id": "OdZIyZwjgsQcOXnmE8X0xy40"
+      },
+      "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",
+      "id": "VJWDivOv3OWy",
+      "metadata": {
+        "id": "VJWDivOv3OWy"
+      },
+      "source": [
+        "# Vertex AI Model Garden - PaliGemma (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_jax_paligemma_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_jax_paligemma_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>\n",
+        "\n",
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to do finetuning PaliGemma with a Vertex AI Custom Training Job, deploying the finetuned model to a Vertex AI Endpoint, and making online predictions.\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "- Prepare data for finetuning.\n",
+        "- Launch a Vertex AI Custom Training Job to finetune PaliGemma, storing the resulting model to a GCS bucket.\n",
+        "- Deploy the finetuned PaliGemma model to a Vertex AI Endpoint.\n",
+        "- Make predictions to the endpoint.\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "2aFHbs1g6Wc-",
+      "metadata": {
+        "id": "2aFHbs1g6Wc-"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "QvQjsmIJ6Y3f",
+      "metadata": {
+        "cellView": "form",
+        "id": "QvQjsmIJ6Y3f"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "# @markdown ### Prerequisites\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. 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\"). If not set, a unique GCS bucket will be created instead.\n",
+        "\n",
+        "# Import the necessary packages\n",
+        "import base64\n",
+        "import json\n",
+        "import os\n",
+        "from datetime import datetime\n",
+        "from io import BytesIO\n",
+        "from typing import Tuple\n",
+        "\n",
+        "import matplotlib.pyplot as plt\n",
+        "import requests\n",
+        "import tensorflow as tf\n",
+        "from google.cloud import aiplatform\n",
+        "from PIL import Image\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\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",
+        "# Cloud Storage bucket for storing the experiment artifacts.\n",
+        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
+        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
+        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "# Create a unique GCS bucket for this notebook, if not specified by the user\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
+        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
+        "else:\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
+        "    bucket_region = shell_output[0].strip().lower()\n",
+        "    if bucket_region != REGION:\n",
+        "        raise ValueError(\n",
+        "            f\"Bucket region {bucket_region} is different from notebook region {REGION}\"\n",
+        "        )\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"paligemma\")\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "# Set up default SERVICE_ACCOUNT\n",
+        "SERVICE_ACCOUNT = None\n",
+        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
+        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
+        "\n",
+        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
+        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "\n",
+        "# The pre-built serving docker images.\n",
+        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-paligemma-train-gpu:20240513_0916_RC00\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-paligemma-serve-gpu:20240513_0916_RC00\"\n",
+        "\n",
+        "pretrained_filename_lookup = {\n",
+        "    \"paligemma-224-float32\": \"pt_224.npz\",\n",
+        "    \"paligemma-448-float32\": \"pt_448.npz\",\n",
+        "    \"paligemma-896-float32\": \"pt_896.npz\",\n",
+        "    \"paligemma-mix-224-float32\": \"mix_224.npz\",\n",
+        "    \"paligemma-mix-448-float32\": \"mix_448.npz\",\n",
+        "}\n",
+        "\n",
+        "\n",
+        "def get_job_name_with_datetime(prefix: str) -> str:\n",
+        "    \"\"\"Gets the job 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: str,\n",
+        "    checkpoint_path: str,\n",
+        "    machine_type: str = \"g2-standard-32\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    resolution: int = 224,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Create a Vertex AI Endpoint and deploy the specified model to the endpoint.\"\"\"\n",
+        "    model_name_with_time = get_job_name_with_datetime(model_name)\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name_with_time}-endpoint\"\n",
+        "    )\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name_with_time,\n",
+        "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
+        "        serving_container_ports=[8080],\n",
+        "        serving_container_predict_route=\"/predict\",\n",
+        "        serving_container_health_route=\"/health\",\n",
+        "        serving_container_environment_variables={\n",
+        "            \"CKPT_PATH\": checkpoint_path,\n",
+        "            \"RESOLUTION\": resolution,\n",
+        "            \"MODEL_ID\": model_name,\n",
+        "        },\n",
+        "    )\n",
+        "    print(\n",
+        "        f\"Deploying {model_name_with_time} 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=1800,\n",
+        "        service_account=SERVICE_ACCOUNT,\n",
+        "        enable_access_logging=True,\n",
+        "        min_replica_count=1,\n",
+        "        sync=True,\n",
+        "    )\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "def download_image(url: str) -> Image.Image:\n",
+        "    \"\"\"Downloads an image from the specified URL.\"\"\"\n",
+        "    response = requests.get(url)\n",
+        "    return Image.open(BytesIO(response.content))\n",
+        "\n",
+        "\n",
+        "def resize_image(image: Image.Image, new_width: int = 1000) -> Image.Image:\n",
+        "    width, height = image.size\n",
+        "    print(f\"original input image size: {width}, {height}\")\n",
+        "    new_height = int(height * new_width / width)\n",
+        "    new_img = image.resize((new_width, new_height))\n",
+        "    print(f\"resized input image size: {new_width}, {new_height}\")\n",
+        "    return new_img\n",
+        "\n",
+        "\n",
+        "def image_to_base64(image: Image.Image, format=\"JPEG\") -> str:\n",
+        "    \"\"\"Converts an image to a base64 string.\"\"\"\n",
+        "    buffer = BytesIO()\n",
+        "    image.save(buffer, format=format)\n",
+        "    image_str = base64.b64encode(buffer.getvalue()).decode(\"utf-8\")\n",
+        "    return image_str\n",
+        "\n",
+        "\n",
+        "def caption_predict(\n",
+        "    endpoint: aiplatform.Endpoint,\n",
+        "    image: Image.Image = None,\n",
+        "    language_code: str = \"en\",\n",
+        "    new_width: int = 1000,\n",
+        ") -> str:\n",
+        "    \"\"\"Predicts a caption for a given image using an Endpoint.\"\"\"\n",
+        "    # Resize and convert image to base64 string.\n",
+        "    resized_image = resize_image(image, new_width)\n",
+        "    resized_image_base64 = image_to_base64(resized_image)\n",
+        "\n",
+        "    # Format caption prompt\n",
+        "    caption_prompt = f\"caption {language_code}\\n\"\n",
+        "\n",
+        "    instances = [\n",
+        "        {\n",
+        "            \"prompt\": caption_prompt,\n",
+        "            \"image\": resized_image_base64,\n",
+        "        },\n",
+        "    ]\n",
+        "    response = endpoint.predict(instances=instances)\n",
+        "    return response.predictions[0].get(\"response\")\n",
+        "\n",
+        "\n",
+        "def get_quota(project_id: str, region: str, resource_id: str) -> int:\n",
+        "    \"\"\"Returns the quota for a resource in a region. Returns -1 if can not figure out the quota.\"\"\"\n",
+        "    service_endpoint = \"aiplatform.googleapis.com\"\n",
+        "    quota_list_output = !gcloud alpha services quota list --service=$service_endpoint  --consumer=projects/$project_id --filter=\"$service_endpoint/$resource_id\" --format=json\n",
+        "    # Use '.s' on the command output because it is an SList type.\n",
+        "    quota_data = json.loads(quota_list_output.s)\n",
+        "    if len(quota_data) == 0 or \"consumerQuotaLimits\" not in quota_data[0]:\n",
+        "        return -1\n",
+        "    if (\n",
+        "        len(quota_data[0][\"consumerQuotaLimits\"]) == 0\n",
+        "        or \"quotaBuckets\" not in quota_data[0][\"consumerQuotaLimits\"][0]\n",
+        "    ):\n",
+        "        return -1\n",
+        "    all_regions_data = quota_data[0][\"consumerQuotaLimits\"][0][\"quotaBuckets\"]\n",
+        "    for region_data in all_regions_data:\n",
+        "        if (\n",
+        "            region_data.get(\"dimensions\")\n",
+        "            and region_data[\"dimensions\"][\"region\"] == region\n",
+        "        ):\n",
+        "            if \"effectiveLimit\" in region_data:\n",
+        "                return int(region_data[\"effectiveLimit\"])\n",
+        "            else:\n",
+        "                return 0\n",
+        "    return -1\n",
+        "\n",
+        "\n",
+        "def get_resource_id(accelerator_type: str, is_for_training: bool) -> str:\n",
+        "    \"\"\"Returns the resource id for a given accelerator type and the use case.\n",
+        "    Args:\n",
+        "      accelerator_type: The accelerator type.\n",
+        "      is_for_training: Whether the resource is used for training. Set false\n",
+        "      for serving use case.\n",
+        "    Returns:\n",
+        "      The resource id.\n",
+        "    \"\"\"\n",
+        "    training_accelerator_map = {\n",
+        "        \"NVIDIA_TESLA_V100\": \"custom_model_training_nvidia_v100_gpus\",\n",
+        "        \"NVIDIA_L4\": \"custom_model_training_nvidia_l4_gpus\",\n",
+        "        \"NVIDIA_TESLA_A100\": \"custom_model_training_nvidia_a100_gpus\",\n",
+        "    }\n",
+        "    serving_accelerator_map = {\n",
+        "        \"NVIDIA_TESLA_V100\": \"custom_model_serving_nvidia_v100_gpus\",\n",
+        "        \"NVIDIA_L4\": \"custom_model_serving_nvidia_l4_gpus\",\n",
+        "        \"NVIDIA_TESLA_A100\": \"custom_model_serving_nvidia_a100_gpus\",\n",
+        "    }\n",
+        "    if is_for_training:\n",
+        "        if accelerator_type in training_accelerator_map:\n",
+        "            return training_accelerator_map[accelerator_type]\n",
+        "        else:\n",
+        "            raise ValueError(\n",
+        "                f\"Could not find accelerator type: {accelerator_type} for training.\"\n",
+        "            )\n",
+        "    else:\n",
+        "        if accelerator_type in serving_accelerator_map:\n",
+        "            return serving_accelerator_map[accelerator_type]\n",
+        "        else:\n",
+        "            raise ValueError(\n",
+        "                f\"Could not find accelerator type: {accelerator_type} for serving.\"\n",
+        "            )\n",
+        "\n",
+        "\n",
+        "def check_quota(\n",
+        "    project_id: str,\n",
+        "    region: str,\n",
+        "    accelerator_type: str,\n",
+        "    accelerator_count: int,\n",
+        "    is_for_training: bool,\n",
+        "):\n",
+        "    \"\"\"Checks if the project and the region has the required quota.\"\"\"\n",
+        "    resource_id = get_resource_id(accelerator_type, is_for_training)\n",
+        "    quota = get_quota(project_id, region, resource_id)\n",
+        "    quota_request_instruction = (\n",
+        "        \"Either use \"\n",
+        "        \"a different region or request additional quota. Follow \"\n",
+        "        \"instructions here \"\n",
+        "        \"https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota\"\n",
+        "        \" to check quota in a region or request additional quota for \"\n",
+        "        \"your project.\"\n",
+        "    )\n",
+        "    if quota == -1:\n",
+        "        raise ValueError(\n",
+        "            f\"\"\"Quota not found for: {resource_id} in {region}.\n",
+        "            {quota_request_instruction}\"\"\"\n",
+        "        )\n",
+        "    if quota < accelerator_count:\n",
+        "        raise ValueError(\n",
+        "            f\"\"\"Quota not enough for {resource_id} in {region}:\n",
+        "            {quota} < {accelerator_count}.\n",
+        "            {quota_request_instruction}\"\"\"\n",
+        "        )"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "Jo9Iv7WZGQub",
+      "metadata": {
+        "cellView": "form",
+        "id": "Jo9Iv7WZGQub"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Access PaliGemma models on Vertex AI for GPU based serving\n",
+        "\n",
+        "\n",
+        "# @markdown Accept the model agreement to access the models:\n",
+        "# @markdown 1. Open the [PaliGemma model card](https://console.cloud.google.com/vertex-ai/publishers/google/model-garden/363) from [Vertex AI Model Garden](https://cloud.google.com/model-garden).\n",
+        "# @markdown 1. 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",
+        "# @markdown 1. After accepting the agreement of PaliGemma, a `gs://` URI containing PaliGemma pretrained models will be shared.\n",
+        "# @markdown 1. Paste the link in the `VERTEX_AI_MODEL_GARDEN_PALIGEMMA` field below.\n",
+        "# @markdown 1. The PaliGemma models will be copied into `BUCKET_URI`.\n",
+        "# @markdown The file transfer can take anywhere from 15 minutes to 30 minutes.\n",
+        "VERTEX_AI_MODEL_GARDEN_PALIGEMMA = \"gs://\"  # @param {type:\"string\", isTemplate:true}\n",
+        "assert (\n",
+        "    VERTEX_AI_MODEL_GARDEN_PALIGEMMA and VERTEX_AI_MODEL_GARDEN_PALIGEMMA != \"gs://\"\n",
+        "), \"Click the agreement of PaliGemma in Vertex AI Model Garden, and get the GCS path of PaliGemma model artifacts.\"\n",
+        "print(\n",
+        "    \"Copying PaliGemma model artifacts from\",\n",
+        "    VERTEX_AI_MODEL_GARDEN_PALIGEMMA,\n",
+        "    \"to \",\n",
+        "    MODEL_BUCKET,\n",
+        ")\n",
+        "\n",
+        "! gsutil -m cp -R $VERTEX_AI_MODEL_GARDEN_PALIGEMMA/* $MODEL_BUCKET\n",
+        "\n",
+        "assert (\n",
+        "    os.system(f\"gsutil ls {MODEL_BUCKET}\") == 0\n",
+        "), f\"MODEL_BUCKET does not exist: {MODEL_BUCKET}.\"\n",
+        "model_path_prefix = MODEL_BUCKET"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "XdVYFARuNsuG",
+      "metadata": {
+        "id": "XdVYFARuNsuG"
+      },
+      "source": [
+        "## Finetune with Vertex AI Custom Training Jobs"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "zSjxLZegONqH",
+      "metadata": {
+        "cellView": "form",
+        "id": "zSjxLZegONqH"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Data preparation\n",
+        "\n",
+        "# @markdown The dataset file format is to be image-string pairs stored in a jsonl file.\n",
+        "# @markdown The value for `\"image\"` can be a GCS path or a URL.\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown {\"image\": \"gs://bucket-name/image.jpg\", \"prefix\": \"What animal is this?\", \"suffix\": \"cat\"}\n",
+        "# @markdown {\"image\": \"https://google.com/image.jpg\", \"prefix\": \"What drink is this?\", \"suffix\": \"soda\"}\n",
+        "# @markdown ```\n",
+        "\n",
+        "dataset_gcs_uri = \"gs://longcap100/data_train90.jsonl\"  # @param {type: \"string\"}\n",
+        "\n",
+        "# @markdown [Optional] You can optionally specify the image fields in the JSONL file to use the\n",
+        "# @markdown filename and fill in the `dataset_image_dir` with the location where the images are stored.\n",
+        "dataset_image_dir = \"\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "YMteCx6OS_bq"
+      },
+      "source": [
+        "### Inference examples before finetuning\n",
+        "\n",
+        "The images below are some of the examples of inference results of the pretrained\n",
+        " `paligemma-224-float32` checkpoint.\n",
+        "\n",
+        "| Image | URI | Caption |\n",
+        "|-----|-----|-----|\n",
+        "| <img src=\"https://storage.googleapis.com/longcap100/91.jpeg\" width=\"200\" >  | gs://longcap100/91.jpeg | the beauty of the sleeve |\n",
+        "| <img src=\"https://storage.googleapis.com/longcap100/92.jpeg\" width=\"200\" >   | gs://longcap100/92.jpeg | how to wear a maxi dress for summer |\n",
+        "| <img src=\"https://storage.googleapis.com/longcap100/93.jpeg\" width=\"200\" >  | gs://longcap100/93.jpeg | a red blazer and black bag , a key piece of the week 's fashion . |\n",
+        "| <img src=\"https://storage.googleapis.com/longcap100/94.jpeg\" width=\"200\" >   | gs://longcap100/94.jpeg | how to wear boyfriend jeans like a fashion blogger |\n",
+        "| <img src=\"https://storage.googleapis.com/longcap100/95.jpeg\" width=\"200\" >   | gs://longcap100/95.jpeg | this graphic sweatshirt is a must have for your wardrobe . |\n",
+        "| <img src=\"https://storage.googleapis.com/longcap100/96.jpeg\" width=\"200\" >   | gs://longcap100/96.jpeg | person in a long shot of our model |\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "k_ET-ZVpWFVk"
+      },
+      "source": [
+        "### Finetune"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "jitdsEHDNqQk",
+      "metadata": {
+        "cellView": "form",
+        "id": "jitdsEHDNqQk"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Run\n",
+        "# @markdown Use the Vertex AI SDK to create and run the custom training jobs. With the default setting, you can train about 30 steps per minute with the batch size of 4.\n",
+        "\n",
+        "model_variant = \"pt\"  # @param [\"mix\", \"pt\"]\n",
+        "model_resolution = 224  # @param [224, 448, 896]\n",
+        "model_precision_type = \"float32\"  # Only float32 is supported.\n",
+        "\n",
+        "if model_variant == \"mix\":\n",
+        "    model_name_prefix = \"paligemma-mix\"\n",
+        "else:\n",
+        "    model_name_prefix = \"paligemma\"\n",
+        "\n",
+        "base_model_name = f\"{model_name_prefix}-{model_resolution}-{model_precision_type}\"\n",
+        "base_model_filename = pretrained_filename_lookup[base_model_name]\n",
+        "base_model_uri = os.path.join(model_path_prefix, base_model_filename)\n",
+        "\n",
+        "# The accelerator to use.\n",
+        "ACCELERATOR_TYPE = \"NVIDIA_L4\"  # @param [\"NVIDIA_TESLA_V100\", \"NVIDIA_L4\"]\n",
+        "\n",
+        "# Batch size for finetuning.\n",
+        "batch_size = 4  # @param {type:\"integer\"}\n",
+        "# Number of epochs to train.\n",
+        "epochs = 3  # @param {type:\"integer\"}\n",
+        "# Learning rate.\n",
+        "learning_rate = 0.1  # @param{type:\"number\"}\n",
+        "# Text length.\n",
+        "text_length = 512  # @param{type:\"integer\"}\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "\n",
+        "if ACCELERATOR_TYPE == \"NVIDIA_TESLA_V100\":\n",
+        "    machine_type = \"n1-standard-8\"\n",
+        "    accelerator_count = 2\n",
+        "elif ACCELERATOR_TYPE == \"NVIDIA_L4\":\n",
+        "    machine_type = \"g2-standard-24\"\n",
+        "    accelerator_count = 2\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        f\"Cannot automatically determine machine type from {ACCELERATOR_TYPE}.\"\n",
+        "    )\n",
+        "\n",
+        "replica_count = 1\n",
+        "\n",
+        "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",
+        "# Setup training job.\n",
+        "job_name = get_job_name_with_datetime(\"paligemma-finetune\")\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",
+        "\n",
+        "# Designate a GCS folder to store the LORA adapter.\n",
+        "finetune_output_dir_name = get_job_name_with_datetime(\"paligemma-finetune\")\n",
+        "finetune_output_dir = os.path.join(STAGING_BUCKET, finetune_output_dir_name)\n",
+        "\n",
+        "train_args = [\n",
+        "    \"--config=big_vision/configs/proj/paligemma/transfers/vertexai_l4.py\",\n",
+        "    f\"--workdir={finetune_output_dir}\",\n",
+        "    f\"--config.model_init={base_model_uri}\",\n",
+        "    f\"--config.input.data.fname={dataset_gcs_uri}\",\n",
+        "    f\"--config.text_len={text_length}\",\n",
+        "]\n",
+        "\n",
+        "if batch_size:\n",
+        "    train_args.append(f\"--config.input.batch_size={batch_size}\")\n",
+        "if epochs:\n",
+        "    train_args.append(f\"--config.total_epochs={epochs}\")\n",
+        "if learning_rate:\n",
+        "    train_args.append(f\"--config.lr={learning_rate}\")\n",
+        "\n",
+        "if dataset_image_dir:\n",
+        "    train_args.append(f\"--config.input.data.fopen_keys.image={dataset_image_dir}\")\n",
+        "train_job.run(\n",
+        "    args=train_args,\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",
+        ")\n",
+        "\n",
+        "print(\"Checkpoint and log files was saved in: \", finetune_output_dir)\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "WngKdCIzOTJV",
+      "metadata": {
+        "cellView": "form",
+        "id": "WngKdCIzOTJV"
+      },
+      "outputs": [],
+      "source": [
+        "# @title View training loss\n",
+        "# @markdown Metrics will be stored in a file named `big_vision_metrics.txt` in the GCS bucket, including training loss, examples seen, and core hours throughout training.\n",
+        "\n",
+        "# @markdown Run this cell to get and plot the training loss.\n",
+        "\n",
+        "# Get relevant metrics from metrics file.\n",
+        "metrics_path = os.path.join(finetune_output_dir, \"big_vision_metrics.txt\")\n",
+        "steps = []\n",
+        "training_losses = []\n",
+        "with tf.io.gfile.GFile(metrics_path, \"r\") as f:\n",
+        "    for line in f:\n",
+        "        metric = json.loads(line)\n",
+        "        steps.append(metric[\"step\"])\n",
+        "        training_losses.append(metric[\"training_loss\"])\n",
+        "\n",
+        "# Plot training plot\n",
+        "plt.plot(steps, training_losses)\n",
+        "plt.title(\"Steps vs. Training Loss\")\n",
+        "plt.xlabel(\"Steps\")\n",
+        "plt.ylabel(\"Training Loss\")\n",
+        "plt.show()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "uVKr3rggXJd0"
+      },
+      "source": [
+        "## Deployment and prediction"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "toY-WPKDFesF"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "# @markdown This section uploads the finetuned PaliGemma model to Model Registry and deploys it to a Vertex AI Endpoint. It takes approximately 15 minutes to finish.\n",
+        "\n",
+        "# @markdown Note: You cannot use accelerator type `NVIDIA_TESLA_V100` to serve prebuilt or finetuned PaliGemma models with resolution `896`.\n",
+        "\n",
+        "last_checkpoint_path = os.path.join(finetune_output_dir, \"checkpoint.bv-LAST\")\n",
+        "with tf.io.gfile.GFile(last_checkpoint_path, \"r\") as f:\n",
+        "    final_checkpoint_name = \"checkpoint.bv-\" + f.read()\n",
+        "    checkpoint_path = os.path.join(finetune_output_dir, final_checkpoint_name)\n",
+        "\n",
+        "model_name = f\"paligemma-{model_resolution}-{model_precision_type}-custom\"\n",
+        "print(f\"Deploying custom PaliGemma model: {model_name}\")\n",
+        "\n",
+        "# @markdown Select the accelerator type to use to deploy the model:\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_V100\"]\n",
+        "if accelerator_type == \"NVIDIA_L4\":\n",
+        "    machine_type = \"g2-standard-16\"\n",
+        "    accelerator_count = 1\n",
+        "elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
+        "    if model_resolution == 896 and model_precision_type == \"float32\":\n",
+        "        raise ValueError(\n",
+        "            \"NVIDIA_TESLA_V100 is not sufficient. Multi-gpu is not supported for PaLIGemma.\"\n",
+        "        )\n",
+        "    else:\n",
+        "        machine_type = \"n1-highmem-8\"\n",
+        "        accelerator_count = 1\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        f\"Recommended machine settings not found for: {accelerator_type}. To use another another accelerator, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `accelerator_count` to the deploy_model function by clicking `Show Code` and then modifying the code.\"\n",
+        "    )\n",
+        "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",
+        "\n",
+        "# @markdown If you want to use other accelerator types not listed above, then check other Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute. You may need to manually set the `machine_type`, `accelerator_type`, and `accelerator_count` in the code by clicking `Show code` first.\n",
+        "model, endpoint = deploy_model(\n",
+        "    model_name=model_name,\n",
+        "    checkpoint_path=checkpoint_path,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    resolution=model_resolution,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "G_FT3hglXMks"
+      },
+      "source": [
+        "### Image captioning"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "0v3GA8NU8SNp",
+      "metadata": {
+        "cellView": "form",
+        "id": "0v3GA8NU8SNp"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown This section uses the deployed PaliGemma model to caption and describe an image in a chosen language. Check how the caption has changed compared to the examples above.\n",
+        "\n",
+        "# @markdown <img src=\"https://storage.googleapis.com/longcap100/91.jpeg\" width=\"400\" >\n",
+        "\n",
+        "image_url = \"https://storage.googleapis.com/longcap100/91.jpeg\"  # @param {type:\"string\"}\n",
+        "\n",
+        "image = download_image(image_url)\n",
+        "display(image)\n",
+        "\n",
+        "# Make a prediction.\n",
+        "image_base64 = image_to_base64(image)\n",
+        "language_code = \"en\"  # @param {type: \"string\"}\n",
+        "caption = caption_predict(endpoint, image, language_code)\n",
+        "\n",
+        "print(\"Caption: \", caption)\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "ZnHpjzpjUMlH"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown <img src=\"https://storage.googleapis.com/longcap100/92.jpeg\" width=\"400\" >\n",
+        "\n",
+        "image_url = \"https://storage.googleapis.com/longcap100/92.jpeg\"  # @param {type:\"string\"}\n",
+        "\n",
+        "image = download_image(image_url)\n",
+        "display(image)\n",
+        "\n",
+        "# Make a prediction.\n",
+        "image_base64 = image_to_base64(image)\n",
+        "language_code = \"en\"  # @param {type: \"string\"}\n",
+        "caption = caption_predict(endpoint, image, language_code)\n",
+        "\n",
+        "print(\"Caption: \", caption)\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "QSDHB4hyUMqS"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown <img src=\"https://storage.googleapis.com/longcap100/93.jpeg\" width=\"400\" >\n",
+        "\n",
+        "image_url = \"https://storage.googleapis.com/longcap100/93.jpeg\"  # @param {type:\"string\"}\n",
+        "\n",
+        "image = download_image(image_url)\n",
+        "display(image)\n",
+        "\n",
+        "# Make a prediction.\n",
+        "image_base64 = image_to_base64(image)\n",
+        "language_code = \"en\"  # @param {type: \"string\"}\n",
+        "caption = caption_predict(endpoint, image, language_code)\n",
+        "\n",
+        "print(\"Caption: \", caption)\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "rSm8TukGUMu3"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown <img src=\"https://storage.googleapis.com/longcap100/94.jpeg\" width=\"400\" >\n",
+        "\n",
+        "image_url = \"https://storage.googleapis.com/longcap100/94.jpeg\"  # @param {type:\"string\"}\n",
+        "\n",
+        "image = download_image(image_url)\n",
+        "display(image)\n",
+        "\n",
+        "# Make a prediction.\n",
+        "image_base64 = image_to_base64(image)\n",
+        "language_code = \"en\"  # @param {type: \"string\"}\n",
+        "caption = caption_predict(endpoint, image, language_code)\n",
+        "\n",
+        "print(\"Caption: \", caption)\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "GchcOq35VHzP"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown <img src=\"https://storage.googleapis.com/longcap100/95.jpeg\" width=\"400\" >\n",
+        "\n",
+        "image_url = \"https://storage.googleapis.com/longcap100/95.jpeg\"  # @param {type:\"string\"}\n",
+        "\n",
+        "image = download_image(image_url)\n",
+        "display(image)\n",
+        "\n",
+        "# Make a prediction.\n",
+        "image_base64 = image_to_base64(image)\n",
+        "language_code = \"en\"  # @param {type: \"string\"}\n",
+        "caption = caption_predict(endpoint, image, language_code)\n",
+        "\n",
+        "print(\"Caption: \", caption)\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "pHk9MHpUVH8y"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown <img src=\"https://storage.googleapis.com/longcap100/96.jpeg\" width=\"400\" >\n",
+        "\n",
+        "image_url = \"https://storage.googleapis.com/longcap100/96.jpeg\"  # @param {type:\"string\"}\n",
+        "\n",
+        "image = download_image(image_url)\n",
+        "display(image)\n",
+        "\n",
+        "# Make a prediction.\n",
+        "image_base64 = image_to_base64(image)\n",
+        "language_code = \"en\"  # @param {type: \"string\"}\n",
+        "caption = caption_predict(endpoint, image, language_code)\n",
+        "\n",
+        "print(\"Caption: \", caption)\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "IrVZ030i4XMY",
+      "metadata": {
+        "id": "IrVZ030i4XMY"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "YsMpOI1kYjil"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Run\n",
+        "# @markdown Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown and avoid unnecessary continuous charges that may incur.\n",
+        "\n",
+        "# Delete the training job.\n",
+        "train_job.delete()\n",
+        "\n",
+        "# Undeploy model and delete endpoint.\n",
+        "endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "model.delete()\n",
+        "\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_URI"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_jax_paligemma_finetuning.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_jax_stable_diffusion_xl.ipynb

@@ -0,0 +1,415 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "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 - Stable Diffusion XL 1.0 - TPU v5e\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_jax_stable_diffusion_xl.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/community/model_garden/model_garden_jax_stable_diffusion_xl.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/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_jax_stable_diffusion_xl.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>\n",
+        "    (a Python-3 GPU notebook with preinstalled HuggingFace/transformer libraries is recommended)\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3de7470326a2"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to deploy the [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) model on Vertex AI for online prediction.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Deploy the model to a [Vertex AI Endpoint resource](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
+        "- Run online predictions for text-to-image.\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "264c07757582"
+      },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "**NOTE**: \n",
+        "\n",
+        "*  Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands.\n",
+        "*  This Notebook demonstrate how to deploy the model [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) on Vertex AI prediction endpoint with a TPU v5e instance (machine type of `ct5lp-hightpu-1t`). Please ensure you have enough resource quota in region `us-west1`. If not, please follow the [instructions](https://cloud.google.com/vertex-ai/docs/predictions/use-tpu#securing_capacity) to get quota."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ioensNKM8ned"
+      },
+      "source": [
+        "### Setup notebook"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d73ffa0c0b83"
+      },
+      "source": [
+        "#### Colab\n",
+        "Run the following commands for Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2707b02ef5df"
+      },
+      "outputs": [],
+      "source": [
+        "if \"google.colab\" in str(get_ipython()):\n",
+        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()\n",
+        "\n",
+        "# Restart the notebook kernel after installs.\n",
+        "import IPython\n",
+        "\n",
+        "app = IPython.Application.instance()\n",
+        "app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bb7adab99e41"
+      },
+      "source": [
+        "### Setup Google Cloud project\n",
+        "\n",
+        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
+        "\n",
+        "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs.\n",
+        "\n",
+        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User` and `Storage Object Admin` roles for deploying models to Vertex AI endpoint."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6c460088b873"
+      },
+      "source": [
+        "Set following variables for experiments environment:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {
+        "id": "855d6b96f291"
+      },
+      "outputs": [],
+      "source": [
+        "# Cloud project id.\n",
+        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# The region you want to launch jobs in.\n",
+        "REGION = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# The Cloud Storage bucket for storing experiments output. Fill it without the 'gs://' prefix.\n",
+        "GCS_BUCKET = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# The service account for deploying fine tuned model.\n",
+        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e828eb320337"
+      },
+      "source": [
+        "Initialize Vertex AI API:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 2,
+      "metadata": {
+        "id": "12cd25839741"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=GCS_BUCKET)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2cc825514deb"
+      },
+      "source": [
+        "### Define constants"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 3,
+      "metadata": {
+        "id": "b42bd4fa2b2d"
+      },
+      "outputs": [],
+      "source": [
+        "# The pre-built serving docker image. It contains serving scripts and models.\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-diffusers-serve-tpu:20240110_1526_RC00\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0c250872074f"
+      },
+      "source": [
+        "### Define common functions"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 6,
+      "metadata": {
+        "id": "354da31189dc"
+      },
+      "outputs": [],
+      "source": [
+        "import base64\n",
+        "from io import BytesIO\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "from PIL import Image\n",
+        "\n",
+        "\n",
+        "def base64_to_image(image_str):\n",
+        "    \"\"\"Convert base64 encoded string to an image.\"\"\"\n",
+        "    image = Image.open(BytesIO(base64.b64decode(image_str)))\n",
+        "    return image\n",
+        "\n",
+        "\n",
+        "def image_grid(imgs, rows=2, cols=2):\n",
+        "    w, h = imgs[0].size\n",
+        "    grid = Image.new(\n",
+        "        mode=\"RGB\", size=(cols * w + 10 * cols, rows * h), color=(255, 255, 255)\n",
+        "    )\n",
+        "    for i, img in enumerate(imgs):\n",
+        "        grid.paste(img, box=(i % cols * w + 10 * i, i // cols * h))\n",
+        "    return grid\n",
+        "\n",
+        "\n",
+        "def deploy_model(model_id):\n",
+        "    \"\"\"Create a Vertex AI Endpoint and deploy the specified model to the endpoint.\"\"\"\n",
+        "    model_name = model_id + \"-tpu\"\n",
+        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
+        "        serving_container_ports=[8080],\n",
+        "        serving_container_predict_route=\"/predict\",\n",
+        "        serving_container_health_route=\"/health\",\n",
+        "    )\n",
+        "    machine_type = \"ct5lp-hightpu-1t\"\n",
+        "\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=SERVICE_ACCOUNT,\n",
+        "        enable_access_logging=True,\n",
+        "        min_replica_count=1,\n",
+        "        sync=True,\n",
+        "    )\n",
+        "    return model, endpoint"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bf7f82732e61"
+      },
+      "source": [
+        "## Upload and Deploy models"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "1cc26e68d7b0"
+      },
+      "source": [
+        "This section uploads the model to Model Registry and deploys it to a Vertex AI Endpoint resource.\n",
+        "\n",
+        "The model deployment step will take ~30 minutes to complete."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cd7b56421392"
+      },
+      "source": [
+        "### Text-to-image"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "6d331b1ea337"
+      },
+      "source": [
+        "Deploy the stable diffusion xl model for the text-to-image task.\n",
+        "\n",
+        "Once deployed, you can send a batch of text prompts to the endpoint to generated images.\n",
+        "\n",
+        "When deployed on one TPU V5e instance, the averaged inference time of one image is ~3 seconds."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bf55e38815dc"
+      },
+      "outputs": [],
+      "source": [
+        "# Set the model_id to \"stabilityai/stable-diffusion-xl-base-1.0\" to load the OSS pre-trained model.\n",
+        "model, endpoint = deploy_model(\n",
+        "    model_id=\"stabilityai/stable-diffusion-xl-base-1.0\",\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "4ab04da3ec9a"
+      },
+      "outputs": [],
+      "source": [
+        "instances = [\n",
+        "    {\n",
+        "        \"prompt\": \"Photorealistic whale swimming in abyss\",\n",
+        "        \"height\": 1024,\n",
+        "        \"width\": 1024,\n",
+        "    },\n",
+        "    {\n",
+        "        \"prompt\": \"Photorealistic happy dog running\",\n",
+        "        \"height\": 1024,\n",
+        "        \"width\": 1024,\n",
+        "    },\n",
+        "]\n",
+        "response = endpoint.predict(instances=instances)\n",
+        "\n",
+        "images = [\n",
+        "    base64_to_image(prediction.get(\"images\")[0]) for prediction in response.predictions\n",
+        "]\n",
+        "image_grid(images, rows=1)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "af21a3cff1e0"
+      },
+      "source": [
+        "### Clean up resources:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# Undeploy model and delete endpoint.\n",
+        "endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "model.delete()"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "name": "model_garden_jax_stable_diffusion_xl.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_jax_vision_transformer.ipynb

@@ -776,12 +776,18 @@
       },
       "outputs": [],
       "source": [
+        "serving_env = {\n",
+        "    \"MODEL_ID\": \"ViT-JAX-\",\n",
+        "    \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "}\n",
+        "\n",
         "jax_vit_model = aiplatform.Model.upload(\n",
         "    display_name=\"jax_vit\",\n",
         "    artifact_uri=saved_model_dir,\n",
         "    serving_container_image_uri=OPTIMIZED_TF_RUNTIME_IMAGE_URI,\n",
         "    serving_container_args=[],\n",
         "    location=REGION,\n",
+        "    serving_container_environment_variables=serving_env,\n",
         ")\n",
         "\n",
         "jax_vit_endpoint = jax_vit_model.deploy(\n",

notebooks/community/model_garden/model_garden_keras_stable_diffusion.ipynb

@@ -289,14 +289,18 @@
         "\n",
         "# Training constants.\n",
         "TRAINING_JOB_PREFIX = \"train\"\n",
-        "TRAIN_CONTAINER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/keras-train:latest\"\n",
+        "TRAIN_CONTAINER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/keras-train:latest\"\n",
+        ")\n",
         "TRAIN_MACHINE_TYPE = \"a2-highgpu-1g\"\n",
         "TRAIN_ACCELERATOR_TYPE = \"NVIDIA_TESLA_A100\"\n",
         "TRAIN_NUM_GPU = 1\n",
         "RESOLUTION = 512\n",
         "\n",
         "# Prediction constants.\n",
-        "PREDICTION_CONTAINER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/keras-serve:latest\"\n",
+        "PREDICTION_CONTAINER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/keras-serve:latest\"\n",
+        ")\n",
         "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
         "PREDICTION_MACHINE_TYPE = \"n1-standard-8\"\n",
         "DEPLOY_JOB_PREFIX = \"deploy\"\n",
@@ -340,16 +344,19 @@
         "    deploy_model_name = get_job_name_with_datetime(DEPLOY_JOB_PREFIX)\n",
         "    print(\"The deployed job name is: \", deploy_model_name)\n",
         "    serving_env = {\n",
+        "        \"MODEL_ID\": \"keras-stable-diffusion-v1-4-001\",\n",
         "        \"MODEL_PATH\": f\"{model_path}\",\n",
         "        \"IMAGE_WIDTH\": f\"{RESOLUTION}\",\n",
         "        \"IMAGE_HEIGHT\": f\"{RESOLUTION}\",\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
         "    }\n",
         "\n",
         "    endpoint = aiplatform.Endpoint.create(display_name=f\"{deploy_model_name}-endpoint\")\n",
+        "\n",
         "    model = aiplatform.Model.upload(\n",
         "        display_name=deploy_model_name,\n",
         "        serving_container_image_uri=PREDICTION_CONTAINER_URI,\n",
-        "        serving_container_ports=[8501],\n",
+        "        serving_container_ports=[8080],\n",
         "        serving_container_predict_route=\"/predict\",\n",
         "        serving_container_health_route=\"/ping\",\n",
         "        serving_container_environment_variables=serving_env,\n",

notebooks/community/model_garden/model_garden_keras_yolov8.ipynb

@@ -625,11 +625,17 @@
       "source": [
         "upload_job_name = get_job_name_with_datetime(UPLOAD_JOB_PREFIX)\n",
         "\n",
+        "serving_env = {\n",
+        "    \"MODEL_ID\": \"keras-yolov8\",\n",
+        "    \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "}\n",
+        "\n",
         "model = aiplatform.Model.upload(\n",
         "    display_name=upload_job_name,\n",
         "    artifact_uri=model_dir,\n",
         "    serving_container_image_uri=SERVING_CONTAINER_URI,\n",
         "    serving_container_args=SERVING_CONTAINER_ARGS,\n",
+        "    serving_container_environment_variables=serving_env,\n",
         ")\n",
         "\n",
         "print(\"The uploaded model name is: \", upload_job_name)\n",

notebooks/community/model_garden/model_garden_llama_guard_deployment.ipynb

@@ -0,0 +1,503 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "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 Guard\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_llama_guard_deployment.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_llama_guard_deployment.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 downloading and deploying prebuilt [Llama Guard models](https://huggingface.co/meta-llama) with [vLLM](https://github.com/vllm-project/vllm) on GPU, and demonstrates using the Llama Guard model to safeguard LLM inputs and outputs with the Vertex Llama 3.1 API service.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Download and deploy prebuilt Llama Guard models with [vLLM](https://github.com/vllm-project/vllm) on GPU\n",
+        "- Use the Llama Guard models to safeguard LLM inputs and outputs with the Vertex Llama 3.1 API service\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "264c07757582"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YXFGIp1l-qtT"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. 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\"). If not set, a unique GCS bucket will be created instead.\n",
+        "\n",
+        "# Import the necessary packages\n",
+        "\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "import importlib\n",
+        "import os\n",
+        "import re\n",
+        "from datetime import datetime\n",
+        "from typing import Tuple\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\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",
+        "# Cloud Storage bucket for storing the experiment artifacts.\n",
+        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
+        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
+        "\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
+        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
+        "else:\n",
+        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
+        "    bucket_region = shell_output[0].strip().lower()\n",
+        "    if bucket_region != REGION:\n",
+        "        raise ValueError(\n",
+        "            \"Bucket region %s is different from notebook region %s\"\n",
+        "            % (bucket_region, REGION)\n",
+        "        )\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "MODEL_BUCKET = BUCKET_URI\n",
+        "\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "# Gets the default SERVICE_ACCOUNT.\n",
+        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
+        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
+        "\n",
+        "\n",
+        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
+        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "# @markdown # Access Llama Guard models on Vertex AI\n",
+        "# @markdown The original models from Meta are converted into the Hugging Face format for serving in Vertex AI.\n",
+        "# @markdown Accept the model agreement to access the models:\n",
+        "# @markdown 1. Open the [Llama Guard model card](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama-guard) from [Vertex AI Model Garden](https://cloud.google.com/model-garden).\n",
+        "# @markdown 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",
+        "# @markdown 3. After accepting the agreement, a `gs://` URI containing Llama Guard pretrained and finetuned models will be shared.\n",
+        "# @markdown 4. Paste the URI in the `VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD` field below.\n",
+        "# @markdown 5. The Llama Guard models will be copied into `BUCKET_URI`.\n",
+        "\n",
+        "\n",
+        "VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "assert (\n",
+        "    VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD\n",
+        "), \"Please click the agreement in Vertex AI Model Garden at https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama-guard, and get the GCS path of Llama Guard model artifacts.\"\n",
+        "parsed_gcs_url = re.search(\"gs://.*?(?=[ ]|$)\", VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD)\n",
+        "if parsed_gcs_url:\n",
+        "    VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD = parsed_gcs_url.group()\n",
+        "assert VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD.startswith(\n",
+        "    \"gs://\"\n",
+        "), \"VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD is expected to be a GCS URI and must start with `gs://`.\"\n",
+        "print(\n",
+        "    \"Copying Llama Guard model artifacts from\",\n",
+        "    VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD,\n",
+        "    \"to \",\n",
+        "    MODEL_BUCKET,\n",
+        ")\n",
+        "\n",
+        "! gsutil -m cp -R $VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD $MODEL_BUCKET\n",
+        "\n",
+        "# The pre-built serving docker images.\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240726_1329_RC00\"\n",
+        "\n",
+        "\n",
+        "def deploy_model_vllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    service_account: str,\n",
+        "    base_model_id: str = None,\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 = 4096,\n",
+        "    dtype: str = \"auto\",\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",
+        "    if not base_model_id:\n",
+        "        base_model_id = model_id\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\"--dtype={dtype}\",\n",
+        "        \"--disable-log-stats\",\n",
+        "        \"--enforce-eager\",\n",
+        "        \"--disable-custom-all-reduce\",\n",
+        "        \"--enable-chunked-prefill\",\n",
+        "        \"--max-num-seqs=12\",\n",
+        "    ]\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"MODEL_ID\": base_model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    try:\n",
+        "        if HF_TOKEN:\n",
+        "            env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
+        "    except:\n",
+        "        pass\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=1800,\n",
+        "        service_account=service_account,\n",
+        "    )\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
+        "    return model, endpoint"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "z-XybZjtgF9M"
+      },
+      "source": [
+        "## Deploy Llama Guard with vLLM on GPU"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "E8OiHHNNE_wj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "\n",
+        "# @markdown This section uploads prebuilt Llama Guard models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
+        "\n",
+        "# @markdown NVIDIA_L4 GPUs are used for demonstration. The serving efficiency of L4 GPUs is inferior to that of A100 GPUs, but L4 GPUs are nevertheless good serving solutions if you do not have A100 quota.\n",
+        "\n",
+        "# @markdown Set the model to deploy.\n",
+        "\n",
+        "MODEL_ID = \"Llama-Guard-3-8B\"  # @param [\"Llama-Guard-3-8B\"] {allow-input: true, isTemplate: true}\n",
+        "model_id = os.path.join(MODEL_BUCKET, MODEL_ID)\n",
+        "\n",
+        "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_A100\"]\n",
+        "\n",
+        "if accelerator_type == \"NVIDIA_L4\":\n",
+        "    machine_type = \"g2-standard-12\"\n",
+        "    accelerator_count = 1\n",
+        "elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
+        "    machine_type = \"a2-highgpu-1g\"\n",
+        "    accelerator_count = 1\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        f\"Recommended GPU setting not found for: {accelerator_type} and {MODEL_ID}.\"\n",
+        "    )\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",
+        "\n",
+        "gpu_memory_utilization = 0.9\n",
+        "max_model_len = 32768  # Maximum context length.\n",
+        "\n",
+        "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=MODEL_ID),\n",
+        "    model_id=model_id,\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": {
+        "id": "192a021iB_DE"
+      },
+      "source": [
+        "## Use the Llama Guard models to safeguard LLM inputs and outputs with the Vertex Llama 3.1 API service\n",
+        "\n",
+        "We use [meta-llama/Llama-Guard-3-8B](https://huggingface.co/meta-llama/Llama-Guard-3-8B) to safeguard input and output conversations with the [Llama 3.1 405B Instruct model API service on Vertex](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama3-405b-instruct-maas).\n",
+        "\n",
+        "Llama Guard 3 builds on the capabilities introduced with Llama Guard 2, adding three new categories, Defamation, Elections and Code Interpreter Abuse. Additionally this model is multilingual and a new prompt format is introduced, making Llama Guard 3’s prompt format consistent with Llama 3+ Instruct models.\n",
+        "\n",
+        "This section references [LlamaGuard.ipynb](https://colab.research.google.com/drive/16s0tlCSEDtczjPzdIK3jq0Le5LlnSYGf?usp=sharing) from [https://huggingface.co/meta-llama/LlamaGuard-7b](https://huggingface.co/meta-llama/LlamaGuard-7b)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "fHC7INgjB_DF"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install --upgrade --quiet openai"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ajjcGNzhB_DF"
+      },
+      "outputs": [],
+      "source": [
+        "import google.auth\n",
+        "import openai\n",
+        "\n",
+        "# @markdown Set up the Llama 3.1 405B Instruct model API service.\n",
+        "\n",
+        "# Programmatically get an access token\n",
+        "creds, _ = google.auth.default(\n",
+        "    scopes=[\"https://www.googleapis.com/auth/cloud-platform\"]\n",
+        ")\n",
+        "auth_req = google.auth.transport.requests.Request()\n",
+        "creds.refresh(auth_req)\n",
+        "# Note: the credential lives for 1 hour by default (https://cloud.google.com/docs/authentication/token-types#at-lifetime); after expiration, it must be refreshed.\n",
+        "\n",
+        "client = openai.OpenAI(\n",
+        "    base_url=f\"https://us-central1-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{REGION}/endpoints/openapi\",\n",
+        "    api_key=creds.token,\n",
+        ")\n",
+        "LLAMA3_405B_INSTRUCT = \"meta/llama3-405b-instruct-maas\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NvSfBcUUB_DF"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown Define input message in conversation and get output message from model.\n",
+        "\n",
+        "message_role = \"user\"  # @param {type: \"string\"}\n",
+        "message_content = \"What is a car?\"  # @param {type: \"string\"}\n",
+        "\n",
+        "messages = [\n",
+        "    {\n",
+        "        \"role\": message_role,\n",
+        "        \"content\": message_content,\n",
+        "    }\n",
+        "]\n",
+        "print(\"Conversation [turn 1]:\", messages)\n",
+        "\n",
+        "response = client.chat.completions.create(\n",
+        "    model=LLAMA3_405B_INSTRUCT,\n",
+        "    messages=messages,\n",
+        ")\n",
+        "print(\"Response:\", response)\n",
+        "\n",
+        "messages.append(\n",
+        "    {\n",
+        "        \"role\": response.choices[0].message.role,\n",
+        "        \"content\": response.choices[0].message.content,\n",
+        "    }\n",
+        ")\n",
+        "print(\"Conversation [turn 2]:\", messages)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Y7-ym3GlB_DG"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown Use Llama Guard to classify the conversation: safe versus unsafe.\n",
+        "# @markdown Classification is performed on the last turn of the conversation.\n",
+        "# @markdown If the content is safe, the model will return `safe`. If the content is unsafe, the model will return `unsafe` and additionally the list of offending categories as a comma-separated list in a new line.\n",
+        "# @markdown Set `\"@requestFormat\": \"chatCompletions\"` to use the OpenAI chat completions format.\n",
+        "\n",
+        "instances = [\n",
+        "    {\n",
+        "        \"messages\": messages,\n",
+        "        \"@requestFormat\": \"chatCompletions\",\n",
+        "    },\n",
+        "]\n",
+        "response = endpoints[\"vllm_gpu\"].predict(instances=instances)\n",
+        "\n",
+        "prediction = response.predictions[0]\n",
+        "print(prediction)\n",
+        "print(\"Llama Guard prediction:\", prediction[\"choices\"][0][\"message\"][\"content\"])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "956x4r7rsrza"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Delete the models and endpoints\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
+        "\n",
+        "# Undeploy model and delete endpoint.\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "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_URI"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_llama_guard_deployment.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_mammut.ipynb

@@ -0,0 +1,749 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ur8xi4C7S06n"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2024 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden - MaMMUT\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_mammut.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_mammut.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": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates deploying MaMMUT to a Vertex AI Endpoint and making online predictions."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d975e698c9a4"
+      },
+      "source": [
+        "### Objective\n",
+        "\n",
+        "- Deploy MaMMUT to a Vertex AI Endpoint.\n",
+        "- Make predictions to the endpoint including:\n",
+        "  - Answering questions about a given image.\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aed92deeb4a0"
+      },
+      "source": [
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BF1j6f9HApxa"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "DU0WWEDqWJLy"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "# @markdown ### Prerequisites\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. 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\"). If not set, a unique GCS bucket will be created instead.\n",
+        "\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "! pip install -q gradio==4.21.0\n",
+        "\n",
+        "import importlib\n",
+        "import os\n",
+        "from datetime import datetime\n",
+        "from typing import Tuple\n",
+        "\n",
+        "import gradio as gr\n",
+        "import numpy as np\n",
+        "from google.cloud import aiplatform\n",
+        "from PIL import Image\n",
+        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Cloud Storage bucket for storing the experiment artifacts.\n",
+        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
+        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "BUCKET_URI = \"gs://\"  # @param {type: \"string\"}\n",
+        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "\n",
+        "# Create a unique GCS bucket for this notebook, if not specified by the user.\n",
+        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
+        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "else:\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
+        "    bucket_region = shell_output[0].strip().lower()\n",
+        "    if bucket_region != REGION:\n",
+        "        raise ValueError(\n",
+        "            \"Bucket region %s is different from notebook region %s\"\n",
+        "            % (bucket_region, REGION)\n",
+        "        )\n",
+        "\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "! gcloud services enable language.googleapis.com\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "\n",
+        "# Set up default SERVICE_ACCOUNT\n",
+        "SERVICE_ACCOUNT = None\n",
+        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
+        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "\n",
+        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
+        "\n",
+        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "\n",
+        "# The pre-built prediction docker image.\n",
+        "OPTIMIZED_TF_RUNTIME_IMAGE_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.nightly:latest\"\n",
+        ")\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
+        "\n",
+        "def resize_image(image: Image.Image, new_width: int = 512) -> Image.Image:\n",
+        "    width, height = image.size\n",
+        "    new_height = int(height * new_width / width)\n",
+        "    new_image = image.resize((new_width, new_height))\n",
+        "    return new_image\n",
+        "\n",
+        "\n",
+        "def load_image(image_url):\n",
+        "    if image_url.startswith(\"gs://\"):\n",
+        "        local_image_path = \"./images/test_image.jpg\"\n",
+        "        common_util.download_gcs_file_to_local(image_url, local_image_path)\n",
+        "        image = common_util.load_img(local_image_path)\n",
+        "    else:\n",
+        "        image = common_util.download_image(image_url)\n",
+        "    return image\n",
+        "\n",
+        "\n",
+        "def deploy_mammut(\n",
+        "    task: str, machine_type: str, accelerator_type: str, accelerator_count: int\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploy the model to a Vertex endpoint for prediction.\"\"\"\n",
+        "    serving_env = {\n",
+        "        \"MODEL_ID\": \"mammut\",\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    if task == \"vqa\":\n",
+        "        model_dir = \"gs://vertex-model-garden-public-us/mammut/vqa\"\n",
+        "    else:\n",
+        "        model_dir = \"gs://vertex-model-garden-public-us/mammut/retrieval\"\n",
+        "\n",
+        "    upload_job_name = common_util.get_job_name_with_datetime(\n",
+        "        prefix=\"mammut-\" + task + \"-upload\"\n",
+        "    )\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=upload_job_name,\n",
+        "        artifact_uri=model_dir,\n",
+        "        serving_container_image_uri=OPTIMIZED_TF_RUNTIME_IMAGE_URI,\n",
+        "        serving_container_args=[],\n",
+        "        location=REGION,\n",
+        "        serving_container_environment_variables=serving_env,\n",
+        "    )\n",
+        "\n",
+        "    print(\"The uploaded model name is: \", upload_job_name)\n",
+        "\n",
+        "    deploy_model_name = common_util.get_job_name_with_datetime(\n",
+        "        prefix=\"mammut-\" + task + \"-deploy\"\n",
+        "    )\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",
+        "\n",
+        "    endpoint = model.deploy(\n",
+        "        deployed_model_display_name=deploy_model_name,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        min_replica_count=1,\n",
+        "        max_replica_count=1,\n",
+        "    )\n",
+        "\n",
+        "    print(\"The deployed job name is: \", deploy_model_name)\n",
+        "\n",
+        "    endpoint_id = endpoint.name\n",
+        "    print(\"endpoint id is: \", endpoint_id)\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "def predict(\n",
+        "    endpoint: aiplatform.Endpoint,\n",
+        "    image: Image.Image,\n",
+        "    prompt: str,\n",
+        "    new_width: int = 1000,\n",
+        "):\n",
+        "    \"\"\"Generates predictions based on the input image and text using an Endpoint.\"\"\"\n",
+        "    # Resize and convert image to base64 string.\n",
+        "    resized_image = resize_image(image, new_width)\n",
+        "    instances = [\n",
+        "        {\n",
+        "            \"image_bytes\": {\"b64\": common_util.image_to_base64(resized_image)},\n",
+        "            \"text\": prompt,\n",
+        "        }\n",
+        "    ]\n",
+        "\n",
+        "    response = endpoint.predict(instances=instances)\n",
+        "    return response.predictions[0]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "iILhhP3TfO8B"
+      },
+      "source": [
+        "## Run online prediction\n",
+        "\n",
+        "Run online prediction with the TF SavedModel."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Lgq-FTe7wak_"
+      },
+      "source": [
+        "### Visual Question Answering"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "74yqis5ufO8B"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "# @markdown Upload TF SavedModel and deploy it to an endpoint for prediction. This step takes around 15 minutes to finish.\n",
+        "\n",
+        "# @markdown Select the accelerator type to use to deploy the model:\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_V100\"]\n",
+        "# @markdown If you want to use other accelerator types not listed above, then check other Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute. You may need to manually set the `machine_type`, `accelerator_type`, and `accelerator_count` in the code by clicking `Show code` first.\n",
+        "\n",
+        "accelerator_count = 1\n",
+        "if accelerator_type == \"NVIDIA_L4\":\n",
+        "    machine_type = \"g2-standard-4\"\n",
+        "elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
+        "    machine_type = \"n1-standard-4\"\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        f\"Recommended machine settings not found for: {accelerator_type}. To use another another accelerator, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `accelerator_count` to the deploy_model function by clicking `Show Code` and then modifying the code.\"\n",
+        "    )\n",
+        "\n",
+        "models[\"vqa\"], endpoints[\"vqa\"] = deploy_mammut(\n",
+        "    \"vqa\", machine_type, accelerator_type, accelerator_count\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "qxj4Xv_DhHXj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "# @markdown Use the deployed MaMMUT model to answer questions about a given image.\n",
+        "\n",
+        "# @markdown **Note: The first prediction can take up to 2 minutes due to one time JIT compilation of the model. This may cause a timeout error below. If you get a timeout error, then wait for 2 minutes and run the prediction again. You will not get the timeout error after that.**\n",
+        "\n",
+        "# @markdown This section uses images from [pexels.com](https://www.pexels.com/) for demoing purposes. All the images have the following license: https://www.pexels.com/license/.\n",
+        "\n",
+        "# @markdown Images will be resized to a width of 1000 pixels by default since requests made to a Vertex Endpoint are limited to 1.500MB.\n",
+        "\n",
+        "# @markdown ![](https://images.pexels.com/photos/4012966/pexels-photo-4012966.jpeg?w=1260&h=750)\n",
+        "\n",
+        "# @markdown This can be either a Cloud Storage path (gs://\\<image-path\\>) or a public url (http://\\<image-path\\>)\n",
+        "image_url = \"https://images.pexels.com/photos/4012966/pexels-photo-4012966.jpeg\"  # @param {type:\"string\"}\n",
+        "\n",
+        "image = load_image(image_url)\n",
+        "display(image)\n",
+        "\n",
+        "# @markdown You may leave question prompts empty and they will be ignored.\n",
+        "question_prompt_1 = \"Is there a person in the image?\"  # @param {type: \"string\"}\n",
+        "question_prompt_2 = \"What is the person doing in the image?\"  # @param {type: \"string\"}\n",
+        "question_prompt_3 = \"What's the color of the cup?\"  # @param {type: \"string\"}\n",
+        "question_prompt_4 = \"How many laptops are in the image?\"  # @param {type: \"string\"}\n",
+        "\n",
+        "questions_list = [\n",
+        "    question_prompt_1,\n",
+        "    question_prompt_2,\n",
+        "    question_prompt_3,\n",
+        "    question_prompt_4,\n",
+        "]\n",
+        "questions_list = [question for question in questions_list if question]\n",
+        "\n",
+        "for question in questions_list:\n",
+        "    answer = predict(endpoints[\"vqa\"], image, question)\n",
+        "    print(f\"Question: {question}\")\n",
+        "    print(f\"Answer: {answer}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "VNLXg2BxZli7"
+      },
+      "source": [
+        "#### Creating a webpage playground with Gradio"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "wRbUhytaY3Nt"
+      },
+      "outputs": [],
+      "source": [
+        "# @title How to use\n",
+        "\n",
+        "# @markdown **Prerequisites**\n",
+        "# @markdown -  Before you can upload an image to make a prediction, you need to select a Vertex prediction endpoint serving MaMMUT\n",
+        "# @markdown from the endpoint dropdown list that has been deployed in the current project and region.\n",
+        "# @markdown -  If no models have been deployed, you can create a new Vertex prediction\n",
+        "# @markdown endpoint by clicking \"Deploy to Vertex\" in the playground or running the `Deploy` cell above.\n",
+        "# @markdown   * New model deployment takes approximately 15 minutes. You can check the progress at [Vertex Online Prediction](https://console.cloud.google.com/vertex-ai/online-prediction/endpoints).\n",
+        "\n",
+        "# @markdown **How to use**\n",
+        "\n",
+        "# @markdown Just run this cell and a link to the playground formatted as `https://####.gradio.live` will be outputted.\n",
+        "# @markdown This link will take you to the playground in a separate browser tab.\n",
+        "\n",
+        "\n",
+        "def list_mammut_endpoints() -> list[str]:\n",
+        "    \"\"\"Returns all valid prediction endpoints for in the project and region.\"\"\"\n",
+        "    # Gets all the valid endpoints in the project and region.\n",
+        "    endpoints = aiplatform.Endpoint.list(order_by=\"create_time desc\")\n",
+        "    # Filters out the endpoints which do not have a deployed model, and the endpoint is for image generation\n",
+        "    endpoints = list(\n",
+        "        filter(\n",
+        "            lambda endpoint: endpoint.traffic_split\n",
+        "            and \"mammut-vqa\" in endpoint.display_name.lower(),\n",
+        "            endpoints,\n",
+        "        )\n",
+        "    )\n",
+        "\n",
+        "    endpoint_names = list(\n",
+        "        map(\n",
+        "            lambda endpoint: f\"{endpoint.name} - {endpoint.display_name[:40]}\",\n",
+        "            endpoints,\n",
+        "        )\n",
+        "    )\n",
+        "\n",
+        "    if not endpoint_names:\n",
+        "        gr.Warning(\"No prediction endpoints were found. Create an Endpoint first.\")\n",
+        "\n",
+        "    return endpoint_names\n",
+        "\n",
+        "\n",
+        "def deploy_model_handler() -> None:\n",
+        "    gr.Info(\"Starting model deployment.\")\n",
+        "    model, endpoint = deploy_mammut(\"vqa\", \"g2-standard-4\", \"NVIDIA_L4\", 1)\n",
+        "    gr.Info(f\"Deploying model ID: {model.name}, endpoint ID: {endpoint.name}\")\n",
+        "\n",
+        "\n",
+        "def get_endpoint(endpoint_name: str) -> aiplatform.Endpoint:\n",
+        "    \"\"\"Returns a Vertex endpoint for the given endpoint_name.\"\"\"\n",
+        "    endpoint_id = endpoint_name.split(\" - \")[0]\n",
+        "    endpoint = aiplatform.Endpoint(\n",
+        "        f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
+        "    )\n",
+        "    return endpoint\n",
+        "\n",
+        "\n",
+        "def predict_handler(\n",
+        "    endpoint_name: str,\n",
+        "    image: Image.Image,\n",
+        "    prompt: str,\n",
+        ") -> str:\n",
+        "    if not endpoint_name:\n",
+        "        raise gr.Error(\"Select (or deploy) a model first!\")\n",
+        "    if not image:\n",
+        "        raise gr.Error(\"You must upload an image!\")\n",
+        "    endpoint = get_endpoint(endpoint_name)\n",
+        "    return predict(endpoint, image, prompt)\n",
+        "\n",
+        "\n",
+        "tip_text = r\"\"\"\n",
+        "<b> Tips: </b>\n",
+        "1. Select a Vertex prediction endpoint with a deployed MaMMUT model or click `Deploy to Vertex` to deploy MaMMUT to Vertex.\n",
+        "2. New model deployment takes approximately 15 minutes. You can check the progress by examining the output section of the notebook cell that runs this playground. Your endpoint will show up at [Vertex Online Prediction](https://console.cloud.google.com/vertex-ai/online-prediction/endpoints) once the deployment is done.\n",
+        "3. After the model deployment is complete, click `Refresh Endpoints list` to view the new endpoint in the dropdown list.\n",
+        "4. Note: The first prediction can take up to 2 minutes due to one time JIT compilation of the model. This may cause a timeout error below. If you get a timeout error, then wait for 2 minutes and run the prediction again. You will not get the timeout error after that.\n",
+        "\"\"\"\n",
+        "\n",
+        "css = \"\"\"\n",
+        ".gradio-container {\n",
+        "  width: 85% !important\n",
+        "}\n",
+        "\"\"\"\n",
+        "with gr.Blocks(\n",
+        "    css=css, theme=gr.themes.Default(primary_hue=\"orange\", secondary_hue=\"blue\")\n",
+        ") as demo:\n",
+        "    gr.Markdown(\"# Model Garden Playground for MaMMUT\")\n",
+        "    with gr.Row(equal_height=True):\n",
+        "        with gr.Column(scale=3):\n",
+        "            gr.Markdown(tip_text)\n",
+        "        with gr.Column(scale=2):\n",
+        "            with gr.Row():\n",
+        "                endpoint_name = gr.Dropdown(\n",
+        "                    scale=7,\n",
+        "                    label=\"Select a model previously deployed on Vertex (Click inside the input box below)\",\n",
+        "                    choices=list_mammut_endpoints(),\n",
+        "                    value=None,\n",
+        "                )\n",
+        "                refresh_button = gr.Button(\n",
+        "                    \"Refresh Endpoints list\",\n",
+        "                    scale=1,\n",
+        "                    variant=\"primary\",\n",
+        "                    min_width=10,\n",
+        "                )\n",
+        "            with gr.Row():\n",
+        "                deploy_model_button = gr.Button(\n",
+        "                    \"Deploy a new model\",\n",
+        "                    scale=1,\n",
+        "                    variant=\"primary\",\n",
+        "                    min_width=10,\n",
+        "                )\n",
+        "    with gr.Row(equal_height=True):\n",
+        "        with gr.Column(scale=1):\n",
+        "            image_input = gr.Image(\n",
+        "                show_label=True,\n",
+        "                type=\"pil\",\n",
+        "                label=\"Upload\",\n",
+        "                visible=True,\n",
+        "                height=400,\n",
+        "            )\n",
+        "            with gr.Group():\n",
+        "                text_input_box = gr.Textbox(label=\"Question\", lines=1)\n",
+        "                submit_button = gr.Button(\"Answer\", variant=\"primary\")\n",
+        "        with gr.Column(scale=1):\n",
+        "            image_output = gr.Image(label=\"Image response:\", visible=False)\n",
+        "            text_output = gr.Textbox(label=\"Text response:\")\n",
+        "\n",
+        "    refresh_button.click(\n",
+        "        fn=lambda: gr.update(choices=list_mammut_endpoints()),\n",
+        "        outputs=[endpoint_name],\n",
+        "    )\n",
+        "    deploy_model_button.click(\n",
+        "        deploy_model_handler,\n",
+        "        outputs=[],\n",
+        "    )\n",
+        "    submit_button.click(\n",
+        "        fn=predict_handler,\n",
+        "        inputs=[\n",
+        "            endpoint_name,\n",
+        "            image_input,\n",
+        "            text_input_box,\n",
+        "        ],\n",
+        "        outputs=[text_output],\n",
+        "    )\n",
+        "show_debug_logs = True  # @param {type: \"boolean\"}\n",
+        "demo.queue()\n",
+        "demo.launch(\n",
+        "    share=True, inline=False, inbrowser=True, debug=show_debug_logs, show_error=True\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "97jEsBSfwm-3"
+      },
+      "source": [
+        "### Retrieval and Multimodal Embeddings"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "oJemp5enwl2g"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "# @markdown Upload TF SavedModel and deploy it to an endpoint for prediction. This step takes around 15 minutes to finish.\n",
+        "\n",
+        "# @markdown Select the accelerator type to use to deploy the model:\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_V100\"]\n",
+        "# @markdown If you want to use other accelerator types not listed above, then check other Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute. You may need to manually set the `machine_type`, `accelerator_type`, and `accelerator_count` in the code by clicking `Show code` first.\n",
+        "\n",
+        "accelerator_count = 1\n",
+        "if accelerator_type == \"NVIDIA_L4\":\n",
+        "    machine_type = \"g2-standard-4\"\n",
+        "elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
+        "    machine_type = \"n1-standard-4\"\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        f\"Recommended machine settings not found for: {accelerator_type}. To use another another accelerator, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `accelerator_count` to the deploy_model function by clicking `Show Code` and then modifying the code.\"\n",
+        "    )\n",
+        "\n",
+        "models[\"retrieval\"], endpoints[\"retrieval\"] = deploy_mammut(\n",
+        "    \"retrieval\", machine_type, accelerator_type, accelerator_count\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "fO0-YrOZi34k"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Image-Text Retrieval\n",
+        "# @markdown Given an image, use the deployed MaMMUT model to find the best matching text out of multiple options based on similarity scores of their embeddings. This example uses only 5 text options but you can modify the example to retrieve over as many text examples as needed.\n",
+        "\n",
+        "# @markdown **Note: The first prediction can take up to 2 minutes due to one time JIT compilation of the model. This may cause a timeout error below. If you get a timeout error, then wait for 2 minutes and run the prediction again. You will not get the timeout error after that.**\n",
+        "\n",
+        "# @markdown This section uses images from [pexels.com](https://www.pexels.com/) for demoing purposes. All the images have the following license: https://www.pexels.com/license/.\n",
+        "\n",
+        "# @markdown Images will be resized to a width of 1000 pixels by default since requests made to a Vertex Endpoint are limited to 1.500MB.\n",
+        "\n",
+        "# @markdown ![](https://images.pexels.com/photos/20427316/pexels-photo-20427316/free-photo-of-a-moped-parked-in-front-of-a-blue-door.jpeg?auto=compress&cs=tinysrgb&w=630&h=375&dpr=2)\n",
+        "\n",
+        "# @markdown This can be either a Cloud Storage path (gs://\\<image-path\\>) or a public url (http://\\<image-path\\>)\n",
+        "image_url = \"https://images.pexels.com/photos/20427316/pexels-photo-20427316/free-photo-of-a-moped-parked-in-front-of-a-blue-door.jpeg?auto=compress&cs=tinysrgb&w=630&h=375&dpr=2\"  # @param {type:\"string\"}\n",
+        "\n",
+        "image = load_image(image_url)\n",
+        "display(image)\n",
+        "\n",
+        "text_1 = \"A tennis player about to serve.\"  # @param {type: \"string\"}\n",
+        "text_2 = \"Green broccolis and fruit in a bowl on a table.\"  # @param {type: \"string\"}\n",
+        "text_3 = \"A moped parked in front of a blue door.\"  # @param {type: \"string\"}\n",
+        "text_4 = \"A baguette with some ham in it.\"  # @param {type: \"string\"}\n",
+        "text_5 = \"Three zebras in a dry land with some bush.\"  # @param {type: \"string\"}\n",
+        "\n",
+        "text_list = [text_1, text_2, text_3, text_4, text_5]\n",
+        "text_list = [text for text in text_list if text]\n",
+        "\n",
+        "image_embeddings = []\n",
+        "text_embeddings = []\n",
+        "for text in text_list:\n",
+        "    prediction = predict(endpoints[\"retrieval\"], image, text)\n",
+        "    image_embeddings.append(np.array(prediction[\"normalized_image_embedding\"]))\n",
+        "    text_embeddings.append(np.array(prediction[\"normalized_text_embedding\"]))\n",
+        "\n",
+        "# predictions = predict(endpoint, image, text_list)\n",
+        "# image_embeddings = [np.array(prediction[\"normalized_image_embedding\"]) for prediction in predictions]\n",
+        "\n",
+        "# text_embeddings = [np.array(prediction[\"normalized_text_embedding\"]) for prediction in predictions]\n",
+        "\n",
+        "image_embeddings = np.vstack(image_embeddings)\n",
+        "text_embeddings = np.vstack(text_embeddings)\n",
+        "similarity = np.matmul(image_embeddings, text_embeddings.T)\n",
+        "argmax_indices = np.argmax(similarity, axis=-1)\n",
+        "argmax = argmax_indices[0]\n",
+        "print(f\"The text that's most similar to the image is: {text_list[argmax]}\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "o2modPjMFYrm"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Text-Image Retrieval\n",
+        "# @markdown Given a text description, use the deployed MaMMUT model to find the best matching image out of multiple options based on similarity scores of their embeddings. This example uses only 5 image options but you can modify the example to retrieve over as many image examples as needed.\n",
+        "\n",
+        "# @markdown This section uses images from [pexels.com](https://www.pexels.com/) for demoing purposes. All the images have the following license: https://www.pexels.com/license/.\n",
+        "\n",
+        "# @markdown Images will be resized to a width of 1000 pixels by default since requests made to a Vertex Endpoint are limited to 1.500MB.\n",
+        "\n",
+        "text = \"A view of the city with many red roofs.\"  # @param {type: \"string\"}\n",
+        "\n",
+        "# @markdown Image URLs can be either a Cloud Storage path (gs://\\<image-path\\>) or a public url (http://\\<image-path\\>)\n",
+        "\n",
+        "image_url_1 = \"https://images.pexels.com/photos/4012966/pexels-photo-4012966.jpeg?w=1260&h=750\"  # @param {type:\"string\"}\n",
+        "# @markdown ![](https://images.pexels.com/photos/4012966/pexels-photo-4012966.jpeg?w=1260&h=750)\n",
+        "\n",
+        "image_url_2 = \"https://images.pexels.com/photos/24427993/pexels-photo-24427993/free-photo-of-a-group-of-cherries-arranged-in-a-row-on-a-white-wall.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=1\"  # @param {type:\"string\"}\n",
+        "# @markdown ![](https://images.pexels.com/photos/24427993/pexels-photo-24427993/free-photo-of-a-group-of-cherries-arranged-in-a-row-on-a-white-wall.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=1)\n",
+        "\n",
+        "image_url_3 = \"https://images.pexels.com/photos/20427316/pexels-photo-20427316/free-photo-of-a-moped-parked-in-front-of-a-blue-door.jpeg?auto=compress&cs=tinysrgb&w=630&h=375&dpr=2\"  # @param {type:\"string\"}\n",
+        "# @markdown ![](https://images.pexels.com/photos/20427316/pexels-photo-20427316/free-photo-of-a-moped-parked-in-front-of-a-blue-door.jpeg?auto=compress&cs=tinysrgb&w=630&h=375&dpr=2)\n",
+        "\n",
+        "image_url_4 = \"https://images.pexels.com/photos/18592009/pexels-photo-18592009/free-photo-of-a-view-of-the-city-with-many-red-roofs.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=1\"  # @param {type:\"string\"}\n",
+        "# @markdown ![](https://images.pexels.com/photos/18592009/pexels-photo-18592009/free-photo-of-a-view-of-the-city-with-many-red-roofs.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=1)\n",
+        "\n",
+        "image_url_5 = \"https://images.pexels.com/photos/1006293/pexels-photo-1006293.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=2\"  # @param {type:\"string\"}\n",
+        "# @markdown ![](https://images.pexels.com/photos/1006293/pexels-photo-1006293.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=2)\n",
+        "\n",
+        "image_url_list = [image_url_1, image_url_2, image_url_3, image_url_4, image_url_5]\n",
+        "image_url_list = [image_url for image_url in image_url_list if image_url]\n",
+        "\n",
+        "images = [load_image(image_url) for image_url in image_url_list]\n",
+        "\n",
+        "text_embeddings = []\n",
+        "image_embeddings = []\n",
+        "for image in images:\n",
+        "    prediction = predict(retrieval_endpoint, image, text)\n",
+        "    image_embeddings.append(np.array(prediction[\"normalized_image_embedding\"]))\n",
+        "    text_embeddings.append(np.array(prediction[\"normalized_text_embedding\"]))\n",
+        "\n",
+        "image_embeddings = np.vstack(image_embeddings)\n",
+        "text_embeddings = np.vstack(text_embeddings)\n",
+        "similarity = np.matmul(image_embeddings, text_embeddings.T)\n",
+        "argmax_indices = np.argmax(similarity, axis=0)\n",
+        "argmax = argmax_indices[0]\n",
+        "print(f\"The image that's most similar to the text is: {image_url_list[argmax]}\")\n",
+        "display(images[argmax])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3aD4PW3d1bG5"
+      },
+      "source": [
+        "## Clean up resources\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "sx_vKniMq9ZX"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Run\n",
+        "\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
+        "\n",
+        "# Delete endpoint resource.\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete model resource.\n",
+        "for model in models.values():\n",
+        "    model.delete()\n",
+        "\n",
+        "# Delete Cloud Storage objects that were created.\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_URI"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_mammut.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_mediapipe_face_stylizer.ipynb

@@ -317,7 +317,7 @@
       "outputs": [],
       "source": [
         "TRAINING_JOB_DISPLAY_NAME = \"mediapipe_face_stylizer_%s\" % now\n",
-        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
+        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
         "TRAINING_MACHINE_TYPE = \"n1-highmem-16\"\n",
         "TRAINING_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
         "TRAINING_ACCELERATOR_COUNT = 2"

notebooks/community/model_garden/model_garden_mediapipe_gesture_recognition.ipynb

@@ -318,7 +318,7 @@
       "outputs": [],
       "source": [
         "TRAINING_JOB_DISPLAY_NAME = \"mediapipe_gesture_recognizer_%s\" % now\n",
-        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
+        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
         "TRAINING_MACHINE_TYPE = \"n1-highmem-16\"\n",
         "TRAINING_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
         "TRAINING_ACCELERATOR_COUNT = 2"

notebooks/community/model_garden/model_garden_mediapipe_image_classification.ipynb

@@ -318,7 +318,7 @@
       "outputs": [],
       "source": [
         "TRAINING_JOB_DISPLAY_NAME = \"mediapipe_image_classifier_%s\" % now\n",
-        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
+        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
         "TRAINING_MACHINE_TYPE = \"n1-highmem-16\"\n",
         "TRAINING_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
         "TRAINING_ACCELERATOR_COUNT = 2"

notebooks/community/model_garden/model_garden_mediapipe_image_generation.ipynb

@@ -440,9 +440,7 @@
       "outputs": [],
       "source": [
         "# Path to the training data folder.\n",
-        "training_data_path = (\n",
-        "    \"gs://mediapipe-tasks/image_generator/teapot\"  # @param {type:\"string\"}\n",
-        ")\n",
+        "training_data_path = \"gs://mediapipe-tasks/image_generator/teapot\"  # @param {type:\"string\"}\n",
         "# An instance description of the training data.\n",
         "training_data_prompt = \"A monadikos teapot\"  # @param {type:\"string\"}"
       ]
@@ -823,10 +821,11 @@
       "source": [
         "serving_env = {\n",
         "    \"TASK\": \"text-to-image-lora\",\n",
-        "    \"BASE_MODEL_ID\": \"runwayml/stable-diffusion-v1-5\",\n",
+        "    \"MODEL_ID\": \"runwayml/stable-diffusion-v1-5\",\n",
         "    \"FINETUNED_LORA_MODEL_PATH\": os.path.join(\n",
         "        MODEL_EXPORT_PATH, f\"checkpoint_{deployed_checkpoint}\"\n",
         "    ),\n",
+        "    \"DEPLOY_SOURCE\": \"notebook\",\n",
         "}\n",
         "\n",
         "model = aiplatform.Model.upload(\n",

notebooks/community/model_garden/model_garden_mediapipe_object_detection.ipynb

@@ -318,7 +318,7 @@
       "outputs": [],
       "source": [
         "TRAINING_JOB_DISPLAY_NAME = \"mediapipe_object_detector_%s\" % now\n",
-        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
+        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
         "TRAINING_MACHINE_TYPE = \"n1-highmem-16\"\n",
         "TRAINING_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
         "TRAINING_ACCELERATOR_COUNT = 2"

notebooks/community/model_garden/model_garden_mediapipe_text_classification.ipynb

@@ -317,7 +317,7 @@
       "outputs": [],
       "source": [
         "TRAINING_JOB_DISPLAY_NAME = \"mediapipe_text_classifier_%s\" % now\n",
-        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
+        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/mediapipe-train\"\n",
         "TRAINING_MACHINE_TYPE = \"n1-highmem-16\"\n",
         "TRAINING_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
         "TRAINING_ACCELERATOR_COUNT = 2"

notebooks/community/model_garden/model_garden_movinet_action_recognition.ipynb

@@ -265,7 +265,7 @@
         "# Prediction constants.\n",
         "# You can adjust accelerator types and machine types to get faster predictions.\n",
         "PREDICTION_CONTAINER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/movinet-serve\"\n",
-        "PREDICTION_PORT = 8501\n",
+        "PREDICTION_PORT = 8080\n",
         "PREDICTION_ACCELERATOR_COUNT = 1\n",
         "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_T4\"\n",
         "PREDICTION_MACHINE_TYPE = \"n1-standard-4\"\n",
@@ -791,6 +791,7 @@
       "outputs": [],
       "source": [
         "serving_env = {\n",
+        "    \"MODEL_ID\": \"tfvision-movinet-var\",\n",
         "    \"MODEL_PATH\": container_args[\"export_path\"],\n",
         "    \"BATCH_SIZE\": 1,\n",
         "    \"NUM_FRAMES\": 32,\n",
@@ -799,6 +800,7 @@
         "    \"OBJECTIVE\": OBJECTIVE,\n",
         "    \"IMAGE_WIDTH\": image_size,\n",
         "    \"IMAGE_HEIGHT\": image_size,\n",
+        "    \"DEPLOY_SOURCE\": \"notebook\",\n",
         "}\n",
         "\n",
         "model = aiplatform.Model.upload(\n",

notebooks/community/model_garden/model_garden_movinet_clip_classification.ipynb

@@ -250,7 +250,7 @@
         "# Prediction constants.\n",
         "# You can adjust accelerator types and machine types to get faster predictions.\n",
         "PREDICTION_CONTAINER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/movinet-serve\"\n",
-        "PREDICTION_PORT = 8501\n",
+        "PREDICTION_PORT = 8080\n",
         "PREDICTION_ACCELERATOR_COUNT = 1\n",
         "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_T4\"\n",
         "PREDICTION_MACHINE_TYPE = \"n1-standard-4\"\n",
@@ -753,12 +753,14 @@
       "outputs": [],
       "source": [
         "serving_env = {\n",
+        "    \"MODEL_ID\": \"tfvision-movinet-vcn\",\n",
         "    \"MODEL_PATH\": container_args[\"export_path\"],\n",
         "    \"BATCH_SIZE\": 1,  # Select a larger batch size to accelerate GPU prediction.\n",
         "    \"NUM_FRAMES\": 32,\n",
         "    \"FPS\": output_fps,\n",
         "    \"OVERLAP_FRAMES\": 24,\n",
         "    \"OBJECTIVE\": OBJECTIVE,\n",
+        "    \"DEPLOY_SOURCE\": \"notebook\",\n",
         "}\n",
         "\n",
         "model = aiplatform.Model.upload(\n",

notebooks/community/model_garden/model_garden_pytorch_autogluon.ipynb

@@ -8,7 +8,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2023 Google LLC\n",
+        "# 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",
@@ -26,30 +26,44 @@
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "JAPoU8Sm5E6e"
+        "id": "99c1c3fc2ca5"
       },
       "source": [
-        "# Train a scikit-learn model with Vertex AI SDK and Bigframes\n",
+        "# Vertex AI Model Garden - AutoGluon\n",
         "\n",
         "<table align=\"left\">\n",
         "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vertex_ai_sdk/remote_training_bigframes_sklearn.ipynb\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_autogluon.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",
+        "\n",
         "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/vertex_ai_sdk/remote_training_bigframes_sklearn.ipynb\">\n",
-        "        <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_autogluon.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/vertex_ai_sdk/remote_training_bigframes_sklearn.ipynb\">\n",
-        "       <img src=\"https://www.gstatic.com/cloud/images/navigation/vertex-ai.svg\" alt=\"Vertex AI logo\">Open in Vertex AI Workbench\n",
+        "  <td>                                                                                               <td>\n",
+        "    <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_autogluon.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>\n",
+        "  </td>\n",
         "</table>"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "24743cf4a1e1"
+      },
+      "source": [
+        "**_NOTE_**: This notebook has been tested in the following environment:\n",
+        "\n",
+        "* Python version = 3.10"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -58,9 +72,7 @@
       "source": [
         "## Overview\n",
         "\n",
-        "This tutorial demonstrates how to train a scikit-learn model using Vertex AI local-to-remote training with Vertex AI SDK and BigQuery Bigframes as the data source.\n",
-        "\n",
-        "Learn more about [bigframes](https://cloud.google.com/bigquery/docs/)."
+        "This notebook demonstrates finetuning a PyTorch based [Autogluon model for tabular data](https://auto.gluon.ai/stable/tutorials/tabular/index.html) on CPU and deploying it on Vertex AI for online prediction."
       ]
     },
     {
@@ -71,42 +83,18 @@
       "source": [
         "### Objective\n",
         "\n",
-        "In this tutorial, you learn to use `Vertex AI SDK` with Bigframes as input data source.\n",
+        "In this tutorial, you learn how to:\n",
         "\n",
+        "- Finetune a PyTorch AutoGluon tabular model.\n",
+        "- Upload the model to [Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction).\n",
+        "- Deploy the model on [Endpoint](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
+        "- Run online predictions for tabular data.\n",
         "\n",
-        "This tutorial uses the following Google Cloud ML services:\n",
+        "This tutorial uses the following Google Cloud ML services and resources:\n",
         "\n",
-        "- `Vertex AI Training`\n",
-        "- `Vertex AI Remote Training`\n",
-        "\n",
-        "\n",
-        "The steps performed include:\n",
-        "\n",
-        "- Initialize a dataframe from a BigQuery table and split the dataset\n",
-        "- Perform transformations as a Vertex AI remote training.\n",
-        "- Train the model remotely and evaluate the model locally\n",
-        "\n",
-        "**Local-to-remote training**\n",
-        "\n",
-        "```\n",
-        "import vertexai\n",
-        "from my_module import MyModelClass\n",
-        "\n",
-        "vertexai.preview.init(remote=True, project=\"my-project\", location=\"my-location\", staging_bucket=\"gs://my-bucket\")\n",
-        "\n",
-        "# Wrap the model class with `vertex_ai.preview.remote`\n",
-        "MyModelClass = vertexai.preview.remote(MyModelClass)\n",
-        "\n",
-        "# Instantiate the class\n",
-        "model = MyModelClass(...)\n",
-        "\n",
-        "# Optional set remote config\n",
-        "model.fit.vertex.remote_config.display_name = \"MyModelClass-remote-training\"\n",
-        "model.fit.vertex.remote_config.staging_bucket = \"gs://my-bucket\"\n",
-        "\n",
-        "# This `fit` call will be executed remotely\n",
-        "model.fit(...)\n",
-        "```"
+        "- Vertex AI Training\n",
+        "- Vertex AI Model Registry\n",
+        "- Vertex AI Online Prediction"
       ]
     },
     {
@@ -117,7 +105,7 @@
       "source": [
         "### Dataset\n",
         "\n",
-        "This tutorial uses the <a href=\"https://scikit-learn.org/stable/auto_examples/datasets/plot_iris_dataset.html\">IRIS dataset</a>, which predicts the iris species."
+        "You can find the details for the [example dataset here](https://auto.gluon.ai/stable/tutorials/tabular/tabular-quick-start.html#example-data)."
       ]
     },
     {
@@ -131,14 +119,9 @@
         "This tutorial uses billable components of Google Cloud:\n",
         "\n",
         "* Vertex AI\n",
-        "* BigQuery\n",
         "* Cloud Storage\n",
         "\n",
-        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing),\n",
-        "[BigQuery pricing](https://cloud.google.com/bigquery/pricing),\n",
-        "and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), \n",
-        "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n",
-        "to generate a cost estimate based on your projected usage."
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
       ]
     },
     {
@@ -149,7 +132,7 @@
       "source": [
         "## Installation\n",
         "\n",
-        "Install the following packages required to execute this notebook. "
+        "Install the following packages required to execute this notebook."
       ]
     },
     {
@@ -160,9 +143,8 @@
       },
       "outputs": [],
       "source": [
-        "# Install the packages\n",
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform[preview]\n",
-        "! pip3 install --upgrade --quiet bigframes"
+        "# Install the packages.\n",
+        "! pip3 install --upgrade google-cloud-aiplatform"
       ]
     },
     {
@@ -171,7 +153,7 @@
         "id": "58707a750154"
       },
       "source": [
-        "### Colab only: Uncomment the following cell to restart the kernel."
+        "### Colab only"
       ]
     },
     {
@@ -182,11 +164,11 @@
       },
       "outputs": [],
       "source": [
-        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
-        "# import IPython\n",
+        "# 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)"
+        "app = IPython.Application.instance()\n",
+        "app.kernel.do_shutdown(True)"
       ]
     },
     {
@@ -203,11 +185,13 @@
         "\n",
         "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
         "\n",
-        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
         "\n",
-        "3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
+        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
         "\n",
-        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
+        "1. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "\n",
+        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User` and `Storage Object Admin` roles for deploying fine tuned model to Vertex AI endpoint.\n"
       ]
     },
     {
@@ -232,7 +216,7 @@
       },
       "outputs": [],
       "source": [
-        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "PROJECT_ID = \"your-project-id\"  # @param {type:\"string\"}\n",
         "\n",
         "# Set the project id\n",
         "! gcloud config set project {PROJECT_ID}"
@@ -253,7 +237,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "region"
+        "id": "twgKk-LsLmX3"
       },
       "outputs": [],
       "source": [
@@ -332,6 +316,20 @@
         "* 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": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Z36ywjGtRey3"
+      },
+      "outputs": [],
+      "source": [
+        "# The service account for deploying fine tuned model.\n",
+        "# The service account looks like:\n",
+        "# '<account_name>@<project>.iam.gserviceaccount.com'\n",
+        "SERVICE_ACCOUNT = \"your-service-account\"  # @param {type:\"string\"}"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -380,7 +378,7 @@
         "id": "960505627ddf"
       },
       "source": [
-        "### Import libraries and define constants"
+        "### Import libraries"
       ]
     },
     {
@@ -391,14 +389,10 @@
       },
       "outputs": [],
       "source": [
-        "import bigframes.pandas as bf\n",
-        "import vertexai\n",
+        "import os\n",
+        "from datetime import datetime\n",
         "\n",
-        "bf.options.bigquery.location = \"us\"  # Dataset is in 'us' not 'us-central1'\n",
-        "bf.options.bigquery.project = PROJECT_ID\n",
-        "\n",
-        "from bigframes.ml.model_selection import \\\n",
-        "    train_test_split as bf_train_test_split"
+        "from google.cloud import aiplatform"
       ]
     },
     {
@@ -407,257 +401,245 @@
         "id": "init_aip:mbsdk,all"
       },
       "source": [
-        "## Initialize Vertex AI SDK for Python\n",
+        "### Initialize Vertex AI SDK for Python\n",
         "\n",
-        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+        "Initialize the Vertex AI SDK for Python for your project."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "init_aip:mbsdk,all"
+        "id": "vS1hQiGuLmX4"
       },
       "outputs": [],
       "source": [
-        "vertexai.init(\n",
-        "    project=PROJECT_ID,\n",
-        "    location=REGION,\n",
-        "    staging_bucket=BUCKET_URI,\n",
+        "staging_bucket = os.path.join(BUCKET_URI, \"autogluon_staging\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=staging_bucket)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2cc825514deb"
+      },
+      "source": [
+        "### Define constants"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b42bd4fa2b2d"
+      },
+      "outputs": [],
+      "source": [
+        "# The pre-built training docker image.\n",
+        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-autogluon-train:20240124_0927_RC00\"\n",
+        "# The pre-built serving docker image.\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-autogluon-serve:20240124_0938_RC00\"\n",
+        "# Serving port.\n",
+        "PORT = 8501"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0c250872074f"
+      },
+      "source": [
+        "### Define common functions\n",
+        "\n",
+        "This section defines functions for:\n",
+        "\n",
+        "- Converting a Cloud Storage path such as `gs://bucket-name` to GCSFuse path format such as `/gcsfuse/bucket-name`.\n",
+        "- Deploy the trained model to Vertex AI Endpoint for prediction."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "XcYUGwr-AJGY"
+      },
+      "outputs": [],
+      "source": [
+        "def gcs_fuse_path(path: str) -> str:\n",
+        "    \"\"\"Try to convert path to gcsfuse path if it starts with gs:// else do not modify it.\"\"\"\n",
+        "    path = path.strip()\n",
+        "    if path.startswith(\"gs://\"):\n",
+        "        return \"/gcs/\" + path[5:]\n",
+        "    return path\n",
+        "\n",
+        "\n",
+        "def deploy_model(model_path):\n",
+        "    \"\"\"Deploy the model to Vertex AI Endpoint for prediction.\"\"\"\n",
+        "    model_name = \"autogluon\"\n",
+        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
+        "    serving_env = {\n",
+        "        \"model_path\": model_path,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "    # Since the model_id is a GCS path, use artifact_uri to pass it\n",
+        "    # to the serving docker.\n",
+        "    artifact_uri = model_path\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
+        "        serving_container_ports=[PORT],\n",
+        "        serving_container_predict_route=\"/predict\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=serving_env,\n",
+        "        artifact_uri=artifact_uri,\n",
+        "    )\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=\"n1-highmem-16\",\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=SERVICE_ACCOUNT,\n",
+        "    )\n",
+        "    return model, endpoint"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aCpLmWPMpJQ8"
+      },
+      "source": [
+        "## Finetune with AutoGluon\n",
+        "\n",
+        "Create and run the training job with the model-garden PyTorch AutoGluon training docker using the Vertex AI SDK. The training uses one CPU and runs for around 3 mins once the training job begins."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aec22792ee84"
+      },
+      "outputs": [],
+      "source": [
+        "# Set up training docker arguments.\n",
+        "\n",
+        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
+        "JOB_NAME = \"pytorch_autogluon\" + TIMESTAMP\n",
+        "\n",
+        "finetuning_workdir = os.path.join(BUCKET_URI, JOB_NAME)\n",
+        "train_data_path = (\n",
+        "    \"https://raw.githubusercontent.com/mli/ag-docs/main/knot_theory/train.csv\"\n",
+        ")\n",
+        "# The column id to predict.\n",
+        "label = \"signature\"\n",
+        "\n",
+        "# We are using the\n",
+        "docker_args_list = [\n",
+        "    \"--train_data_path\",\n",
+        "    train_data_path,\n",
+        "    \"--label\",\n",
+        "    label,\n",
+        "    \"--model_save_path\",\n",
+        "    f\"{gcs_fuse_path(finetuning_workdir)}\",\n",
+        "]\n",
+        "print(docker_args_list)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2ELphfgj1f3Q"
+      },
+      "outputs": [],
+      "source": [
+        "# Create and run the training job.\n",
+        "# Click on the generated link in the output under \"View backing custom job:\" to see your run in the Cloud Console.\n",
+        "container_uri = TRAIN_DOCKER_URI\n",
+        "job = aiplatform.CustomContainerTrainingJob(\n",
+        "    display_name=JOB_NAME,\n",
+        "    container_uri=container_uri,\n",
+        ")\n",
+        "model = job.run(\n",
+        "    args=docker_args_list,\n",
+        "    base_output_dir=f\"{finetuning_workdir}\",\n",
+        "    replica_count=1,\n",
+        "    machine_type=\"n1-highmem-16\",\n",
         ")"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "105334524e96"
+        "id": "iILhhP3TfO8B"
       },
       "source": [
-        "## Prepare the dataset\n",
+        "## Run online prediction\n",
         "\n",
-        "Now load the Iris dataset and split the data into train and test sets."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "b44cdc4e03f1"
-      },
-      "outputs": [],
-      "source": [
-        "df = bf.read_gbq(\"bigquery-public-data.ml_datasets.iris\")\n",
-        "\n",
-        "species_categories = {\n",
-        "    \"versicolor\": 0,\n",
-        "    \"virginica\": 1,\n",
-        "    \"setosa\": 2,\n",
-        "}\n",
-        "df[\"species\"] = df[\"species\"].map(species_categories)\n",
-        "\n",
-        "# Assign an index column name\n",
-        "index_col = \"index\"\n",
-        "df.index.name = index_col"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "9cb8616b1997"
-      },
-      "outputs": [],
-      "source": [
-        "feature_columns = df[[\"sepal_length\", \"sepal_width\", \"petal_length\", \"petal_width\"]]\n",
-        "label_columns = df[[\"species\"]]\n",
-        "train_X, test_X, train_y, test_y = bf_train_test_split(\n",
-        "    feature_columns, label_columns, test_size=0.2\n",
-        ")\n",
-        "\n",
-        "print(\"X_train size: \", train_X.size)\n",
-        "print(\"X_test size: \", test_X.size)"
+        "Run online prediction with the trained model."
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "8306545fcc57"
+        "id": "XswgX6JqRwFK"
       },
       "source": [
-        "## Feature transformation\n",
-        "\n",
-        "Next, you do feature transformations on the data using the Vertex AI remote training service.\n",
-        "\n",
-        "First, you re-initialize Vertex AI to enable remote training."
+        "Upload the trained model and deploy it to an endpoint for prediction. This step takes around 20 minutes to finish."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "55e701c31036"
+        "id": "74yqis5ufO8B"
       },
       "outputs": [],
       "source": [
-        "# Switch to remote mode for training\n",
-        "vertexai.preview.init(remote=True)"
+        "model, endpoint = deploy_model(model_path=finetuning_workdir)"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "4a0e9d59b273"
+        "id": "iiozz1aVR7Pe"
       },
       "source": [
-        "### Execute remote job for fit_transform() on training data\n",
-        "\n",
-        "Next, indicate that the `StandardScalar` class is to be executed remotely. Then set up the data transform and call the `fit_transform()` method is executed remotely."
+        "Send the prediction request for the query data. The expected `signature` label for this example query data is `-2`. You can also send comma separated multiple queries."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "90333089d362"
+        "id": "qxj4Xv_DhHXj"
       },
       "outputs": [],
       "source": [
-        "from sklearn.preprocessing import StandardScaler\n",
-        "\n",
-        "# Wrap classes to enable Vertex remote execution\n",
-        "StandardScaler = vertexai.preview.remote(StandardScaler)\n",
-        "\n",
-        "# Instantiate transformer\n",
-        "transformer = StandardScaler()\n",
-        "\n",
-        "# Execute transformer on Vertex (train_X is bigframes.dataframe.DataFrame, X_train is np.array)\n",
-        "X_train = transformer.fit_transform(train_X)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6bf95574c907"
-      },
-      "source": [
-        "### Remote transform on test data"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "da6eea22a89a"
-      },
-      "outputs": [],
-      "source": [
-        "# Execute transformer on Vertex (test_X is bigframes.dataframe.DataFrame, X_test is np.array)\n",
-        "X_test = transformer.transform(test_X)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ddf906c886e4"
-      },
-      "source": [
-        "## Remote training\n",
-        "\n",
-        "First, train the scikit-learn model as a remote training job:\n",
-        "\n",
-        "- Set LogisticRegression for the remote training job.\n",
-        "- Invoke LogisticRegression locally which will launch the remote training job."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "c7b0116fa60c"
-      },
-      "outputs": [],
-      "source": [
-        "from sklearn.linear_model import LogisticRegression\n",
-        "\n",
-        "# Wrap classes to enable Vertex remote execution\n",
-        "LogisticRegression = vertexai.preview.remote(LogisticRegression)\n",
-        "\n",
-        "# Instantiate model, warm_start=True for uptraining\n",
-        "model = LogisticRegression(warm_start=True)\n",
-        "\n",
-        "# Train model on Vertex\n",
-        "model.fit(train_X, train_y)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ffe1d5903bcb"
-      },
-      "source": [
-        "## Remote prediction\n",
-        "\n",
-        "Obtain predictions from the trained model."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "d00ce35920fa"
-      },
-      "outputs": [],
-      "source": [
-        "# Remote evaluation\n",
-        "vertexai.preview.init(remote=True)\n",
-        "\n",
-        "predictions = model.predict(test_X)\n",
-        "\n",
-        "print(f\"Remote predictions: {predictions}\")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "a8cd6cbd4403"
-      },
-      "source": [
-        "## Local evaluation\n",
-        "\n",
-        "Score model results locally."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "dc105dafdfb9"
-      },
-      "outputs": [],
-      "source": [
-        "# User must convert bigframes to pandas dataframe for local evaluation\n",
-        "train_X_pd = train_X.to_pandas().reset_index(drop=True)\n",
-        "train_y_pd = train_y.to_pandas().reset_index(drop=True)\n",
-        "\n",
-        "test_X_pd = test_X.to_pandas().reset_index(drop=True)\n",
-        "test_y_pd = test_y.to_pandas().reset_index(drop=True)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "25fec549de69"
-      },
-      "outputs": [],
-      "source": [
-        "# Switch to local mode for testing\n",
-        "vertexai.preview.init(remote=False)\n",
-        "\n",
-        "# Evaluate model's accuracy score\n",
-        "print(f\"Train accuracy: {model.score(train_X_pd, train_y_pd)}\")\n",
-        "\n",
-        "print(f\"Test accuracy: {model.score(test_X_pd, test_y_pd)}\")"
+        "instances = [\n",
+        "    {\n",
+        "        \"Unnamed: 0\": 70746,\n",
+        "        \"chern_simons\": 0.0905302166938781,\n",
+        "        \"cusp_volume\": 12.226321765565215,\n",
+        "        \"hyperbolic_adjoint_torsion_degree\": 0,\n",
+        "        \"hyperbolic_torsion_degree\": 10,\n",
+        "        \"injectivity_radius\": 0.5077560544013977,\n",
+        "        \"longitudinal_translation\": 10.685555458068848,\n",
+        "        \"meridinal_translation_imag\": 1.1441915035247805,\n",
+        "        \"meridinal_translation_real\": -0.5191566348075867,\n",
+        "        \"short_geodesic_imag_part\": -2.7606005668640137,\n",
+        "        \"short_geodesic_real_part\": 1.0155121088027954,\n",
+        "        \"Symmetry_0\": 0.0,\n",
+        "        \"Symmetry_D3\": 0.0,\n",
+        "        \"Symmetry_D4\": 0.0,\n",
+        "        \"Symmetry_D6\": 0.0,\n",
+        "        \"Symmetry_D8\": 0.0,\n",
+        "        \"Symmetry_Z/2 + Z/2\": 1.0,\n",
+        "        \"volume\": 11.393224716186523,\n",
+        "    },\n",
+        "]\n",
+        "predictions = endpoint.predict(instances=instances).predictions\n",
+        "print(predictions)"
       ]
     },
     {
@@ -682,9 +664,13 @@
       },
       "outputs": [],
       "source": [
-        "import os\n",
+        "# Delete endpoint resource.\n",
+        "endpoint.delete(force=True)\n",
         "\n",
-        "# Delete Cloud Storage objects that were created\n",
+        "# Delete model resource.\n",
+        "model.delete()\n",
+        "\n",
+        "# Delete Cloud Storage objects that were created.\n",
         "delete_bucket = False\n",
         "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
         "    ! gsutil -m rm -r $BUCKET_URI"
@@ -693,8 +679,7 @@
   ],
   "metadata": {
     "colab": {
-      "collapsed_sections": [],
-      "name": "remote_training_bigframes_sklearn.ipynb",
+      "name": "model_garden_pytorch_autogluon.ipynb",
       "toc_visible": true
     },
     "kernelspec": {

notebooks/community/model_garden/model_garden_pytorch_bart_large_cnn.ipynb

@@ -235,6 +235,7 @@
         "    serving_env = {\n",
         "        \"MODEL_ID\": model_id,\n",
         "        \"TASK\": task,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
         "    }\n",
         "    # If the model_id is a GCS path, use artifact_uri to pass it to serving docker.\n",
         "    artifact_uri = model_id if model_id.startswith(\"gs://\") else None\n",

notebooks/community/model_garden/model_garden_pytorch_biogpt_serve.ipynb

@@ -231,7 +231,7 @@
       "outputs": [],
       "source": [
         "# The pre-built training and serving docker images.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-biogpt-serve:latest\""
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-biogpt-serve\""
       ]
     },
     {
@@ -275,10 +275,11 @@
         "    \"\"\"Deploys trained models into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
         "    serving_env = {\n",
-        "        \"BASE_MODEL_ID\": model_id,\n",
+        "        \"MODEL_ID\": model_id,\n",
         "        \"TASK\": task,\n",
         "        \"MAX_LENGTH\": max_length,\n",
         "        \"NUM_RETURN_SEQUENCES\": num_return_sequences,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
         "    }\n",
         "    # If the model_id is a GCS path, use artifact_uri to pass it to serving docker.\n",
         "    artifact_uri = model_id if model_id.startswith(\"gs://\") else None\n",

notebooks/community/model_garden/model_garden_pytorch_biomedclip.ipynb

@@ -278,7 +278,7 @@
         "    \"\"\"\n",
         "    !wget -O image.jpg $url\n",
         "    !base64 image.jpg > image.txt\n",
-        "    return open(\"image.txt\", \"r\").read()\n",
+        "    return open(\"image.txt\").read()\n",
         "\n",
         "\n",
         "def deploy_model(\n",
@@ -297,6 +297,7 @@
         "        \"MODEL\": model_id,\n",
         "        \"TASK\": task,\n",
         "        \"PRECISION\": precision,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
         "    }\n",
         "    # If the model_id is a GCS path, use artifact_uri to pass it to serving docker.\n",
         "    artifact_uri = model_id if model_id.startswith(\"gs://\") else None\n",

notebooks/community/model_garden/model_garden_pytorch_blip2.ipynb

@@ -273,6 +273,7 @@
         "    serving_env = {\n",
         "        \"MODEL_ID\": model_id,\n",
         "        \"TASK\": task,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
         "    }\n",
         "    # If the model_id is a GCS path, use artifact_uri to pass it to serving docker.\n",
         "    artifact_uri = model_id if model_id.startswith(\"gs://\") else None\n",

notebooks/community/model_garden/model_garden_pytorch_blip_image_captioning.ipynb

@@ -272,6 +272,7 @@
         "    serving_env = {\n",
         "        \"MODEL_ID\": model_id,\n",
         "        \"TASK\": task,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
         "    }\n",
         "    # If the model_id is a GCS path, use artifact_uri to pass it to serving docker.\n",
         "    artifact_uri = model_id if model_id.startswith(\"gs://\") else None\n",

notebooks/community/model_garden/model_garden_pytorch_blip_vqa.ipynb

@@ -272,6 +272,7 @@
         "    serving_env = {\n",
         "        \"MODEL_ID\": model_id,\n",
         "        \"TASK\": task,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
         "    }\n",
         "    # If the model_id is a GCS path, use artifact_uri to pass it to serving docker.\n",
         "    artifact_uri = model_id if model_id.startswith(\"gs://\") else None\n",

notebooks/community/model_garden/model_garden_pytorch_clip.ipynb

@@ -272,6 +272,7 @@
         "    serving_env = {\n",
         "        \"MODEL_ID\": model_id,\n",
         "        \"TASK\": task,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
         "    }\n",
         "    # If the model_id is a GCS path, use artifact_uri to pass it to serving docker.\n",
         "    artifact_uri = model_id if model_id.startswith(\"gs://\") else None\n",

notebooks/community/model_garden/model_garden_pytorch_codellama_evaluation.ipynb

@@ -0,0 +1,434 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ZwuKo6wnznjo"
+      },
+      "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": "RYqGQYXF0WaA"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Code LLaMA Evaluation\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_codellama_evaluation.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>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_codellama_evaluation.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",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "TlqQD2DM0lQS"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to evaluate the Code LLaMA models using EleutherAI's [Language Model Evaluation Harness (lm-evaluation-harness)](https://github.com/EleutherAI/lm-evaluation-harness) with Vertex CustomJob. Please reference the peak GPU memory usgaes for serving and adjust the machine type, accelerator type and accelerator count accordingly.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Evaluate Code LLaMA models\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0nC-ZtQY02Pw"
+      },
+      "source": [
+        "## Run the notebook"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "GNT-d4681YW-"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. 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\"). If not set, a unique GCS bucket will be created instead.\n",
+        "\n",
+        "import os\n",
+        "import sys\n",
+        "from datetime import datetime\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\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",
+        "# Cloud Storage bucket for storing the experiment artifacts.\n",
+        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
+        "# prefer using your own GCS bucket, please change the value yourself below.\n",
+        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    # Create a unique GCS bucket for this notebook if not specified\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
+        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
+        "else:\n",
+        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
+        "    bucket_region = shell_output[0].strip().lower()\n",
+        "    if bucket_region != REGION:\n",
+        "        raise ValueError(\n",
+        "            \"Bucket region %s is different from notebook region %s\"\n",
+        "            % (bucket_region, REGION)\n",
+        "        )\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "# Gets the default BUCKET_URI and SERVICE_ACCOUNT if they were not specified by the user.\n",
+        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
+        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "\n",
+        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
+        "\n",
+        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_URI\n",
+        "\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "! gcloud services enable language.googleapis.com\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "MODEL_BUCKET = os.path.join(STAGING_BUCKET, \"model\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
+        "\n",
+        "\n",
+        "def get_job_name_with_datetime(prefix: str) -> str:\n",
+        "    \"\"\"Gets the job 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",
+        "# The pre-built serving docker images.\n",
+        "EVAL_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-lm-evaluation-harness:20231011_0934_RC00\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "I7EvfUgRZPtd"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Access pretrained Code LLaMA models\n",
+        "\n",
+        "# @markdown The original models from Meta are converted into the HuggingFace format for serving in Vertex AI.\n",
+        "\n",
+        "# @markdown Accept the model agreement to access the models:\n",
+        "# @markdown 1. Open the [Code LLaMA model card](https://console.cloud.google.com/vertex-ai/publishers/google/model-garden/137).\n",
+        "# @markdown 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",
+        "# @markdown 3. A Cloud Storage bucket (starting with ‘gs://’) containing Code LLaMA pretrained and finetuned models will be shared under the “Documentation” section and its “Get started” subsection.\n",
+        "\n",
+        "# This path will be shared once click the agreement in Code LLaMA model card\n",
+        "# as described in the `Access pretrained Code LLaMA models` section.\n",
+        "VERTEX_AI_MODEL_GARDEN_CODE_LLAMA = \"\"  # @param {type: \"string\"}\n",
+        "assert (\n",
+        "    VERTEX_AI_MODEL_GARDEN_CODE_LLAMA\n",
+        "), \"Please click the agreement of Code LLaMA in Vertex AI Model Garden, and get the GCS path of Code LLaMA model artifacts.\"\n",
+        "print(\n",
+        "    \"Copying Code LLaMA model artifacts from\",\n",
+        "    VERTEX_AI_MODEL_GARDEN_CODE_LLAMA,\n",
+        "    \"to \",\n",
+        "    MODEL_BUCKET,\n",
+        ")\n",
+        "! gsutil -m cp -R $VERTEX_AI_MODEL_GARDEN_CODE_LLAMA/* $MODEL_BUCKET"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Yy1ldWUHy7By"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Evaluate Code LLaMA models\n",
+        "\n",
+        "# @markdown This section demonstrates how to evaluate the Code LLaMA models using EleutherAI's [Language Model Evaluation Harness (lm-evaluation-harness)](https://github.com/EleutherAI/lm-evaluation-harness) with Vertex CustomJob. Please reference the peak GPU memory usgaes for serving and adjust the machine type, accelerator type and accelerator count accordingly.\n",
+        "\n",
+        "# @markdown This example uses the dataset [GSM8K](https://arxiv.org/abs/2110.14168). All supported tasks are listed in [this task table](https://github.com/EleutherAI/lm-evaluation-harness/blob/master/docs/task_table.md).\n",
+        "\n",
+        "eval_dataset = \"gsm8k\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# @markdown Set the model name.\n",
+        "model_name = \"CodeLlama-7b-Instruct-hf\"  # @param [\"CodeLlama-7b-hf\", \"CodeLlama-7b-Python-hf\", \"CodeLlama-7b-Instruct-hf\", \"CodeLlama-13b-hf\", \"CodeLlama-13b-Python-hf\", \"CodeLlama-13b-Instruct-hf\", \"CodeLlama-34b-hf\", \"CodeLlama-34b-Python-hf\", \"CodeLlama-34b-Instruct-hf\", \"CodeLlama-70b-hf\", \"CodeLlama-70b-Python-hf\", \"CodeLlama-70b-Instruct-hf\"]\n",
+        "model_id = os.path.join(MODEL_BUCKET, model_name)\n",
+        "print(model_id)\n",
+        "\n",
+        "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
+        "\n",
+        "accelerator_type = \"NVIDIA_TESLA_V100\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_V100\", \"NVIDIA_TESLA_A100\"]\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "# Find Vertex AI supported accelerators and regions in:\n",
+        "#  https://cloud.google.com/vertex-ai/docs/training/configure-compute\n",
+        "\n",
+        "# Setup evaluation job.\n",
+        "job_name = get_job_name_with_datetime(prefix=\"code-llama-eval\")\n",
+        "eval_output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
+        "eval_output_dir_gcsfuse = eval_output_dir.replace(\"gs://\", \"/gcs/\")\n",
+        "model_id_gcsfuse = model_id.replace(\"gs://\", \"/gcs/\")\n",
+        "\n",
+        "if \"7b\" in model_id_gcsfuse:\n",
+        "    # Sets A100 (40G) to evaluate 7B models.\n",
+        "    if accelerator_type == \"NVIDIA_TESLA_A100\":\n",
+        "        machine_type = \"a2-highgpu-1g\"\n",
+        "        accelerator_count = 1\n",
+        "    # Sets 1 V100 (16G) to evaluate 7B models.\n",
+        "    elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
+        "        machine_type = \"n1-standard-8\"\n",
+        "        accelerator_count = 1\n",
+        "    # Sets 1 L4 (24G) to evaluate 7B models.\n",
+        "    elif accelerator_type == \"NVIDIA_L4\":\n",
+        "        machine_type = \"g2-standard-8\"\n",
+        "        accelerator_count = 1\n",
+        "    else:\n",
+        "        raise ValueError(\n",
+        "            f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
+        "        )\n",
+        "elif \"13b\" in model_id_gcsfuse:\n",
+        "    # Sets A100 (40G) to evaluate 13B models.\n",
+        "    if accelerator_type == \"NVIDIA_TESLA_A100\":\n",
+        "        machine_type = \"a2-highgpu-1g\"\n",
+        "        accelerator_count = 1\n",
+        "    # Sets 2 V100 (16G) to evaluate 13B models.\n",
+        "    elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
+        "        machine_type = \"n1-standard-16\"\n",
+        "        accelerator_count = 2\n",
+        "    # Sets 2 L4 (24G) to evaluate 13B models.\n",
+        "    elif accelerator_type == \"NVIDIA_L4\":\n",
+        "        machine_type = \"g2-standard-24\"\n",
+        "        accelerator_count = 2\n",
+        "    else:\n",
+        "        raise ValueError(\n",
+        "            f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
+        "        )\n",
+        "elif \"34b\" in model_id_gcsfuse:\n",
+        "    # Sets 2 A100 (40G) to evaluate 34B models.\n",
+        "    if accelerator_type == \"NVIDIA_TESLA_A100\":\n",
+        "        machine_type = \"a2-highgpu-2g\"\n",
+        "        accelerator_count = 2\n",
+        "    # Sets 8 V100 (16G) to evaluate 34B models.\n",
+        "    elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
+        "        machine_type = \"n1-standard-32\"\n",
+        "        accelerator_count = 8\n",
+        "    # Sets 4 L4 (24G) to evaluate 34B models.\n",
+        "    elif accelerator_type == \"NVIDIA_L4\":\n",
+        "        machine_type = \"g2-standard-48\"\n",
+        "        accelerator_count = 4\n",
+        "    else:\n",
+        "        raise ValueError(\n",
+        "            f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
+        "        )\n",
+        "elif \"70b\" in model_id_gcsfuse:\n",
+        "    # Sets 4 A100 (40G) to evaluate 70B models.\n",
+        "    if accelerator_type == \"NVIDIA_TESLA_A100\":\n",
+        "        machine_type = \"a2-highgpu-4g\"\n",
+        "        accelerator_count = 4\n",
+        "    # Sets 8 L4 (24G) to evaluate 70B models.\n",
+        "    elif accelerator_type == \"NVIDIA_L4\":\n",
+        "        machine_type = \"g2-standard-96\"\n",
+        "        accelerator_count = 8\n",
+        "    else:\n",
+        "        raise ValueError(\n",
+        "            f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
+        "        )\n",
+        "\n",
+        "replica_count = 1\n",
+        "\n",
+        "# Setup evaluation job.\n",
+        "job_name = get_job_name_with_datetime(prefix=\"code-llama-eval\")\n",
+        "eval_output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
+        "eval_output_dir_gcsfuse = eval_output_dir.replace(\"gs://\", \"/gcs/\")\n",
+        "model_id_gcsfuse = model_id.replace(\"gs://\", \"/gcs/\")\n",
+        "\n",
+        "# Prepare evaluation command that runs the evaluation harness.\n",
+        "# Set `trust_remote_code = True` because evaluating the model requires\n",
+        "# executing code from the model repository.\n",
+        "# Set `use_accelerate = True` to enable evaluation across multiple GPUs.\n",
+        "eval_command = [\n",
+        "    \"python\",\n",
+        "    \"main.py\",\n",
+        "    \"--model\",\n",
+        "    \"hf-causal-experimental\",\n",
+        "    \"--model_args\",\n",
+        "    f\"pretrained={model_id_gcsfuse},trust_remote_code=True,use_accelerate=True,device_map_option=auto\",\n",
+        "    \"--tasks\",\n",
+        "    f\"{eval_dataset}\",\n",
+        "    \"--output_path\",\n",
+        "    f\"{eval_output_dir_gcsfuse}\",\n",
+        "]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "ZZIqS9vkzPXY"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Submit evaluation CustomJob\n",
+        "\n",
+        "# @markdown Pass evaluation arguments and launch job.\n",
+        "\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": machine_type,\n",
+        "            \"accelerator_type\": accelerator_type,\n",
+        "            \"accelerator_count\": accelerator_count,\n",
+        "        },\n",
+        "        \"replica_count\": replica_count,\n",
+        "        \"disk_spec\": {\n",
+        "            \"boot_disk_size_gb\": 500,\n",
+        "        },\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": EVAL_DOCKER_URI,\n",
+        "            \"command\": eval_command,\n",
+        "            \"args\": [],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "eval_job = aiplatform.CustomJob(\n",
+        "    display_name=job_name,\n",
+        "    worker_pool_specs=worker_pool_specs,\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": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "1AnAEuOAzTe9"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Fetch and print evaluation results\n",
+        "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}\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "W8yhtUxTzVCF"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Clean up resources\n",
+        "\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_URI\n",
+        "    # Uncomment below to delete all artifacts\n",
+        "    # !gsutil -m rm -r $STAGING_BUCKET $MODEL_BUCKET\n",
+        "\n",
+        "\n",
+        "# Delete evaluation job.\n",
+        "eval_job.delete()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_pytorch_codellama_evaluation.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_pytorch_controlnet.ipynb

@@ -8,7 +8,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2023 Google LLC\n",
+        "# 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",
@@ -33,23 +33,16 @@
         "\n",
         "<table align=\"left\">\n",
         "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_controlnet.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\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_controlnet.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>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_controlnet.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br>\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
-        "  <td>\n",
-        "    <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_controlnet.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>\n",
-        "    (a Python-3 CPU notebook is recommended)\n",
-        "  </td>\n",
         "</table>"
       ]
     },
@@ -61,11 +54,10 @@
       "source": [
         "## Overview\n",
         "\n",
-        "This notebook demonstrates finetuning the [ControlNet](https://huggingface.co/lllyasviel/ControlNet) with the [fusing/fill50k](https://huggingface.co/datasets/fusing/fill50k) dataset and deploying the model on Vertex AI for online prediction.\n",
+        "This notebook demonstrates deploying the [ControlNet](https://huggingface.co/lllyasviel/ControlNet) model on Vertex AI for online prediction.\n",
         "\n",
         "### Objective\n",
         "\n",
-        "- Finetune the ControlNet model.\n",
         "- Upload the model to [Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction).\n",
         "- Deploy the model on [Endpoint](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
         "- Run online predictions for text-guided-image-to-image.\n",
@@ -86,220 +78,92 @@
         "id": "264c07757582"
       },
       "source": [
-        "## Setup environment\n",
-        "\n",
-        "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "d73ffa0c0b83"
-      },
-      "source": [
-        "### Colab only"
+        "## Run the notebook"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
+        "cellView": "form",
         "id": "2707b02ef5df"
       },
       "outputs": [],
       "source": [
-        "!pip3 install --upgrade google-cloud-aiplatform"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "b60a4d7100bf"
-      },
-      "outputs": [],
-      "source": [
-        "from google.colab import auth as google_auth\n",
+        "# @title Setup Google Cloud project\n",
         "\n",
-        "google_auth.authenticate_user()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "fb671e75ca7b"
-      },
-      "source": [
-        "### Install dependencies"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "dc8ee367fb42"
-      },
-      "outputs": [],
-      "source": [
-        "# Install gdown for downloading example training images.\n",
-        "!pip install gdown\n",
-        "# Install libs for generating conditioning images for ControlNet.\n",
-        "!pip install opencv-python"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "5244aac3d929"
-      },
-      "source": [
-        "Restart the notebook kernel after installs."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "567212ff53a6"
-      },
-      "outputs": [],
-      "source": [
-        "import IPython\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
         "\n",
-        "app = IPython.Application.instance()\n",
-        "app.kernel.do_shutdown(True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "bb7adab99e41"
-      },
-      "source": [
-        "### Setup Google Cloud project\n",
+        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. 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\"). If not set, a unique GCS bucket will be created instead.\n",
         "\n",
-        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
-        "\n",
-        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
-        "\n",
-        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
-        "\n",
-        "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs.\n",
-        "\n",
-        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User` and `Storage Object Admin` roles for deploying fine tuned model to Vertex AI endpoint."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6c460088b873"
-      },
-      "source": [
-        "Fill following variables for experiments environment:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "855d6b96f291"
-      },
-      "outputs": [],
-      "source": [
-        "# Cloud project id.\n",
-        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# The region you want to launch jobs in.\n",
-        "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# The Cloud Storage bucket for storing experiments output. Fill it without the 'gs://' prefix.\n",
-        "GCS_BUCKET = \"\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# The service account you created in step-5 above, it's like \"<account_name>@<project>.iam.gserviceaccount.com\"\n",
-        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "e828eb320337"
-      },
-      "source": [
-        "Initialize Vertex AI API:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "12cd25839741"
-      },
-      "outputs": [],
-      "source": [
-        "from google.cloud import aiplatform\n",
-        "\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=GCS_BUCKET)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "2cc825514deb"
-      },
-      "source": [
-        "### Define constants"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "b42bd4fa2b2d"
-      },
-      "outputs": [],
-      "source": [
-        "# The pre-built training docker image. It contains training scripts and models.\n",
-        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-train:latest\"\n",
-        "\n",
-        "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\""
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0c250872074f"
-      },
-      "source": [
-        "### Define common functions"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "354da31189dc"
-      },
-      "outputs": [],
-      "source": [
         "import base64\n",
         "import os\n",
+        "import sys\n",
+        "import uuid\n",
         "from datetime import datetime\n",
         "from io import BytesIO\n",
         "\n",
         "import cv2\n",
         "import numpy as np\n",
         "import requests\n",
+        "from google.cloud import aiplatform\n",
         "from PIL import Image\n",
         "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
         "\n",
-        "def create_job_name(prefix):\n",
-        "    user = os.environ.get(\"USER\")\n",
-        "    now = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
-        "    job_name = f\"{prefix}-{user}-{now}\"\n",
-        "    return job_name\n",
+        "# Get the default region for launching jobs.\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
+        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
+        "\n",
+        "# Cloud Storage bucket for storing the experiment artifacts.\n",
+        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
+        "# prefer using your own GCS bucket, please change the value yourself below.\n",
+        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "BUCKET_URI = \"gs://\"  # @param {type: \"string\"}\n",
+        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "\n",
+        "# Create a unique GCS bucket for this notebook, if not specified by the user.\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
+        "else:\n",
+        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
+        "    bucket_region = shell_output[0].strip().lower()\n",
+        "    if bucket_region != REGION:\n",
+        "        raise ValueError(\n",
+        "            \"Bucket region %s is different from notebook region %s.\"\n",
+        "            % (bucket_region, REGION)\n",
+        "        )\n",
+        "\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "# Set up the default SERVICE_ACCOUNT.\n",
+        "SERVICE_ACCOUNT = None\n",
+        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
+        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "\n",
+        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
+        "\n",
+        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
         "\n",
         "\n",
+        "# The pre-built serving docker image. It contains serving scripts and models.\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
+        "\n",
+        "\n",
+        "# Define common functions.\n",
         "def download_image(url):\n",
         "    response = requests.get(url)\n",
         "    return Image.open(BytesIO(response.content))\n",
@@ -319,9 +183,9 @@
         "\n",
         "def image_grid(imgs, rows=2, cols=2):\n",
         "    w, h = imgs[0].size\n",
-        "    grid = Image.new(\"RGB\", size=(cols * w, rows * h))\n",
+        "    grid = Image.new(\"RGB\", size=(cols * w, rows * h), color=(255, 255, 255))\n",
         "    for i, img in enumerate(imgs):\n",
-        "        grid.paste(img, box=(i % cols * w, i // cols * h))\n",
+        "        grid.paste(img, box=(i % cols * w + 10 * i, i // cols * h))\n",
         "    return grid\n",
         "\n",
         "\n",
@@ -340,6 +204,7 @@
         "    serving_env = {\n",
         "        \"MODEL_ID\": model_id,\n",
         "        \"TASK\": task,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
         "    }\n",
         "    model = aiplatform.Model.upload(\n",
         "        display_name=model_name,\n",
@@ -351,8 +216,8 @@
         "    )\n",
         "    model.deploy(\n",
         "        endpoint=endpoint,\n",
-        "        machine_type=\"n1-standard-8\",\n",
-        "        accelerator_type=\"NVIDIA_TESLA_V100\",\n",
+        "        machine_type=\"g2-standard-8\",\n",
+        "        accelerator_type=\"NVIDIA_L4\",\n",
         "        accelerator_count=1,\n",
         "        deploy_request_timeout=1800,\n",
         "        service_account=SERVICE_ACCOUNT,\n",
@@ -360,288 +225,88 @@
         "    return model, endpoint"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "e70e3519ff8b"
-      },
-      "source": [
-        "## Finetune with fill50k dataset"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "0dc65d8f0689"
-      },
-      "source": [
-        "This section uses the [fusing/fill50k](https://huggingface.co/datasets/fusing/fill50k) dataset to finetune the ControlNet model.\n",
-        "\n",
-        "The job will run on 1 A100 GPU and take ~7 hours to finish 1 epoch of training.\n",
-        "\n",
-        "The ControlNet model will be saved after the finetuning job finishs and it can be loaded to run inference later."
-      ]
-    },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "65467b361315"
+        "cellView": "form",
+        "id": "a9bbe0f7c237"
       },
       "outputs": [],
       "source": [
-        "# The pre-trained stable diffusion model to be loaded.\n",
-        "stable_diffusion_model_id = \"runwayml/stable-diffusion-v1-5\"\n",
-        "# The datase id to be loaded.\n",
-        "dataset_id = \"fusing/fill50k\"\n",
-        "# The output path.\n",
-        "output_dir = f\"/gcs/{GCS_BUCKET}/controlnet/output\"\n",
-        "# The training steps. If provided, it overrides num_train_epochs. Set it to None to use num_train_epochs only.\n",
-        "max_train_steps = 10  # @param {type:\"integer\"}\n",
-        "# The training epochs. Set it to a bigger number, like 10, to make training converge better.\n",
-        "num_train_epochs = 1  # @param {type:\"integer\"}\n",
+        "# @title Upload and deploy model\n",
         "\n",
-        "# Worker pool spec.\n",
-        "machine_type = \"a2-highgpu-1g\"\n",
-        "num_nodes = 1\n",
-        "gpu_type = \"NVIDIA_TESLA_A100\"\n",
-        "num_gpus = 1\n",
+        "# @markdown This step deploys the pre-trained [lllyasviel/sd-controlnet-canny](https://huggingface.co/lllyasviel/sd-controlnet-canny) model for the text-guided image-to-image task.\n",
         "\n",
-        "# Setup training job.\n",
-        "job_name = create_job_name(\"controlnet\")\n",
-        "job = aiplatform.CustomContainerTrainingJob(\n",
-        "    display_name=job_name,\n",
-        "    container_uri=TRAIN_DOCKER_URI,\n",
-        ")\n",
+        "# @markdown The model deployment step will take ~15 minutes to complete.\n",
         "\n",
-        "# Pass training arguments and launch job.\n",
-        "# See https://github.com/huggingface/diffusers/blob/main/examples/controlnet/train_controlnet.py\n",
-        "# for a full list of training arguments.\n",
-        "model = job.run(\n",
-        "    args=[\n",
-        "        f\"--num_machines={num_nodes}\",\n",
-        "        f\"--num_processes={num_gpus}\",\n",
-        "        \"--machine_rank=0\",\n",
-        "        \"--mixed_precision=no\",\n",
-        "        \"--gpu_ids=all\",\n",
-        "        \"--same_network\",\n",
-        "        \"--dynamo_backend=no\",\n",
-        "        \"controlnet/train_controlnet.py\",\n",
-        "        \"--tracker_project_name=train_controlnet\",\n",
-        "        f\"--pretrained_model_name_or_path={stable_diffusion_model_id}\",\n",
-        "        f\"--output_dir={output_dir}\",\n",
-        "        f\"--dataset_name={dataset_id}\",\n",
-        "        f\"--max_train_steps={max_train_steps}\",\n",
-        "        f\"--num_train_epochs={num_train_epochs}\",\n",
-        "        \"--resolution=512\",\n",
-        "        \"--learning_rate=1e-5\",\n",
-        "        \"--train_batch_size=2\",\n",
-        "        \"--checkpointing_steps=50000\",\n",
-        "        \"--checkpoints_total_limit=1\",\n",
-        "    ],\n",
-        "    replica_count=num_nodes,\n",
-        "    machine_type=machine_type,\n",
-        "    accelerator_type=gpu_type,\n",
-        "    accelerator_count=num_gpus,\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "bf7f82732e61"
-      },
-      "source": [
-        "## Upload and Deploy models"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "1cc26e68d7b0"
-      },
-      "source": [
-        "This section uploads the model to Model Registry and deploys it on the Endpoint.\n",
-        "\n",
-        "The model deployment step will take ~15 minutes to complete."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "cd7b56421392"
-      },
-      "source": [
-        "### Pre-trained canny model"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "6d331b1ea337"
-      },
-      "source": [
-        "Deploy the pre-trained [lllyasviel/sd-controlnet-canny](https://huggingface.co/lllyasviel/sd-controlnet-canny) model for the text-guided image-to-image task. When deployed on one V100 GPU, the average inference time of a request is ~15 seconds."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "bf55e38815dc"
-      },
-      "outputs": [],
-      "source": [
         "model, endpoint = deploy_model(\n",
         "    model_id=\"lllyasviel/sd-controlnet-canny\", task=\"controlnet\"\n",
         ")"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "80b3fd2ace09"
-      },
-      "source": [
-        "NOTE: The model weights will be downloaded 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."
-      ]
-    },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "4ab04da3ec9a"
+        "cellView": "form",
+        "id": "9c1c39133dd1"
       },
       "outputs": [],
       "source": [
-        "init_image = download_image(\n",
-        "    \"https://huggingface.co/takuma104/controlnet_dev/resolve/main/gen_compare/output_images/diffusers/output_bird_canny_1.png\"\n",
-        ")\n",
-        "display(init_image)\n",
-        "image = canny(init_image)\n",
-        "display(image)\n",
+        "# @title Predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompt and canny image.\n",
+        "\n",
+        "# @markdown When deployed on one L4 GPU (the default machine type), the averaged inference time of a request is ~15 seconds.\n",
+        "\n",
+        "# @markdown You may adjust the parameters below to achieve best image quality.\n",
+        "\n",
+        "prompt = \"bird\"  # @param {type: \"string\"}\n",
+        "image = \"https://huggingface.co/takuma104/controlnet_dev/resolve/main/gen_compare/output_images/diffusers/output_bird_canny_1.png\"  # @param {type: \"string\"}\n",
+        "num_inference_steps = 25  # @param {type:\"number\"}\n",
+        "\n",
+        "init_image = download_image(image)\n",
+        "canny_image = canny(init_image)\n",
         "\n",
         "instances = [\n",
         "    {\n",
-        "        \"prompt\": \"bird\",\n",
-        "        \"image\": image_to_base64(image),\n",
+        "        \"prompt\": prompt,\n",
+        "        \"image\": image_to_base64(canny_image),\n",
+        "        \"num_inference_steps\": num_inference_steps,\n",
         "    },\n",
         "]\n",
         "response = endpoint.predict(instances=instances)\n",
         "images = [base64_to_image(image) for image in response.predictions]\n",
-        "display(images[0])"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "af21a3cff1e0"
-      },
-      "source": [
-        "Clean up resources:"
+        "new_image = images[0]\n",
+        "\n",
+        "image_grid([init_image, canny_image, new_image], rows=1, cols=3)"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "911406c1561e"
+        "cellView": "form",
+        "id": "7b827b2370bd"
       },
       "outputs": [],
       "source": [
+        "# @title Clean up resources\n",
+        "\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
+        "\n",
         "# Undeploy model and delete endpoint.\n",
         "endpoint.delete(force=True)\n",
         "\n",
         "# Delete models.\n",
-        "model.delete()"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "c1e51f764a60"
-      },
-      "source": [
-        "### Custom finetuned fill50k model"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "fa686a54047c"
-      },
-      "source": [
-        "Deploy the finetuned fill50k model above for the text-guided image-to-image task. When deployed on one V100 GPU, the averaged inference time of a request is ~15 seconds."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "65e32356fbd1"
-      },
-      "outputs": [],
-      "source": [
-        "model, endpoint = deploy_model(\n",
-        "    model_id=f\"gs://{GCS_BUCKET}/controlnet/output\", task=\"controlnet\"\n",
-        ")"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "80b3fd2ace09"
-      },
-      "source": [
-        "NOTE: The model weights will be downloaded 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."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "83a50fd4a1ed"
-      },
-      "outputs": [],
-      "source": [
-        "init_image = download_image(\n",
-        "    \"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_training/conditioning_image_1.png\"\n",
-        ")\n",
-        "display(init_image)\n",
+        "model.delete()\n",
         "\n",
-        "instances = [\n",
-        "    {\n",
-        "        \"prompt\": \"red circle with green background\",\n",
-        "        \"image\": image_to_base64(init_image, format=\"PNG\"),\n",
-        "    },\n",
-        "]\n",
-        "response = endpoint.predict(instances=instances)\n",
-        "images = [base64_to_image(image) for image in response.predictions]\n",
-        "display(images[0])"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "ed3795d474b9"
-      },
-      "source": [
-        "Clean up resources:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "b53b883257b4"
-      },
-      "outputs": [],
-      "source": [
-        "# Undeploy model and delete endpoint.\n",
-        "endpoint.delete(force=True)\n",
-        "\n",
-        "# Delete models.\n",
-        "model.delete()"
+        "# Delete bucket.\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_NAME"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_pytorch_detectron2.ipynb

@@ -301,6 +301,7 @@
         "        \"MODEL_PTH_FILE\": model_pth_file,\n",
         "        \"CONFIG_YAML_FILE\": model_cfg_yaml_file,\n",
         "        \"TEST_THRESHOLD\": test_threshold,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
         "    }\n",
         "\n",
         "    model = aiplatform.Model.upload(\n",
@@ -475,7 +476,7 @@
         "    )\n",
         "    try:\n",
         "        font = ImageFont.truetype(\"arial.ttf\", 24)\n",
-        "    except IOError:\n",
+        "    except OSError:\n",
         "        font = ImageFont.load_default()\n",
         "\n",
         "    # If the total height of the display strings added to the top of the bounding\n",