update: post-Next

* feat: RLHF bulk infer

* fix: review comments

* fix: review comment

* fix: review comment

* update: text-bison@001

* fix: revert change

* fix: colab endpoint

.cloud-build/cleanup/cleanup.py

@@ -1,5 +1,14 @@
 from typing import List
 from ratemate import RateLimit
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--dry_run",
+                    type=bool,
+                    default=False)
+args = parser.parse_args()
+
+
 from resource_cleanup_manager import (
     DatasetResourceCleanupManager,
     ModelResourceCleanupManager,
@@ -7,6 +16,14 @@ from resource_cleanup_manager import (
     ResourceCleanupManager,
     MatchingEngineIndexEndpointResourceCleanupManager,
     MatchingEngineIndexResourceCleanupManager,
+    FeatureStoreCleanupManager,
+    PipelineJobCleanupManager,
+    TrainingJobCleanupManager,
+    HyperparameterTuningCleanupManager,
+    BatchPredictionJobCleanupManager,
+    ExperimentCleanupManager,
+    BucketCleanupManager,
+    ArtifactRegistryCleanupManager
 )
 
 rate_limit = RateLimit(max_count=25, per=60, greedy=False)
@@ -23,7 +40,6 @@ def run_cleanup_managers(managers: List[ResourceCleanupManager], is_dry_run: boo
             try:
                 if not manager.is_deletable(resource):
                     continue
-
                 if is_dry_run:
                     resource_name = manager.resource_name(resource)
                     print(f"Will delete '{type_name}': {resource_name}")
@@ -36,9 +52,7 @@ def run_cleanup_managers(managers: List[ResourceCleanupManager], is_dry_run: boo
         print("")
 
 
-is_dry_run = False
-
-if is_dry_run:
+if args.dry_run:
     print("Starting cleanup in dry run mode...")
 
 # List of all cleanup managers
@@ -48,6 +62,14 @@ managers: List[ResourceCleanupManager] = [
     ModelResourceCleanupManager(),  # ModelResourceCleanupManager must follow EndpointResourceCleanupManager due to deployed models blocking model deletion.
     MatchingEngineIndexEndpointResourceCleanupManager(),
     MatchingEngineIndexResourceCleanupManager(),
+    FeatureStoreCleanupManager(),
+    PipelineJobCleanupManager(),
+    TrainingJobCleanupManager(),
+    HyperparameterTuningCleanupManager(),
+    BatchPredictionJobCleanupManager(),
+    ExperimentCleanupManager(), # Experiment missing _resource_noun
+    BucketCleanupManager(),
+    ArtifactRegistryCleanupManager()
 ]
 
-run_cleanup_managers(managers=managers, is_dry_run=is_dry_run)
+run_cleanup_managers(managers=managers, is_dry_run=args.dry_run)

.cloud-build/cleanup/resource_cleanup_manager.py

@@ -1,8 +1,18 @@
+'''
+READ FIRST BEFORE MAKING CHANGES
+-  Create a convention for resources created from vertex-ai-samples GH. We already have one IIRC
+- Only delete those objects as part of our clean-up script.
+- Don't run any tests on python-docs-samples-tests project, especially ones that affect resources created outside of our purview
+- Add --dry-run option to the clean-up script. This option will just output the list of resources the script will delete instead of actually deleting the resources.
+- Have a larger conversation in DEE before touching any resources that were not created as part of vertex-ai-samples
+'''
+import os
 import abc
 from typing import Any, Type
 
 from google.cloud import aiplatform
 from google.cloud.aiplatform import base
+from google.cloud import storage
 from proto.datetime_helpers import DatetimeWithNanoseconds
 
 # If a resource was updated within this number of seconds, do not delete.
@@ -69,7 +79,7 @@ class VertexAIResourceCleanupManager(ResourceCleanupManager):
     def delete(self, resource):
         resource.delete()
 
-    def get_seconds_since_modification(self, resource: Any) -> bool:
+    def get_seconds_since_modification(self, resource: Any) -> float:
         update_time = resource.update_time
         current_time = DatetimeWithNanoseconds.now(tz=update_time.tzinfo)
         return (current_time - update_time).total_seconds()
@@ -97,13 +107,10 @@ class EndpointResourceCleanupManager(VertexAIResourceCleanupManager):
     vertex_ai_resource = aiplatform.Endpoint
 
     def delete(self, resource):
-        # TODO: Remove this once https://github.com/googleapis/python-aiplatform/issues/1441 is fixed
-        resource._sync_gca_resource()
         for deployed_model_id in [
             models.id for models in resource._gca_resource.deployed_models
         ]:
             resource._undeploy(deployed_model_id=deployed_model_id)
-
         resource.delete(force=True)
 
 
@@ -117,3 +124,137 @@ class MatchingEngineIndexResourceCleanupManager(VertexAIResourceCleanupManager):
 
 class MatchingEngineIndexEndpointResourceCleanupManager(VertexAIResourceCleanupManager):
     vertex_ai_resource = aiplatform.MatchingEngineIndexEndpoint
+
+    def delete(self, resource):
+        resource.undeploy_all()
+        resource.delete(force=True)
+
+class FeatureStoreCleanupManager(VertexAIResourceCleanupManager):
+    vertex_ai_resource = aiplatform.Featurestore
+
+    def resource_name(self, resource: Any) -> str:
+        return resource.name
+
+class PipelineJobCleanupManager(VertexAIResourceCleanupManager):
+    vertex_ai_resource = aiplatform.PipelineJob
+
+class TrainingJobCleanupManager(VertexAIResourceCleanupManager):
+    vertex_ai_resource = aiplatform.training_jobs._CustomTrainingJob
+
+    job_types = [
+            aiplatform.AutoMLImageTrainingJob,
+            aiplatform.AutoMLTextTrainingJob,
+            aiplatform.AutoMLTabularTrainingJob,
+            aiplatform.AutoMLVideoTrainingJob,
+            aiplatform.AutoMLForecastingTrainingJob,
+            aiplatform.CustomJob,
+            aiplatform.CustomTrainingJob,
+            aiplatform.CustomContainerTrainingJob,
+            aiplatform.CustomPythonPackageTrainingJob
+    ]
+
+    def list(self) -> Any:
+        return [
+            job
+            for job_type in self.job_types
+            for job in job_type.list()
+        ]
+
+class HyperparameterTuningCleanupManager(VertexAIResourceCleanupManager):
+    vertex_ai_resource = aiplatform.HyperparameterTuningJob
+
+
+class BatchPredictionJobCleanupManager(VertexAIResourceCleanupManager):
+    vertex_ai_resource = aiplatform.BatchPredictionJob
+
+class ExperimentCleanupManager(VertexAIResourceCleanupManager):
+    vertex_ai_resource = aiplatform.Experiment
+
+    @property
+    def type_name(self) -> str:
+        return "Experiment"
+
+    def resource_name(self, resource: Any) -> str:
+        return resource.name
+
+    def get_seconds_since_modification(self, resource: Any) -> float:
+        update_time = resource._metadata_context.update_time
+        current_time = DatetimeWithNanoseconds.now()
+        return float(current_time.timestamp() - update_time.timestamp())
+
+class BucketCleanupManager(ResourceCleanupManager):
+    vertex_ai_resource = storage.bucket.Bucket
+
+    def list(self) -> Any:
+        storage_client = storage.Client()
+        return list(storage_client.list_buckets())
+
+    def delete(self, resource):
+        try:
+            resource.delete(force=True)
+        except Exception as e:
+            print(e)
+
+    @property
+    def type_name(self) -> str:
+        return "Bucket"
+
+    def get_seconds_since_modification(self, resource: Any) -> float:
+        # Bucket has no last_update property, only time created
+        created_time = resource.time_created
+        current_time = DatetimeWithNanoseconds.now()
+        return float(current_time.timestamp() - created_time.timestamp())
+
+    def resource_name(self, resource: Any) -> str:
+        return resource.name
+
+    def is_deletable(self, resource: Any) -> bool:
+        time_difference = self.get_seconds_since_modification(resource)
+
+        if not self.resource_name(resource).startswith('your-bucket-name'):
+            print(f"Skipping '{resource}' not a Vertex AI notebook bucket")
+            return False
+
+        # Check that it wasn't created too recently, to prevent race conditions
+        if time_difference <= RESOURCE_UPDATE_BUFFER_IN_SECONDS:
+            print(
+                f"Skipping '{resource}' due to update_time being '{time_difference}', which is less than '{RESOURCE_UPDATE_BUFFER_IN_SECONDS}'."
+            )
+            return False
+        return True
+
+class ArtifactRegistryCleanupManager(ResourceCleanupManager):
+    vertex_ai_resource = "Artifact Registry"
+
+    def list(self) -> Any:
+        import subprocess
+
+        result = subprocess.run(["gcloud artifacts repositories list --location=us-central1"], 
+                                shell=True, capture_output=True, text=True)
+
+        ret = []
+        lines = result.stdout.split('\n')[2:]
+        for line in lines:
+            repo = line.split(' ')[0]
+            if repo.startswith("my-docker-repo"):
+                ret.append(repo)
+
+        return ret
+
+    def delete(self, resource):
+        os.system(f"! gcloud artifacts repositories delete {resource} --location=us-central1")
+
+    @property
+    def type_name(self) -> str:
+        return "ArtifactRepository"
+
+    def resource_name(self, resource: Any) -> str:
+        return resource
+
+    # delete repository regardless of age
+    def get_seconds_since_modification(self, resource: Any) -> float:
+        return RESOURCE_UPDATE_BUFFER_IN_SECONDS + 1
+    
+    def is_deleteable(self, resource: Any) -> bool:
+        return True
+

.cloud-build/execute_changed_notebooks_cli.py

@@ -17,6 +17,7 @@
 
 import argparse
 import pathlib
+import os
 
 import execute_changed_notebooks_helper
 
@@ -39,6 +40,19 @@ parser.add_argument(
     help="The path to the file that has newline-limited folders of notebooks that should be tested.",
     required=True,
 )
+parser.add_argument(
+    "--test_percent",
+    type=int,
+    help="The percent of notebooks to be tested (between 1 and 100).",
+    required=False,
+    default=100,
+)
+parser.add_argument(
+    "--build_id",
+    type=str,
+    help="The build id (which may be a Cloud Build job specific or user explicit.",
+    required=True
+)
 parser.add_argument(
     "--base_branch",
     help="The base git branch to diff against to find changed files.",
@@ -107,24 +121,60 @@ parser.add_argument(
     default=True,
     help="Should run notebooks in parallel.",
 )
+parser.add_argument(
+    "--concurrent_notebooks",
+    type=int,
+    help="Maximum number of parallel notebook executions per minute",
+    default=10,
+    required=False,
+)
+parser.add_argument(
+    "--dry_run",
+    type=str2bool,
+    default=False,
+    help="Dry run for testing - no execution",
+)
 
 args = parser.parse_args()
 
-notebooks = execute_changed_notebooks_helper.get_changed_notebooks(
+changed_notebooks = execute_changed_notebooks_helper.get_changed_notebooks(
     test_paths_file=args.test_paths_file,
     base_branch=args.base_branch,
 )
 
-execute_changed_notebooks_helper.process_and_execute_notebooks(
-    notebooks=notebooks,
-    container_uri=args.container_uri,
-    staging_bucket=args.staging_bucket,
-    artifacts_bucket=args.artifacts_bucket,
-    should_parallelize=args.should_parallelize,
-    timeout=args.timeout,
-    variable_project_id=args.variable_project_id,
-    variable_region=args.variable_region,
-    variable_service_account=args.variable_service_account,
-    variable_vpc_network=args.variable_vpc_network,
-    private_pool_id=args.private_pool_id,
+
+results_bucket = f"{args.artifacts_bucket}"
+# artifacts_bucket may get set by trigger to a full gs:// folder path
+if results_bucket.startswith("gs://"):
+    results_bucket = results_bucket[5:]
+results_bucket = results_bucket.split('/')[0]
+results_file = f"build_results/{args.build_id}.json"
+
+if args.test_percent == 100:
+    notebooks = changed_notebooks
+    accumulative_results = {}
+else:
+    accumulative_results = execute_changed_notebooks_helper.load_results(results_bucket, results_file)
+
+    notebooks = [changed_notebook for changed_notebook in changed_notebooks if execute_changed_notebooks_helper.select_notebook(changed_notebook, accumulative_results, args.test_percent)]
+
+if args.dry_run:
+    print("Dry run ...\n")
+    for notebook in notebooks:
+        print(f"Would execute: {notebook}")
+else:
+    execute_changed_notebooks_helper.process_and_execute_notebooks(
+        notebooks=notebooks,
+        container_uri=args.container_uri,
+        staging_bucket=args.staging_bucket,
+        artifacts_bucket=args.artifacts_bucket,
+        results_file=results_file,
+        should_parallelize=args.should_parallelize,
+        timeout=args.timeout,
+        variable_project_id=args.variable_project_id,
+        variable_region=args.variable_region,
+        variable_service_account=args.variable_service_account,
+        variable_vpc_network=args.variable_vpc_network,
+        private_pool_id=args.private_pool_id,
+        concurrent_notebooks=args.concurrent_notebooks,
 )

.cloud-build/execute_changed_notebooks_helper.py

@@ -21,18 +21,21 @@ import json
 import git
 import operator
 import os
+import io
+import json
 import pathlib
 import re
 import subprocess
+import random
+from google.cloud import storage
 import utils
-from typing import List, Optional
+from typing import List, Optional, Dict, Any
 from utils import util
 
 import execute_notebook_helper
 import execute_notebook_remote
 import nbformat
 from google.cloud.devtools.cloudbuild_v1.types import BuildOperationMetadata
-from ratemate import RateLimit
 from tabulate import tabulate
 from utils import NotebookProcessors, util
 
@@ -40,6 +43,9 @@ from utils import NotebookProcessors, util
 WORKER_TIMEOUT_BUFFER_IN_SECONDS: int = 60 * 60
 PYTHON_VERSION = "3.9"  # Set default python version
 
+# rolling time window for accumulating build results for selecting notebooks
+MAX_RESULTS_AGE_SECONDS: int = (60 * 60) * 24 * 60  # 60 days
+
 
 def format_timedelta(delta: datetime.timedelta) -> str:
     """Formats a timedelta duration to [N days] %H:%M:%S format"""
@@ -65,7 +71,9 @@ def format_timedelta(delta: datetime.timedelta) -> str:
 @dataclasses.dataclass
 class NotebookExecutionResult:
     name: str
+    path: str
     duration: datetime.timedelta
+    start_time: datetime.datetime
     is_pass: bool
     log_url: str
     output_uri: str
@@ -81,6 +89,75 @@ class NotebookExecutionResult:
             return None
 
 
+def load_results(results_bucket: str,
+                 results_file: str) -> Dict[str, Any]:
+    '''
+    Load accumulated notebook test results
+    '''
+
+    print("Loading existing accumulative results ...")
+    accumulative_results = {}
+    try:
+        client = storage.Client()
+        bucket = client.bucket(results_bucket)
+
+        build_results_dir = os.path.dirname(results_file)
+        blobs = client.list_blobs(results_bucket, prefix=build_results_dir)
+        for blob in blobs:
+            time_created = blob.time_created.replace(tzinfo=None)
+            if (datetime.datetime.now().replace(tzinfo=None) - time_created).total_seconds() > MAX_RESULTS_AGE_SECONDS:
+                continue
+
+            content = util.download_blob_into_memory(results_bucket, blob.name, download_as_text=True)
+
+            try:
+                build_results = json.loads(content)
+            except:
+                continue  # skip corrupted build results files
+            for notebook in build_results:
+                if notebook in accumulative_results:
+                    accumulative_results[notebook]['passed'] += build_results[notebook]['passed']
+                    accumulative_results[notebook]['failed'] += build_results[notebook]['failed']
+                else:
+                    accumulative_results[notebook] = build_results[notebook]
+
+        print(accumulative_results)
+    except Exception as e:
+        print(e)
+
+    # If there are no accumulative results, an empty dict is returned
+    return accumulative_results
+
+def select_notebook(changed_notebook: str,
+                    accumulative_results: Dict[str, Any],
+                    test_percent: int) -> bool:
+    '''
+    Algorithm to randomly select a notebook, but weight the propbability of selected based on past failures
+    '''
+
+    if changed_notebook in accumulative_results:
+        pass_count = accumulative_results[changed_notebook]['passed']
+        fail_count = accumulative_results[changed_notebook]['failed']
+    else:
+        pass_count = 1
+        fail_count = 0
+
+    inferred_failure_rate = fail_count / (pass_count + fail_count)
+
+    # If failure rate is high, the chance of testing should be higher
+    should_test_due_to_failure = random.uniform(0, 1) <= inferred_failure_rate
+
+    # Additionally, only test a percentage of these
+    should_test_due_to_random_subset = random.uniform(0, 1) <= (test_percent / 100)
+
+    if should_test_due_to_failure or should_test_due_to_random_subset:
+        print(f"Selected: {changed_notebook}, {should_test_due_to_failure}, {should_test_due_to_random_subset}")
+        return True
+    else:
+        print(f"Not Selected: {changed_notebook}, pass {pass_count}, fail {fail_count}")
+        return False
+
+
 def _process_notebook(
     notebook_path: str,
     variable_project_id: str,
@@ -156,7 +233,6 @@ def _create_tag(filepath: str) -> str:
     return tag
 
 
-rate_limit = RateLimit(max_count=50, per=60, greedy=True)
 
 
 def process_and_execute_notebook(
@@ -172,7 +248,6 @@ def process_and_execute_notebook(
     notebook: str,
     should_get_tail_logs: bool = False,
 ) -> NotebookExecutionResult:
-    rate_limit.wait()  # wait before creating the task
 
     print(f"Running notebook: {notebook}")
 
@@ -191,7 +266,9 @@ def process_and_execute_notebook(
 
     result = NotebookExecutionResult(
         name=tag,
+        path=notebook,
         duration=datetime.timedelta(seconds=0),
+        start_time=datetime.datetime.now(),
         is_pass=False,
         output_uri=notebook_output_uri,
         log_url="",
@@ -201,7 +278,6 @@ def process_and_execute_notebook(
     )
 
     # TODO: Handle cases where multiple notebooks have the same name
-    time_start = datetime.datetime.now()
     operation = None
     try:
         # Get the python version for running the notebook if specified
@@ -245,11 +321,12 @@ def process_and_execute_notebook(
         result.logs_bucket = operation_metadata.build.logs_bucket
 
         # Block and wait for the result
-        operation_result = operation.result()
+        operation_result = operation.result(timeout=timeout_in_seconds)
 
-        result.duration = datetime.datetime.now() - time_start
+        result.duration = datetime.datetime.now() - result.start_time
         result.is_pass = True
         print(f"{notebook} PASSED in {format_timedelta(result.duration)}.")
+
     except Exception as error:
         result.error_message = str(error)
 
@@ -268,7 +345,7 @@ def process_and_execute_notebook(
             except Exception as error:
                 result.error_message = str(error)
 
-        result.duration = datetime.datetime.now() - time_start
+        result.duration = datetime.datetime.now() - result.start_time
         result.is_pass = False
 
         print(
@@ -336,12 +413,44 @@ def get_changed_notebooks(
 
     return notebooks
 
+def _save_results(results: List[NotebookExecutionResult],
+                  artifacts_bucket: str,
+                  results_file: str):
+
+    artifacts_bucket = artifacts_bucket.replace("gs://", "").split('/')[0]
+
+    print("Updating build results ...")
+    build_results = {}
+    for result in results:
+        if result.is_pass:
+            pass_count = 1
+            fail_count = 0
+        else:
+            pass_count = 0
+            fail_count = 1
+        build_results[result.path] = {
+                'duration': result.duration.total_seconds(),
+                'start_time': str(result.start_time),
+                'passed': pass_count,
+                'failed': fail_count
+        }
+        print(f"adding {result.path}")
+
+    print("Saving accumulative results ...")
+    content = json.dumps(build_results)
+
+    client = storage.Client()
+    bucket = client.get_bucket(artifacts_bucket)
+    bucket.blob(str(results_file)).upload_from_string(content, 'text/json')
+
+
 
 def process_and_execute_notebooks(
     notebooks: List[str],
     container_uri: str,
     staging_bucket: str,
     artifacts_bucket: str,
+    results_file: str,
     should_parallelize: bool,
     timeout: int,
     variable_project_id: str,
@@ -349,6 +458,7 @@ def process_and_execute_notebooks(
     variable_service_account: str,
     variable_vpc_network: Optional[str] = None,
     private_pool_id: Optional[str] = None,
+    concurrent_notebooks: Optional[int] = 10,
 ):
     """
     Run the notebooks that exist under the folders defined in the test_paths_file.
@@ -369,6 +479,8 @@ def process_and_execute_notebooks(
             Required. The GCS staging bucket to write source code to.
         artifacts_bucket (str):
             Required. The GCS staging bucket to write executed notebooks to.
+        results_file (str):
+            Required: The path to the artifacts bucket to save results 
         variable_project_id (str):
             Required. The value for PROJECT_ID to inject into notebooks.
         variable_region (str):
@@ -377,6 +489,7 @@ def process_and_execute_notebooks(
             Required. Should run notebooks in parallel using a thread pool as opposed to in sequence.
         timeout (str):
             Required. Timeout string according to https://cloud.google.com/build/docs/build-config-file-schema#timeout.
+        concurrent_notebooks (int): Max number of notebooks per minute to run in parallel.
     """
 
     # Calculate deadline
@@ -393,7 +506,9 @@ def process_and_execute_notebooks(
             print(
                 "Running notebooks in parallel, so no logs will be displayed. Please wait..."
             )
-            with concurrent.futures.ThreadPoolExecutor(max_workers=100) as executor:
+
+
+            with concurrent.futures.ThreadPoolExecutor(max_workers=concurrent_notebooks) as executor:
                 print(f"Max workers: {executor._max_workers}")
 
                 notebook_execution_results = list(
@@ -471,7 +586,7 @@ def process_and_execute_notebooks(
             print("=" * 100)
 
             build_id = results_sorted[0].build_id
-            logs_bucket_name = (results_sorted[0].logs_bucket).removeprefix("gs://")
+            logs_bucket_name = (results_sorted[0].logs_bucket).replace("gs://", "")
             log_file_name = f"log-{build_id}.txt"
 
             log_contents = util.download_blob_into_memory(
@@ -489,6 +604,10 @@ def process_and_execute_notebooks(
             else:
                 print(log_contents)
 
+        _save_results(results_sorted, 
+                      artifacts_bucket, 
+                      results_file)
+
         print("\n=== END RESULTS===\n")
 
         total_notebook_duration = functools.reduce(

.cloud-build/notebook-execution-test-cloudbuild.yaml

@@ -36,7 +36,7 @@ steps:
     - -c
     - |
       . workspace/env/bin/activate &&
-      python3 .cloud-build/execute_changed_notebooks_cli.py --test_paths_file "${_TEST_PATHS_FILE}" --base_branch "${_FORCED_BASE_BRANCH}" --container_uri ${_PYTHON_IMAGE} --staging_bucket ${_GCS_STAGING_BUCKET}  --artifacts_bucket ${_GCS_STAGING_BUCKET}/executed_notebooks/PR_${_PR_NUMBER}/BUILD_${BUILD_ID} --variable_project_id ${PROJECT_ID} --variable_region ${_GCP_REGION} --variable_service_account ${_GCP_SERVICE_ACCOUNT} --variable_vpc_network "${_GPC_VPC_NETWORK_NAME}" `if [ ! -z "${_PRIVATE_POOL_NAME}" ]; then echo "--private_pool_id ${_PRIVATE_POOL_NAME}"; fi`
+      python3 .cloud-build/execute_changed_notebooks_cli.py --test_paths_file "${_TEST_PATHS_FILE}" --base_branch "${_FORCED_BASE_BRANCH}" --container_uri ${_PYTHON_IMAGE} --staging_bucket ${_GCS_STAGING_BUCKET}  --artifacts_bucket ${_GCS_STAGING_BUCKET}/executed_notebooks/PR_${_PR_NUMBER}/BUILD_${BUILD_ID} --variable_project_id ${PROJECT_ID} --variable_region ${_GCP_REGION} --variable_service_account ${_GCP_SERVICE_ACCOUNT} --variable_vpc_network "${_GPC_VPC_NETWORK_NAME}" `if [ ! -z "${_PRIVATE_POOL_NAME}" ]; then echo "--private_pool_id ${_PRIVATE_POOL_NAME}"; fi` --build_id ${BUILD_ID} --test_percent=${_TEST_PERCENT} --concurrent_notebooks=${_CONCURRENT_NOTEBOOKS}
     env:
     - 'IS_TESTING=1'
 timeout: 86400s

.cloud-build/requirements.txt

@@ -3,11 +3,15 @@ numpy
 jupyter
 nbconvert
 papermill
-pandas
 matplotlib
 tabulate
 google-cloud-aiplatform
 google-cloud-storage
 google-cloud-build
+google-cloud-storage
 ratemate
-GitPython
+GitPython
+tqdm
+fsspec
+pandas
+

.cloud-build/test_failing.txt

@@ -0,0 +1,40 @@
+notebooks/official/training/pytorch_gcs_data_training.ipynb
+notebooks/official/custom/custom_training_tensorboard_profiler.ipynb
+notebooks/official/custom/custom-tabular-bq-managed-dataset.ipynb
+notebooks/official/ml_metadata/sdk-metric-parameter-tracking-for-locally-trained-models.ipynb
+notebooks/official/ml_metadata/sdk-metric-parameter-tracking-for-custom-jobs.ipynb
+notebooks/official/tabnet/tabnet_vertex_tutorial.ipynb
+notebooks/official/tabnet/get_started_with_tabnet.ipynb
+notebooks/official/pipelines/google_cloud_pipeline_components_automl_text.ipynb
+notebooks/official/pipelines/multicontender_vs_champion_deployment_method.ipynb
+notebooks/official/pipelines/google_cloud_pipeline_components_automl_images.ipynb
+notebooks/official/pipelines/rapid_prototyping_bqml_automl.ipynb
+notebooks/official/pipelines/challenger_vs_blessed_deployment_method.ipynb
+notebooks/official/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings.ipynb
+notebooks/official/matching_engine/sdk_matching_engine_for_indexing.ipynb
+notebooks/official/matching_engine/sdk_matching_engine_create_text_to_image_embeddings.ipynb
+notebooks/official/explainable_ai/sdk_custom_image_classification_online_explain.ipynb
+notebooks/official/explainable_ai/xai_image_classification_feature_attributions.ipynb
+notebooks/official/explainable_ai/sdk_custom_image_classification_batch_explain.ipynb
+notebooks/official/tabular_workflows/tabnet_on_vertex_pipelines.ipynb
+notebooks/official/model_registry/get_started_with_model_registry.ipynb
+notebooks/official/model_registry/bqml_vertexai_model_registry.ipynb
+notebooks/official/sdk/SDK_Custom_Training_Python_Package_Managed_Text_Dataset_Tensorflow_Serving_Container.ipynb
+notebooks/official/model_monitoring/batch_prediction_model_monitoring.ipynb
+notebooks/official/model_monitoring/get_started_with_model_monitoring_setup.ipynb
+notebooks/official/model_monitoring/get_started_with_model_monitoring_custom.ipynb
+notebooks/official/model_monitoring/get_started_with_model_monitoring_custom_tf_serving.ipynb
+notebooks/official/model_monitoring/model_monitoring.ipynb
+notebooks/official/tensorboard/tensorboard_profiler_custom_training_with_prebuilt_container.ipynb
+notebooks/official/tensorboard/tensorboard_hyperparameter_tuning_with_hparams.ipynb
+notebooks/official/tensorboard/tensorboard_profiler_custom_training.ipynb
+notebooks/official/model_evaluation/custom_tabular_regression_model_evaluation.ipynb
+notebooks/official/model_evaluation/custom_tabular_classification_model_evaluation.ipynb
+notebooks/official/model_evaluation/automl_video_classification_model_evaluation.ipynb
+notebooks/official/experiments/comparing_local_trained_models.ipynb
+notebooks/official/automl/automl_image_classification_online_online_prediction.ipynb
+notebooks/official/automl/automl-text-classification.ipynb
+notebooks/official/automl/sdk_automl_video_object_tracking_batch.ipynb
+notebooks/official/feature_store/sdk-feature-store-pandas.ipynb
+notebooks/official/prediction/custom_batch_prediction_feature_filter.ipynb
+notebooks/official/prediction/pytorch_image_classification_with_prebuilt_serving_containers.ipynb

.cloud-build/test_passing.txt

@@ -0,0 +1,80 @@
+notebooks/official/training/hyperparameter_tuning_tensorflow.ipynb
+notebooks/official/training/get_started_with_vertex_distributed_training.ipynb
+notebooks/official/training/hyperparameter_tuning_xgboost.ipynb
+notebooks/official/training/multi_node_ddp_gloo_vertex_training_with_custom_container.ipynb
+notebooks/official/training/distributed_hyperparameter_tuning.ipynb
+notebooks/official/training/pytorch-text-sentiment-classification-custom-train-deploy.ipynb
+notebooks/official/training/xgboost_data_parallel_training_on_cpu_using_dask.ipynb
+notebooks/official/training/multi_node_ddp_nccl_vertex_training_with_custom_container.ipynb
+notebooks/official/bigquery_ml/get_started_with_bqml_training.ipynb
+notebooks/official/bigquery_ml/bqml-online-prediction.ipynb
+notebooks/official/custom/custom_training_container_and_model_registry.ipynb
+notebooks/official/custom/sdk-custom-image-classification-online.ipynb
+notebooks/official/custom/sdk-custom-image-classification-batch.ipynb
+notebooks/official/custom/SDK_FBProphet_Forecasting_Online.ipynb
+notebooks/official/custom/get_started_vertex_training_xgboost.ipynb
+notebooks/official/custom/get_started_with_vertex_endpoint_and_shared_vm.ipynb
+notebooks/official/custom/SDK_Custom_Container_Prediction.ipynb
+notebooks/official/reduction_server/pytorch_distributed_training_reduction_server.ipynb
+notebooks/official/tabnet/ai-explanations-tabnet-algorithm.ipynb
+notebooks/official/vizier/get_started_vertex_vizier.ipynb
+notebooks/official/vizier/gapic-vizier-multi-objective-optimization.ipynb
+notebooks/official/pipelines/get_started_with_hpt_pipeline_components.ipynb
+notebooks/official/pipelines/google_cloud_pipeline_components_automl_tabular.ipynb
+notebooks/official/pipelines/custom_tabular_train_batch_pred_bq_pipeline.ipynb
+notebooks/official/pipelines/metrics_viz_run_compare_kfp.ipynb
+notebooks/official/pipelines/google_cloud_pipeline_components_model_upload_predict_evaluate.ipynb
+notebooks/official/pipelines/google_cloud_pipeline_components_model_train_upload_deploy.ipynb
+notebooks/official/pipelines/get_started_with_machine_management.ipynb
+notebooks/official/pipelines/custom_model_training_and_batch_prediction.ipynb
+notebooks/official/pipelines/control_flow_kfp.ipynb
+notebooks/official/pipelines/lightweight_functions_component_io_kfp.ipynb
+notebooks/official/pipelines/google_cloud_pipeline_components_bqml_text.ipynb
+notebooks/official/pipelines/pipelines_intro_kfp.ipynb
+notebooks/official/pipelines/automl_tabular_classification_beans.ipynb
+notebooks/official/pipelines/google_cloud_pipeline_components_dataproc_tabular.ipynb
+notebooks/official/explainable_ai/sdk_automl_tabular_classification_online_explain.ipynb
+notebooks/official/explainable_ai/sdk_custom_tabular_regression_online_explain.ipynb
+notebooks/official/explainable_ai/sdk_automl_tabular_binary_classification_batch_explain.ipynb
+notebooks/official/explainable_ai/sdk_custom_tabular_regression_online_explain_get_metadata.ipynb
+notebooks/official/explainable_ai/sdk_custom_tabular_regression_batch_explain.ipynb
+notebooks/official/tabular_workflows/prophet_on_vertex_pipelines.ipynb
+notebooks/official/tabular_workflows/wide_and_deep_on_vertex_pipelines.ipynb
+notebooks/official/sdk/SDK_AutoML_Video_Classification.ipynb
+notebooks/official/model_monitoring/get_started_with_model_monitoring_automl.ipynb
+notebooks/official/model_monitoring/get_started_with_model_monitoring_automl_image_batch.ipynb
+notebooks/official/model_monitoring/get_started_with_model_monitoring_automl_image_online.ipynb
+notebooks/official/model_monitoring/get_started_with_model_monitoring_xgboost.ipynb
+notebooks/official/tensorboard/tensorboard_custom_training_with_custom_container.ipynb
+notebooks/official/tensorboard/tensorboard_custom_training_with_prebuilt_container.ipynb
+notebooks/official/tensorboard/tensorboard_vertex_ai_pipelines_integration.ipynb
+notebooks/official/model_evaluation/automl_text_classification_model_evaluation.ipynb
+notebooks/official/model_evaluation/get_started_with_custom_model_evaluation_import.ipynb
+notebooks/official/model_evaluation/automl_tabular_classification_model_evaluation.ipynb
+notebooks/official/model_evaluation/automl_tabular_regression_model_evaluation.ipynb
+notebooks/official/experiments/get_started_with_vertex_experiments.ipynb
+notebooks/official/experiments/comparing_pipeline_runs.ipynb
+notebooks/official/experiments/get_started_with_vertex_experiments_autologging.ipynb
+notebooks/official/experiments/build_model_experimentation_lineage_with_prebuild_code.ipynb
+notebooks/official/experiments/delete_outdated_tensorboard_experiments.ipynb
+notebooks/official/automl/sdk_automl_tabular_regression_batch_bq.ipynb
+notebooks/official/automl/sdk_automl_text_sentiment_analysis_online.ipynb
+notebooks/official/automl/sdk_automl_text_entity_extraction_online.ipynb
+notebooks/official/automl/sdk_automl_forecasting_hierarchical_batch.ipynb
+notebooks/official/automl/automl_text_entity_extraction_batch_prediction.ipynb
+notebooks/official/automl/automl_image_classification_batch_prediction.ipynb
+notebooks/official/automl/automl_text_sentiment_analysis_batch_prediction.ipynb
+notebooks/official/automl/sdk_automl_tabular_regression_online_bq.ipynb
+notebooks/official/automl/get_started_automl_training.ipynb
+notebooks/official/automl/automl-tabular-classification.ipynb
+notebooks/official/automl/automl_image_object_detection_export_edge.ipynb
+notebooks/official/automl/sdk_automl_image_object_detection_batch.ipynb
+notebooks/official/automl/automl_tabular_on_vertex_pipelines.ipynb
+notebooks/official/automl/sdk_automl_video_classification_batch.ipynb
+notebooks/official/automl/sdk_automl_video_action_recognition_batch.ipynb
+notebooks/official/automl/sdk_automl_tabular_forecasting_batch.ipynb
+notebooks/official/automl/automl_image_object_detection_online_prediction.ipynb
+notebooks/official/automl/automl_forecasting_bqml_arima_plus_comparison.ipynb
+notebooks/official/datasets/get_started_bq_datasets.ipynb
+notebooks/official/datasets/get_started_with_data_labeling.ipynb
+notebooks/official/feature_store/feature_store_streaming_ingestion_sdk.ipynb

.cloud-build/tests/flaky.py

@@ -0,0 +1,46 @@
+# grep PASSED tests.txt | cut -c 10-100 >passed.txt
+
+import os
+
+repo_dir = '/home/jupyter/vertex-ai-samples/'
+repo_dir_len = len(repo_dir)
+official_dir = repo_dir + 'notebooks/official'
+
+entries = os.scandir(official_dir)
+folders = []
+for entry in entries:
+    if entry.is_dir():
+        folders.append(entry.path)
+
+# Passing
+with open('passed.txt', 'r') as pass_file:
+    notebook_names = pass_file.readlines()
+        
+notebooks = []
+for folder in folders:
+    entries = os.scandir(folder)
+    for entry in entries:
+        for notebook in notebook_names:
+            if entry.name == notebook.rstrip():
+                notebooks.append(entry.path[repo_dir_len:])
+                
+with open('passing_tests.txt', 'w') as f:
+    for notebook in notebooks:
+        f.write(notebook + '\n')
+        
+        
+# Failing
+with open('failed.txt', 'r') as fail_file:
+    notebook_names = fail_file.readlines()
+        
+notebooks = []
+for folder in folders:
+    entries = os.scandir(folder)
+    for entry in entries:
+        for notebook in notebook_names:
+            if entry.name == notebook.rstrip():
+                notebooks.append(entry.path[repo_dir_len:])
+                
+with open('failing_tests.txt', 'w') as f:
+    for notebook in notebooks:
+        f.write(notebook + '\n')

.cloud-build/tests/select_notebooks_test.py

@@ -0,0 +1,33 @@
+import sys
+
+from execute_changed_notebooks_helper import (load_results, select_notebook)
+
+
+def test_load_results():
+    bucket: str = "cloud-build-notebooks-presubmit"
+    bucket_file: str = "build_results"
+
+    accum = load_results(bucket, bucket_file)
+
+    print(accum)
+
+    assert len(accum) > 0
+
+def test_select_notebook():
+    bucket: str = "cloud-build-notebooks-presubmit"
+    bucket_file: str = "build_results"
+
+    accum = load_results(bucket, bucket_file)
+
+    n_select = 0
+    n_notselect = 0
+    for notebook in accum:
+        if select_notebook(notebook, accum, 50):
+            n_select += 1
+        else:
+            n_notselect += 1
+
+    print(f"SELECTED {n_select}, NOT SELECTED {n_notselect}")
+
+    assert n_select > 0
+    assert n_notselect > 0

.cloud-build/utils/build_results_viewer.py

@@ -0,0 +1,22 @@
+'''
+Viewer for the weekly regression testing of the official notebooks
+
+Cloud Storage location: gs://cloud-build-notebooks-presubmit/build_results/
+'''
+import argparse
+import json
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--file', dest='file',
+                    default='build.json', type=str, help='build results file')
+import json
+
+with open('build.json', 'r') as f:
+    results = json.load(f)
+
+for item in results.items():
+    if item[1]['passed']:
+        print(f"{item[0]},PASSED")
+    else:
+        print(f"{item[0]},FAILED")

.cloud-build/web-index/notebook-index-cloudbuild.yaml

@@ -0,0 +1,23 @@
+steps:
+  # Fetch full repo for diff purposes
+  - name: gcr.io/cloud-builders/git
+    args: [fetch, --unshallow, --quiet]
+  # Create a virtual environment
+  - name: ${_PYTHON_IMAGE}
+    entrypoint: /bin/sh
+    args:
+    - -c
+    - python3 -m venv workspace/env
+  # Install Python dependencies and run testing script
+  - name: ${_PYTHON_IMAGE}
+    entrypoint: /bin/sh
+    args:
+    - -c
+    - |
+      . workspace/env/bin/activate &&
+      python3 notebooks/notebook_template_review.py --web --title --steps --desc --linkback --notebook-dir=notebooks/official >web.html
+artifacts:
+  objects:
+          location: gs://${_GCS_ARTIFACTS_BUCKET}/webdoc
+          paths: ['web.html']
+timeout: 86400s

.github/workflows/linter/requirements.txt

@@ -2,9 +2,9 @@ git+https://github.com/tensorflow/docs
 ipython
 jupyter
 nbconvert
-black==22.10.0
-pyupgrade==2.38.4
-isort==5.10.1
-flake8==4.0.1
-nbqa==1.5.3
+black==23.3.0
+pyupgrade==3.7.0
+isort==5.12.0
+flake8==6.0.0
+nbqa==1.7.0
 

.github/workflows/linter/run_linter.sh

@@ -84,7 +84,7 @@ if [ ${#notebooks[@]} -gt 0 ]; then
                 # python3 -m nbqa black "$notebook" --check
                 # BLACK_RTN=$?
                 echo "Running pyupgrade..."
-                python3 -m nbqa pyupgrade "$notebook"
+                python3 -m nbqa pyupgrade --exit-zero-even-if-changed "$notebook"
                 PYUPGRADE_RTN=$?
                 echo "Running isort..."
                 python3 -m nbqa isort "$notebook" --check
@@ -97,7 +97,7 @@ if [ ${#notebooks[@]} -gt 0 ]; then
                 python3 -m nbqa black "$notebook"
                 BLACK_RTN=$?
                 echo "Running pyupgrade..."
-                python3 -m nbqa pyupgrade "$notebook"
+                python3 -m nbqa pyupgrade --exit-zero-even-if-changed "$notebook"
                 PYUPGRADE_RTN=$?
                 echo "Running isort..."
                 python3 -m nbqa isort "$notebook"

CONTRIBUTING.md

@@ -44,12 +44,10 @@ Finally, run this code block to check for errors. Each step will attempt to
 automatically fix any issues. If the fixes can't be performed automatically,
 then you will need to manually address them before submitting your PR.
 
+Note: For official, only submit one notebook per PR.
+
 ```shell
-nbqa black "$notebook"
-nbqa pyupgrade "$notebook"
-nbqa isort "$notebook"
-nbqa flake8 "$notebook" --extend-ignore=W391,E501,F821,E402,F404,W503,E203,E722,W293,W291
-python3 -m tensorflow_docs.tools.nbfmt --remove_outputs "$notebook"
+docker run -v ${PWD}:/setup/app gcr.io/cloud-devrel-public-resources/notebook_linter:latest your_notebook
 ```
 
 ## Code Reviews

community-content/CODEOWNERS

@@ -8,3 +8,16 @@
 /cpr-examples @samthrasher
 /Train_tabular_models_with_many_frameworks_and_import_to_Vertex_AI_using_Pipelines @Ark-kun
 /pipeline_components @Ark-kun
+/pipeline_components/image_ml_model_training @lakeyk
+/prediction_featurestore_integration @googleapis/vertex-prediction-team
+/vertex_vision_model_garden/model_oss/util @weigary
+/vertex_vision_model_garden/model_oss/diffusers @weigary
+/vertex_vision_model_garden/model_oss/keras @dstnluong-google
+/vertex_vision_model_garden/model_oss/transformers @dstnluong-google
+/vertex_vision_model_garden/model_oss/pic2word @jismailyan-google
+/vertex_vision_model_garden/model_oss/open_clip @lydhr
+/vertex_vision_model_garden/model_oss/movinet @KCFindstr
+/vertex_vision_model_garden/model_oss/data_converter @KCFindstr
+/vertex_vision_model_garden/model_oss/peft @weigary
+/vertex_vision_model_garden/model_oss/lm-evaluation-harness @kathyyu-google
+

community-content/Train_tabular_models_with_many_frameworks_and_import_to_Vertex_AI_using_Pipelines/Train_tabular_classification_logistic_regression_model_using_Scikit_learn_and_import_to_Vertex_AI/pipeline.py

@@ -6,8 +6,8 @@ download_from_gcs_op = components.load_component_from_url("https://raw.githubuse
 select_columns_using_Pandas_on_CSV_data_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/pandas/Select_columns/in_CSV_format/component.yaml")
 fill_all_missing_values_using_Pandas_on_CSV_data_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/pandas/Fill_all_missing_values/in_CSV_format/component.yaml")
 binarize_column_using_Pandas_on_CSV_data_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/pandas/Binarize_column/in_CSV_format/component.yaml")
-train_logistic_regression_model_using_scikit_learn_from_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/ML_frameworks/Scikit_learn/Train_logistic_regression_model/from_CSV/component.yaml")
-upload_Scikit_learn_pickle_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/1f5cf6e06409b704064b2086c0a705e4e6b4fcde/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_Scikit-learn_pickle_model/component.yaml")
+train_logistic_regression_model_using_scikit_learn_from_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/1f5cf6e06409b704064b2086c0a705e4e6b4fcde/community-content/pipeline_components/ML_frameworks/Scikit_learn/Train_logistic_regression_model/from_CSV/component.yaml")
+upload_Scikit_learn_pickle_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_Scikit-learn_pickle_model/component.yaml")
 deploy_model_to_endpoint_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Deploy_to_endpoint/component.yaml")
 
 # %% Pipeline definition

community-content/pipeline_components/image_ml_model_training/load_image_classification_model/component.yaml

@@ -0,0 +1,112 @@
+name: Load image classification model from tfhub
+description: |
+  Loads specified model from TFHub, creates layer to receive additional (3 channel) imagery data.
+  Args:
+    class_names (Sequence[str]):
+        Sequence of strings of categories for classification corresponding to input data.
+    loaded_model_path (str):
+        Output path for the loaded model.
+    image_size_path (str):
+        Output path for the model expected image size.
+    model_name (Optional[str]):
+        Name of the pre-trained image classification model to load from TFHub.
+        Eligible model_name:
+        - efficientnetv2-s
+        - efficientnetv2-m
+        - efficientnetv2-l
+        - efficientnetv2-s-21k
+        - efficientnetv2-m-21k
+        - efficientnetv2-l-21k
+        - efficientnetv2-xl-21k
+        - efficientnetv2-b0-21k
+        - efficientnetv2-b1-21k
+        - efficientnetv2-b2-21k
+        - efficientnetv2-b3-21k
+        - efficientnetv2-s-21k-ft1k
+        - efficientnetv2-m-21k-ft1k
+        - efficientnetv2-l-21k-ft1k
+        - efficientnetv2-xl-21k-ft1k
+        - efficientnetv2-b0-21k-ft1k
+        - efficientnetv2-b1-21k-ft1k
+        - efficientnetv2-b2-21k-ft1k
+        - efficientnetv2-b3-21k-ft1k
+        - efficientnetv2-b0
+        - efficientnetv2-b1
+        - efficientnetv2-b2
+        - efficientnetv2-b3
+        - efficientnet_b0
+        - efficientnet_b1
+        - efficientnet_b2
+        - efficientnet_b3
+        - efficientnet_b4
+        - efficientnet_b5
+        - efficientnet_b6
+        - efficientnet_b7
+        - bit_s-r50x1
+        - inception_v3
+        - inception_resnet_v2
+        - resnet_v1_50
+        - resnet_v1_101
+        - resnet_v1_152
+        - resnet_v2_50
+        - resnet_v2_101
+        - resnet_v2_152
+        - nasnet_large
+        - nasnet_mobile
+        - pnasnet_large
+        - mobilenet_v2_100_224
+        - mobilenet_v2_130_224
+        - mobilenet_v2_140_224
+        - mobilenet_v3_small_100_224
+        - mobilenet_v3_small_075_224
+        - mobilenet_v3_large_100_224
+        - mobilenet_v3_large_075_224
+    dropout_rate (Optional[float]):
+        Fraction of input units to drop in the last layer. Value should be between 0.0 and 1.0.
+    trainable (Optional[bool]):
+        If true fine tuning will be performed on entire Hub model. If false only additional
+        layers will be trained.
+    l2_regularization_penalty (Optional[float]):
+        l2 regularization penalty.
+inputs:
+- {name: class_names, type: 'typing.List[str]', description: List of class names corresponding
+    to the input image data}
+- {name: model_name, type: String, description: Name of the TFHub model to load, default: efficientnetv2-xl-21k,
+  optional: true}
+- {name: dropout_rate, type: Float, description: Dropout rate, default: '0.2', optional: true}
+- name: trainable
+  type: Boolean
+  description: True if fine tuning should be performed
+  default: "True"
+  optional: true
+- {name: l2_regularization_penalty, type: Float, description: Regularization penalty,
+  default: '0.0001', optional: true}
+outputs:
+- {name: loaded_model_path, type: TensorflowSavedModel, description: Output path for
+    the loaded model}
+- {name: image_size_path, type: HeightWidth}
+implementation:
+  container:
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
+    # command is a list of strings (command-line arguments).
+    # The YAML language has two syntaxes for lists and you can use either of them.
+    # Here we use the "flow syntax" - comma-separated strings inside square brackets.
+    command: [
+      python3,
+      # Path of the program inside the container
+      /pipelines/component/src/loading_component.py,
+      --loaded-model-path,
+      {outputPath: loaded_model_path},
+      --class-names,
+      {inputValue: class_names},
+      --model-name,
+      {inputValue: model_name},
+      --dropout-rate,
+      {inputValue: dropout_rate},
+      --trainable,
+      {inputValue: trainable},
+      --l2-regularization-penalty,
+      {inputValue: l2_regularization_penalty},
+      --image-size-path,
+      {outputPath: image_size_path},
+    ]

community-content/pipeline_components/image_ml_model_training/pipeline.py

@@ -0,0 +1,62 @@
+# python3 -m pip install "kfp<2.0.0" "google-cloud-aiplatform>=1.16.0" --upgrade --quiet
+from kfp import components
+from kfp.v2 import dsl
+
+# %% Loading components
+upload_Tensorflow_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url('https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_Tensorflow_model/component.yaml')
+deploy_model_to_endpoint_op = components.load_component_from_url('https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Deploy_to_endpoint/component.yaml')
+transcode_imagedataset_tfrecord_from_csv_op = components.load_component_from_url('https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/community-content/pipeline_components/image_ml_model_training/transcode_tfrecord_image_dataset_from_csv/component.yaml')
+load_image_classification_model_from_tfhub_op = components.load_component_from_url('https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/b5b65198a6c2ffe8c0fa2aa70127e3325752df68/community-content/pipeline_components/image_ml_model_training/load_image_classification_model/component.yaml')
+preprocess_image_data_op = components.load_component_from_url('https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/community-content/pipeline_components/image_ml_model_training/preprocess_image_data/component.yaml')
+train_tensorflow_image_classification_model_op = components.load_component_from_url('https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/community-content/pipeline_components/image_ml_model_training/train_image_classification_model/component.yaml')
+
+
+# %% Pipeline definition
+def image_classification_pipeline():
+    class_names = ['daisy', 'dandelion', 'roses', 'sunflowers', 'tulips']
+    csv_image_data_path = 'gs://cloud-samples-data/ai-platform/flowers/flowers.csv'
+    deploy_model = False
+
+    image_data = dsl.importer(
+        artifact_uri=csv_image_data_path, artifact_class=dsl.Dataset).output
+
+    image_tfrecord_data = transcode_imagedataset_tfrecord_from_csv_op(
+        csv_image_data_path=image_data,
+        class_names=class_names
+    ).outputs['tfrecord_image_data_path']
+
+    loaded_model_outputs = load_image_classification_model_from_tfhub_op(
+        class_names=class_names,
+    ).outputs
+
+    preprocessed_data = preprocess_image_data_op(
+        image_tfrecord_data,
+        height_width_path=loaded_model_outputs['image_size_path'],
+    ).outputs
+
+    trained_model = (train_tensorflow_image_classification_model_op(
+        preprocessed_training_data_path = preprocessed_data['preprocessed_training_data_path'],
+        preprocessed_validation_data_path = preprocessed_data['preprocessed_validation_data_path'],
+        model_path=loaded_model_outputs['loaded_model_path']).
+                   set_cpu_limit('96').
+                   set_memory_limit('128G').
+                   add_node_selector_constraint('cloud.google.com/gke-accelerator', 'NVIDIA_TESLA_A100').
+                   set_gpu_limit('8').
+                   outputs['trained_model_path'])
+
+    vertex_model_name = upload_Tensorflow_model_to_Google_Cloud_Vertex_AI_op(
+        model=trained_model,
+    ).outputs['model_name']
+
+    # Deploying the model might incur additional costs over time
+    if deploy_model:
+        vertex_endpoint_name = deploy_model_to_endpoint_op(
+            model_name=vertex_model_name,
+        ).outputs['endpoint_name']
+
+pipeline_func = image_classification_pipeline
+
+# %% Pipeline submission
+if __name__ == '__main__':
+    from google.cloud import aiplatform
+    aiplatform.PipelineJob.from_pipeline_func(pipeline_func=pipeline_func).submit()

community-content/pipeline_components/image_ml_model_training/preprocess_image_data/component.yaml

@@ -0,0 +1,57 @@
+name: Preprocess image data
+description: |
+  Preprocess the image data and split between train and validation.
+  Args:
+    input_data_path (str):
+        Input path for the TFRecord image data. Data will be formatted as 'label' (encoded image
+        label), and 'image_raw' (the binary string of the image data).
+    height_width_path (str):
+        Path to square height and width to resize images to. File should contain single float value.
+        Value is dependent on training model.
+    preprocessed_training_data_path (str):
+        Output path for the TFRecord training data. Data will be formatted as 'label' (encoded image
+        label), and 'image_raw' (the binary string of the image data).
+    preprocessed_validation_data_path (str):
+        Output path for the TFRecord validation data. Data will be formatted as 'label' (encoded
+        image label), and 'image_raw' (the binary string of the image data).
+    validation_split (Optional[float]):
+        Fraction of data that will make up validation dataset. Value should be between 0.0 and 1.0.
+    seed (Optional[int]):
+        The global random seed to ensure the system gets a unique random sequence
+        that is deterministic (https://www.tensorflow.org/api_docs/python/tf/random/set_seed).
+inputs:
+- {name: input_data_path, type: ImageDatasetTFRecord, description: 'Input path for
+    the TFRecord image data,'}
+- {name: height_width_path, type: HeightWidth, description: 'Path to square height and width to
+    resize images to,'}
+- {name: validation_split, type: Float, description: 'Fraction of data that will make
+    up validation dataset,', default: '0.2', optional: true}
+- {name: seed, type: Integer, description: Random seed, default: '0', optional: true}
+outputs:
+- {name: preprocessed_training_data_path, type: ImageDatasetTFRecord, description: 'Output
+    path for the training data,'}
+- {name: preprocessed_validation_data_path, type: ImageDatasetTFRecord, description: 'Output
+    path for the validation data,'}
+implementation:
+  container:
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
+    # command is a list of strings (command-line arguments).
+    # The YAML language has two syntaxes for lists and you can use either of them.
+    # Here we use the "flow syntax" - comma-separated strings inside square brackets.
+    command: [
+      python3,
+      # Path of the program inside the container
+      /pipelines/component/src/preprocessing_component.py,
+      --input-data-path,
+      {inputPath: input_data_path},
+      --height-width-path,
+      {inputPath: height_width_path},
+      --validation-split,
+      {inputValue: validation_split},
+      --seed,
+      {inputValue: seed},
+      --preprocessed-training-data-path,
+      {outputPath: preprocessed_training_data_path},
+      --preprocessed-validation-data-path,
+      {outputPath: preprocessed_validation_data_path},
+    ]

community-content/pipeline_components/image_ml_model_training/train_image_classification_model/component.yaml

@@ -0,0 +1,90 @@
+name: Train tensorflow image classification model
+description: |
+  Creates a trained image classification TensorFlow model.
+  Args:
+    preprocessed_training_data_path (str):
+        Input path to the TFRecord training data. Data will be formatted as 'label' (encoded image
+        label), and 'image_raw' (the binary string of the image data).
+    preprocessed_validation_data_path (str):
+        Input path to the TFRecord validation data. Data will be formatted as 'label' (encoded
+        image label), and 'image_raw' (the binary string of the image data).
+    model_path (str):
+        Input path to the loaded pre-trained model.
+    trained_model_path (str):
+        Output path to save the trained model to.
+    optimizer_name (Optional[str]):
+        Name of the tf.keras optimizer. Available optimizers are listed at
+        https://keras.io/api/optimizers/
+    optimizer_parameters (Optional[Dict[str, str]]):
+        Optimizer parameters.
+    loss_function_name (Optional[str]):
+        Name of the loss function.
+    loss_function_parameters (Optional[Dict[str, str]]):
+        Loss function parameters.
+    number_of_epochs (Optional[int]):
+        Number of training iterations over data.
+    metric_names (Optional[Sequence[str]]):
+        List of tf.keras.metrics to be evaluated by the model during training and testing. Available
+        metrics are listed at https://keras.io/api/metrics/.
+    seed Optional(int):
+        The global random seed to ensure the system gets a unique random sequence
+        that is deterministic (https://www.tensorflow.org/api_docs/python/tf/random/set_seed).
+inputs:
+- {name: preprocessed_training_data_path, type: ImageDatasetTFRecord, description: 'Input
+    path for the training data,'}
+- {name: preprocessed_validation_data_path, type: ImageDatasetTFRecord, description: 'Input
+    path for the validation data,'}
+- {name: model_path, type: TensorflowSavedModel, description: 'Input path for the
+    model,'}
+- {name: optimizer_name, type: String, description: 'Name of the optimizer,', default: SGD,
+  optional: true}
+- {name: optimizer_parameters, type: 'typing.Dict[str, str]', description: 'Optimizer
+    parameters,', default: '{}', optional: true}
+- {name: loss_function_name, type: String, description: 'Name of the loss function,',
+  default: CategoricalCrossentropy, optional: true}
+- {name: loss_function_parameters, type: 'typing.Dict[str, str]', description: 'Loss
+    function parameters,', default: '{}', optional: true}
+- {name: number_of_epochs, type: Integer, description: 'Number of epochs,', default: '10',
+  optional: true}
+- {name: metric_names, type: 'typing.List[str]', description: 'List of metrics to
+    use,', default: '["accuracy"]', optional: true}
+- {name: seed, type: Integer, description: 'Random seed,', default: '0', optional: true}
+- {name: batch_size, type: Integer, description: Batch size, default: '16', optional: true}
+outputs:
+- {name: trained_model_path, type: TensorflowSavedModel, description: 'Output path
+    for the saved model,'}
+implementation:
+  container:
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
+    # command is a list of strings (command-line arguments).
+    # The YAML language has two syntaxes for lists and you can use either of them.
+    # Here we use the "flow syntax" - comma-separated strings inside square brackets.
+    command: [
+      python3,
+      # Path of the program inside the container
+      /pipelines/component/src/training_component.py,
+      --preprocessed-training-data-path,
+      {inputPath: preprocessed_training_data_path},
+      --preprocessed-validation-data-path,
+      {inputPath: preprocessed_validation_data_path},
+      --model-path,
+      {inputPath: model_path},
+      --trained-model-path,
+      {outputPath: trained_model_path},
+      --optimizer-name,
+      {inputValue: optimizer_name},
+      --loss-function-name,
+      {inputValue: loss_function_name},
+      --number-of-epochs,
+      {inputValue: number_of_epochs},
+      --seed,
+      {inputValue: seed},
+      --batch-size,
+      {inputValue: batch_size},
+      --metric-names,
+      {inputValue: metric_names},
+      --optimizer-parameters,
+      {inputValue: optimizer_parameters},
+      --loss-function-parameters,
+      {inputValue: loss_function_parameters},
+    ]

community-content/pipeline_components/image_ml_model_training/transcode_tfrecord_image_dataset_from_csv/component.yaml

@@ -0,0 +1,37 @@
+name: Transcode imagedataset tfrecord from csv
+description: |
+  Transcodes CSV Data into TFRecord file of TFExamples.
+  Args:
+    csv_image_data_path (str):
+        Path to the CSV image data. Data must include 'image_filepath' (Path to image file) and
+        'image_label' (output for a prediction) fields.
+    class_names (Sequence[str]):
+        Sequence of strings of categories for classification corresponding to input data.
+    tfrecord_image_data_path (str):
+        Output path for the TFRecord image data. Data will be formatted as 'label' (encoded image
+        label), and 'image_raw' (the binary string of the image data).
+inputs:
+- {name: csv_image_data_path, type: ImageDatasetCSV, description: Input path for the
+    CSV image data}
+- {name: class_names, type: 'typing.List[str]', description: List of class names corresponding
+    to the input image data}
+outputs:
+- {name: tfrecord_image_data_path, type: ImageDatasetTFRecord, description: Output
+    path for the TFRecord image data}
+implementation:
+  container:
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
+    # command is a list of strings (command-line arguments).
+    # The YAML language has two syntaxes for lists and you can use either of them.
+    # Here we use the "flow syntax" - comma-separated strings inside square brackets.
+    command: [
+      python3,
+      # Path of the program inside the container
+      /pipelines/component/src/transcoding_csv_component.py,
+      --csv-image-data-path,
+      {inputPath: csv_image_data_path},
+      --tfrecord-image-data-path,
+      {outputPath: tfrecord_image_data_path},
+      --class-names,
+      {inputValue: class_names},
+    ]

community-content/pipeline_components/image_ml_model_training/transcode_tfrecord_image_dataset_from_jsonl/component.yaml

@@ -0,0 +1,39 @@
+name: Transcode imagedataset tfrecord from jsonlines
+description: |
+  Transcodes JSONL Data into TFRecord file of TFExamples.
+  Args:
+    jsonl_image_data_path (str):
+        Input path for the JSONL image data
+        Path to the JSONL image data. Each line corresponds to a JSON input describing an image.
+        Schema follows AutoML image classification JSONL format
+        https://cloud.google.com/vertex-ai/docs/image-data/classification/prepare-data#json-lines.
+    class_names (Sequence[str]):
+        Sequence of strings of categories for classification corresponding to input data.
+    tfrecord_image_data_path (str):
+        Output path for the TFRecord image data. Data will be formatted as 'label' (encoded image
+        label), and 'image_raw' (the binary string of the image data).
+inputs:
+- {name: jsonl_image_data_path, type: ImageDatasetJsonLines, description: Input path
+    for the JSONL image data}
+- {name: class_names, type: 'typing.List[str]', description: List of class names corresponding
+    to the input image data}
+outputs:
+- {name: tfrecord_image_data_path, type: ImageDatasetTFRecord, description: Output
+    path for the TFRecord image data}
+implementation:
+  container:
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
+    # command is a list of strings (command-line arguments).
+    # The YAML language has two syntaxes for lists and you can use either of them.
+    # Here we use the "flow syntax" - comma-separated strings inside square brackets.
+    command: [
+      python3,
+      # Path of the program inside the container
+      /pipelines/component/src/transcoding_jsonl_component.py,
+      --jsonl-image-data-path,
+      {inputPath: jsonl_image_data_path},
+      --tfrecord-image-data-path,
+      {outputPath: tfrecord_image_data_path},
+      --class-names,
+      {inputValue: class_names},
+    ]

community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/requirements.txt

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

community-content/pytorch_image_classification_single_gpu_with_vertex_sdk_and_torchserve/trainer/requirements.txt

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

community-content/pytorch_pre_built_images_deployment/pytorch-text-classification-vertex-ai-deploy.ipynb

@@ -31,17 +31,7 @@
       "source": [
         "# Deploying a PyTorch Text Classification Model on [Vertex AI](https://cloud.google.com/vertex-ai)\n",
         "\n",
-        "**This is an Experimental release**, covered by the Pre-GA Offerings Terms of your Google Cloud Platform [Terms of Service](https://cloud.google.com/terms).\n",
-        "\n",
-        "Experiments are focused on validating a prototype and are not guaranteed to be released. They are not intended for production use or covered by any SLA, support obligation, or deprecation policy and might be subject to backward-incompatible changes.\n",
-        "\n",
-        "**Kindly drop us a note before you run any scale tests.**\n",
-        "\n",
-        "**Do not hesitate to contact vertexai-prediction-preview-feedback@google.com if you have any questions or run into any issues.**\n",
-        "\n",
-        "The usage of the product is free during the Experimental release period: you will still incur charges for other GCP products usage, such as storage.\n",
-        "\n",
-        "The projects need to be added to the allowlist in order to deploy PyTorch models using Vertex AI Prediction pre-built PyTorch images. If you are interested in the feature, please send an email to vertexai-prediction-preview-feedback@google.com to provide your project numbers OR project ids."
+        "**Kindly reach out to Vertex AI before you run any scale tests or you have any questions.**\n"
       ]
     },
     {

community-content/vertex_vision_model_garden/.pylintrc

@@ -0,0 +1,4 @@
+[MASTER]
+
+generated-members=get_concrete_function,cv2.*
+ignored-modules=tensorflow,google.cloud

community-content/vertex_vision_model_garden/model_oss/data_converter/common_lib.py

@@ -0,0 +1,623 @@
+"""Library with functions to use for data conversion."""
+
+import json
+import os
+import random
+from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union
+import uuid
+
+from absl import logging
+import apache_beam as beam
+import cv2
+import numpy as np
+import pandas as pd
+import PIL
+from PIL import Image
+import tensorflow as tf
+import yaml
+
+from util import constants
+from util import fileutils
+from apache_beam.options import pipeline_options
+
+REFORMATTED_CSV_SUFFIX = '-reformatted.csv'
+
+LABEL_MAP_NAME = 'label_map.yaml'
+
+_SPLIT_RATIO_ERROR_THRESHOLD = 1e-5
+# Internal constant. Only for distinguishing rows without ML use.
+ML_USE_UNASSIGNED = 'unassigned'
+ALL_ML_USES = (
+    constants.ML_USE_TRAINING,
+    constants.ML_USE_VALIDATION,
+    constants.ML_USE_TEST,
+    ML_USE_UNASSIGNED,
+)
+COLUMN_NAME_ML_USE = 'ml_use'
+COLUMN_NAME_GCS_FILE_PATH = 'gcs_file_path'
+COLUMN_NAME_LABEL = 'label'
+COLUMN_NAME_START_SEC = 'start_sec'
+COLUMN_NAME_END_SEC = 'end_sec'
+# Output filenames
+TRAIN_TFRECORD_NAME = 'train.tfrecord'
+VALIDATION_TFRECORD_NAME = 'val.tfrecord'
+TEST_TFRECORD_NAME = 'test.tfrecord'
+# Jsonl keys
+JSON_GCS_URI_KEY = 'imageGcsUri'
+JSON_RESOURCE_LABEL_KEY = 'dataItemResourceLabels'
+JSON_ML_USE_KEY = 'aiplatform.googleapis.com/ml_use'
+# I/O parameters
+READ_CHUNK_SIZE = 1024 * 1024 * 1024  # 1GB
+
+
+class WriteToTFRecord(beam.DoFn):
+  """DoFn to write TF examples to sharded TF record files."""
+
+  def __init__(
+      self,
+      output_prefix: str,
+      num_shards: int,
+      convert_fn: Callable[[Dict[str, Any]], tf.train.Example],
+  ):
+    self.output_prefix = output_prefix
+    self.num_shards = num_shards
+    self.writer: list[tf.io.TFRecordWriter] = []
+    self.sharded_files: list[str] = []
+    self.convert_fn = convert_fn
+    self.success_counter = beam.metrics.Metrics.counter(
+        self.__class__.__name__, 'Success'
+    )
+    self.failure_counter = beam.metrics.Metrics.counter(
+        self.__class__.__name__, 'Failure'
+    )
+
+  def start_bundle(self):
+    logging.info('Start writing TF Record to %s.', self.output_prefix)
+    unique_str = uuid.uuid4().hex
+    for i in range(self.num_shards):
+      uri = f'{self.output_prefix}-{i}-{unique_str}'
+      self.sharded_files.append(uri)
+      self.writer.append(tf.io.TFRecordWriter(uri))
+
+  def process(self, data: Dict[str, Any]) -> Iterable[Tuple[int, str]]:
+    try:
+      example = self.convert_fn(data)
+      data = example.SerializeToString()
+      idx = hash(data) % self.num_shards
+      self.writer[idx].write(data)
+      self.success_counter.inc()
+      yield (idx, self.sharded_files[idx])
+    # pylint: disable-next=broad-exception-caught
+    except Exception as err:
+      logging.error('Failed to process %s', data)
+      logging.exception(err)
+      self.failure_counter.inc()
+
+  def finish_bundle(self):
+    logging.info('Finish writing TF Record to %s.', self.output_prefix)
+    for writer in self.writer:
+      writer.close()
+    self.writer = []
+
+
+def convert_to_feature(
+    value: Union[List[Union[int, float, bytes]], int, float, bytes],
+    value_type: Optional[str] = None,
+) -> tf.train.Feature:
+  """Converts the given python object to a tf.train.Feature.
+
+  This is copied from tensorflow_models/official/vision/data/tfrecord_lib.py.
+
+  Args:
+    value: int, float, bytes or a list of them.
+    value_type: optional, if specified, forces the feature to be of the given
+      type. Otherwise, type is inferred automatically. Can be one of ['bytes',
+      'int64', 'float', 'bytes_list', 'int64_list', 'float_list']
+
+  Returns:
+    feature: A tf.train.Feature object.
+  """
+
+  if value_type is None:
+    element = value[0] if isinstance(value, list) else value
+
+    if isinstance(element, bytes):
+      value_type = 'bytes'
+
+    elif isinstance(element, (int, np.integer)):
+      value_type = 'int64'
+
+    elif isinstance(element, (float, np.floating)):
+      value_type = 'float'
+
+    else:
+      raise ValueError(
+          'Cannot convert type {} to feature'.format(type(element))
+      )
+
+    if isinstance(value, list):
+      value_type = value_type + '_list'
+
+  if value_type == 'int64':
+    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
+
+  elif value_type == 'int64_list':
+    value = np.asarray(value).astype(np.int64).reshape(-1)
+    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))
+
+  elif value_type == 'float':
+    return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))
+
+  elif value_type == 'float_list':
+    value = np.asarray(value).astype(np.float32).reshape(-1)
+    return tf.train.Feature(float_list=tf.train.FloatList(value=value))
+
+  elif value_type == 'bytes':
+    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
+
+  elif value_type == 'bytes_list':
+    return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))
+
+  else:
+    raise ValueError('Unknown value_type parameter - {}'.format(value_type))
+
+
+def convert_to_string_feature(
+    value: str, encoding: str = 'utf-8'
+) -> tf.train.Feature:
+  """Returns a bytes_list from an encoded string."""
+  return convert_to_feature(value.encode(encoding))
+
+
+def convert_to_list_string_feature(
+    lst: list[str], encoding: str = 'utf-8'
+) -> tf.train.Feature:
+  """Returns a bytes_list from a list of encoded strings."""
+  return convert_to_feature([value.encode(encoding) for value in lst])
+
+
+def create_ml_use_array_with_split(
+    total_size: int,
+    split_ratio: Sequence[float],
+) -> list[str]:
+  """Create randomized list of 'training', 'validation', 'test'.
+
+  The list of will be of length total_size with ratios according to train_size,
+  validation_size, and test_size.
+
+  Args:
+    total_size: Length of sequence to return
+    split_ratio: Proportions to split into 'training', 'validation', and 'test'
+
+  Returns:
+    List containing 'training', 'validation', and 'test'
+  """
+  train_size, validation_size, _ = split_ratio
+  num_train = round(train_size * total_size)
+  num_validation = round(validation_size * total_size)
+  num_test = total_size - num_train - num_validation
+  ml_use_row = (
+      [constants.ML_USE_TRAINING] * num_train
+      + [constants.ML_USE_VALIDATION] * num_validation
+      + [constants.ML_USE_TEST] * num_test
+  )
+  random.shuffle(ml_use_row)
+  return ml_use_row
+
+
+def format_ml_use_column(df: pd.DataFrame):
+  df[COLUMN_NAME_ML_USE].replace(
+      # We need to support non-standard ML uses other than documented ones,
+      # since they are used by some existing datasets.
+      [r'(?i)^train(ing)?$', r'(?i)^test$', r'(?i)^validat(ion|e)$'],
+      [
+          constants.ML_USE_TRAINING,
+          constants.ML_USE_TEST,
+          constants.ML_USE_VALIDATION,
+      ],
+      inplace=True,
+      regex=True,
+  )
+
+
+def insert_missing_ml_use(df: pd.DataFrame) -> None:
+  """For every row that does not have ml_use as the first column, insert a column containing 'unassigned' to the front.
+
+  Args:
+    df: The DataFrame to process. The first column should be 'ml_use'.
+  """
+  df[COLUMN_NAME_ML_USE].fillna(ML_USE_UNASSIGNED, inplace=True)
+  rows_to_fill = ~df[COLUMN_NAME_ML_USE].isin(ALL_ML_USES)
+  df.loc[rows_to_fill] = df[rows_to_fill].shift(
+      axis=1, fill_value=ML_USE_UNASSIGNED
+  )
+
+
+def replace_unassigned_ml_use(
+    ml_uses: List[str],
+    split_ratio: Sequence[float],
+):
+  """Replace `unassigned` in ml_uses with `training`, `validation`, and `test` with ratios according to split_ratio.
+
+  Args:
+    ml_uses: List of ml_use string values.
+    split_ratio: Proportions to split into `training`, `validation`, and `test`.
+  """
+  unassigned_indices = [
+      i for i, ml_use in enumerate(ml_uses) if ml_use == ML_USE_UNASSIGNED
+  ]
+  ml_use_arr = create_ml_use_array_with_split(
+      len(unassigned_indices), split_ratio
+  )
+  for unassigned_index, ml_use in zip(unassigned_indices, ml_use_arr):
+    ml_uses[unassigned_index] = ml_use
+
+
+def merge_seq_into_dicts(
+    key: str, values: Sequence[Any], dicts: Sequence[Dict[Any, Any]]
+):
+  """Merges a list of values into a list of dicts, inserted with the given key.
+
+  Args:
+    key: Key to insert or overwrite in the dictionary.
+    values: A list of values to insert.
+    dicts: A list of dictionaries. Each value will be inserted into the
+      corresponding dictionary. The original value will be overwritten if the
+      key already existed.
+
+  Raises:
+    ValueError: The values and dicts have different lengths.
+  """
+  if len(values) != len(dicts):
+    raise ValueError(
+        f'Length of values and dicts must match, got {len(values)} and'
+        f' {len(dicts)}'
+    )
+  for val, d in zip(values, dicts):
+    d[key] = val
+
+
+def drop_invalid_rows(df: pd.DataFrame) -> int:
+  """Drops DataFrame rows missing the gcs_file_path column or the label column.
+
+  Args:
+    df: The DataFrame to process in place.
+
+  Returns:
+    The number of rows dropped.
+  """
+  original_rows = df.shape[0]
+  df.dropna(subset=[COLUMN_NAME_GCS_FILE_PATH, COLUMN_NAME_LABEL], inplace=True)
+  dropped_num = original_rows - df.shape[0]
+  if dropped_num > 0:
+    df.reset_index(drop=True, inplace=True)
+  return dropped_num
+
+
+def check_split_ratio(split_ratio: Sequence[float]):
+  """Checks if the give split ratio is valid.
+
+  Args:
+    split_ratio: Proportions to split into 'training', 'validation', and 'test'
+
+  Raises:
+    ValueError: Must have valid entries, correct length, and sum to 1.
+  """
+  if len(split_ratio) != 3:
+    raise ValueError('split_ratio must contain exactly 3 values.')
+  if abs(sum(split_ratio) - 1) > _SPLIT_RATIO_ERROR_THRESHOLD:
+    raise ValueError('split_ratio must sum to 1.')
+  if not all([0 <= val <= 1 for val in split_ratio]):
+    raise ValueError('Entries of split_ratio must be in the range [0, 1].')
+
+
+def check_num_shard(num_shard: Sequence[int]):
+  """Checks if the number of shards is valid.
+
+  Args:
+    num_shard: The number of shards for each tfrecord.
+
+  Raises:
+    ValueError: Must have valid entries and correct length.
+  """
+  if len(num_shard) != 3:
+    raise ValueError('num_shard must contain exactly 3 values.')
+  if not all([val >= 1 for val in num_shard]):
+    raise ValueError('Shards must be at least 1.')
+
+
+def create_label_map_yaml(meta_data_path: str, output_dir: str) -> None:
+  """Generate label_map.yaml from meta_data.yaml.
+
+  Args:
+    meta_data_path: Path to a meta_data.yaml file.
+    output_dir: Directory to output label_map.yaml.
+  """
+  tf.io.gfile.copy(
+      meta_data_path, os.path.join(output_dir, LABEL_MAP_NAME), overwrite=True
+  )
+
+
+def reformat_bbox(
+    bbox: Sequence[int], img_width: int, img_height: int
+) -> Tuple[float, float, float, float]:
+  """Converts XYWH unnormalized bounding box with to a normalized XYXY bounding box.
+
+  Args:
+    bbox: Relative bounding box with unnormalized coordinates as [x, y, width,
+      height].
+    img_width: Image's pixel width.
+    img_height: Image's pixel height.
+
+  Returns:
+    Absolute bounding box with normalized coordinates as
+      [xmin, ymin, xmax, ymax].
+  """
+  x, y, width, height = bbox
+  xmin = x / img_width
+  ymin = y / img_height
+  xmax = (x + width) / img_width
+  ymax = (y + height) / img_height
+  return xmin, ymin, xmax, ymax
+
+
+def encode_image(
+    filepath: str,
+    output_shape: Optional[Sequence[int]] = None,
+    image_format: str = 'png',
+) -> Tuple[bytes, Sequence[int]]:
+  """Encodes an image at the given path.
+
+  Args:
+    filepath: Path to the image.
+    output_shape: The output shape of the image, (height, width).
+    image_format: The format of the output image.
+
+  Returns:
+    The encoded image data in bytes and the shape of the image, (height, width).
+
+  Raises:
+    IOError: The image file is corrupt.
+  """
+  filepath = fileutils.force_gcs_fuse_path(filepath)
+  with open(filepath, 'rb') as f:
+    # If an output_shape is specified, resize the image and set data to the new
+    # bytes.
+    try:
+      img = Image.open(f)
+    except PIL.UnidentifiedImageError as e:
+      raise IOError(f'Failed to open {filepath}') from e
+
+    try:
+      if output_shape is not None:
+        rgb_img = img.resize((output_shape[1], output_shape[0])).convert('RGB')
+      else:
+        rgb_img = img.convert('RGB')
+      rgb_img = np.array(rgb_img)
+
+      _, data = cv2.imencode(f'.{image_format}', rgb_img)
+      data = data.tobytes()
+      return data, rgb_img.shape
+    except cv2.error as e:
+      raise IOError(f'Failed to encode {filepath}') from e
+    finally:
+      img.close()
+
+
+def encode_video(
+    filepath: str,
+    start_sec: float,
+    end_sec: float,
+    output_fps: int = 5,
+    output_shape: Optional[Sequence[int]] = None,
+    image_format: str = 'jpg',
+) -> Sequence[bytes]:
+  """Encodes a video clip at the given path with start and end timestamps.
+
+  Args:
+    filepath: Path to the video.
+    start_sec: Start timestamp of the video clip in seconds.
+    end_sec: End timestamp of the video clip in seconds.
+    output_fps: The output frame rate per second.
+    output_shape: The output shape of each frame, (height, width).
+    image_format: The format of the encoded frames.
+
+  Returns:
+    A list of the encoded frames data in bytes.
+
+  Raises:
+    IOError if the video file is corrupt.
+  """
+  filepath = fileutils.force_gcs_fuse_path(filepath)
+  video = None
+
+  try:
+    video = cv2.VideoCapture(filepath)
+    frames = []
+    frame_interval = 1 / output_fps
+    total_frames = video.get(cv2.CAP_PROP_FRAME_COUNT)
+    original_fps = video.get(cv2.CAP_PROP_FPS)
+    if not original_fps:
+      # 0 or None indicates the video is invalid
+      raise IOError(f'Failed to load {filepath}')
+    video_length = total_frames / original_fps
+    start_sec = max(start_sec, 0)
+    end_sec = min(end_sec, video_length)
+    for t in np.arange(start_sec, end_sec, frame_interval):
+      frame_idx = min(total_frames - 1, round(t * original_fps))
+      video.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
+      ret, frame = video.read()
+      if not ret:
+        raise IOError(f'Failed to load {filepath} at frame {frame_idx}')
+      if output_shape is not None:
+        frame = cv2.resize(frame, (output_shape[1], output_shape[0]))
+      _, data = cv2.imencode(f'.{image_format}', frame)
+      frames.append(data.tobytes())
+  except cv2.error as e:
+    raise IOError(f'Failed to load {filepath}') from e
+  finally:
+    if video:
+      video.release()
+  return frames
+
+
+def create_label_map(
+    labels: Sequence[str],
+) -> Tuple[Sequence[int], Dict[int, str]]:
+  """Creates a label map from a sequence of label strings.
+
+  Args:
+    labels: The sequence of labels to create label map from. Must not contain
+      invalid values, which means data without labels should be filtered first.
+
+  Returns:
+    The integer labels and the mapping from integers to the original strings.
+  """
+  inverse_label_map: Dict[str, int] = dict()
+  num_labels = 0
+  for label in labels:
+    if label not in inverse_label_map:
+      num_labels += 1
+      inverse_label_map[label] = num_labels
+  int_labels = [inverse_label_map[label] for label in labels]
+  label_map = {value: key for key, value in inverse_label_map.items()}
+  return int_labels, label_map
+
+
+def write_label_map(output_file: str, label_map: Dict[int, str]) -> None:
+  """Writes a label map to the output file, which can be a GCS uri."""
+  with tf.io.gfile.GFile(output_file, 'w') as f:
+    yaml.dump({'label_map': label_map}, f)
+
+
+def detectron_json_to_image_rows(input_json: str) -> list[Dict[str, Any]]:
+  """Converts a Detectron JSON file to a list of image rows.
+
+  Args:
+    input_json: A path to a Detectron JSON or JSONL file.
+
+  Returns:
+    A list of dictionaries, where each dictionary contains Detectron format
+    entry.
+
+  Raises:
+    ValueError: If the input JSON is invalid.
+  """
+
+  image_rows = []
+  with tf.io.gfile.GFile(input_json, 'r') as f:
+    for line in f:
+      json_data = json.loads(line)
+      if isinstance(json_data, dict):
+        image_rows.append(json_data)
+      elif isinstance(json_data, list):
+        image_rows.extend(json_data)
+      else:
+        raise ValueError(
+            'The input JSON is invalid. Dict or list is expected, but got '
+            f'{type(json_data)}.'
+        )
+  return image_rows
+
+
+def coco_json_to_image_rows(
+    input_json: str,
+) -> List[Dict[str, Any]]:
+  """Converts a COCO JSON file to a list of image rows.
+
+  Args:
+    input_json: A path to a COCO JSON or JSONL file.
+
+  Returns:
+    A list of dictionaries, where each dictionary contains COCO format entry.
+
+  Raises:
+    ValueError: If the input JSON is invalid.
+  """
+
+  with tf.io.gfile.GFile(input_json, 'r') as f:
+    coco_json = json.load(f)
+  if 'annotations' not in coco_json:
+    raise ValueError('"annotations" is not in the dataset.')
+  if 'images' not in coco_json:
+    raise ValueError('"images" is not in the dataset.')
+
+  images = coco_json['images']
+  return images
+
+
+def partition_by_ml_use(element: Dict[str, Any], num_partitions: int) -> int:
+  """Beam partition function to split data by ml_use."""
+  del num_partitions
+  try:
+    partition = ALL_ML_USES.index(element[COLUMN_NAME_ML_USE])
+  except Exception as e:
+    raise ValueError(f'Invalid ML use: {element[COLUMN_NAME_ML_USE]}') from e
+  return partition
+
+
+def run_beam_pipeline(pipeline: Any) -> None:
+  """Runs a beam pipeline. Works in both internal and docker environment."""
+  options = pipeline_options.PipelineOptions([
+      '--runner=FlinkRunner',
+      '--faster_copy',
+      '--max_parallelism', '8',
+  ])
+  p = beam.Pipeline(options=options)
+  pipeline(p)
+  result = p.run()
+  result.wait_until_finish()
+  for counter in result.metrics().query()['counters']:
+    logging.info('%s counter: %s.', counter.key.metric.name, counter)
+  logging.info('Completing beam pipeline.')
+
+
+def beam_convert_tfexamples(
+    root: beam.Pipeline,
+    data_list: Sequence[Dict[str, Any]],
+    convert_fn: Callable[[Dict[str, Any]], tf.train.Example],
+    output_dir: str,
+    num_shards: Sequence[int],
+) -> None:
+  """Constructs beam pipelines to convert train, val, test TF Examples."""
+  names = [TRAIN_TFRECORD_NAME, VALIDATION_TFRECORD_NAME, TEST_TFRECORD_NAME]
+  split_data = (
+      root
+      | 'Create PCollection' >> beam.Create(data_list)
+      | 'Data split' >> beam.Partition(partition_by_ml_use, 3)
+  )
+  for i in range(3):
+    ml_use: str = ALL_ML_USES[i]
+    num_shard = num_shards[i]
+    output_prefix = os.path.join(output_dir, names[i])
+    _ = (
+        split_data[i]
+        | f'Convert {ml_use} TF Examples'
+        >> beam.ParDo(WriteToTFRecord(output_prefix, num_shard, convert_fn))
+        | f'Group {ml_use} TF Record files' >> beam.GroupBy(lambda x: x[0])
+        | f'Merge {ml_use} TF Record files'
+        >> beam.Map(merge_tfrecords_func(output_prefix, num_shard))
+    )
+
+
+def merge_tfrecords_func(output_prefix: str, num_shard: int) -> ...:
+  """Returns a function to merge sharded worker output into expected shards."""
+  output_prefix = fileutils.force_gcs_fuse_path(output_prefix)
+
+  def merge_tfrecords(worker_output: Tuple[int, Sequence[Tuple[int, str]]]):
+    idx = worker_output[0]
+    files: Sequence[str] = np.unique([x[1] for x in worker_output[1]])
+    output_file = f'{output_prefix}-{idx:05d}-of-{num_shard:05d}'
+    with open(output_file, 'wb') as f:
+      for file in files:
+        logging.info('Merging %s.', file)
+        file = fileutils.force_gcs_fuse_path(file)
+        with open(file, 'rb') as fin:
+          while True:
+            data = fin.read(READ_CHUNK_SIZE)
+            if not data:
+              break
+            f.write(data)
+        os.remove(file)
+
+  return merge_tfrecords

community-content/vertex_vision_model_garden/model_oss/data_converter/convert_coco_label_map_iod.py

@@ -0,0 +1,111 @@
+r"""Converts COCO labels as yamls for model garden playground (IOD).
+"""
+
+import os
+import urllib.request
+
+from absl import app
+from absl import flags
+import tensorflow as tf
+import yaml
+
+from object_detection.utils import label_map_util
+
+_CONVERT_LABEL_TYPE_COCO_80 = 'coco_80'
+_CONVERT_LABEL_TYPE_COCO_91 = 'coco_91'
+
+_CONVERT_LABEL_TYPE = flags.DEFINE_enum(
+    'convert_label_type',
+    None,
+    [
+        _CONVERT_LABEL_TYPE_COCO_80,
+        _CONVERT_LABEL_TYPE_COCO_91,
+    ],
+    'Different types of label type conversion.',
+    required=True,
+)
+
+_TEMPORARY_PATH = flags.DEFINE_string(
+    'temporary_path',
+    None,
+    'The tempory path.',
+    required=True,
+)
+
+_OUTPUT_YAML_FILEPATH = flags.DEFINE_string(
+    'output_yaml_filepath',
+    None,
+    'The output yaml filepath.',
+    required=True,
+)
+
+
+def convert_coco_label_map_91(
+    output_yaml_filepath: str,
+) -> None:
+  """Converts coco label map 91."""
+  input_proto_filepath = 'https://raw.githubusercontent.com/tensorflow/models/master/research/object_detection/data/mscoco_label_map.pbtxt'
+  local_input_proto_filepath = os.path.join(
+      _TEMPORARY_PATH.value, 'mscoco_label_map.pbtxt'
+  )
+  with open(local_input_proto_filepath, 'w') as writer:
+    contents = (
+        urllib.request.urlopen(input_proto_filepath).read().decode('utf-8')
+    )
+    writer.write(contents)
+
+  label_map = label_map_util.load_labelmap(local_input_proto_filepath)
+  label_map_dict = label_map_util.get_label_map_dict(
+      label_map, use_display_name=True
+  )
+  swapped_label_map_dict = {v: k for k, v in label_map_dict.items()}
+  print(swapped_label_map_dict)
+
+  # Saves new label maps as yamls.
+  with tf.io.gfile.GFile(output_yaml_filepath, 'w') as writer:
+    writer.write(yaml.dump(swapped_label_map_dict))
+
+
+def convert_coco_label_map_80(
+    output_yaml_filepath: str,
+) -> None:
+  """Converts coco label map 80."""
+  # Loads label maps from texts.
+  input_text_filepath = 'https://gist.githubusercontent.com/AruniRC/7b3dadd004da04c80198557db5da4bda/raw/2f10965ace1e36c4a9dca76ead19b744f5eb7e88/ms_coco_classnames.txt'
+  local_input_text_filepath = os.path.join(
+      _TEMPORARY_PATH.value, 'ms_coco_classnames.txt'
+  )
+  with open(local_input_text_filepath, 'w') as writer:
+    contents = (
+        urllib.request.urlopen(input_text_filepath).read().decode('utf-8')
+    )
+    writer.write(contents)
+  with open(local_input_text_filepath, 'r') as file:
+    content = file.read()
+    label_map = yaml.safe_load(content)
+
+  # Removes background in label maps.
+  new_label_map = {}
+  for k, v in label_map.items():
+    if k == 0:
+      continue
+    new_label_map[k - 1] = v
+  print(new_label_map)
+  # Saves new label maps as yamls.
+  with tf.io.gfile.GFile(output_yaml_filepath, 'w') as writer:
+    writer.write(yaml.dump(new_label_map))
+
+
+def main(_) -> None:
+  if _CONVERT_LABEL_TYPE.value == _CONVERT_LABEL_TYPE_COCO_80:
+    convert_coco_label_map_80(_OUTPUT_YAML_FILEPATH.value)
+  elif _CONVERT_LABEL_TYPE.value == _CONVERT_LABEL_TYPE_COCO_91:
+    convert_coco_label_map_91(
+        _OUTPUT_YAML_FILEPATH.value,
+    )
+  else:
+    print('Not supported convert label type: ', _CONVERT_LABEL_TYPE.value)
+
+
+if __name__ == '__main__':
+  app.run(main)

community-content/vertex_vision_model_garden/model_oss/data_converter/convert_imagenet_label_map.py

@@ -0,0 +1,86 @@
+r"""Converts ImageNet label texts as yamls for model garden playground.
+
+# ImageNet1K will have label maps with background.
+"""
+
+import urllib.request
+from absl import app
+from absl import flags
+import tensorflow as tf
+import yaml
+
+
+_INPUT_TEXT_FILEPATH = flags.DEFINE_string(
+    'input_text_filepath',
+    None,
+    'The input text filepath.',
+    required=True,
+)
+
+_ADD_BACKGROUND_LABEL = flags.DEFINE_boolean(
+    'add_background_label',
+    None,
+    'Whether or not add background labels.',
+    required=True,
+)
+
+_ADD_IDS = flags.DEFINE_boolean(
+    'add_ids',
+    None,
+    'Whether or not add ids.',
+    required=True,
+)
+
+_OUTPUT_YAML_FILEPATH = flags.DEFINE_string(
+    'output_yaml_filepath',
+    None,
+    'The output yaml filepath.',
+    required=True,
+)
+
+
+def convert_imagenet_label_map_from_text_to_yaml(
+    input_text_filepath: str,
+    add_background_label: bool,
+    add_ids: bool,
+    output_yaml_filepath: str,
+) -> None:
+  """Converts imagenet label map from text to yamls."""
+  label_map = {}
+
+  # Shifts all keys by 1, and add 0 as 'background'.
+  if add_background_label:
+    label_map = yaml.safe_load(
+        urllib.request.urlopen(input_text_filepath).read()
+    )
+    new_label_map = {}
+    for key, value in label_map.items():
+      new_label_map[key + 1] = value
+    new_label_map[0] = 'background'
+    label_map = new_label_map
+
+  # Adds maps from id to each line.
+  if add_ids:
+    lines = urllib.request.urlopen(input_text_filepath).readlines()
+    current_id = 0
+    for line in lines:
+      label_map[current_id] = line.decode('ascii').strip()
+      print(label_map[current_id])
+      current_id += 1
+
+  # Saves new label maps as yamls.
+  with tf.io.gfile.GFile(output_yaml_filepath, 'w') as writer:
+    writer.write(yaml.dump(label_map))
+
+
+def main(_) -> None:
+  convert_imagenet_label_map_from_text_to_yaml(
+      _INPUT_TEXT_FILEPATH.value,
+      _ADD_BACKGROUND_LABEL.value,
+      _ADD_IDS.value,
+      _OUTPUT_YAML_FILEPATH.value,
+  )
+
+
+if __name__ == '__main__':
+  app.run(main)

community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_icn_lib.py

@@ -0,0 +1,199 @@
+"""Converts ICN CSV/JSONL files to TFRecord with apache beam."""
+
+import json
+from os import path
+from typing import Any, Dict, Sequence, Union, cast
+
+from absl import logging
+import apache_beam as beam
+import pandas as pd
+import tensorflow as tf
+
+from data_converter import common_lib
+
+
+_COLUMN_NAMES = [
+    common_lib.COLUMN_NAME_ML_USE,
+    common_lib.COLUMN_NAME_GCS_FILE_PATH,
+    common_lib.COLUMN_NAME_LABEL,
+]
+_JSON_GCS_URI_KEY = 'imageGcsUri'
+_JSON_CLASS_ANNOTATION_KEY = 'classificationAnnotation'
+_JSON_RESOURCE_LABEL_KEY = 'dataItemResourceLabels'
+_JSON_CLASS_NAME_KEY = 'displayName'
+_JSON_ML_USE_KEY = 'aiplatform.googleapis.com/ml_use'
+
+
+def build_tf_example(element: Dict[str, Union[str, int]]) -> tf.train.Example:
+  """Builds a TF Example from an image uri and label.
+
+  Args:
+    element: A dict with the keys gcs_file_path and label.
+
+  Returns:
+    The created TF Example.
+  """
+  image_uri = cast(str, element[common_lib.COLUMN_NAME_GCS_FILE_PATH])
+  label = cast(int, element[common_lib.COLUMN_NAME_LABEL])
+  image_bytes, shape = common_lib.encode_image(image_uri, image_format='jpeg')
+  features = tf.train.Features(
+      feature={
+          'image/encoded': common_lib.convert_to_feature(image_bytes),
+          'image/format': common_lib.convert_to_string_feature('jpeg'),
+          'image/height': common_lib.convert_to_feature(shape[0]),
+          'image/width': common_lib.convert_to_feature(shape[1]),
+          'image/class/label': common_lib.convert_to_feature(label),
+      },
+  )
+  return tf.train.Example(features=features)
+
+
+def _run_convert_pipeline(
+    output_dir: str, df: pd.DataFrame, num_shards: Sequence[int]
+) -> None:
+  """Starts a Beam pipeline to write DataFrame as TF Records.
+
+  Args:
+    output_dir: TF Records output directory.
+    df: DataFrame to convert from.
+    num_shards: Number of shards for train/validation/test TFRecord files.
+  """
+  images_list = df.to_dict('records')
+
+  def pipeline(root: beam.Pipeline):
+    common_lib.beam_convert_tfexamples(
+        root,
+        images_list,
+        build_tf_example,
+        output_dir,
+        num_shards,
+    )
+
+  common_lib.run_beam_pipeline(pipeline)
+
+
+def _convert_df_to_tfrecord(
+    df: pd.DataFrame,
+    output_dir: str,
+    split_ratio: Sequence[float],
+    num_shard: Sequence[int],
+) -> None:
+  """Converts a DataFrame into three separate tfrecords for training, validation, and testing into output_dir.
+
+  Args:
+    df: DataFrame to convert.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    split_ratio: List specifying the training, validation, and testing splits
+      for unassigned TFRecords.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+  """
+  # Replaces ml_use with common_lib string constants for consistency.
+  common_lib.format_ml_use_column(df)
+  common_lib.insert_missing_ml_use(df)
+
+  # Ignores invalid rows.
+  dropped_row_num = common_lib.drop_invalid_rows(df)
+  if dropped_row_num > 0:
+    logging.warning('Ignored %d invalid rows.', dropped_row_num)
+
+  common_lib.replace_unassigned_ml_use(
+      df[common_lib.COLUMN_NAME_ML_USE], split_ratio
+  )
+
+  # Converts labels to integers as required by training.
+  new_labels, label_map = common_lib.create_label_map(
+      df[common_lib.COLUMN_NAME_LABEL]
+  )
+  df[common_lib.COLUMN_NAME_LABEL] = new_labels
+  label_map_path = path.join(output_dir, common_lib.LABEL_MAP_NAME)
+  logging.info('Writing label map to %s.', label_map_path)
+  common_lib.write_label_map(label_map_path, label_map)
+
+  _run_convert_pipeline(output_dir, df, num_shard)
+
+
+def convert_csv_to_tfrecord(
+    input_csv: str,
+    output_dir: str,
+    split_ratio: Sequence[float] = (0.8, 0.1, 0.1),
+    num_shard: Sequence[int] = (10, 10, 10),
+) -> None:
+  """Parses input_csv file into three separate tfrecords for training, validation, and testing into output_dir.
+
+  The csv format is shown in
+    https://cloud.google.com/vertex-ai/docs/image-data/classification/prepare-data#csv.
+
+  If an ml_use column is not provided, one will be created.
+
+  label_map.yaml containing the label map will be placed in output_dir.
+
+  Args:
+    input_csv: Name of the csv file.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    split_ratio: List specifying the training, validation, and testing splits
+      for unassigned TFRecords.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+  """
+  with tf.io.gfile.GFile(input_csv, 'r') as f:
+    df: pd.DataFrame = pd.read_csv(
+        f, header=None, names=_COLUMN_NAMES, on_bad_lines='warn'
+    )
+
+  _convert_df_to_tfrecord(df, output_dir, split_ratio, num_shard)
+
+
+def convert_jsonl_to_tfrecord(
+    input_jsonl: str,
+    output_dir: str,
+    split_ratio: Sequence[float] = (0.8, 0.1, 0.1),
+    num_shard: Sequence[int] = (10, 10, 10),
+) -> None:
+  """Parses input_jsonl file into three separate tfrecords for training, validation, and testing into output_dir.
+
+  The JSONL format is shown in
+    https://cloud.google.com/vertex-ai/docs/image-data/classification/prepare-data#json-lines.
+
+  If an ml_use column is not provided, one will be created.
+
+  label_map.yaml containing the label map will be placed in output_dir.
+
+  Args:
+    input_jsonl: Name of the JSONL file.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    split_ratio: List specifying the training, validation, and testing splits
+      for unassigned TFRecords.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+  """
+  df_rows = []
+  with tf.io.gfile.GFile(input_jsonl, 'r') as f:
+    lines = f.read().rstrip().splitlines()
+
+  for i, line in enumerate(lines, 1):
+    try:
+      item: Dict[str, Any] = json.loads(line)
+
+      gcs_uri = item.get(_JSON_GCS_URI_KEY)
+      label = item.get(_JSON_CLASS_ANNOTATION_KEY, {}).get(_JSON_CLASS_NAME_KEY)
+      if not gcs_uri or not label:
+        logging.warning('Invalid JSON at line %d, skipped.', i)
+        continue
+
+      ml_use = item.get(_JSON_RESOURCE_LABEL_KEY, {}).get(
+          _JSON_ML_USE_KEY, common_lib.ML_USE_UNASSIGNED
+      )
+    except (json.JSONDecodeError, AttributeError):
+      logging.warning('Invalid JSON at line %d, skipped.', i)
+      continue
+
+    df_rows.append([ml_use, gcs_uri, label])
+
+  df = pd.DataFrame(
+      data=df_rows,
+      columns=[
+          common_lib.COLUMN_NAME_ML_USE,
+          common_lib.COLUMN_NAME_GCS_FILE_PATH,
+          common_lib.COLUMN_NAME_LABEL,
+      ],
+  )
+
+  _convert_df_to_tfrecord(df, output_dir, split_ratio, num_shard)

community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_iod_lib.py

@@ -0,0 +1,430 @@
+"""Converts IOD dataset files to TFRecord with apache beam."""
+
+import collections
+import json
+from os import path
+from typing import Any, Dict, Sequence
+
+from absl import logging
+import apache_beam as beam
+import pandas as pd
+import tensorflow as tf
+
+from data_converter import common_lib
+from util import constants
+
+COLUMN_NAME_LABEL_INT = 'label_int'
+_COLUMN_NAME_XMIN = 'X_MIN'
+_COLUMN_NAME_YMIN = 'Y_MIN'
+_COLUMN_NAME_XMAX = 'X_MAX'
+_COLUMN_NAME_YMAX = 'Y_MAX'
+COLUMN_NAMES = [
+    common_lib.COLUMN_NAME_ML_USE,
+    common_lib.COLUMN_NAME_GCS_FILE_PATH,
+    common_lib.COLUMN_NAME_LABEL,
+    _COLUMN_NAME_XMIN,
+    _COLUMN_NAME_YMIN,
+    'XMAX_NOT_USED',
+    'YMIN_NOT_USED',
+    _COLUMN_NAME_XMAX,
+    _COLUMN_NAME_YMAX,
+    'XMIN_NOT_USED',
+    'YMAX_NOT_USED',
+]
+_BOUNDING_BOX_COLUMNS = [
+    _COLUMN_NAME_XMIN,
+    _COLUMN_NAME_YMIN,
+    _COLUMN_NAME_XMAX,
+    _COLUMN_NAME_YMAX,
+]
+_JSON_BBOX_ANNOTATIONS_KEY = 'boundingBoxAnnotations'
+_JSON_DISPLAY_NAME_KEY = 'displayName'
+_JSON_X_MIN_KEY = 'xMin'
+_JSON_X_MAX_KEY = 'xMax'
+_JSON_Y_MIN_KEY = 'yMin'
+_JSON_Y_MAX_KEY = 'yMax'
+
+
+def build_tf_example(image_row: Dict[str, Any]) -> tf.train.Example:
+  """Builds a TF Example from an image row.
+
+  Args:
+    image_row: A dictionary containing information about the image, such as its
+      GCS uri, labels, and bounding box coordinates.
+
+  Returns:
+    A tf.train.Example containing the encoded image and optionally a
+      bounding box and label.
+  """
+  image_uri = image_row[common_lib.COLUMN_NAME_GCS_FILE_PATH]
+  image_bytes, shape = common_lib.encode_image(image_uri, image_format='jpeg')
+  feature = {
+      'image/encoded': common_lib.convert_to_feature(image_bytes),
+      'image/format': common_lib.convert_to_string_feature('jpeg'),
+      'image/height': common_lib.convert_to_feature(shape[0]),
+      'image/width': common_lib.convert_to_feature(shape[1]),
+      'image/source_id': common_lib.convert_to_string_feature(image_uri),
+      'image/object/bbox/xmin': common_lib.convert_to_feature(
+          image_row[_COLUMN_NAME_XMIN]
+      ),
+      'image/object/bbox/ymin': common_lib.convert_to_feature(
+          image_row[_COLUMN_NAME_YMIN]
+      ),
+      'image/object/bbox/xmax': common_lib.convert_to_feature(
+          image_row[_COLUMN_NAME_XMAX]
+      ),
+      'image/object/bbox/ymax': common_lib.convert_to_feature(
+          image_row[_COLUMN_NAME_YMAX]
+      ),
+      'image/object/class/text': common_lib.convert_to_list_string_feature(
+          image_row[common_lib.COLUMN_NAME_LABEL]
+      ),
+      'image/object/class/label': common_lib.convert_to_feature(
+          image_row[COLUMN_NAME_LABEL_INT]
+      ),
+  }
+  return tf.train.Example(features=tf.train.Features(feature=feature))
+
+
+def _run_convert_pipeline(
+    output_dir: str,
+    image_rows: Sequence[Dict[str, Any]],
+    num_shards: Sequence[int],
+) -> None:
+  """Starts a Beam pipeline to write DataFrame as TF Records.
+
+  Args:
+    output_dir: TF Records output directory.
+    image_rows: Contains all necessary information to create a TF Example.
+    num_shards: Number of shards for train/validation/test TFRecord files.
+  """
+
+  def pipeline(root: beam.Pipeline):
+    common_lib.beam_convert_tfexamples(
+        root,
+        image_rows,
+        build_tf_example,
+        output_dir,
+        num_shards,
+    )
+
+  common_lib.run_beam_pipeline(pipeline)
+
+
+def _convert_df_to_tfrecord(
+    df: pd.DataFrame,
+    output_dir: str,
+    split_ratio: Sequence[float],
+    num_shard: Sequence[int],
+) -> None:
+  """Converts a DataFrame into three separate tfrecords for training, validation, and testing into output_dir.
+
+  Args:
+    df: DataFrame to convert.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    split_ratio: List specifying the training, validation, and testing splits
+      for unassigned TFRecords.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+  """
+  # Replaces ml_use with common_lib string constants for consistency.
+  common_lib.format_ml_use_column(df)
+  common_lib.insert_missing_ml_use(df)
+
+  # Specify bounding box columns to be numeric.
+  df[_BOUNDING_BOX_COLUMNS] = df[_BOUNDING_BOX_COLUMNS].apply(pd.to_numeric)
+
+  # Ignores invalid rows.
+  dropped_row_num = common_lib.drop_invalid_rows(df)
+  dropped_row_num += drop_rows_without_bbox(df)
+  if dropped_row_num > 0:
+    logging.warning('Ignored %d invalid rows.', dropped_row_num)
+
+  # Converts labels to integers as required by training.
+  int_labels, label_map = common_lib.create_label_map(
+      df[common_lib.COLUMN_NAME_LABEL]
+  )
+  df[COLUMN_NAME_LABEL_INT] = int_labels
+  label_map_path = path.join(output_dir, common_lib.LABEL_MAP_NAME)
+  logging.info('Writing label map to %s.', label_map_path)
+  common_lib.write_label_map(label_map_path, label_map)
+
+  image_rows = _condense_bounding_boxes(df.to_dict(orient='records'))
+  ml_uses = [row[common_lib.COLUMN_NAME_ML_USE] for row in image_rows]
+  common_lib.replace_unassigned_ml_use(ml_uses, split_ratio)
+  common_lib.merge_seq_into_dicts(
+      common_lib.COLUMN_NAME_ML_USE, ml_uses, image_rows
+  )
+
+  _run_convert_pipeline(output_dir, image_rows, num_shard)
+
+
+def _condense_bounding_boxes(
+    image_rows: Sequence[Dict[str, Any]]
+) -> Sequence[Dict[str, Any]]:
+  """Gather all the bounding boxes in an image and put them in the same dictionary.
+
+  Args:
+    image_rows: List of dictionaries, each containing information about the
+      image, such as its GCS uri, labels, and bounding box coordinates.
+
+  Returns:
+    List of dictionaries such that each contains all the bounding boxes for a
+    given gcs_file_path.
+
+  Raises:
+    RuntimeError: This is raised when the input data contains images that have
+      annotations in different ml_use classes.
+  """
+  output = {}
+  for image_row in image_rows:
+    ml_use = image_row[common_lib.COLUMN_NAME_ML_USE]
+    gcs_file_path = image_row[common_lib.COLUMN_NAME_GCS_FILE_PATH]
+    label = image_row[common_lib.COLUMN_NAME_LABEL]
+    xmin = image_row[_COLUMN_NAME_XMIN]
+    ymin = image_row[_COLUMN_NAME_YMIN]
+    xmax = image_row[_COLUMN_NAME_XMAX]
+    ymax = image_row[_COLUMN_NAME_YMAX]
+    label_int = image_row[COLUMN_NAME_LABEL_INT]
+    if gcs_file_path in output:
+      d = output[gcs_file_path]
+      if ml_use != common_lib.ML_USE_UNASSIGNED:
+        if d[common_lib.COLUMN_NAME_ML_USE] == common_lib.ML_USE_UNASSIGNED:
+          d[common_lib.COLUMN_NAME_ML_USE] = ml_use
+        elif ml_use != d[common_lib.COLUMN_NAME_ML_USE]:
+          raise RuntimeError(
+              f'Image {gcs_file_path} can only be placed in one of'
+              f' training/validation/test. It is currently in {ml_use} and'
+              f' {d[common_lib.COLUMN_NAME_ML_USE]}.'
+          )
+      d[common_lib.COLUMN_NAME_LABEL].append(label)
+      d[_COLUMN_NAME_XMIN].append(xmin)
+      d[_COLUMN_NAME_YMIN].append(ymin)
+      d[_COLUMN_NAME_XMAX].append(xmax)
+      d[_COLUMN_NAME_YMAX].append(ymax)
+      d[COLUMN_NAME_LABEL_INT].append(label_int)
+    else:
+      output[gcs_file_path] = {
+          common_lib.COLUMN_NAME_ML_USE: ml_use,
+          common_lib.COLUMN_NAME_GCS_FILE_PATH: gcs_file_path,
+          common_lib.COLUMN_NAME_LABEL: [label],
+          _COLUMN_NAME_XMIN: [xmin],
+          _COLUMN_NAME_YMIN: [ymin],
+          _COLUMN_NAME_XMAX: [xmax],
+          _COLUMN_NAME_YMAX: [ymax],
+          COLUMN_NAME_LABEL_INT: [label_int],
+      }
+  return list(output.values())
+
+
+def convert_csv_to_tfrecord(
+    input_csv: str,
+    output_dir: str,
+    split_ratio: Sequence[float] = (0.8, 0.1, 0.1),
+    num_shard: Sequence[int] = (10, 10, 10),
+) -> None:
+  """Parses input_csv file into three separate tfrecords for training, validation, and testing into output_dir.
+
+  The csv format is shown in
+    https://cloud.google.com/vertex-ai/docs/image-data/object-detection/prepare-data#csv.
+
+  If an ml_use column is not provided, one will be created.
+
+  label_map.yaml containing the label map will be placed in output_dir.
+
+  Args:
+    input_csv: Name of the csv file.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    split_ratio: List specifying the train, validation, and test splits for
+      unassigned TFRecords.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+  """
+  with tf.io.gfile.GFile(input_csv, 'r') as f:
+    df: pd.DataFrame = pd.read_csv(
+        f, header=None, names=COLUMN_NAMES, on_bad_lines='warn'
+    )
+
+  _convert_df_to_tfrecord(df, output_dir, split_ratio, num_shard)
+
+
+def drop_rows_without_bbox(df: pd.DataFrame) -> int:
+  """Drops DataFrame rows without bounding_boxes.
+
+  Args:
+    df: The DataFrame to process in place.
+
+  Returns:
+    The number of rows dropped.
+  """
+  invalid_rows = df.index[~(df[_BOUNDING_BOX_COLUMNS].notnull().all(axis=1))]
+  dropped_num = len(invalid_rows)
+  if dropped_num > 0:
+    invalid_df = df.loc[invalid_rows].to_dict(orient='records')
+    for entry in invalid_df:
+      logging.warning('Skipping entry due to missing bounding box: %s.', entry)
+    df.drop(invalid_rows, inplace=True)
+    df.reset_index(drop=True, inplace=True)
+  return dropped_num
+
+
+def convert_coco_json_categories_to_label_map(
+    categories: Sequence[Dict[str, Any]]
+) -> Dict[int, str]:
+  return {category['id']: category['name'] for category in categories}
+
+
+def convert_coco_json_to_tfrecord(
+    input_coco_json: str,
+    output_dir: str,
+    split_ratio: Sequence[float] = (0.8, 0.1, 0.1),
+    num_shard: Sequence[int] = (10, 10, 10),
+) -> None:
+  """Parses input_csv file into three separate tfrecords for training, validation, and testing into output_dir.
+
+  The COCO json format is shown here: https://cocodataset.org/#format-data.
+
+  label_map.yaml containing the label map will be placed in output_dir.
+
+  Args:
+    input_coco_json: Name of coco json file.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    split_ratio: List specifying the train, validation, and test splits for
+      dataset.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+  """
+  with tf.io.gfile.GFile(input_coco_json, 'r') as f:
+    coco_json = json.load(f)
+  # Writes label map from coco json categories.
+  label_map = convert_coco_json_categories_to_label_map(
+      coco_json[constants.COCO_JSON_CATEGORIES]
+  )
+  label_map_path = path.join(output_dir, common_lib.LABEL_MAP_NAME)
+  logging.info('Writes label map to %s.', label_map_path)
+  common_lib.write_label_map(label_map_path, label_map)
+
+  img_to_anns = collections.defaultdict(list)
+  imgs = {}
+  if constants.COCO_JSON_ANNOTATIONS in coco_json:
+    for ann in coco_json[constants.COCO_JSON_ANNOTATIONS]:
+      img_to_anns[ann[constants.COCO_JSON_ANNOTATION_IMAGE_ID]].append(ann)
+
+  if constants.COCO_JSON_IMAGES in coco_json:
+    for img in coco_json[constants.COCO_JSON_IMAGES]:
+      imgs[img[constants.COCO_JSON_IMAGE_ID]] = img
+
+  df_rows = []
+
+  for image_id, annotations in img_to_anns.items():
+    img = imgs[image_id]
+    for ann in annotations:
+      xmin, ymin, xmax, ymax = common_lib.reformat_bbox(
+          ann[constants.COCO_ANNOTATION_BBOX],
+          img[constants.COCO_JSON_IMAGE_WIDTH],
+          img[constants.COCO_JSON_IMAGE_HEIGHT],
+      )
+      df_rows.append([
+          common_lib.ML_USE_UNASSIGNED,
+          img[constants.COCO_JSON_IMAGE_COCO_URL],
+          label_map[ann[constants.COCO_JSON_ANNOTATION_CATEGORY_ID]],
+          xmin,
+          ymin,
+          xmax,
+          ymin,
+          xmax,
+          ymax,
+          xmin,
+          ymax,
+          ann[constants.COCO_JSON_ANNOTATION_CATEGORY_ID],
+      ])
+  df = pd.DataFrame(
+      data=df_rows,
+      columns=COLUMN_NAMES + [COLUMN_NAME_LABEL_INT],
+  )
+
+  # Replaces ml_use with common_lib string constants for consistency.
+  common_lib.format_ml_use_column(df)
+  common_lib.insert_missing_ml_use(df)
+
+  # Species bounding box columns to be numeric.
+  df[_BOUNDING_BOX_COLUMNS] = df[_BOUNDING_BOX_COLUMNS].apply(pd.to_numeric)
+
+  # Ignores invalid rows.
+  dropped_row_num = common_lib.drop_invalid_rows(df)
+  dropped_row_num += drop_rows_without_bbox(df)
+  if dropped_row_num > 0:
+    logging.warning('Ignored %d invalid rows.', dropped_row_num)
+
+  image_rows = _condense_bounding_boxes(df.to_dict(orient='records'))
+  ml_uses = [row[common_lib.COLUMN_NAME_ML_USE] for row in image_rows]
+  common_lib.replace_unassigned_ml_use(ml_uses, split_ratio)
+  common_lib.merge_seq_into_dicts(
+      common_lib.COLUMN_NAME_ML_USE, ml_uses, image_rows
+  )
+
+  _run_convert_pipeline(output_dir, image_rows, num_shard)
+
+
+def convert_jsonl_to_tfrecord(
+    input_jsonl: str,
+    output_dir: str,
+    split_ratio: Sequence[float] = (0.8, 0.1, 0.1),
+    num_shard: Sequence[int] = (10, 10, 10),
+) -> None:
+  """Parses input_jsonl file into three separate tfrecords for training, validation, and testing into output_dir.
+
+  The JSONL format is shown in
+    https://cloud.google.com/vertex-ai/docs/image-data/object-detection/prepare-data#json-lines.
+
+  If an ml_use column is not provided, one will be created.
+
+  label_map.yaml containing the label map will be placed in output_dir.
+
+  Args:
+    input_jsonl: Name of the JSONL file.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    split_ratio: List specifying the training, validation, and testing splits
+      for unassigned TFRecords.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+  """
+  df_rows = []
+  with tf.io.gfile.GFile(input_jsonl, 'r') as f:
+    lines = f.read().rstrip().splitlines()
+
+  for i, line in enumerate(lines, start=1):
+    try:
+      item: Dict[str, Any] = json.loads(line)
+    except (json.JSONDecodeError, AttributeError):
+      logging.warning('Invalid JSON at line %d skipped.', i)
+      continue
+
+    gcs_uri = item.get(common_lib.JSON_GCS_URI_KEY)
+    if not gcs_uri:
+      logging.warning(
+          'Invalid JSON at line %d skipped. Missing gcs_uri_key.', i
+      )
+      continue
+    ml_use = item.get(common_lib.JSON_RESOURCE_LABEL_KEY, {}).get(
+        common_lib.JSON_ML_USE_KEY, common_lib.ML_USE_UNASSIGNED
+    )
+
+    for bbox in item.get(_JSON_BBOX_ANNOTATIONS_KEY, []):
+      label = bbox.get(_JSON_DISPLAY_NAME_KEY)
+      xmin = bbox.get(_JSON_X_MIN_KEY)
+      ymin = bbox.get(_JSON_Y_MIN_KEY)
+      xmax = bbox.get(_JSON_X_MAX_KEY)
+      ymax = bbox.get(_JSON_Y_MAX_KEY)
+
+      df_rows.append([ml_use, gcs_uri, label, xmin, ymin, xmax, ymax])
+
+  df = pd.DataFrame(
+      data=df_rows,
+      columns=[
+          common_lib.COLUMN_NAME_ML_USE,
+          common_lib.COLUMN_NAME_GCS_FILE_PATH,
+          common_lib.COLUMN_NAME_LABEL,
+          _COLUMN_NAME_XMIN,
+          _COLUMN_NAME_YMIN,
+          _COLUMN_NAME_XMAX,
+          _COLUMN_NAME_YMAX,
+      ],
+  )
+  _convert_df_to_tfrecord(df, output_dir, split_ratio, num_shard)

community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_isg_lib.py

@@ -0,0 +1,328 @@
+"""Python script to convert different file formats for ISG to tfrecords."""
+
+import hashlib
+import os
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
+
+from absl import logging
+import apache_beam as beam
+from apache_beam.io import tfrecordio
+import cv2
+import numpy as np
+from pycocotools import coco
+import tensorflow as tf
+import yaml
+
+from data_converter import common_lib
+from util import constants
+from util import fileutils
+
+_IMAGE_FORMAT = 'PNG'
+
+
+def build_tf_example(
+    image_info: dict[str, Union[str, int]],
+    segmentation_image: List[List[int]],
+    output_shape: Optional[Tuple[int, int]] = None,
+) -> tf.train.Example:
+  """Encodes an image and its segmentation mask into a tf.train.Example.
+
+  Args:
+    image_info: A dictionary containing information about the image, such as its
+      file name, height, and width.
+    segmentation_image: 2D image in list of lists having category ids.
+    output_shape: The desired output shape of the image. If None, the original
+      image shape will be used.
+
+  Returns:
+    A tf.train.Example containing the encoded image and segmentation mask.
+
+  Raises:
+    IOError: If image cannot be found in the path.
+  """
+  file_name = image_info[constants.COCO_JSON_FILE_NAME]
+  height = int(image_info[constants.COCO_JSON_IMAGE_HEIGHT])
+  width = int(image_info[constants.COCO_JSON_IMAGE_WIDTH])
+
+  segmentation_image = np.expand_dims(
+      np.asarray(segmentation_image, dtype=np.int32), axis=-1
+  )
+  _, encoded_seg = cv2.imencode(f'.{_IMAGE_FORMAT.lower()}', segmentation_image)
+  encoded_seg = encoded_seg.tobytes()
+
+  encoded_img, _ = common_lib.encode_image(
+      image_info[constants.COCO_JSON_IMAGE_COCO_URL],
+      output_shape=output_shape,
+      image_format=_IMAGE_FORMAT.lower(),
+  )
+
+  key = hashlib.sha256(encoded_img).hexdigest()
+
+  return tf.train.Example(
+      features=tf.train.Features(
+          feature={
+              'image/height': common_lib.convert_to_feature(height),
+              'image/width': common_lib.convert_to_feature(width),
+              'image/filename': common_lib.convert_to_string_feature(file_name),
+              'image/sha256': common_lib.convert_to_string_feature(key),
+              'image/encoded': common_lib.convert_to_feature(encoded_img),
+              'image/format': common_lib.convert_to_string_feature(
+                  _IMAGE_FORMAT
+              ),
+              'image/segmentation/class/encoded': common_lib.convert_to_feature(
+                  encoded_seg
+              ),
+              'image/segmentation/class/format': (
+                  common_lib.convert_to_string_feature(_IMAGE_FORMAT)
+              ),
+              'image/segmentation/class/height': common_lib.convert_to_feature(
+                  height
+              ),
+              'image/segmentation/class/width': common_lib.convert_to_feature(
+                  width
+              ),
+          }
+      )
+  )
+
+
+class AcquireTFExampleDoFn(beam.DoFn):
+  """Beam DoFn to build TF Examples from a single row of image_info data."""
+
+  # These tags will be used to tag the outputs of this DoFn.
+  output_tag_train = constants.ML_USE_TRAINING
+  output_tag_validation = constants.ML_USE_VALIDATION
+  output_tag_test = constants.ML_USE_TEST
+
+  valid_ml_use_set = set(
+      [output_tag_train, output_tag_validation, output_tag_test]
+  )
+
+  def __init__(self, output_shape: Optional[Tuple[int, int]] = None):
+    self.acquired_examples_counter = beam.metrics.Metrics.counter(
+        self.__class__.__name__, 'Success'
+    )
+    self.failure_counter = beam.metrics.Metrics.counter(
+        self.__class__.__name__, 'Failure'
+    )
+    self.output_shape = output_shape
+
+  def process(
+      self,
+      row: Tuple[str, Dict[str, Union[str, int]], List[List[int]]],
+  ) -> Iterator[tf.train.Example]:
+    ml_use, image_info, annotation_info = row
+    if ml_use not in self.valid_ml_use_set:
+      logging.warning('ml_use invalid: %s', ml_use)
+      self.failure_counter.inc()
+      return
+
+    try:
+      tf_example = build_tf_example(
+          image_info, annotation_info, self.output_shape
+      )
+    except IOError as e:
+      logging.warning('Failed to build TF Example: %s', e)
+      self.failure_counter.inc()
+    else:
+      self.acquired_examples_counter.inc()
+      yield beam.pvalue.TaggedOutput(ml_use, tf_example)
+
+
+def _define_data_conversion_pipeline(
+    root: beam.Pipeline,
+    ml_use_rows: List[str],
+    image_rows: List[Dict[str, Union[str, int]]],
+    segmentation_rows: List[List[List[int]]],
+    output_dir: str,
+    output_shape: Optional[Tuple[int, int]],
+    num_shard_list: List[int],
+):
+  """Define a data conversion pipeline.
+
+  Args:
+    root: A Beam pipeline.
+    ml_use_rows: List containing the ml_use.
+    image_rows: List of dictionaries containing information about the image,
+      such as its file name, height, and width.
+    segmentation_rows: List of 2D images of integers representing segmentation
+      masks.
+    output_dir: Directory where the output TFRecords will be written.
+    output_shape: Desired output shape of the image. If None, the original image
+      shape will be used.
+    num_shard_list: Number of shards to write to each output TFRecord.
+
+  Returns:
+    A Beam pipeline.
+  """
+  train, validation, test = (
+      root
+      | 'Load ml use and image rows to beam'
+      >> beam.Create(zip(ml_use_rows, image_rows, segmentation_rows))
+      | 'Build TF Examples'
+      >> beam.ParDo(AcquireTFExampleDoFn(output_shape)).with_outputs(
+          AcquireTFExampleDoFn.output_tag_train,
+          AcquireTFExampleDoFn.output_tag_validation,
+          AcquireTFExampleDoFn.output_tag_test,
+      )
+  )
+
+  # Save each split to TFRecord.
+  _ = train | 'Save train split to TFRecord' >> tfrecordio.WriteToTFRecord(
+      os.path.join(output_dir, common_lib.TRAIN_TFRECORD_NAME),
+      coder=beam.coders.ProtoCoder(tf.train.Example),
+      num_shards=num_shard_list[0],
+  )
+  _ = (
+      validation
+      | 'Save validation split to TFRecord'
+      >> tfrecordio.WriteToTFRecord(
+          os.path.join(output_dir, common_lib.VALIDATION_TFRECORD_NAME),
+          coder=beam.coders.ProtoCoder(tf.train.Example),
+          num_shards=num_shard_list[1],
+      )
+  )
+  _ = test | 'Save test split to TFRecord' >> tfrecordio.WriteToTFRecord(
+      os.path.join(output_dir, common_lib.TEST_TFRECORD_NAME),
+      coder=beam.coders.ProtoCoder(tf.train.Example),
+      num_shards=num_shard_list[2],
+  )
+
+
+def _image_info_to_segmentation_image(
+    img: Dict[str, Any],
+    coco_dataset: coco.COCO,
+    label_id_by_category_id: Dict[int, int],
+) -> List[List[int]]:
+  """Convert image information to a segmentation image.
+
+  Args:
+      img: The image information.
+      coco_dataset: The COCO dataset.
+      label_id_by_category_id: The mapping from label id used for training to
+        category_id defined in dataset.
+
+  Returns:
+      The segmentation image.
+
+  Raises:
+    ValueError: If the mask size does not match the image or if a pixel has
+      multiple labels.
+  """
+  seg_img = np.zeros(
+      shape=(
+          img[constants.COCO_JSON_IMAGE_HEIGHT],
+          img[constants.COCO_JSON_IMAGE_WIDTH],
+      ),
+      dtype=np.int32,
+  )
+  for ann in coco_dataset.imgToAnns[img[constants.COCO_JSON_IMAGE_ID]]:
+    new_category_id = ann[constants.COCO_JSON_ANNOTATION_CATEGORY_ID]
+    binary_mask = coco_dataset.annToMask(ann)
+    if seg_img.shape != binary_mask.shape:
+      raise ValueError(
+          'Binary mask does not have the same shape as image. image_id:'
+          f' {img["id"]}'
+      )
+    boolean_mask = binary_mask == 1
+    if (seg_img[boolean_mask] != 0).any():
+      raise ValueError(
+          'Error: Some pixels have more than one label in image_id:'
+          f' {img["id"]}.'
+      )
+    seg_img[boolean_mask] = label_id_by_category_id[new_category_id]
+
+  return seg_img.tolist()
+
+
+def get_input_rows(
+    coco_dataset: coco.COCO,
+    split_ratio: List[float],
+    label_id_by_category_id: Dict[int, int],
+) -> Tuple[List[str], List[Dict[str, Union[str, int]]], List[List[List[int]]]]:
+  """Get input rows for training and validation.
+
+  Args:
+      coco_dataset: The COCO dataset.
+      split_ratio: The split ratio for training and validation.
+      label_id_by_category_id: The mapping from label id used for training to
+        category_id defined in dataset.
+
+  Returns:
+      - A list of ml_use strings.
+      - A list of image informations.
+      - A list of segmentation images for the corresponding images.
+  """
+  image_rows = coco_dataset.dataset[constants.COCO_JSON_IMAGES]
+
+  segmentation_rows = [
+      _image_info_to_segmentation_image(
+          img, coco_dataset, label_id_by_category_id
+      )
+      for img in image_rows
+  ]
+
+  ml_use_rows = common_lib.create_ml_use_array_with_split(
+      len(image_rows), split_ratio
+  )
+  return ml_use_rows, image_rows, segmentation_rows
+
+
+def beam_build_tfrecord_from_coco_json(
+    input_json: str,
+    output_dir: str,
+    split_ratio: List[float],
+    num_shard_list: List[int],
+    output_shape: Optional[Tuple[int, int]] = None,
+) -> None:
+  """Builds TFRecord files from COCO dataset.
+
+  The output file names are `_TRAIN_TFRECORD_NAME`, `_VALIDATION_TFRECORD_NAME`,
+  and `_TEST_TFRECORD_NAME`.
+
+  Args:
+    input_json: Path to a COCO JSON or JSONL file.
+    output_dir: Directory to output the TFRecord files.
+    split_ratio: List of how to split entries to train, validation, and test
+      TFRecords.
+    num_shard_list: List of the number of shards for each TFRecord file.
+    output_shape: The desired output shape of the image. If None, the original
+      image shape will be used.
+  """
+  # `coco` cannot access gcs uri. Use gcsfuse, it is faster.
+  input_json = fileutils.force_gcs_fuse_path(input_json)
+  coco_dataset = coco.COCO(input_json)
+
+  label_map = {}
+  label_id_by_category_id = {}
+  for idx, category in enumerate(
+      coco_dataset.dataset[constants.COCO_JSON_CATEGORIES], start=1
+  ):
+    label_map[idx] = category[constants.COCO_JSON_CATEGORY_NAME]
+    label_id_by_category_id[category[constants.COCO_JSON_CATEGORY_ID]] = idx
+  label_map_path = os.path.join(output_dir, common_lib.LABEL_MAP_NAME)
+  logging.info('Writing label map to %s.', label_map_path)
+  common_lib.write_label_map(label_map_path, label_map)
+
+  with tf.io.gfile.GFile(
+      os.path.join(output_dir, 'label_id_by_category_id.yaml'), 'w'
+  ) as f:
+    yaml.dump(label_id_by_category_id, f)
+
+  ml_use_rows, image_rows, segmentation_rows = get_input_rows(
+      coco_dataset, split_ratio, label_id_by_category_id
+  )
+
+  def pipeline(root):
+    _define_data_conversion_pipeline(
+        root,
+        ml_use_rows,
+        image_rows,
+        segmentation_rows,
+        output_dir,
+        output_shape,
+        num_shard_list,
+    )
+
+  logging.info('Beginning beam pipeline to acquire tfrecords.')
+  common_lib.run_beam_pipeline(pipeline)

community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_main.py

@@ -0,0 +1,166 @@
+r"""Python script to convert user input data to training docker format.
+
+
+Note: the training format is designed to be tfrecord as in the design doc.
+If there are training efficiency issues for pytorch algorithms, we will also
+support pytorch formats as well.
+"""
+
+from absl import app
+from absl import flags
+from absl import logging
+from data_converter import common_lib
+from data_converter import data_converter_icn_lib
+from data_converter import data_converter_iod_lib
+from data_converter import data_converter_isg_lib
+from data_converter import data_converter_vcn_lib
+from util import constants
+
+
+_INPUT_FILE_PATH = flags.DEFINE_string(
+    'input_file_path',
+    None,
+    'Input file path.',
+    required=True,
+)
+_INPUT_FILE_TYPE = flags.DEFINE_enum(
+    'input_file_type',
+    None,
+    [
+        constants.INPUT_FILE_TYPE_CSV,
+        constants.INPUT_FILE_TYPE_JSONL,
+        constants.INPUT_FILE_TYPE_COCO_JSON,
+    ],
+    'Input file type.',
+    required=True,
+)
+_OBJECTIVE = flags.DEFINE_enum(
+    'objective',
+    None,
+    [
+        constants.OBJECTIVE_IMAGE_CLASSIFICATION,
+        constants.OBJECTIVE_IMAGE_OBJECT_DETECTION,
+        constants.OBJECTIVE_IMAGE_SEGMENTATION,
+        constants.OBJECTIVE_VIDEO_CLASSIFICATION,
+    ],
+    'The objective of this training job.',
+    required=True,
+)
+_OUTPUT_DIR = flags.DEFINE_string(
+    'output_dir',
+    None,
+    'The output directory for converted data and label map files.',
+    required=True,
+)
+_SPLIT_RATIO = flags.DEFINE_list(
+    'split_ratio',
+    '0.8,0.1,0.1',
+    'Proportion of data to split into train/validation/test.',
+)
+_NUM_SHARD = flags.DEFINE_list(
+    'num_shard', '10,10,10', 'The number of shards for train/validation/test.'
+)
+_OUTPUT_FPS = flags.DEFINE_integer(
+    'output_fps', 5, 'For videos only. The output frames rate per second.'
+)
+
+
+def main(_) -> None:
+  logging.info(
+      (
+          'Start data converter on: %s (type: %s) with split: %s for %s'
+          ' (shard=%s), and output to %s.'
+      ),
+      _INPUT_FILE_PATH.value,
+      _INPUT_FILE_TYPE.value,
+      _SPLIT_RATIO.value,
+      _OBJECTIVE.value,
+      _NUM_SHARD.value,
+      _OUTPUT_DIR.value,
+  )
+  split_ratio = list(map(float, _SPLIT_RATIO.value))
+  num_shard = list(map(int, _NUM_SHARD.value))
+  common_lib.check_split_ratio(split_ratio)
+  common_lib.check_num_shard(num_shard)
+  if (
+      _OBJECTIVE.value == constants.OBJECTIVE_IMAGE_OBJECT_DETECTION
+      and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_CSV
+  ):
+    data_converter_iod_lib.convert_csv_to_tfrecord(
+        _INPUT_FILE_PATH.value,
+        _OUTPUT_DIR.value,
+        split_ratio,
+        num_shard,
+    )
+  elif (
+      _OBJECTIVE.value == constants.OBJECTIVE_IMAGE_OBJECT_DETECTION
+      and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_JSONL
+  ):
+    data_converter_iod_lib.convert_jsonl_to_tfrecord(
+        _INPUT_FILE_PATH.value, _OUTPUT_DIR.value, split_ratio, num_shard
+    )
+  elif (
+      _OBJECTIVE.value == constants.OBJECTIVE_IMAGE_OBJECT_DETECTION
+      and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_COCO_JSON
+  ):
+    data_converter_iod_lib.convert_coco_json_to_tfrecord(
+        _INPUT_FILE_PATH.value,
+        _OUTPUT_DIR.value,
+        split_ratio,
+        num_shard,
+    )
+  elif _OBJECTIVE.value == constants.OBJECTIVE_IMAGE_SEGMENTATION:
+    data_converter_isg_lib.beam_build_tfrecord_from_coco_json(
+        _INPUT_FILE_PATH.value,
+        _OUTPUT_DIR.value,
+        split_ratio,
+        num_shard,
+    )
+  elif (
+      _OBJECTIVE.value == constants.OBJECTIVE_IMAGE_CLASSIFICATION
+      and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_CSV
+  ):
+    data_converter_icn_lib.convert_csv_to_tfrecord(
+        _INPUT_FILE_PATH.value,
+        _OUTPUT_DIR.value,
+        split_ratio,
+        num_shard,
+    )
+  elif (
+      _OBJECTIVE.value == constants.OBJECTIVE_IMAGE_CLASSIFICATION
+      and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_JSONL
+  ):
+    data_converter_icn_lib.convert_jsonl_to_tfrecord(
+        _INPUT_FILE_PATH.value, _OUTPUT_DIR.value, split_ratio, num_shard
+    )
+  elif (
+      _OBJECTIVE.value == constants.OBJECTIVE_VIDEO_CLASSIFICATION
+      and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_CSV
+  ):
+    data_converter_vcn_lib.convert_csv_to_tfrecord(
+        _INPUT_FILE_PATH.value,
+        _OUTPUT_DIR.value,
+        _OUTPUT_FPS.value,
+        split_ratio,
+        num_shard,
+    )
+  elif (
+      _OBJECTIVE.value == constants.OBJECTIVE_VIDEO_CLASSIFICATION
+      and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_JSONL
+  ):
+    data_converter_vcn_lib.convert_jsonl_to_tfrecord(
+        _INPUT_FILE_PATH.value,
+        _OUTPUT_DIR.value,
+        _OUTPUT_FPS.value,
+        split_ratio,
+        num_shard,
+    )
+  else:
+    raise NotImplementedError(
+        f'File format {_INPUT_FILE_TYPE.value} is not supported for'
+        f' {_OBJECTIVE.value}.'
+    )
+
+
+if __name__ == '__main__':
+  app.run(main)

community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_vcn_lib.py

@@ -0,0 +1,289 @@
+"""Converts VCN CSV/JSONL files to TFRecord with apache beam."""
+
+import json
+from os import path
+from typing import Any, Dict, Iterator, Sequence, Union, cast
+
+from absl import logging
+import apache_beam as beam
+from apache_beam.io import tfrecordio
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+
+from data_converter import common_lib
+from util import constants
+
+
+_COLUMN_NAMES = [
+    common_lib.COLUMN_NAME_ML_USE,
+    common_lib.COLUMN_NAME_GCS_FILE_PATH,
+    common_lib.COLUMN_NAME_LABEL,
+    common_lib.COLUMN_NAME_START_SEC,
+    common_lib.COLUMN_NAME_END_SEC,
+]
+_JSON_GCS_URI_KEY = 'videoGcsUri'
+_JSON_CLASS_ANNOTATION_KEY = 'timeSegmentAnnotations'
+_JSON_CLASS_NAME_KEY = 'displayName'
+_JSON_START_TIME_KEY = 'startTime'
+_JSON_END_TIME_KEY = 'endTime'
+_JSON_RESOURCE_LABEL_KEY = 'dataItemResourceLabels'
+_JSON_ML_USE_KEY = 'aiplatform.googleapis.com/ml_use'
+
+
+def build_tf_example(
+    video_uri: str,
+    label: int,
+    start_sec: float,
+    end_sec: float,
+    output_fps: int,
+) -> tf.train.SequenceExample:
+  """Builds a TF Example from a video clip.
+
+  Args:
+    video_uri: GCS URI to the video file.
+    label: Class label as an integer.
+    start_sec: Start timestamp of the video clip in seconds.
+    end_sec: End timestamp of the video clip in seconds.
+    output_fps: The output frame rate per second.
+
+  Returns:
+    The created TF Example.
+  """
+  frame_bytes = common_lib.encode_video(
+      video_uri, start_sec, end_sec, output_fps, image_format='jpg'
+  )
+  seq_example = tf.train.SequenceExample()
+  seq_example.context.feature['clip/label/index'].int64_list.value[:] = [label]
+  for frame in frame_bytes:
+    seq_example.feature_lists.feature_list.get_or_create(
+        'image/encoded'
+    ).feature.add().bytes_list.value[:] = [frame]
+
+  return seq_example
+
+
+class AcquireTFExampleDoFn(beam.DoFn):
+  """Beam DoFn to build TF Examples from a DataFrame row dict for VCN."""
+
+  def __init__(self, output_fps: int):
+    self._success_counter = beam.metrics.Metrics.counter(
+        self.__class__.__name__, 'Success'
+    )
+    self._failure_counter = beam.metrics.Metrics.counter(
+        self.__class__.__name__, 'Failure'
+    )
+    self._output_fps = output_fps
+
+  def process(
+      self, element: Dict[str, Union[float, int, str]]
+  ) -> Iterator[tf.train.SequenceExample]:
+    ml_use: str = cast(str, element[common_lib.COLUMN_NAME_ML_USE])
+    video_uri: str = cast(str, element[common_lib.COLUMN_NAME_GCS_FILE_PATH])
+
+    try:
+      label: int = int(element[common_lib.COLUMN_NAME_LABEL])
+      start_sec: float = float(element[common_lib.COLUMN_NAME_START_SEC])
+      end_sec: float = float(element[common_lib.COLUMN_NAME_END_SEC])
+
+      tf_example = build_tf_example(
+          video_uri,
+          label,
+          start_sec,
+          end_sec,
+          self._output_fps,
+      )
+      self._success_counter.inc()
+      yield beam.pvalue.TaggedOutput(ml_use, tf_example)
+    except (ValueError, IOError) as err:
+      logging.error('Failed to process %s', video_uri)
+      logging.exception(err)
+      self._failure_counter.inc()
+
+
+def _run_convert_pipeline(
+    output_dir: str,
+    df: pd.DataFrame,
+    num_shards: Sequence[int],
+    output_fps: int,
+) -> None:
+  """Starts a Beam pipeline to write DataFrame as TF Records.
+
+  Args:
+    output_dir: TF Records output directory.
+    df: DataFrame to convert from.
+    num_shards: Number of shards for train/validation/test TFRecord files.
+    output_fps: The output frame rate per second.
+  """
+  clip_list = df.to_dict('records')
+
+  def pipeline(root):
+    train, val, test = (
+        root
+        | 'Create PCollection' >> beam.Create(clip_list)
+        | 'Convert to TF Example'
+        >> beam.ParDo(AcquireTFExampleDoFn(output_fps)).with_outputs(
+            constants.ML_USE_TRAINING,
+            constants.ML_USE_VALIDATION,
+            constants.ML_USE_TEST,
+        )
+    )
+    _ = train | 'Save train TF Record' >> tfrecordio.WriteToTFRecord(
+        path.join(output_dir, common_lib.TRAIN_TFRECORD_NAME),
+        coder=beam.coders.ProtoCoder(tf.train.Example),
+        num_shards=num_shards[0],
+    )
+    _ = val | 'Save val TF Record' >> tfrecordio.WriteToTFRecord(
+        path.join(output_dir, common_lib.VALIDATION_TFRECORD_NAME),
+        coder=beam.coders.ProtoCoder(tf.train.Example),
+        num_shards=num_shards[1],
+    )
+    _ = test | 'Save test TF Record' >> tfrecordio.WriteToTFRecord(
+        path.join(output_dir, common_lib.TEST_TFRECORD_NAME),
+        coder=beam.coders.ProtoCoder(tf.train.Example),
+        num_shards=num_shards[2],
+    )
+
+  common_lib.run_beam_pipeline(pipeline)
+
+
+def _convert_df_to_tfrecord(
+    df: pd.DataFrame,
+    output_dir: str,
+    split_ratio: Sequence[float],
+    num_shard: Sequence[int],
+    output_fps: int,
+) -> None:
+  """Converts a DataFrame into three separate tfrecords for training, validation, and testing into output_dir.
+
+  Args:
+    df: DataFrame to convert.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    split_ratio: List specifying the training, validation, and testing splits
+      for unassigned TFRecords.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+    output_fps: The output frame rate per second.
+  """
+  # Replaces ml_use with common_lib string constants for consistency.
+  common_lib.format_ml_use_column(df)
+  common_lib.insert_missing_ml_use(df)
+
+  # Ignores invalid rows.
+  dropped_row_num = common_lib.drop_invalid_rows(df)
+  if dropped_row_num > 0:
+    logging.warning('Ignored %d invalid rows.', dropped_row_num)
+
+  common_lib.replace_unassigned_ml_use(
+      df[common_lib.COLUMN_NAME_ML_USE], split_ratio
+  )
+
+  # Converts labels to integers as required by training.
+  new_labels, label_map = common_lib.create_label_map(
+      df[common_lib.COLUMN_NAME_LABEL]
+  )
+  df[common_lib.COLUMN_NAME_LABEL] = new_labels
+  label_map_path = path.join(output_dir, common_lib.LABEL_MAP_NAME)
+  logging.info('Writing label map to %s.', label_map_path)
+  common_lib.write_label_map(label_map_path, label_map)
+
+  # Missing start / end times are treated as 0, inf, respectively.
+  df[common_lib.COLUMN_NAME_START_SEC].fillna(0, inplace=True)
+  df[common_lib.COLUMN_NAME_END_SEC].fillna(np.inf, inplace=True)
+
+  _run_convert_pipeline(output_dir, df, num_shard, output_fps)
+
+
+def convert_csv_to_tfrecord(
+    input_csv: str,
+    output_dir: str,
+    output_fps: int,
+    split_ratio: Sequence[float] = (0.8, 0.1, 0.1),
+    num_shard: Sequence[int] = (10, 10, 10),
+) -> None:
+  """Parses input_csv file into three separate tfrecords for training, validation, and testing into output_dir.
+
+  The csv format is shown in
+    https://cloud.google.com/vertex-ai/docs/video-data/classification/prepare-data#csv
+
+  If an ml_use column is not provided, one will be created.
+
+  label_map.yaml containing the label map will be placed in output_dir.
+
+  Args:
+    input_csv: Name of the csv file.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    output_fps: The output frame rate per second.
+    split_ratio: List specifying the training, validation, and testing splits
+      for unassigned TFRecords.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+  """
+  with tf.io.gfile.GFile(input_csv, 'r') as f:
+    df: pd.DataFrame = pd.read_csv(
+        f, header=None, names=_COLUMN_NAMES, on_bad_lines='warn'
+    )
+
+  _convert_df_to_tfrecord(df, output_dir, split_ratio, num_shard, output_fps)
+
+
+def convert_jsonl_to_tfrecord(
+    input_jsonl: str,
+    output_dir: str,
+    output_fps: int,
+    split_ratio: Sequence[float] = (0.8, 0.1, 0.1),
+    num_shard: Sequence[int] = (10, 10, 10),
+) -> None:
+  """Parses input_jsonl file into three separate tfrecords for training, validation, and testing into output_dir.
+
+  The JSONL format is shown in
+    https://cloud.google.com/vertex-ai/docs/video-data/classification/prepare-data#jsonl.
+
+  If an ml_use column is not provided, one will be created.
+
+  label_map.yaml containing the label map will be placed in output_dir.
+
+  Args:
+    input_jsonl: Name of the JSONL file.
+    output_dir: The directory to save TFRecords and label_map.yaml.
+    output_fps: The output frame rate per second.
+    split_ratio: List specifying the training, validation, and testing splits
+      for unassigned TFRecords.
+    num_shard: Number of shards for train/validation/test TFRecord files.
+  """
+  df_rows = []
+  with tf.io.gfile.GFile(input_jsonl, 'r') as f:
+    lines = f.read().rstrip().splitlines()
+
+  for i, line in enumerate(lines, 1):
+    try:
+      item: Dict[str, Any] = json.loads(line)
+
+      gcs_uri = item.get(_JSON_GCS_URI_KEY)
+      if not gcs_uri:
+        logging.warning('Invalid JSON at line %d, skipped.', i)
+        continue
+
+      annotations = item.get(_JSON_CLASS_ANNOTATION_KEY, [])
+      ml_use = item.get(_JSON_RESOURCE_LABEL_KEY, {}).get(
+          _JSON_ML_USE_KEY, common_lib.ML_USE_UNASSIGNED
+      )
+
+      for j, annotation in enumerate(annotations):
+        label = annotation.get(_JSON_CLASS_NAME_KEY)
+        if not label:
+          logging.warning('Invalid annotation #%d at line %d, skipped.', j, i)
+          continue
+        # The example in external documentation uses strings like "1.0s", so we
+        # need to remove the "s" suffix.
+        start_time = annotation.get(_JSON_START_TIME_KEY, '0').removesuffix('s')
+        end_time = annotation.get(_JSON_END_TIME_KEY, 'inf').removesuffix('s')
+        df_rows.append([ml_use, gcs_uri, label, start_time, end_time])
+    except (json.JSONDecodeError, AttributeError):
+      logging.warning('Invalid JSON at line %d, skipped.', i)
+      continue
+
+  df = pd.DataFrame(
+      data=df_rows,
+      columns=_COLUMN_NAMES,
+  )
+
+  _convert_df_to_tfrecord(df, output_dir, split_ratio, num_shard, output_fps)

community-content/vertex_vision_model_garden/model_oss/data_converter/dockerfile/base.Dockerfile

@@ -0,0 +1,50 @@
+FROM python:3.9
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Install basic libs.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        cmake \
+        curl \
+        wget \
+        sudo \
+        gnupg \
+        python3-opencv \
+        lsb-release \
+        ca-certificates \
+        build-essential \
+        git \
+        vim \
+        screen \
+        libportaudio2 \
+        libusb-1.0-0-dev \
+        openjdk-17-jre
+
+# Add gcsfuse distribution URL as a package source and import its public key.
+RUN echo "deb https://packages.cloud.google.com/apt gcsfuse-`lsb_release -c -s` main" | sudo tee /etc/apt/sources.list.d/gcsfuse.list
+RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -
+
+# Install gcsfuse.
+RUN apt-get update && apt-get install -y --no-install-recommends gcsfuse
+
+# Install google cloud SDK.
+RUN wget -q https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-359.0.0-linux-x86_64.tar.gz
+RUN tar xzf google-cloud-sdk-359.0.0-linux-x86_64.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 required libs.
+RUN pip install --upgrade pip
+RUN pip install pyyaml==5.4.1
+RUN pip install pycocotools==2.0.6
+RUN pip install opencv-python-headless==4.7.0.72
+RUN pip install numpy==1.24.2
+RUN pip install pandas==1.5.3
+RUN pip install Pillow==9.4.0
+RUN pip install apache-beam[gcp]==2.45.0
+RUN pip install object-detection==0.0.3
+RUN pip install google-cloud-storage==1.42.3
+RUN pip install gcsfs==2021.10.1
+RUN pip install pylint==2.17.2

community-content/vertex_vision_model_garden/model_oss/data_converter/dockerfile/run.Dockerfile

@@ -0,0 +1,23 @@
+FROM gcr.io/automl-migration-test/automl-vision-data-converter-base:latest
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+COPY model_oss/data_converter /automl_vision/data_converter
+COPY model_oss/util /automl_vision/util
+
+WORKDIR /automl_vision
+ENV PYTHONPATH "${PYTHONPATH}:/automl_vision"
+
+# Run pylint to validate code.
+COPY .pylintrc /automl_vision/.pylintrc
+RUN find . -type f -name "*.py" | xargs pylint --rcfile=./.pylintrc --errors-only
+
+ENTRYPOINT ["python3","data_converter/data_converter_main.py"]
+
+CMD ["--input_file_path=YOUR_INPUT_FILE",\
+"--input_file_type=csv",\
+"--objective=iod",\
+"--output_dir=YOUR_OUTPUT_DIR",\
+"--num_shard=10,10,10",\
+"--split_ratio=0.8,0.1,0.1"]

community-content/vertex_vision_model_garden/model_oss/diffusers/dockerfile/serve.Dockerfile

@@ -0,0 +1,74 @@
+# Dockerfile for Diffuser Serving.
+#
+# To build:
+# docker build -f model_oss/diffusers/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/torchserve:0.7.0-gpu
+
+USER root
+
+ENV infer_port=7080
+ENV mng_port=7081
+ENV model_name="diffusers_serving"
+ENV PATH="/home/model-server/:${PATH}"
+
+# Install libraries.
+ENV PIP_ROOT_USER_ACTION=ignore
+RUN python3 -m pip install --upgrade pip
+RUN pip install torch==1.13.1
+RUN pip install torchvision==0.14.1
+RUN pip install transformers==4.27.4
+RUN pip install datasets==2.9.0
+RUN pip install accelerate==0.17.0
+RUN pip install triton==2.0.0.dev20221120
+RUN pip install xformers==0.0.16
+RUN pip install google-cloud-storage==2.7.0
+RUN pip install imageio[ffmpeg]==2.31.0
+RUN pip install absl-py==1.4.0
+
+# Copy LICENSE file
+RUN apt-get update && apt-get install wget
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Install diffusers from main branch source code with a pinned commit.
+RUN git clone --depth 1 --branch v0.18.1 https://github.com/huggingface/diffusers.git
+WORKDIR diffusers
+RUN pip install -e .
+
+# Copy model artifacts.
+COPY model_oss/diffusers/handler.py /home/model-server/handler.py
+COPY model_oss/util/ /home/model-server/util/
+ENV PYTHONPATH /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}" >> /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.
+CMD ["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_vision_model_garden/model_oss/diffusers/dockerfile/train.Dockerfile

@@ -0,0 +1,47 @@
+# Dockerfile for Diffuser Training.
+#
+# To build:
+# docker build -f model_oss/diffusers/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}
+
+# Base on pytorch-cuda image.
+FROM pytorch/pytorch:1.13.0-cuda11.6-cudnn8-runtime
+
+# Install tools.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        git \
+        vim
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Install libraries.
+RUN pip install torchvision==0.14.1
+RUN pip install transformers==4.26.1
+RUN pip install datasets==2.9.0
+RUN pip install accelerate==0.17.0
+RUN pip install triton==2.0.0.dev20221120
+RUN pip install xformers==0.0.16
+RUN pip install Jinja2==3.1.2
+RUN pip install ftfy==6.1.1
+RUN pip install cloudml-hypertune==0.1.0.dev6
+RUN pip install tensorboard==2.12.0
+
+# Install diffusers from main branch source code with a pinned commit.
+RUN git clone --depth 1 --branch v0.18.1 https://github.com/huggingface/diffusers.git
+WORKDIR diffusers
+RUN pip install -e .
+
+# Switch to diffusers examples folder.
+WORKDIR examples
+
+# Config accelerate.
+COPY model_oss/diffusers/train.sh train.sh
+
+# Generate accelerate config at the beginning of docker run.
+ENTRYPOINT ["/bin/bash", "train.sh"]

community-content/vertex_vision_model_garden/model_oss/diffusers/handler.py

@@ -0,0 +1,256 @@
+"""Custom handler for huggingface/diffusers models."""
+
+# pylint: disable=g-importing-member
+# pylint: disable=logging-fstring-interpolation
+
+import base64
+import io
+import logging
+import os
+from typing import Any, List, Sequence, Tuple
+
+from diffusers import ControlNetModel
+from diffusers import DiffusionPipeline
+from diffusers import DPMSolverMultistepScheduler
+from diffusers import EulerAncestralDiscreteScheduler
+from diffusers import StableDiffusionControlNetPipeline
+from diffusers import StableDiffusionImg2ImgPipeline
+from diffusers import StableDiffusionInpaintPipeline
+from diffusers import StableDiffusionInstructPix2PixPipeline
+from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionUpscalePipeline
+from diffusers import TextToVideoZeroPipeline
+from diffusers import UniPCMultistepScheduler
+import imageio
+import numpy as np
+from PIL import Image
+import torch
+from ts.torch_handler.base_handler import BaseHandler
+
+from util import constants
+from util import fileutils
+from util import image_format_converter
+from video_util import video_format_converter
+
+STABLE_DIFFUSION_MODEL = "runwayml/stable-diffusion-v1-5"
+
+# Tasks
+TEXT_TO_IMAGE = "text-to-image"
+IMAGE_TO_IMAGE = "image-to-image"
+IMAGE_INPAINTING = "image-inpainting"
+INSTRUCT_PIX2PIX = "instruct-pix2pix"
+CONTROLNET = "controlnet"
+CONDITIONED_SUPER_RES = "conditioned-super-res"
+TEXT_TO_VIDEO_ZERO_SHOT = "text-to-video-zero-shot"
+TEXT_TO_VIDEO = "text-to-video"
+
+
+def frames_to_video_bytes(frames: Sequence[np.ndarray], fps: int) -> bytes:
+  images = [Image.fromarray(array) for array in frames]
+  io_obj = io.BytesIO()
+  imageio.mimsave(io_obj, images, format=".mp4", fps=fps)
+  return io_obj.getvalue()
+
+
+class DiffusersHandler(BaseHandler):
+  """Custom handler for TIMM models."""
+
+  def initialize(self, context: Any):
+    """Custom initialize."""
+
+    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.model_id = os.environ["MODEL_ID"]
+    if self.model_id.startswith(constants.GCS_URI_PREFIX):
+      gcs_path = self.model_id[len(constants.GCS_URI_PREFIX) :]
+      local_model_dir = os.path.join(constants.LOCAL_MODEL_DIR, gcs_path)
+      logging.info(f"Download {self.model_id} to {local_model_dir}")
+      fileutils.download_gcs_dir_to_local(self.model_id, local_model_dir)
+      self.model_id = local_model_dir
+
+    self.task = os.environ.get("TASK", TEXT_TO_IMAGE)
+    logging.info(f"Using task:{self.task}, model:{self.model_id}")
+
+    if self.task == TEXT_TO_IMAGE:
+      pipeline = StableDiffusionPipeline.from_pretrained(
+          self.model_id, torch_dtype=torch.float16
+      )
+      pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
+          pipeline.scheduler.config
+      )
+      pipeline = pipeline.to(self.map_location)
+      # Reduce memory footprint.
+      pipeline.enable_attention_slicing()
+    elif self.task == IMAGE_TO_IMAGE:
+      pipeline = StableDiffusionImg2ImgPipeline.from_pretrained(
+          self.model_id, torch_dtype=torch.float16
+      )
+      pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
+          pipeline.scheduler.config
+      )
+      pipeline = pipeline.to(self.map_location)
+      # Reduce memory footprint.
+      pipeline.enable_attention_slicing()
+    elif self.task == IMAGE_INPAINTING:
+      pipeline = StableDiffusionInpaintPipeline.from_pretrained(
+          self.model_id, torch_dtype=torch.float16
+      )
+      pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
+          pipeline.scheduler.config
+      )
+      pipeline = pipeline.to(self.map_location)
+      # Reduce memory footprint.
+      pipeline.enable_attention_slicing()
+    elif self.task == INSTRUCT_PIX2PIX:
+      pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+          self.model_id, torch_dtype=torch.float16
+      )
+      pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
+          pipeline.scheduler.config
+      )
+      pipeline = pipeline.to(self.map_location)
+      # Reduce memory footprint.
+      pipeline.enable_attention_slicing()
+    elif self.task == CONTROLNET:
+      controlnet = ControlNetModel.from_pretrained(
+          self.model_id, torch_dtype=torch.float16
+      )
+      pipeline = StableDiffusionControlNetPipeline.from_pretrained(
+          STABLE_DIFFUSION_MODEL,
+          controlnet=controlnet,
+          torch_dtype=torch.float16,
+      )
+      pipeline.scheduler = UniPCMultistepScheduler.from_config(
+          pipeline.scheduler.config
+      )
+      pipeline.enable_xformers_memory_efficient_attention()
+      pipeline.enable_model_cpu_offload()
+      pipeline = pipeline.to(self.map_location)
+      # Reduce memory footprint.
+      pipeline.enable_attention_slicing()
+    elif self.task == CONDITIONED_SUPER_RES:
+      pipeline = StableDiffusionUpscalePipeline.from_pretrained(
+          self.model_id, torch_dtype=torch.float16
+      )
+      pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
+          pipeline.scheduler.config
+      )
+      # This is necessary to 4x upscale >=256x256 input images with V100.
+      logging.info("Enable xformers memory efficient attention for inference.")
+      pipeline.enable_xformers_memory_efficient_attention()
+      pipeline = pipeline.to(self.map_location)
+      # Reduce memory footprint.
+      pipeline.enable_attention_slicing()
+    elif self.task == TEXT_TO_VIDEO_ZERO_SHOT:
+      pipeline = TextToVideoZeroPipeline.from_pretrained(
+          STABLE_DIFFUSION_MODEL, torch_dtype=torch.float16
+      )
+      # Memory optimization.
+      pipeline.enable_xformers_memory_efficient_attention()
+      pipeline.enable_model_cpu_offload()
+      pipeline = pipeline.to(self.map_location)
+    elif self.task == TEXT_TO_VIDEO:
+      pipeline = DiffusionPipeline.from_pretrained(
+          self.model_id, torch_dtype=torch.float16, variant="fp16"
+      )
+      pipeline.enable_model_cpu_offload()
+      # Memory optimization.
+      pipeline.enable_vae_slicing()
+      pipeline.scheduler = DPMSolverMultistepScheduler.from_config(
+          pipeline.scheduler.config
+      )
+    else:
+      raise ValueError(f"Invalid TASK: {self.task}")
+
+    self.pipeline = pipeline
+    self.initialized = True
+    logging.info("Handler initialization done.")
+
+  def preprocess(self, data: Any) -> Tuple[Any, Any, Any]:
+    """Preprocess input data."""
+    prompts = [item["prompt"] for item in data]
+    images = None
+    mask_images = None
+
+    if "image" in data[0]:
+      images = [
+          image_format_converter.base64_to_image(item["image"]) for item in data
+      ]
+    if "mask_image" in data[0]:
+      mask_images = [
+          image_format_converter.base64_to_image(item["mask_image"])
+          for item in data
+      ]
+    return prompts, images, mask_images
+
+  def inference(self, data: Any, *args, **kwargs) -> List[Image.Image]:
+    """Run the inference."""
+    prompts, images, mask_images = data
+    if self.task == TEXT_TO_IMAGE:
+      predicted_images = self.pipeline(prompt=prompts).images
+    elif self.task == IMAGE_TO_IMAGE:
+      predicted_images = self.pipeline(prompt=prompts, image=images).images
+    elif self.task == IMAGE_INPAINTING:
+      predicted_images = self.pipeline(
+          prompt=prompts, image=images, mask_image=mask_images
+      ).images
+    elif self.task == INSTRUCT_PIX2PIX:
+      predicted_images = self.pipeline(prompt=prompts, image=images).images
+    elif self.task == CONTROLNET:
+      predicted_images = self.pipeline(
+          prompt=prompts, image=images, num_inference_steps=20
+      ).images
+    elif self.task == CONDITIONED_SUPER_RES:
+      predicted_images = self.pipeline(
+          prompt=prompts, image=images, num_inference_steps=20
+      ).images
+    elif self.task == TEXT_TO_VIDEO_ZERO_SHOT:
+      # For each given prompt, generate a short video.
+      # The pipeline doesn't support multiple prompts in one run yet.
+      videos = []
+      for prompt in prompts:
+        numpy_arrays = self.pipeline(prompt=prompt).images
+        numpy_arrays = [(i * 255).astype("uint8") for i in numpy_arrays]
+        videos.append(
+            frames_to_video_bytes(numpy_arrays, fps=4)
+        )
+      return videos
+    elif self.task == TEXT_TO_VIDEO:
+      predicted_images = np.asarray(self.pipeline(prompt=prompts).frames)
+      # For multiple prompts, the model concatenates video frames, i.e. the
+      # output shape is (num_frames, height, width * len(prompts), channels).
+      # Therefore we need to split the output into different videos.
+      predicted_images = np.array_split(predicted_images, len(prompts), axis=2)
+      videos = [
+          frames_to_video_bytes(images, fps=8)
+          for images in predicted_images
+      ]
+      return videos
+    else:
+      raise ValueError(f"Invalid TASK: {self.task}")
+    return predicted_images
+
+  def postprocess(self, data: Any) -> List[str]:
+    """Convert the images to base64 string."""
+    outputs = []
+    for prediction in data:
+      if isinstance(prediction, bytes):
+        # This is the video bytes.
+        outputs.append(base64.b64encode(prediction).decode("utf-8"))
+      else:
+        outputs.append(image_format_converter.image_to_base64(prediction))
+    return outputs
+
+
+# pylint: enable=logging-fstring-interpolation

community-content/vertex_vision_model_garden/model_oss/diffusers/train.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+
+# Setup accelerate config before running trainer.
+python -c "from accelerate.utils import write_basic_config; write_basic_config(mixed_precision='fp16')"
+
+accelerate launch "$@"

community-content/vertex_vision_model_garden/model_oss/keras/dockerfile/serve.Dockerfile

@@ -0,0 +1,118 @@
+# Dockerfile for basic serving dockers with Keras.
+#
+# To build:
+# docker build -f model_oss/keras/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 tensorflow/tensorflow:2.12.0-gpu
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# This is added to fix docker build error related to Nvidia key update.
+RUN rm -f /etc/apt/sources.list.d/cuda.list
+RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -
+
+# 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 \
+        vim \
+        screen \
+        libtcmalloc-minimal4
+
+
+# Install google cloud SDK.
+RUN wget -q https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-359.0.0-linux-x86_64.tar.gz
+RUN tar xzf google-cloud-sdk-359.0.0-linux-x86_64.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 required libs.
+RUN pip install --upgrade pip
+RUN pip install cloud-tpu-client==0.10
+RUN pip install pyyaml==5.4.1
+RUN pip install fsspec==2021.10.1
+RUN pip install gcsfs==2021.10.1
+RUN pip install tensorflow-text==2.11.0
+RUN pip install pyglove==0.1.0
+RUN pip install cloudml-hypertune==0.1.0.dev6
+RUN pip install pylint==2.17.2
+RUN pip install keras-cv==0.4.0
+RUN pip install tensorflow-datasets==4.8.3
+RUN pip install protobuf==3.20.3
+RUN pip install Pillow==9.5.0
+RUN pip install flask==2.3.2
+RUN pip install waitress==2.1.2
+
+# Installs Reduction Server NCCL plugin.
+RUN echo "deb https://packages.cloud.google.com/apt google-fast-socket main" | tee /etc/apt/sources.list.d/google-fast-socket.list \
+&&  curl -s -L https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - \
+&&  apt update && apt install -y google-reduction-server
+
+# Downloading gcloud package
+RUN curl https://dl.google.com/dl/cloudsdk/release/google-cloud-sdk.tar.gz > /tmp/google-cloud-sdk.tar.gz
+
+# Installing the package
+RUN mkdir -p /usr/local/gcloud \
+  && tar -C /usr/local/gcloud -xvf /tmp/google-cloud-sdk.tar.gz \
+  && /usr/local/gcloud/google-cloud-sdk/install.sh
+
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Adding the package path to local
+ENV PATH $PATH:/usr/local/gcloud/google-cloud-sdk/bin
+
+
+ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=cpp
+
+# Lower the memory fragmentation, and speed up the training.
+# https://github.com/tensorflow/tensorflow/issues/44176#issuecomment-783768033
+ENV LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4
+
+# Enable userspace DNS cache
+ENV GCS_RESOLVE_REFRESH_SECS=60
+ENV GCS_REQUEST_CONNECTION_TIMEOUT_SECS=300
+ENV GCS_METADATA_REQUEST_TIMEOUT_SECS=300
+ENV GCS_READ_REQUEST_TIMEOUT_SECS=300
+ENV GCS_WRITE_REQUEST_TIMEOUT_SECS=600
+# Each opened GCS file takes GCS_READ_CACHE_BLOCK_SIZE_MB of RAM, reduce the
+# value from the default 64MB to 8MB to decrease memory footprint.
+ENV GCS_READ_CACHE_BLOCK_SIZE_MB=8
+
+EXPOSE 8501
+
+WORKDIR /usr/local/lib/python3.8/dist-packages/official/vision
+
+COPY model_oss/keras /automl_vision/keras
+COPY model_oss/util /automl_vision/util
+
+WORKDIR /automl_vision
+
+ENV PYTHONPATH "${PYTHONPATH}:/automl_vision/util"
+ENV MODEL_PATH ""
+ENV IMAGE_WIDTH "512"
+ENV IMAGE_HEIGHT "512"
+
+COPY model_oss/keras/serve.py ./app.py
+
+# Run pylint to validate code.
+COPY .pylintrc /automl_vision/.pylintrc
+RUN find . -type f -name "*.py" | xargs pylint --rcfile=./.pylintrc --errors-only
+
+ENTRYPOINT ["flask","run"]
+CMD ["--host=0.0.0.0", "--port=8501"]

community-content/vertex_vision_model_garden/model_oss/keras/dockerfile/train.Dockerfile

@@ -0,0 +1,111 @@
+# Dockerfile for basic training dockers with Keras.
+#
+# To build:
+# docker build -f model_oss/keras/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 tensorflow/tensorflow:2.12.0-gpu
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# This is added to fix docker build error related to Nvidia key update.
+RUN rm -f /etc/apt/sources.list.d/cuda.list
+RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -
+
+# 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 \
+        vim \
+        screen \
+        libtcmalloc-minimal4
+
+
+# Install google cloud SDK.
+RUN wget -q https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-359.0.0-linux-x86_64.tar.gz
+RUN tar xzf google-cloud-sdk-359.0.0-linux-x86_64.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 required libs.
+RUN pip install --upgrade pip
+RUN pip install cloud-tpu-client==0.10
+RUN pip install pyyaml==5.4.1
+RUN pip install fsspec==2021.10.1
+RUN pip install gcsfs==2021.10.1
+RUN pip install tensorflow-text==2.11.0
+RUN pip install pyglove==0.1.0
+RUN pip install cloudml-hypertune==0.1.0.dev6
+RUN pip install pylint==2.17.2
+RUN pip install keras-cv==0.4.0
+RUN pip install tensorflow-datasets==4.8.3
+RUN pip install tensorflow-estimator==2.12.0
+RUN pip install tensorflow-gcs-config==2.12.0
+RUN pip install tensorflow-hub==0.13.0
+RUN pip install tensorflow-io-gcs-filesystem==0.32.0
+RUN pip install tensorflow-metadata==1.13.1
+RUN pip install tensorflow-probability==0.19.0
+RUN pip install tensorboard==2.12.2
+RUN pip install tensorboard-data-server==0.7.0
+RUN pip install tensorboard-plugin-wit==1.8.1
+RUN pip install protobuf==3.20.3
+RUN pip install pandas==1.5.3
+RUN pip install pandas-datareader==0.10.0
+RUN pip install pandas-gbq==0.17.9
+RUN pip install pycocotools==2.0.6
+
+# Installs Reduction Server NCCL plugin.
+RUN echo "deb https://packages.cloud.google.com/apt google-fast-socket main" | tee /etc/apt/sources.list.d/google-fast-socket.list \
+&&  curl -s -L https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - \
+&&  apt update && apt install -y google-reduction-server
+
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=cpp
+
+# Lower the memory fragmentation, and speed up the training.
+# https://github.com/tensorflow/tensorflow/issues/44176#issuecomment-783768033
+ENV LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4
+
+# Enable userspace DNS cache
+ENV GCS_RESOLVE_REFRESH_SECS=60
+ENV GCS_REQUEST_CONNECTION_TIMEOUT_SECS=300
+ENV GCS_METADATA_REQUEST_TIMEOUT_SECS=300
+ENV GCS_READ_REQUEST_TIMEOUT_SECS=300
+ENV GCS_WRITE_REQUEST_TIMEOUT_SECS=600
+# Each opened GCS file takes GCS_READ_CACHE_BLOCK_SIZE_MB of RAM, reduce the
+# value from the default 64MB to 8MB to decrease memory footprint.
+ENV GCS_READ_CACHE_BLOCK_SIZE_MB=8
+
+WORKDIR /usr/local/lib/python3.8/dist-packages/official/vision
+
+COPY model_oss/keras /automl_vision/keras
+COPY model_oss/util /automl_vision/util
+
+WORKDIR /automl_vision
+
+# Keras stable diffusion training codes set width and height as RESOLUTION.
+ENV RESOLUTION "512"
+
+ENV PYTHONPATH "${PYTHONPATH}:/automl_vision/util"
+
+# Run pylint to validate code.
+COPY .pylintrc /automl_vision/.pylintrc
+RUN find . -type f -name "*.py" | xargs pylint --rcfile=./.pylintrc --errors-only
+
+ENTRYPOINT ["python3","keras/train.py"]

community-content/vertex_vision_model_garden/model_oss/keras/serve.py

@@ -0,0 +1,184 @@
+r"""Servers Keras Stable Diffusion models.
+
+python serve.py --model_path=<model path in gcs>
+
+curl -d \
+'{"prompt":"Hello Kitty"}' \
+-H "Content-Type: application/json" \
+-X POST http://localhost:8501/predict
+"""
+
+import base64
+import io
+import json
+import os
+from typing import List, Tuple
+
+from absl import app
+# The docker builds could not find flask and waitress.
+# pylint: disable=import-error
+from flask import Flask
+from flask import request
+from flask import Response
+import keras_cv
+from PIL import Image
+from waitress import serve
+
+from util import constants
+from util import fileutils
+
+
+flask_app = Flask(__name__)
+
+stable_diffusion_model = None
+
+
+model_path = os.environ.get('MODEL_PATH', '')
+if model_path.startswith(constants.GCS_URI_PREFIX):
+  print('Downloading models from gcs to local.')
+  os.makedirs(constants.LOCAL_MODEL_DIR, exist_ok=True)
+  fileutils.download_gcs_dir_to_local(
+      os.path.dirname(model_path), constants.LOCAL_MODEL_DIR
+  )
+  model_path = os.path.join(
+      constants.LOCAL_MODEL_DIR, os.path.basename(model_path)
+  )
+
+image_width = int(os.environ.get('IMAGE_WIDTH', 512))
+image_height = int(os.environ.get('IMAGE_HEIGHT', 512))
+
+print('image_width=', image_width, 'image_height=', image_height)
+print('Create Keras stable diffusion models.')
+stable_diffusion_model = keras_cv.models.StableDiffusion(
+    img_width=image_width,
+    img_height=image_height,
+    jit_compile=True,
+)
+
+if model_path:
+  # We just reload the weights of the fine-tuned diffusion model.
+  print('Initialize finetuned models from: ', model_path)
+  stable_diffusion_model.diffusion_model.load_weights(model_path)
+
+
+def error(message: str) -> str:
+  """Returns a JSON representing an error response."""
+  return json.dumps({
+      'success': False,
+      'error': message,
+  })
+
+
+def check_key_in_json(content: str, keys: List[str]) -> str:
+  for key in keys:
+    if key not in content:
+      return error('No {} in request {}.'.format(key, content))
+  return None
+
+
+def validate_json_key(json_key_string: str) -> Tuple[str, bool]:
+  try:
+    json_key = json.loads(json_key_string)
+  except (ValueError, TypeError):
+    return (error('Invalid key found in request'), False)
+  return (json_key, True)
+
+
+# The health check route is required for docker deployment in google cloud.
+@flask_app.route('/ping')
+def ping() -> Response:
+  """Health checks."""
+  return Response(status=200)
+
+
+# The return should be `Response` for docker deployment in google cloud.
+@flask_app.route('/predict', methods=['GET', 'POST'])
+def predict_model() -> Response:
+  """Predictions."""
+  if request.method == 'POST':
+    contents = request.get_json(force=True)
+
+    print('The input contents are:', contents)
+    batch_size = 1
+    num_steps = 25
+    seed = 1234
+    if 'parameters' in contents:
+      parameters = contents['parameters']
+      if 'batch_size' in parameters:
+        batch_size = int(parameters['batch_size'])
+      if 'num_steps' in parameters:
+        num_steps = int(parameters['num_steps'])
+      if 'seed' in parameters:
+        seed = int(parameters['seed'])
+    print('batch_size=', batch_size, 'num_steps=', num_steps, 'seed=', seed)
+    if batch_size < 1:
+      return Response(
+          response=error('The batch size must be a positive integar.'),
+          status=200,
+          mimetype='text/plain',
+      )
+    if num_steps < 1:
+      return Response(
+          response=error('The num steps must be a positive integar.'),
+          status=200,
+          mimetype='text/plain',
+      )
+    predictions = []
+    for content in contents['instances']:
+      print('Processing:', content)
+      prompt = content['prompt']
+      generated_image_array = stable_diffusion_model.text_to_image(
+          prompt=prompt,
+          batch_size=batch_size,
+          num_steps=num_steps,
+          seed=seed,
+      )
+
+      generated_image_bytes_array = []
+      for i in range(batch_size):
+        generated_image = Image.fromarray(generated_image_array[i])
+        # Converts the image to a base64-encoded string.
+        buffered_image = io.BytesIO()
+        generated_image.save(buffered_image, format='JPEG')
+        generated_image_bytes = base64.b64encode(
+            buffered_image.getvalue()
+        ).decode('utf-8')
+        generated_image_bytes_array.append(generated_image_bytes)
+      prediction = {
+          'prompt': prompt,
+          'predicted_image': generated_image_bytes_array,
+      }
+      predictions.append(prediction)
+
+    return Response(
+        response=json.dumps({
+            'success': True,
+            'predictions': predictions,
+        }),
+        status=200,
+        mimetype='text/plain',
+    )
+  else:
+    return Response(
+        response=json.dumps({
+            'success': True,
+            'isalive': stable_diffusion_model is not None,
+        }),
+        status=200,
+        mimetype='text/plain',
+    )
+
+
+def serve_main(unused_argv):
+  """The main function to serve Keras models."""
+  del unused_argv
+  # This is used when running locally only. When deploying to Google App
+  # Engine, a webserver process such as Gunicorn will serve the app.
+  # # Debug deployment.
+  # flask_app.run(host='0.0.0.0', port=8501, debug=True)
+  # Prod deployment.
+  serve(flask_app, host='0.0.0.0', port=8501)
+
+
+if __name__ == '__main__':
+  app.run(serve_main)

community-content/vertex_vision_model_garden/model_oss/keras/train.py

@@ -0,0 +1,363 @@
+"""Train Keras Stable Diffusion.
+
+Most the codes below are from
+https://keras.io/examples/generative/finetune_stable_diffusion/.
+"""
+import os
+
+from absl import app
+from absl import flags
+from absl import logging
+import keras_cv
+# pylint: disable=g-importing-member
+from keras_cv.models.stable_diffusion.clip_tokenizer import SimpleTokenizer
+from keras_cv.models.stable_diffusion.diffusion_model import DiffusionModel
+from keras_cv.models.stable_diffusion.image_encoder import ImageEncoder
+from keras_cv.models.stable_diffusion.noise_scheduler import NoiseScheduler
+from keras_cv.models.stable_diffusion.text_encoder import TextEncoder
+import numpy as np
+# The docker builds could not find pandas.
+# pylint: disable=import-error
+import pandas as pd
+import tensorflow as tf
+from tensorflow import keras
+import tensorflow.experimental.numpy as tnp
+
+from util import constants
+from util import fileutils
+
+_INPUT_CSV_PATH = flags.DEFINE_string(
+    'input_csv_path',
+    None,
+    'The input csv path.',
+    required=True,
+)
+
+_USE_MP = flags.DEFINE_bool(
+    'use_mp',
+    True,
+    'Enable mixed-precision training if the underlying GPU has tensor cores.',
+)
+
+_EPOCHS = flags.DEFINE_integer('epochs', 1, 'The number of epochs.')
+
+_OUTPUT_MODEL_DIR = flags.DEFINE_string(
+    'output_model_dir',
+    None,
+    'The output model dir.',
+    required=True,
+)
+
+# These hyperparameters defaults come from this tutorial by Hugging Face:
+# https://huggingface.co/docs/diffusers/training/text2image
+_LEARNING_RATE = flags.DEFINE_float(
+    'learning_rate', 1e-5, 'The learning rate parameter for AdamW optimizer.'
+)
+
+_BETA_1 = flags.DEFINE_float(
+    'beta_1', 0.9, 'The beta_1 parameter for AdamW optimizer.'
+)
+
+_BETA_2 = flags.DEFINE_float(
+    'beta_2', 0.999, 'The beta_2 parameter for AdamW optimizer.'
+)
+
+_WEIGHT_DECAY = flags.DEFINE_float(
+    'weight_decay', 1e-2, 'The weight decay parameter for AdamW optimizer.'
+)
+
+_EPSILON = flags.DEFINE_float(
+    'epsilon', 1e-08, 'The epsilon parameter for AdamW optimizer.'
+)
+
+RESOLUTION = int(os.environ.get('RESOLUTION', 512))
+
+# The padding token and maximum prompt length are specific to the text encoder.
+# If you're using a different text encoder be sure to change them accordingly.
+PADDING_TOKEN = 49407
+MAX_PROMPT_LENGTH = 77
+
+AUTO = tf.data.AUTOTUNE
+POS_IDS = tf.convert_to_tensor([list(range(MAX_PROMPT_LENGTH))], dtype=tf.int32)
+
+
+augmenter = keras.Sequential(
+    layers=[
+        keras_cv.layers.CenterCrop(RESOLUTION, RESOLUTION),
+        keras_cv.layers.RandomFlip(),
+        tf.keras.layers.Rescaling(scale=1.0 / 127.5, offset=-1),
+    ]
+)
+text_encoder = TextEncoder(MAX_PROMPT_LENGTH)
+
+
+def process_image(image_path, tokenized_text):
+  image = tf.io.read_file(image_path)
+  image = tf.io.decode_png(image, 3)
+  image = tf.image.resize(image, (RESOLUTION, RESOLUTION))
+  return image, tokenized_text
+
+
+def apply_augmentation(image_batch, token_batch):
+  return augmenter(image_batch), token_batch
+
+
+def run_text_encoder(image_batch, token_batch):
+  return (
+      image_batch,
+      token_batch,
+      text_encoder([token_batch, POS_IDS], training=False),
+  )
+
+
+def prepare_dict(image_batch, token_batch, encoded_text_batch):
+  return {
+      'images': image_batch,
+      'tokens': token_batch,
+      'encoded_text': encoded_text_batch,
+  }
+
+
+def prepare_dataset(image_paths, tokenized_texts, batch_size=1):
+  dataset = tf.data.Dataset.from_tensor_slices((image_paths, tokenized_texts))
+  dataset = dataset.shuffle(batch_size * 10)
+  dataset = dataset.map(process_image, num_parallel_calls=AUTO).batch(
+      batch_size
+  )
+  dataset = dataset.map(apply_augmentation, num_parallel_calls=AUTO)
+  dataset = dataset.map(run_text_encoder, num_parallel_calls=AUTO)
+  dataset = dataset.map(prepare_dict, num_parallel_calls=AUTO)
+  return dataset.prefetch(AUTO)
+
+
+def prepare_training_dataset(dataset_csv):
+  """Prepares training datasets."""
+  if dataset_csv.startswith(constants.GCS_URI_PREFIX):
+    if not os.path.exists(constants.LOCAL_DATA_DIR):
+      os.makedirs(constants.LOCAL_DATA_DIR)
+    logging.info(
+        'Start to download data from %s to %s.',
+        os.path.dirname(dataset_csv),
+        constants.LOCAL_DATA_DIR,
+    )
+    fileutils.download_gcs_dir_to_local(
+        os.path.dirname(dataset_csv), constants.LOCAL_DATA_DIR
+    )
+    data_frame = pd.read_csv(
+        os.path.join(constants.LOCAL_DATA_DIR, os.path.basename(dataset_csv))
+    )
+    data_frame['image_path'] = data_frame['image_path'].apply(
+        lambda x: os.path.join(constants.LOCAL_DATA_DIR, x)
+    )
+  else:
+    # Keeps the following codes for experiments with
+    # https://keras.io/examples/generative/finetune_stable_diffusion/.
+    data_path = tf.keras.utils.get_file(origin=dataset_csv, untar=True)
+    data_frame = pd.read_csv(os.path.join(data_path, 'data.csv'))
+    data_frame['image_path'] = data_frame['image_path'].apply(
+        lambda x: os.path.join(data_path, x)
+    )
+  data_frame.head()
+
+  # Load the tokenizer.
+  tokenizer = SimpleTokenizer()
+
+  #  Method to tokenize and pad the tokens.
+  def process_text(caption):
+    tokens = tokenizer.encode(caption)
+    tokens = tokens + [PADDING_TOKEN] * (MAX_PROMPT_LENGTH - len(tokens))
+    return np.array(tokens)
+
+  # Collate the tokenized captions into an array.
+  tokenized_texts = np.empty((len(data_frame), MAX_PROMPT_LENGTH))
+
+  all_captions = list(data_frame['caption'].values)
+  for i, caption in enumerate(all_captions):
+    tokenized_texts[i] = process_text(caption)
+
+  # Prepare the dataset.
+  training_dataset = prepare_dataset(
+      np.array(data_frame['image_path']), tokenized_texts, batch_size=4
+  )
+
+  return training_dataset
+
+
+class Trainer(tf.keras.Model):
+  """The trainer for Keras Stable Diffusion."""
+
+  # Reference:
+  # https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image.py
+
+  def __init__(
+      self,
+      diffusion_model,
+      vae,
+      noise_scheduler,
+      use_mixed_precision=False,
+      max_grad_norm=1.0,
+      **kwargs,
+  ):
+    super().__init__(**kwargs)
+
+    self.diffusion_model = diffusion_model
+    self.vae = vae
+    self.noise_scheduler = noise_scheduler
+    self.max_grad_norm = max_grad_norm
+
+    self.use_mixed_precision = use_mixed_precision
+    self.vae.trainable = False
+
+  def train_step(self, inputs):
+    images = inputs['images']
+    encoded_text = inputs['encoded_text']
+    batch_size = tf.shape(images)[0]
+
+    with tf.GradientTape() as tape:
+      # Project image into the latent space and sample from it.
+      latents = self.sample_from_encoder_outputs(
+          self.vae(images, training=False)
+      )
+      # Know more about the magic number here:
+      # https://keras.io/examples/generative/fine_tune_via_textual_inversion/
+      latents = latents * 0.18215
+
+      # Sample noise that we'll add to the latents.
+      noise = tf.random.normal(tf.shape(latents))
+
+      # Sample a random timestep for each image.
+      timesteps = tnp.random.randint(
+          0, self.noise_scheduler.train_timesteps, (batch_size,)
+      )
+
+      # Add noise to the latents according to the noise magnitude at each
+      # timestep (this is the forward diffusion process).
+      noisy_latents = self.noise_scheduler.add_noise(
+          tf.cast(latents, noise.dtype), noise, timesteps
+      )
+
+      # Get the target for loss depending on the prediction type
+      # just the sampled noise for now.
+      target = noise  # noise_schedule.predict_epsilon == True
+
+      # Predict the noise residual and compute loss.
+      # pylint: disable=unnecessary-lambda
+      timestep_embedding = tf.map_fn(
+          lambda t: self.get_timestep_embedding(t), timesteps, dtype=tf.float32
+      )
+      timestep_embedding = tf.squeeze(timestep_embedding, 1)
+      model_pred = self.diffusion_model(
+          [noisy_latents, timestep_embedding, encoded_text], training=True
+      )
+      loss = self.compiled_loss(target, model_pred)
+      if self.use_mixed_precision:
+        loss = self.optimizer.get_scaled_loss(loss)
+
+    # Update parameters of the diffusion model.
+    trainable_vars = self.diffusion_model.trainable_variables
+    gradients = tape.gradient(loss, trainable_vars)
+    if self.use_mixed_precision:
+      gradients = self.optimizer.get_unscaled_gradients(gradients)
+    gradients = [tf.clip_by_norm(g, self.max_grad_norm) for g in gradients]
+    self.optimizer.apply_gradients(zip(gradients, trainable_vars))
+
+    return {m.name: m.result() for m in self.metrics}
+
+  def get_timestep_embedding(self, timestep, dim=320, max_period=10000):
+    half = dim // 2
+    log_max_preiod = tf.math.log(tf.cast(max_period, tf.float32))
+    # The docker builds could not support unary `-`.
+    # pylint: disable=invalid-unary-operand-type
+    freqs = tf.math.exp(
+        -log_max_preiod * tf.range(0, half, dtype=tf.float32) / half
+    )
+    args = tf.convert_to_tensor([timestep], dtype=tf.float32) * freqs
+    embedding = tf.concat([tf.math.cos(args), tf.math.sin(args)], 0)
+    embedding = tf.reshape(embedding, [1, -1])
+    return embedding
+
+  def sample_from_encoder_outputs(self, outputs):
+    mean, logvar = tf.split(outputs, 2, axis=-1)
+    logvar = tf.clip_by_value(logvar, -30.0, 20.0)
+    std = tf.exp(0.5 * logvar)
+    sample = tf.random.normal(tf.shape(mean), dtype=mean.dtype)
+    return mean + std * sample
+
+  def save_weights(
+      self, filepath, overwrite=True, save_format=None, options=None
+  ):
+    # Overriding this method will allow us to use the `ModelCheckpoint`
+    # callback directly with this trainer class. In this case, it will
+    # only checkpoint the `diffusion_model` since that's what we're training
+    # during fine-tuning.
+    self.diffusion_model.save_weights(
+        filepath=filepath,
+        overwrite=overwrite,
+        save_format=save_format,
+        options=options,
+    )
+
+
+def main(_) -> None:
+  # _INPUT_CSV_PATH and _OUTPUT_MODEL_DIR should have the format as
+  # gs://<bucket_name>/<object_name>.
+  if _INPUT_CSV_PATH.value:
+    if not _INPUT_CSV_PATH.value.startswith(constants.GCS_URI_PREFIX):
+      raise ValueError('The input csv path should be a gcs path like gs://<>')
+  if _OUTPUT_MODEL_DIR.value:
+    if not _OUTPUT_MODEL_DIR.value.startswith(constants.GCS_URI_PREFIX):
+      raise ValueError('The output model dir should be a gcs path like gs://<>')
+
+  if _USE_MP.value:
+    keras.mixed_precision.set_global_policy('mixed_float16')
+
+  image_encoder = ImageEncoder(RESOLUTION, RESOLUTION)
+  diffusion_ft_trainer = Trainer(
+      diffusion_model=DiffusionModel(RESOLUTION, RESOLUTION, MAX_PROMPT_LENGTH),
+      # Remove the top layer from the encoder, which cuts off the variance and
+      # only returns the mean.
+      vae=tf.keras.Model(
+          image_encoder.input,
+          image_encoder.layers[-2].output,
+      ),
+      noise_scheduler=NoiseScheduler(),
+      use_mixed_precision=_USE_MP.value,
+  )
+
+  optimizer = tf.keras.optimizers.experimental.AdamW(
+      learning_rate=_LEARNING_RATE.value,
+      weight_decay=_WEIGHT_DECAY.value,
+      beta_1=_BETA_1.value,
+      beta_2=_BETA_2.value,
+      epsilon=_EPSILON.value,
+  )
+  diffusion_ft_trainer.compile(optimizer=optimizer, loss='mse')
+
+  training_dataset = prepare_training_dataset(_INPUT_CSV_PATH.value)
+
+  # Note: gcsfuse does not work for Keras. We saves the trained models locally
+  # first, and then copy to gcs storages.
+  if not os.path.exists(constants.LOCAL_MODEL_DIR):
+    os.makedirs(constants.LOCAL_MODEL_DIR)
+  # The default saved model is in HDF5.
+  ckpt_path = os.path.join(constants.LOCAL_MODEL_DIR, 'saved_model.h5')
+  ckpt_callback = tf.keras.callbacks.ModelCheckpoint(
+      ckpt_path,
+      save_weights_only=True,
+      monitor='loss',
+      mode='min',
+  )
+  diffusion_ft_trainer.fit(
+      training_dataset, epochs=_EPOCHS.value, callbacks=[ckpt_callback]
+  )
+
+  # Copies the files in constants.LOCAL_MODEL_DIR to output_model_dir.
+  fileutils.upload_local_dir_to_gcs(
+      constants.LOCAL_MODEL_DIR, _OUTPUT_MODEL_DIR.value
+  )
+
+  return
+
+
+if __name__ == '__main__':
+  app.run(main)

community-content/vertex_vision_model_garden/model_oss/lm-evaluation-harness/dockerfile/eval.Dockerfile

@@ -0,0 +1,40 @@
+# Dockerfile for lm-evaluation-harness evaluation.
+#
+# To build:
+# docker build -f model_oss/lm-evaluation-harness/dockerfile/eval.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.0.0-cuda11.7-cudnn8-devel
+
+USER root
+
+# Install tools.
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update
+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
+
+# 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 pip install google-cloud-storage==2.7.0
+RUN pip install absl-py==1.4.0
+
+# Install lm-evaluation-harness
+RUN git clone https://github.com/EleutherAI/lm-evaluation-harness
+WORKDIR lm-evaluation-harness
+# Pin version up to date 08/08/2023
+RUN git reset --hard b952a206de210b72b1bf750fbab38c26121e0dc0
+# Edit tokenizer loading function to avoid using fast tokenizer for OpenLLaMA
+RUN sed -i '355 i\        use_fast = not pretrained.startswith("openlm-research/open_llama")' lm_eval/models/huggingface.py
+RUN sed -i '360 i\            use_fast=use_fast,' lm_eval/models/huggingface.py
+# Install from source while including the sentencepiece dependency
+RUN pip install -e ".[sentencepiece]"

community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/base.Dockerfile

@@ -0,0 +1,64 @@
+FROM tensorflow/build:2.12-python3.9
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# This is added to fix docker build error related to Nvidia key update.
+RUN rm -f /etc/apt/sources.list.d/cuda.list
+RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -
+
+# 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 \
+        vim \
+        libtcmalloc-minimal4
+
+
+# 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 required libs.
+RUN pip install --upgrade pip
+RUN pip install cloud-tpu-client==0.10
+RUN pip install pyyaml==6.0
+RUN pip install fsspec==2023.4.0
+RUN pip install gcsfs==2023.4.0
+RUN pip install tf-models-official==2.12.0
+RUN pip install cloudml-hypertune==0.1.0.dev6
+RUN pip install pylint==2.17.3
+
+# Installs Reduction Server NCCL plugin.
+RUN echo "deb https://packages.cloud.google.com/apt google-fast-socket main" | tee /etc/apt/sources.list.d/google-fast-socket.list \
+&&  curl -s -L https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - \
+&&  apt update && apt install -y google-reduction-server
+
+ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=cpp
+
+# Lower the memory fragmentation, and speed up the training.
+# https://github.com/tensorflow/tensorflow/issues/44176#issuecomment-783768033
+ENV LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4
+
+# Enable userspace DNS cache
+ENV GCS_RESOLVE_REFRESH_SECS=60
+ENV GCS_REQUEST_CONNECTION_TIMEOUT_SECS=300
+ENV GCS_METADATA_REQUEST_TIMEOUT_SECS=300
+ENV GCS_READ_REQUEST_TIMEOUT_SECS=300
+ENV GCS_WRITE_REQUEST_TIMEOUT_SECS=600
+# Each opened GCS file takes GCS_READ_CACHE_BLOCK_SIZE_MB of RAM, reduce the
+# value from the default 64MB to 8MB to decrease memory footprint.
+ENV GCS_READ_CACHE_BLOCK_SIZE_MB=8

community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/model_export.Dockerfile

@@ -0,0 +1,13 @@
+FROM gcr.io/automl-migration-test/movinet-base:latest
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+RUN wget https://raw.githubusercontent.com/tensorflow/models/954dd73bffd43174bd3ca26a4a34abebe4147570/official/projects/movinet/tools/export_saved_model.py \
+  -O /usr/local/lib/python3.9/dist-packages/official/projects/movinet/tools/export_saved_model.py
+
+WORKDIR /automl_vision
+
+ENV PYTHONPATH "${PYTHONPATH}:/automl_vision/util"
+
+ENTRYPOINT ["python3", "-m", "official.projects.movinet.tools.export_saved_model"]

community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/serve.Dockerfile

@@ -0,0 +1,18 @@
+FROM gcr.io/automl-migration-test/movinet-base:latest
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+RUN pip install flask==2.3.2
+RUN pip install waitress==2.1.2
+
+RUN mkdir -p /automl_vision/movinet/serving
+COPY model_oss/movinet/serving /automl_vision/movinet/serving
+COPY model_oss/util /automl_vision/util
+
+WORKDIR /automl_vision
+
+ENV PYTHONPATH "${PYTHONPATH}:/automl_vision/util"
+
+ENTRYPOINT ["flask", "--app", "movinet.serving.serving_main", "run"]
+CMD ["--host=0.0.0.0", "--port=8501"]

community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/train.Dockerfile

@@ -0,0 +1,18 @@
+FROM gcr.io/automl-migration-test/movinet-base:latest
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+RUN mkdir -p /automl_vision/movinet
+COPY model_oss/movinet/*.py /automl_vision/movinet/
+COPY model_oss/util /automl_vision/util
+
+WORKDIR /automl_vision
+
+ENV PYTHONPATH "${PYTHONPATH}:/automl_vision/util"
+
+# Run pylint to validate code.
+COPY .pylintrc /automl_vision/.pylintrc
+RUN find . -type f -name "*.py" | xargs pylint --rcfile=./.pylintrc --errors-only
+
+ENTRYPOINT ["python3","movinet/train.py"]

community-content/vertex_vision_model_garden/model_oss/movinet/serving/serving_main.py

@@ -0,0 +1,142 @@
+"""Main executable for MoViNet online / batch predictions."""
+
+from collections.abc import Sequence
+import json
+import os
+
+from absl import app
+from absl import logging
+import flask
+import tensorflow as tf
+import waitress
+
+from movinet.serving import video_serving_lib
+from util import constants
+
+
+flask_app = flask.Flask(__name__)
+logging.set_verbosity(logging.INFO)
+
+movinet_model = None
+
+_BATCH_SIZE = int(os.environ.get('BATCH_SIZE', '1'))
+_NUM_FRAMES = int(os.environ.get('NUM_FRAMES', '32'))
+_FPS = float(os.environ.get('FPS', '5'))
+_OVERLAP_FRAMES = int(os.environ.get('OVERLAP_FRAMES', '24'))
+_OBJECTIVE = os.environ.get(
+    'OBJECTIVE', constants.OBJECTIVE_VIDEO_CLASSIFICATION
+).lower()
+
+# VAR parameters.
+_CONFIDENCE_THRESHOLD = float(os.environ.get('CONFIDENCE_THRESHOLD', '0.5'))
+_MIN_GAP_TIME = float(os.environ.get('MIN_GAP_TIME', '1.5'))
+
+
+def load_movinet_model() -> None:
+  model_path = os.environ.get('MODEL_PATH')
+
+  if not model_path:
+    raise app.UsageError('Missing MODEL_PATH environment variable.')
+
+  # We just reload the weights of the fine-tuned diffusion model.
+  logging.info('Initialize finetuned models from: %s', model_path)
+  global movinet_model
+  movinet_model = tf.saved_model.load(model_path)
+
+
+load_movinet_model()
+
+
+def error(message: str) -> str:
+  """Returns a JSON representing an error response."""
+  return json.dumps({
+      'success': False,
+      'error': message,
+  })
+
+
+# The health check route is required for docker deployment in google cloud.
+@flask_app.route('/ping')
+def ping() -> flask.Response:
+  """Health checks."""
+  return flask.Response(status=200)
+
+
+# The return should be `Response` for docker deployment in google cloud.
+@flask_app.route('/predict', methods=['GET', 'POST'])
+def predict_model() -> flask.Response:
+  """Predictions."""
+  if flask.request.method == 'POST':
+    contents = flask.request.get_json(force=True)
+
+    logging.info('The input contents are: %s', contents)
+    instances = contents.get('instances', [])
+
+    try:
+      predictions = []
+      for instance in instances:
+        executor = video_serving_lib.parse_request(instance)
+        prediction = executor.get_prediction(
+            movinet_model,
+            _BATCH_SIZE,
+            _FPS,
+            _NUM_FRAMES,
+            _OVERLAP_FRAMES,
+            _OBJECTIVE,
+        )
+        if _OBJECTIVE == constants.OBJECTIVE_VIDEO_CLASSIFICATION:
+          prediction = video_serving_lib.postprocess_vcn(prediction)
+        elif _OBJECTIVE == constants.OBJECTIVE_VIDEO_ACTION_RECOGNITION:
+          prediction = video_serving_lib.postprocess_var(
+              executor.windows, prediction, _CONFIDENCE_THRESHOLD, _MIN_GAP_TIME
+          )
+        predictions.append(prediction)
+    except ValueError as e:
+      return flask.Response(
+          error(str(e)), status=500, mimetype='application/json'
+      )
+
+    return flask.Response(
+        response=json.dumps({
+            'success': True,
+            'predictions': predictions,
+        }),
+        status=200,
+        mimetype='application/json',
+    )
+  else:
+    return flask.Response(
+        response=json.dumps({
+            'success': True,
+            'isalive': movinet_model is not None,
+        }),
+        status=200,
+        mimetype='application/json',
+    )
+
+
+def main(argv: Sequence[str]) -> None:
+  if len(argv) > 1:
+    raise app.UsageError('Too many command-line arguments.')
+  # This is used when running locally only. When deploying to Google App
+  # Engine, a webserver process such as Gunicorn will serve the app.
+  # # Debug deployment.
+  # flask_app.run(host='0.0.0.0', port=8501, debug=True)
+  # Prod deployment.
+  if _OBJECTIVE not in [
+      constants.OBJECTIVE_VIDEO_CLASSIFICATION,
+      constants.OBJECTIVE_VIDEO_ACTION_RECOGNITION,
+  ]:
+    raise app.UsageError('Objective must be vcn or var.')
+  logging.info(
+      'Env: batch_size: %s, num_frames: %s, fps: %s, overlap_frames: %s',
+      _BATCH_SIZE,
+      _NUM_FRAMES,
+      _FPS,
+      _OVERLAP_FRAMES,
+  )
+  waitress.serve(flask_app, host='0.0.0.0', port=8501)
+
+
+if __name__ == '__main__':
+  app.run(main)

community-content/vertex_vision_model_garden/model_oss/movinet/serving/video_serving_lib.py

@@ -0,0 +1,462 @@
+"""Lib for handling video prediction requests.
+
+The VCN inference algorithm is as follows:
+1. Find all video frames within the given clip according to the sampling FPS.
+2. Create possibly overlapping sliding windows according to the num_frames and
+overlap_frames parameters. The last window might have a larger overlap if it
+doesn't exactly fit.
+3. Run model inference on each sliding window and compute softmax to obtain
+probabilities.
+4. Average the probabilities over all sliding windows.
+
+The VAR inference algorithm is very similar to VCN, with a few differences:
+1. The last sliding window is discarded if it does not exactly fit.
+2. Instead of averaging, the postprocessing consists of temporal nonmaximal
+suppression and removing background and low-confidence labels.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import os
+from typing import Any, Dict, Optional, Sequence, Union, cast
+
+from absl import logging
+import cv2
+import numpy as np
+import tensorflow as tf
+
+from util import constants
+from util import fileutils
+
+
+_JSON_LABEL_KEY = 'label'
+_JSON_GCS_URI_KEY = 'content'
+_JSON_CONFIDENCE_KEY = 'confidence'
+_JSON_START_TIME_KEY = 'timeSegmentStart'
+_JSON_END_TIME_KEY = 'timeSegmentEnd'
+_BACKGROUND_LABEL = 0
+_JSON_REQUIRED_KEYS = [
+    _JSON_GCS_URI_KEY,
+    _JSON_START_TIME_KEY,
+    _JSON_END_TIME_KEY,
+]
+_IMAGE_WIDTH = int(os.environ.get('IMAGE_WIDTH', '172'))
+_IMAGE_HEIGHT = int(os.environ.get('IMAGE_HEIGHT', '172'))
+
+
+@dataclasses.dataclass
+class DetectionOutput:
+  timestamp: float
+  label: int
+  confidence: float
+
+  def to_json_obj(self) -> Dict[str, Union[int, float]]:
+    """Encodes self as a dict for JSON serialization."""
+    return {
+        _JSON_LABEL_KEY: self.label,
+        _JSON_START_TIME_KEY: self.timestamp,
+        _JSON_END_TIME_KEY: self.timestamp,
+        _JSON_CONFIDENCE_KEY: self.confidence,
+    }
+
+
+def create_detection_output(
+    timestamp: float, predictions: np.ndarray
+) -> DetectionOutput:
+  label = np.argmax(predictions).item()
+  confidence: float = predictions[label].item()
+  return DetectionOutput(timestamp, label, confidence)
+
+
+class SlidingWindow:
+  """Represents a sliding window with start / end timestamps."""
+
+  def __init__(self, fps: float, frames: Sequence[int]):
+    if not frames:
+      raise ValueError('Sliding window cannot be empty.')
+    self.frames = frames
+    self.start_time = frames[0] / fps
+    self.end_time = frames[-1] / fps
+    self.frame_data: list[Optional[np.ndarray]] = []
+    self.clear_frame_data()
+
+  def load_cache_from(self, other: SlidingWindow) -> int:
+    """Loads cache from another sliding window if possible."""
+    cache_count = 0
+    for i, frame in enumerate(self.frames):
+      try:
+        other_idx = other.frames.index(frame)
+        self.frame_data[i] = other.frame_data[other_idx]
+        cache_count += 1
+      except ValueError:
+        # Cache miss.
+        pass
+    return cache_count
+
+  def load_frames(self, video: Any) -> Sequence[np.ndarray]:
+    """Loads frames of this sliding window from a video."""
+    for i, frame in enumerate(self.frames):
+      if self.frame_data[i] is None:
+        video.set(cv2.CAP_PROP_POS_FRAMES, frame)
+        ret, frame = video.read()
+        if not ret:
+          raise IOError(f'Failed to read video at frame {frame}.')
+        self.frame_data[i] = cv2.resize(frame, (_IMAGE_WIDTH, _IMAGE_HEIGHT))
+    return cast(Sequence[np.ndarray], self.frame_data)
+
+  def clear_frame_data(self) -> None:
+    """Clears frame data of this sliding window to reduce memory usage."""
+    self.frame_data: list[Optional[np.ndarray]] = [None] * len(self)
+
+  def __len__(self) -> int:
+    return len(self.frames)
+
+  @property
+  def middle_timestamp(self) -> float:
+    return (self.start_time + self.end_time) / 2
+
+
+def _get_sliding_windows(
+    frames: Sequence[int],
+    original_fps: float,
+    window_size: int,
+    overlap: int,
+    flush_last_window: bool,
+) -> Sequence[SlidingWindow]:
+  """Computes a list of sliding windows from frames.
+
+  Args:
+    frames: A list of frame indices.
+    original_fps: Frames per second of the original video.
+    window_size: Number of frames in a single window.
+    overlap: Number of overlapping frames in adjacent windows.
+    flush_last_window: Where to flush the last window if there are not enough
+      frames left.
+
+  Returns:
+    A list of sliding windows, each has a list of frame indices. The last two
+    windows might have a larger overlap if the last window does not exactly fit
+    and flush_last_window is set to True.
+
+  Raises:
+    ValueError: Arguments are invalid.
+  """
+  if window_size <= overlap:
+    raise ValueError(f'Window size {window_size} <= overlap {overlap}')
+  total_frames = len(frames)
+  windows: list[SlidingWindow] = []
+  for i in range(0, total_frames, window_size - overlap):
+    if i == 0 or i + window_size <= total_frames:
+      windows.append(SlidingWindow(original_fps, frames[i : i + window_size]))
+    elif i + overlap < total_frames and flush_last_window:
+      # Some frames in this window are not covered by the previous window.
+      windows.append(
+          SlidingWindow(
+              original_fps, frames[total_frames - window_size : total_frames]
+          )
+      )
+  return windows
+
+
+def _sample_frame_indices(
+    start_time: float,
+    end_time: float,
+    original_fps: float,
+    sample_fps: float,
+    max_frames: int,
+    padding_left: int = 0,
+    padding_right: int = 0,
+) -> Sequence[int]:
+  """Samples frames from start_time to end_time by sample_fps.
+
+  Args:
+    start_time: Start timestamp in seconds.
+    end_time: End timestamp in seconds.
+    original_fps: Frames per second of the original video.
+    sample_fps: Number of frames to sample per second.
+    max_frames: Total number of frames in the video.
+    padding_left: Padding to add to the start in frames. Padded frames will be
+      duplicates of the first frame.
+    padding_right: Padding to add to the end in frames. Padded frames will be
+      duplicates of the last frame.
+
+  Returns:
+    A list of sampled frame indices.
+  """
+  ret = [
+      min(max_frames - 1, round(t * original_fps))
+      for t in np.arange(start_time, end_time, 1 / sample_fps)
+  ]
+  if ret:
+    ret = [ret[0]] * padding_left + ret + [ret[-1]] * padding_right
+  return ret
+
+
+class VideoPredictionExecutor:
+  """Represents a Video prediction request with a video clip."""
+
+  def __init__(self, gcs_uri: str, start_time: float, end_time: float):
+    self._gcs_uri = gcs_uri
+    self._start_time = start_time
+    self._end_time = end_time
+    self.windows: Sequence[SlidingWindow] = []
+    self._last_window: SlidingWindow = None
+
+  def _read_frames_from_window(
+      self, video: Any, new_window: SlidingWindow
+  ) -> Sequence[np.ndarray]:
+    """Reads video frames from the new window.
+
+    Args:
+      video: Video loaded with cv2.
+      new_window: A list of sorted frame indices in the new window.
+
+    Returns:
+      Frame data from the video as a list of numpy arrays.
+
+    Raises:
+      IOError: Failed to read video.
+    """
+    # Caches frames as much as possible.
+    if self._last_window is not None:
+      cache_count = new_window.load_cache_from(self._last_window)
+      logging.info('Cached %d frames.', cache_count)
+      self._last_window.clear_frame_data()
+    self._last_window = new_window
+    return new_window.load_frames(video)
+
+  def _predict(
+      self, model: Any, video: Any, batched_windows: Sequence[SlidingWindow]
+  ) -> np.ndarray:
+    """Run model inference on specific frames of a video.
+
+    Args:
+      model: MoViNet model.
+      video: Video loaded with cv2.
+      batched_windows: A batch of sliding windows to predict. Each element is an
+        integer frame index. Must have equal number of frames in each window.
+
+    Returns:
+      Prediction results.
+
+    Raises:
+      ValueError: Batched windows are not sorted, or do not have equal number of
+        frames in each window.
+      IOError: Failed to read video.
+    """
+    if any(
+        (
+            len(window) != len(batched_windows[0])
+            for window in batched_windows[1:]
+        )
+    ):
+      raise ValueError(
+          'Batched windows do not have equal number of frames in each window.'
+      )
+    batch = []
+    logging.info('Loading video frames...')
+    for window in batched_windows:
+      logging.info('Predict frames: %s', window.frames)
+      frames = self._read_frames_from_window(video, window)
+      batch.append(frames)
+    input_tensor = tf.convert_to_tensor(batch, dtype=tf.float32) / 255.0
+    logging.info('Predict: Input tensor shape %s', input_tensor.shape)
+    predictions = model({'image': input_tensor})
+    logging.info('Running softmax on predictions...')
+    predictions = tf.nn.softmax(predictions, axis=1)
+    return predictions.numpy()
+
+  def get_prediction(
+      self,
+      model: Any,
+      batch_size: int,
+      fps: float,
+      num_frames: int,
+      overlap_frames: int,
+      objective: str,
+  ) -> Sequence[np.ndarray]:
+    """Predicts the video clip with the model.
+
+    Args:
+      model: The loaded MoViNet model.
+      batch_size: Batch size for prediction.
+      fps: Video sampling FPS.
+      num_frames: Number of frames in a single predictions. If the model is
+        exported with a fixed input shape, this must match its num_frames
+        dimension.
+      overlap_frames: Number of overlapping frames of consecutive sliding
+        windows.
+      objective: A string `vcn` or `var`.
+
+    Returns:
+      A list of floats as the prediction response.
+
+    Raises:
+      IOError: The video fails to load.
+      ValueError: Some arguments are invalid.
+    """
+    if objective not in [
+        constants.OBJECTIVE_VIDEO_CLASSIFICATION,
+        constants.OBJECTIVE_VIDEO_ACTION_RECOGNITION,
+    ]:
+      raise ValueError(f'{objective} objective is not supported.')
+
+    # cv2 expects a local path so we need to download the video from GCS.
+    local_file_path = fileutils.generate_tmp_path(
+        os.path.splitext(self._gcs_uri)[1]
+    )
+    logging.info('Downloading %s to %s...', self._gcs_uri, local_file_path)
+    fileutils.download_gcs_file_to_local(self._gcs_uri, local_file_path)
+    logging.info('Download %s complete.', self._gcs_uri)
+
+    # Loads video.
+    video = cv2.VideoCapture(local_file_path)
+    total_frames = video.get(cv2.CAP_PROP_FRAME_COUNT)
+    original_fps = video.get(cv2.CAP_PROP_FPS)
+    if not original_fps:
+      # 0 or None indicates the video is invalid.
+      raise IOError(f'Failed to load {self._gcs_uri}.')
+    video_length = total_frames / original_fps
+    self._start_time = max(0, self._start_time)
+    self._end_time = min(video_length, self._end_time)
+    padding = (
+        (num_frames // 2)
+        if objective == constants.OBJECTIVE_VIDEO_ACTION_RECOGNITION
+        else 0
+    )
+
+    # Computes sliding windows.
+    frame_indices = _sample_frame_indices(
+        self._start_time,
+        self._end_time,
+        original_fps,
+        fps,
+        total_frames,
+        padding,
+        padding,
+    )
+    logging.info('Frame indices: %s', frame_indices)
+    self.windows = _get_sliding_windows(
+        frame_indices,
+        original_fps,
+        num_frames,
+        overlap_frames,
+        objective != 'var',
+    )
+    if not self.windows:
+      raise ValueError(
+          f'No sliding windows found from {self._start_time} to'
+          f' {self._end_time}.'
+      )
+    self._last_window = None
+
+    # Runs inference.
+    predictions = []
+    for i in range(0, len(self.windows), batch_size):
+      predictions.extend(
+          self._predict(model, video, self.windows[i : i + batch_size])
+      )
+    return predictions
+
+
+def parse_request(req_json: Any) -> VideoPredictionExecutor:
+  """Parses VideoPredictionExecutor from request JSON object.
+
+  Args:
+    req_json: Request JSON object.
+
+  Returns:
+    Parsed VideoPredictionExecutor.
+
+  Raises:
+    ValueError: Request JSON object is invalid.
+  """
+  for key in _JSON_REQUIRED_KEYS:
+    if key not in req_json:
+      raise ValueError(f'{key} not found in {req_json}.')
+  gcs_uri = req_json[_JSON_GCS_URI_KEY]
+  start_time = float(req_json[_JSON_START_TIME_KEY].removesuffix('s'))
+  end_time = float(req_json[_JSON_END_TIME_KEY].removesuffix('s'))
+  return VideoPredictionExecutor(gcs_uri, start_time, end_time)
+
+
+def postprocess_vcn(predictions: Sequence[np.ndarray]) -> Sequence[float]:
+  """Aggregates VCN predictions of sliding windows."""
+  return np.mean(predictions, axis=0).tolist()
+
+
+def temporal_nonmaximal_suppression(
+    detections: Sequence[DetectionOutput], min_gap_time: float
+) -> Sequence[DetectionOutput]:
+  """Nonmaximal suppression for key frame detection.
+
+  For consecutive packets of the same label within a pre-defined duration, we
+  only keep the one with the highest confidence score. Such duration can be
+  determined by performing data analysis on users' dataset.
+
+  Args:
+    detections: A list of DetectionOutputs.
+    min_gap_time: Minimum time between consecutive key frames of the same label
+      in seconds.
+
+  Returns:
+    DetectionOutput after nonmaximal suppression sorted in ascending timestamps.
+  """
+  max_label = max([detection.label for detection in detections])
+  prev_detections: list[Optional[DetectionOutput]] = [None] * (max_label + 1)
+  ret: list[DetectionOutput] = []
+  by_time = lambda x: x.timestamp
+  for detection in sorted(detections, key=by_time):
+    prev_detection = prev_detections[detection.label]
+    prev_detections[detection.label] = detection
+    if not prev_detection:
+      continue
+    if detection.timestamp - prev_detection.timestamp > min_gap_time:
+      ret.append(prev_detection)
+      continue
+    detection.confidence = max(detection.confidence, prev_detection.confidence)
+  ret.extend((d for d in prev_detections if d is not None))
+  return sorted(ret, key=by_time)
+
+
+def postprocess_var(
+    windows: Sequence[SlidingWindow],
+    predictions: Sequence[np.ndarray],
+    confidence_threshold: float,
+    min_gap_time: float,
+) -> Sequence[Dict[str, Any]]:
+  """Generates a list of detected keyframes from sliding window predictions.
+
+  Args:
+    windows: Sliding windows.
+    predictions: A list of predictions of sliding windows.
+    confidence_threshold: Only probabilities greater than this threshold will
+      contribute to the final result.
+    min_gap_time: Minimum time between consecutive key frames of the same label
+      in seconds. Used in temporal nonmaximal suppression.
+
+  Returns:
+    A sequence of dictionaries, each item has the following keys:
+    - label: Integer label of the detection result.
+    - timeSegmentStart: Start timestamp in seconds.
+    - timeSegmentEnd: End timestamp in seconds. Always equals timeSegmentStart.
+  """
+  if len(windows) != len(predictions):
+    raise ValueError('Mismatched # of windows with # of predictions.')
+
+  # Creates detection results from windows, filtering out the background label.
+  detections = [
+      create_detection_output(window.middle_timestamp, predictions[i])
+      for i, window in enumerate(windows)
+  ]
+
+  # Temporal nonmaximal suppression.
+  detections = temporal_nonmaximal_suppression(detections, min_gap_time)
+
+  # Filters out ones with low confidence and the background label.
+  return [
+      x.to_json_obj()
+      for x in detections
+      if x.label != _BACKGROUND_LABEL and x.confidence > confidence_threshold
+  ]

community-content/vertex_vision_model_garden/model_oss/movinet/train.py

@@ -0,0 +1,210 @@
+"""Main executable for MoViNet docker."""
+
+import json
+import os
+from typing import Sequence, Any
+
+from absl import app
+from absl import flags
+from absl import logging
+import gin
+import hypertune
+import tensorflow as tf
+
+from util import constants
+from util import hypertune_utils
+from official.common import distribute_utils
+from official.common import flags as tfm_flags
+from official.core import task_factory
+from official.core import train_lib
+from official.core import train_utils
+from official.modeling import performance
+# Import movinet libraries to register the backbone and model into tf.vision
+# model garden factory.
+# pylint: disable=unused-import
+from official.projects.movinet.modeling import movinet
+from official.projects.movinet.modeling import movinet_model
+from official.vision import registry_imports
+# pylint: enable=unused-import
+
+
+FLAGS = flags.FLAGS
+
+_FILE_TYPE_TFRECORD = 'tfrecord'
+
+_LEARNING_RATE = flags.DEFINE_float(
+    'learning_rate', None, 'The learning rate of this training job.'
+)
+
+_NUM_CLASSES = flags.DEFINE_integer(
+    'num_classes', None, 'The number of classes.'
+)
+
+_INIT_CHECKPOINT = flags.DEFINE_string(
+    'init_checkpoint', None, 'The initial checkpoint of this training job.'
+)
+
+_INPUT_TRAIN_DATA_PATH = flags.DEFINE_string(
+    'input_train_data_path', None, 'Input train data path.'
+)
+
+_INPUT_VALIDATION_DATA_PATH = flags.DEFINE_string(
+    'input_validation_data_path', None, 'Input validation data path.'
+)
+
+_GLOBAL_BATCH_SIZE = flags.DEFINE_integer(
+    'global_batch_size', None, 'Global batch size.'
+)
+
+_PREFETCH_BUFFER_SIZE = flags.DEFINE_integer(
+    'prefetch_buffer_size', None, 'Prefetch buffer size.'
+)
+
+_SHUFFLE_BUFFER_SIZE = flags.DEFINE_integer(
+    'shuffle_buffer_size', None, 'Shuffle buffer size.'
+)
+
+_TRAIN_STEPS = flags.DEFINE_integer('train_steps', None, 'Train steps.')
+_LOG_LEVEL = flags.DEFINE_enum(
+    'log_level',
+    'INFO',
+    ['FATAL', 'ERROR', 'WARNING', 'INFO', 'DEBUG'],
+    'Log level.',
+)
+
+
+def parse_params() -> Any:
+  """Parses parameters."""
+  gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_params)
+  params = train_utils.parse_configuration(FLAGS, lock_return=False)
+  if _INIT_CHECKPOINT.value:
+    params.task.init_checkpoint = _INIT_CHECKPOINT.value
+    params.task.init_checkpoint_modules = 'backbone'
+  if _NUM_CLASSES.value:
+    params.task.model.num_classes = _NUM_CLASSES.value
+    params.task.train_data.num_classes = _NUM_CLASSES.value
+    params.task.validation_data.num_classes = _NUM_CLASSES.value
+  # If users set input train/validation data path, we assume the data are
+  # converted from data converter as tfrecord. Users can use tfds by writing
+  # their own config directly, and no need to override this parameter.
+  if _INPUT_TRAIN_DATA_PATH.value:
+    params.task.train_data.input_path = _INPUT_TRAIN_DATA_PATH.value
+    params.task.train_data.file_type = _FILE_TYPE_TFRECORD
+    params.task.train_data.tfds_name = ''
+  if _INPUT_VALIDATION_DATA_PATH.value:
+    params.task.validation_data.input_path = _INPUT_VALIDATION_DATA_PATH.value
+    params.task.validation_data.file_type = _FILE_TYPE_TFRECORD
+    params.task.validation_data.tfds_name = ''
+  if _GLOBAL_BATCH_SIZE.value:
+    params.task.train_data.global_batch_size = _GLOBAL_BATCH_SIZE.value
+    params.task.validation_data.global_batch_size = _GLOBAL_BATCH_SIZE.value
+  if _PREFETCH_BUFFER_SIZE.value:
+    params.task.train_data.prefetch_buffer_size = _PREFETCH_BUFFER_SIZE.value
+    params.task.validation_data.prefetch_buffer_size = (
+        _PREFETCH_BUFFER_SIZE.value
+    )
+  if _SHUFFLE_BUFFER_SIZE.value:
+    params.task.train_data.shuffle_buffer_size = _SHUFFLE_BUFFER_SIZE.value
+  if _TRAIN_STEPS.value:
+    params.trainer.train_steps = _TRAIN_STEPS.value
+  if _LEARNING_RATE.value:
+    logging.info('Updating learning_rate: %s', _LEARNING_RATE.value)
+    # Use `get` method of train_utils.hyperparams.OneOfConfig to get learning
+    # rate config.
+    learning_rate = params.trainer.optimizer_config.learning_rate.get()
+    if hasattr(learning_rate, 'initial_learning_rate'):
+      learning_rate.initial_learning_rate = _LEARNING_RATE.value
+    else:
+      logging.warning('Cannot set learning rate for %s', learning_rate)
+  # Set default params for best checkpoints.
+  params.trainer.best_checkpoint_export_subdir = constants.BEST_CKPT_DIRNAME
+  params.trainer.best_checkpoint_metric_comp = constants.BEST_CKPT_METRIC_COMP
+  params.trainer.best_checkpoint_eval_metric = (
+      constants.VIDEO_CLASSIFICATION_BEST_EVAL_METRIC
+  )
+  return params
+
+
+def main(argv: Sequence[str]) -> None:
+  logging.set_verbosity(_LOG_LEVEL.value)
+  if len(argv) > 1:
+    raise app.UsageError('Too many command-line arguments.')
+  params = parse_params()
+  logging.info('The actual training parameters are:\n%s', params.as_dict())
+  model_dir: str = os.path.join(
+      FLAGS.model_dir,
+      constants.TRIAL_PREFIX + hypertune_utils.get_trial_id_from_environment(),
+  )
+  logging.info('model_dir: %s', model_dir)
+
+  if 'train' in FLAGS.mode:
+    # Pure eval modes do not output yaml files. Otherwise continuous eval job
+    # may race against the train job for writing the same file.
+    train_utils.serialize_config(params, model_dir)
+
+  # Sets mixed_precision policy. Using 'mixed_float16' or 'mixed_bfloat16'
+  # can have significant impact on model speeds by utilizing float16 in case of
+  # GPUs, and bfloat16 in the case of TPUs. loss_scale takes effect only when
+  # dtype is float16
+  if params.runtime.mixed_precision_dtype:
+    performance.set_mixed_precision_policy(params.runtime.mixed_precision_dtype)
+  distribution_strategy = distribute_utils.get_distribution_strategy(
+      distribution_strategy=params.runtime.distribution_strategy,
+      all_reduce_alg=params.runtime.all_reduce_alg,
+      num_gpus=params.runtime.num_gpus,
+      tpu_address=params.runtime.tpu,
+  )
+
+  # Create task and run experiment.
+  with distribution_strategy.scope():
+    task = task_factory.get_task(params.task, logging_dir=model_dir)
+
+  train_lib.run_experiment(
+      distribution_strategy=distribution_strategy,
+      task=task,
+      mode=FLAGS.mode,
+      params=params,
+      model_dir=model_dir,
+  )
+
+  train_utils.save_gin_config(FLAGS.mode, model_dir)
+
+  eval_metric_name = constants.VIDEO_CLASSIFICATION_BEST_EVAL_METRIC
+
+  eval_filepath = os.path.join(
+      model_dir, constants.BEST_CKPT_DIRNAME, constants.BEST_CKPT_EVAL_FILENAME
+  )
+  logging.info('Load eval metrics from: %s.', eval_filepath)
+
+  with tf.io.gfile.GFile(eval_filepath, 'rb') as f:
+    eval_metric_results = json.load(f)
+  logging.info('eval metrics are: %s.', eval_metric_results)
+  if (
+      eval_metric_name in eval_metric_results
+      and constants.BEST_CKPT_STEP_NAME in eval_metric_results
+  ):
+    hp_metric = eval_metric_results[eval_metric_name]
+    hp_step = int(eval_metric_results[constants.BEST_CKPT_STEP_NAME])
+    hpt = hypertune.HyperTune()
+    hpt.report_hyperparameter_tuning_metric(
+        hyperparameter_metric_tag=constants.HP_METRIC_TAG,
+        metric_value=hp_metric,
+        global_step=hp_step,
+    )
+    logging.info(
+        'Send HP metric: %f and steps %d to hyperparameter tuning.',
+        hp_metric,
+        hp_step,
+    )
+  else:
+    logging.info(
+        'Either %s or %s is not included in the evaluation results: %s.',
+        eval_metric_name,
+        constants.BEST_CKPT_STEP_NAME,
+        eval_metric_results,
+    )
+
+
+if __name__ == '__main__':
+  tfm_flags.define_flags()
+  app.run(main)

community-content/vertex_vision_model_garden/model_oss/open_clip/dockerfile/serve.Dockerfile

@@ -0,0 +1,67 @@
+# Dockerfile for basic serving dockers for OpenCLIP.
+#
+# To build:
+# docker build -f model_oss/open_clip/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}
+
+# Switch to this base image for gpu serve.
+FROM pytorch/torchserve:0.7.1-gpu
+
+USER root
+
+# Install tools.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim
+
+# 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="transformers_serving"
+ENV PATH="/home/model-server/:${PATH}"
+
+# Install libraries.
+RUN python3 -m pip install --upgrade pip
+RUN pip install torch==1.13.1
+RUN pip install open_clip_torch==2.20.0
+RUN pip install pillow==9.5.0
+RUN pip install google-cloud-storage==2.7.0
+
+# Copy model artifacts.
+COPY model_oss/open_clip/handler.py /home/model-server/handler.py
+COPY model_oss/util/ /home/model-server/util/
+ENV PYTHONPATH /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}" >> /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.
+CMD ["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_vision_model_garden/model_oss/open_clip/dockerfile/train.Dockerfile

@@ -0,0 +1,53 @@
+# Dockerfile for training dockers with OpenCLIP.
+#
+# To build:
+# docker build -f model_oss/open_clilp/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.0.0-cuda11.7-cudnn8-devel
+
+# Install tools.
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update
+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
+RUN apt-get install -y --no-install-recommends build-essential
+
+ENV PIP_ROOT_USER_ACTION=ignore
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Prepare artifacts.
+WORKDIR /workspace
+RUN git clone --branch main https://github.com/mlfoundations/open_clip.git
+WORKDIR ./open_clip
+RUN git reset --hard 67e5e5ec8741281eb9b30f640c26f91c666308b7
+
+# Install libraries.
+RUN pip install webdataset==0.2.5
+RUN pip install regex==2023.6.3
+RUN pip install ftfy==6.1.1
+RUN pip install pandas==2.0.3
+RUN pip install braceexpand==0.1.7
+RUN pip install huggingface_hub==0.16.4
+RUN pip install transformers==4.31.0
+RUN pip install timm==0.9.2
+RUN pip install fsspec==2023.6.0
+RUN pip install sentencepiece==0.1.99
+RUN pip install protobuf==3.20.3
+RUN pip install tensorboard==2.12.2
+
+# Switch work folder for training.
+WORKDIR ./src

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

@@ -0,0 +1,142 @@
+"""Custom handler for OpenCLIP model."""
+
+# pylint:disable=g-importing-member
+import enum
+import logging
+import os
+from typing import Any, Dict, List
+
+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"
+
+
+# Supported checkpoint&model pairs:
+# https://github.com/mlfoundations/open_clip#pretrained-model-interface
+_DEFAULT_CHECKPOINT = "openai"
+_DEFAULT_MODEL = "RN50"
+_DEFAULT_PRECISION = Precision.AMP
+_ZERO_CLASSIFICATION = "zero-shot-image-classification"
+_FEATURE_EMBEDDING = "feature-embedding"
+_VALID_TASKS = frozenset([_ZERO_CLASSIFICATION, _FEATURE_EMBEDDING])
+
+_IMAGE_KEY = "image"
+_TEXT_KEY = "text"
+_IMAGE_FEATURES_KEY = "image_features"
+_TEXT_FEATURES_KEY = "text_features"
+
+
+class OpenclipHandler(BaseHandler):
+  """Custom handler for OpenCLIP."""
+
+  def initialize(self, context: Any):
+    """Custom initialize."""
+
+    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
+
+    model_name = os.environ.get("MODEL", _DEFAULT_MODEL)
+    precision = os.environ.get("PRECISION", _DEFAULT_PRECISION)
+    checkpoint = os.environ.get("CHECKPOINT", _DEFAULT_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,
+        precision,
+        checkpoint,
+    )
+
+    if checkpoint != _DEFAULT_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.tokenizer = open_clip.get_tokenizer(model_name)
+
+    self.initialized = True
+
+  def preprocess(self, data: Any) -> List[Dict[str, Any]]:
+    """Preprocess input data."""
+    logging.info("preprocessing: %d instances received.", len(data))
+    processed_list = []
+    for item in data:
+      sample = {}
+      if _IMAGE_KEY in item:
+        sample[_IMAGE_KEY] = self.preprocessor(
+            image_format_converter.base64_to_image(item[_IMAGE_KEY])
+        ).unsqueeze(0)
+      if _TEXT_KEY in item:
+        sample[_TEXT_KEY] = self.tokenizer(item[_TEXT_KEY])
+      processed_list.append(sample)
+    return processed_list
+
+  def inference(
+      self, data: List[Dict[str, Any]], *args, **kwargs
+  ) -> List[Dict[str, Any]]:
+    feature_list = []
+    with torch.no_grad(), torch.cuda.amp.autocast():
+      for item in data:
+        sample = {}
+        if _IMAGE_KEY in item:
+          sample[_IMAGE_FEATURES_KEY] = self.model.encode_image(
+              item[_IMAGE_KEY]
+          )
+        if _TEXT_KEY in item:
+          sample[_TEXT_FEATURES_KEY] = self.model.encode_text(item[_TEXT_KEY])
+        feature_list.append(sample)
+    return feature_list
+
+  def postprocess(self, features: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Postprocess the image/text featreus for downstream task."""
+    preds = []
+    if self.task == _FEATURE_EMBEDDING:
+      for item in features:
+        preds.append({k: v.tolist() for k, v in item.items()})
+    elif self.task == _ZERO_CLASSIFICATION:
+      for item in features:
+        image_features = item.get(_IMAGE_FEATURES_KEY, None)
+        text_features = item.get(_TEXT_FEATURES_KEY, None)
+        if image_features is None or text_features is None:
+          raise ValueError(
+              "Missing input for {} task. {} received.".format(
+                  _ZERO_CLASSIFICATION, item.keys()
+              )
+          )
+        image_features /= image_features.norm(dim=-1, keepdim=True)
+        text_features /= text_features.norm(dim=-1, keepdim=True)
+        text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)
+        preds.append(text_probs.tolist())
+
+    return preds

community-content/vertex_vision_model_garden/model_oss/peft/causal_language_modeling_lora.py

@@ -0,0 +1,142 @@
+"""Causal language modeling with LoRA models."""
+
+# pylint: disable=g-importing-member
+
+from datasets import load_dataset
+from peft import get_peft_model
+from peft import LoraConfig
+import torch
+from torch import nn
+import transformers
+from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
+from transformers import BitsAndBytesConfig
+from transformers import TrainingArguments
+from util import constants
+
+
+def finetune_causal_language_modeling(
+    pretrained_model_id: str,
+    dataset_name: str,
+    output_dir: str,
+    precision_mode: str = None,
+    lora_rank: int = 16,
+    lora_alpha: int = 32,
+    lora_dropout: float = 0.05,
+    warmup_steps: int = 10,
+    max_steps: int = 10,
+    learning_rate: float = 2e-4,
+    local_pretrained_model_id: str = None,
+) -> None:
+  """Finetunes causal language modelings."""
+  if precision_mode == constants.PRECISION_MODE_32:
+    model = AutoModelForCausalLM.from_pretrained(
+        local_pretrained_model_id
+        if local_pretrained_model_id
+        else pretrained_model_id,
+        torch_dtype=torch.float32,
+        device_map="auto",
+    )
+  elif precision_mode == constants.PRECISION_MODE_16:
+    model = AutoModelForCausalLM.from_pretrained(
+        local_pretrained_model_id
+        if local_pretrained_model_id
+        else pretrained_model_id,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+    )
+  elif precision_mode == constants.PRECISION_MODE_8:
+    quantization_config = BitsAndBytesConfig(
+        load_in_8bit=True, int8_threshold=0
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        local_pretrained_model_id
+        if local_pretrained_model_id
+        else pretrained_model_id,
+        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(
+        local_pretrained_model_id
+        if local_pretrained_model_id
+        else pretrained_model_id,
+        device_map="auto",
+        torch_dtype=torch.bfloat16,
+        quantization_config=quantization_config,
+    )
+
+  tokenizer = AutoTokenizer.from_pretrained(
+      local_pretrained_model_id
+      if local_pretrained_model_id
+      else pretrained_model_id
+  )
+  if "llama" in pretrained_model_id:
+    tokenizer.pad_token = "[PAD]"
+
+  for param in model.parameters():
+    # Freezes the model - train adapters later.
+    param.requires_grad = False
+    if param.ndim == 1:
+      # Casts the small parameters (e.g. layernorm) to fp32 for stability.
+      param.data = param.data.to(torch.float32)
+
+  # Reduces the number of stored activations.
+  model.gradient_checkpointing_enable()
+  model.enable_input_require_grads()
+
+  class CastOutputToFloat(nn.Sequential):
+
+    def forward(self, x):
+      return super().forward(x).to(torch.float32)
+
+  model.lm_head = CastOutputToFloat(model.lm_head)
+
+  config = LoraConfig(
+      r=lora_rank,
+      lora_alpha=lora_alpha,
+      target_modules=["q_proj", "v_proj"],
+      lora_dropout=lora_dropout,
+      bias="none",
+      task_type="CAUSAL_LM",
+  )
+
+  model = get_peft_model(model, config)
+  model.print_trainable_parameters()
+
+  data = load_dataset(dataset_name)
+  data = data.map(
+      lambda samples: tokenizer(samples["quote"]),
+      batched=True,
+  )
+
+  trainer = transformers.Trainer(
+      model=model,
+      train_dataset=data["train"],
+      args=TrainingArguments(
+          per_device_train_batch_size=4,
+          gradient_accumulation_steps=4,
+          warmup_steps=warmup_steps,
+          max_steps=max_steps,
+          learning_rate=learning_rate,
+          fp16=True,
+          logging_steps=1,
+          output_dir=output_dir,
+          ddp_find_unused_parameters=False,
+      ),
+      data_collator=transformers.DataCollatorForLanguageModeling(
+          tokenizer,
+          mlm=False,
+      ),
+  )
+  # Silence the warnings. Please re-enable for inference!
+  model.config.use_cache = False
+  trainer.train()
+
+  model.save_pretrained(output_dir)

community-content/vertex_vision_model_garden/model_oss/peft/config.properties

@@ -0,0 +1,21 @@
+number_of_netty_threads=32
+job_queue_size=1000
+model_store=/home/model-server/model-store
+workflow_store=/home/model-server/wf-store
+default_response_timeout=1800
+service_envelope=json
+inference_address=http://0.0.0.0:7080
+management_address=http://0.0.0.0:7081
+metrics_address=http://0.0.0.0:7082
+
+models={\
+  "peft_serving": {\
+    "1.0": {\
+        "defaultVersion": true,\
+        "marName": "peft_serving.mar",\
+        "minWorkers": 1,\
+        "maxWorkers": 1,\
+        "batchSize": 1\
+    }\
+  }\
+}

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

@@ -0,0 +1,107 @@
+# Dockerfile for PEFT Serving.
+#
+# To build:
+# docker build -f model_oss/peft/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/torchserve:0.7.0-gpu
+
+USER root
+
+ENV infer_port=7080
+ENV mng_port=7081
+ENV model_name="peft_serving"
+ENV PATH="/home/model-server/:${PATH}"
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim \
+        git \
+        git-lfs
+RUN git lfs install
+
+# Install libraries.
+ENV PIP_ROOT_USER_ACTION=ignore
+RUN python3 -m pip install --upgrade pip
+RUN pip install --upgrade torch==2.0.1
+RUN pip install torchvision==0.15.2
+RUN pip install tokenizers==0.13.3
+RUN pip install accelerate==0.21.0
+RUN pip install sentencepiece==0.1.99
+RUN pip install grpcio-status==1.33.2
+RUN pip install protobuf==3.19.6
+RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/peft.git
+RUN pip install datasets==2.14.4
+RUN pip install triton==2.0.0.dev20221120
+RUN pip install xformers==0.0.20
+RUN pip install google-cloud-storage==2.7.0
+RUN pip install absl-py==1.4.0
+RUN pip install scipy==1.10.1
+RUN pip install evaluate==0.4.0
+RUN pip install scikit-learn==1.2.2
+RUN pip install loralib==0.1.1
+RUN pip install bitsandbytes==0.39.0
+RUN pip install trl==0.4.4
+RUN pip install einops==0.6.1
+
+# Install diffusers from source.
+RUN git clone --depth 1 --branch v0.16.1 https://github.com/huggingface/diffusers.git
+WORKDIR diffusers
+RUN pip install -e .
+WORKDIR /home/model-server
+
+# Install transformers from source.
+RUN git clone --depth 1 --branch v4.31.0 https://github.com/huggingface/transformers.git
+# The patch is used to change the transformers loading model behavior:
+# 1) For models on Huggingface hub: if the model has multiple shards, each shard
+#    will be downloaded separately and get deleted after loading to GPU.
+# 2) For models on local disk: if a model bin file is actually a text file
+#    recording a GCS path, the model file will be downloaded and get deleted
+#    after loading to GPU.
+COPY model_oss/peft/hf_transformers_lazy_download.patch /home/model-server/hf_transformers_lazy_download.patch
+WORKDIR transformers
+RUN git apply /home/model-server/hf_transformers_lazy_download.patch
+RUN pip install -e .
+WORKDIR /home/model-server
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Copy model artifacts.
+COPY model_oss/peft/handler.py /home/model-server/handler.py
+COPY model_oss/peft/config.properties /home/model-server/config.properties
+COPY model_oss/util/ /home/model-server/util/
+ENV PYTHONPATH /home/model-server/
+
+# Expose ports.
+EXPOSE ${infer_port}
+EXPOSE ${mng_port}
+
+# Set environments.
+ENV TASK "causal-language-modeling-lora"
+ENV BASE_MODEL_ID "openlm-research/open_llama_7b"
+ENV PRECISION_LOADING_MODE "float16"
+ENV FINETUNED_LORA_MODEL_PATH ""
+
+
+# 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.
+CMD ["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_vision_model_garden/model_oss/peft/dockerfile/train.Dockerfile

@@ -0,0 +1,111 @@
+# Dockerfile for PEFT Training.
+#
+# To build:
+# docker build -f model_oss/peft/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}
+
+# Builds GPU docker image of PyTorch
+# Uses multi-staged approach to reduce size
+# Stage 1
+# Use base conda image to reduce time
+FROM continuumio/miniconda3:latest AS compile-image
+# Specify py version
+ENV PYTHON_VERSION=3.8
+# Install apt libs - copied from https://github.com/huggingface/accelerate/blob/main/docker/accelerate-gpu/Dockerfile
+RUN apt-get update && \
+    apt-get install -y curl git wget software-properties-common git-lfs && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists*
+
+# Install audio-related libraries
+RUN apt-get update && \
+    apt install -y ffmpeg
+
+RUN apt install -y libsndfile1-dev
+RUN git lfs install
+
+# Create our conda env - copied from https://github.com/huggingface/accelerate/blob/main/docker/accelerate-gpu/Dockerfile
+RUN conda create --name peft python=${PYTHON_VERSION} ipython jupyter pip
+RUN python3 -m pip install --no-cache-dir --upgrade pip
+
+# Below is copied from https://github.com/huggingface/accelerate/blob/main/docker/accelerate-gpu/Dockerfile
+# We don't install pytorch here yet since CUDA isn't available
+# instead we use the direct torch wheel
+ENV PATH /opt/conda/envs/peft/bin:$PATH
+# Activate our bash shell
+RUN chsh -s /bin/bash
+SHELL ["/bin/bash", "-c"]
+# Activate the conda env and install transformers + accelerate from source
+RUN source activate peft
+RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/transformers
+RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/accelerate
+RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/peft#egg=peft[test]
+RUN python3 -m pip install --no-cache-dir bitsandbytes
+
+# Stage 2
+FROM nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04 AS build-image
+COPY --from=compile-image /opt/conda /opt/conda
+ENV PATH /opt/conda/bin:$PATH
+
+# Install apt libs
+RUN apt-get update && \
+    apt-get install -y curl git wget vim && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists*
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+RUN echo "source activate peft" >> ~/.profile
+
+# Install libraries.
+RUN pip install --upgrade torch==2.0.1
+RUN pip install torchvision==0.15.2
+RUN pip install git+https://github.com/huggingface/transformers@de9255de27abfcae4a1f816b904915f0b1e23cd9
+RUN pip install transformers -U
+RUN pip install accelerate==0.21.0
+RUN pip install sentencepiece==0.1.99
+RUN pip install grpcio-status==1.33.2
+RUN pip install protobuf==3.19.6
+RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/peft.git
+RUN pip install datasets==2.9.0
+RUN pip install triton==2.0.0.dev20221120
+RUN pip install xformers==0.0.20
+RUN pip install Jinja2==3.1.2
+RUN pip install ftfy==6.1.1
+RUN pip install cloudml-hypertune==0.1.0.dev6
+RUN pip install tensorboard==2.12.0
+RUN pip install scipy==1.10.1
+RUN pip install evaluate==0.4.0
+RUN pip install scikit-learn==1.2.2
+RUN pip install loralib==0.1.1
+RUN pip install bitsandbytes==0.39.0
+RUN pip install trl==0.4.4
+RUN pip install einops==0.6.1
+RUN pip install google-cloud-storage==2.7.0
+
+RUN git clone --depth 1 --branch v0.16.1 https://github.com/huggingface/diffusers.git
+WORKDIR diffusers
+RUN pip install -e .
+
+# Switch to diffusers examples folder.
+WORKDIR examples
+
+# NOTE: use 'sed' to modify train_text_to_image_lora.py to
+#       fix the bug for accelerator.
+RUN sed -i \
+"s#logging_dir=logging_dir#project_dir=logging_dir#g" \
+text_to_image/train_text_to_image_lora.py
+
+# Config accelerate.
+RUN mkdir -p ./vertex_vision_model_garden_peft/
+COPY model_oss/peft/train.sh ./vertex_vision_model_garden_peft/train.sh
+COPY model_oss/peft/*.py     ./vertex_vision_model_garden_peft/
+COPY model_oss/util /diffusers/examples/util
+ENV PYTHONPATH /diffusers/examples/
+
+# Generate accelerate config at the beginning of docker run.
+ENTRYPOINT ["python3", "vertex_vision_model_garden_peft/main.py"]

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

@@ -0,0 +1,250 @@
+"""Custom handler for huggingface/peft models."""
+
+# pylint: disable=g-importing-member
+# pylint: disable=logging-fstring-interpolation
+
+import logging
+import os
+from typing import Any, List
+
+from absl import logging
+from diffusers import DPMSolverMultistepScheduler
+from diffusers import StableDiffusionPipeline
+from peft import PeftModel
+from PIL import Image
+import torch
+import transformers
+from transformers import AutoModelForCausalLM
+from transformers import AutoModelForSequenceClassification
+from transformers import AutoTokenizer
+from transformers import BitsAndBytesConfig
+from ts.torch_handler.base_handler import BaseHandler
+
+from util import constants
+from util import fileutils
+from util import image_format_converter
+
+# Tasks
+TEXT_TO_IMAGE_LORA = "text-to-image-lora"
+SEQUENCE_CLASSIFICATION_LORA = "sequence-classification-lora"
+CAUSAL_LANGUAGE_MODELING_LORA = "causal-language-modeling-lora"
+INSTRUCT_LORA = "instruct-lora"
+
+# Inference parameters.
+_NUM_INFERENCE_STEPS = 25
+_MAX_LENGTH_DEFAULT = 200
+_TOP_K_DEFAULT = 10
+
+
+class PeftHandler(BaseHandler):
+  """Custom handler for Peft models."""
+
+  def initialize(self, context: Any):
+    """Initializes the handler."""
+    logging.info("Start to initialize the PEFT 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.precision_mode = os.environ.get(
+        "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):
+      fileutils.download_gcs_dir_to_local(
+          self.base_model_id,
+          constants.LOCAL_BASE_MODEL_DIR,
+          skip_hf_model_bin=True,
+      )
+      self.base_model_id = constants.LOCAL_BASE_MODEL_DIR
+    self.finetuned_lora_model_path = os.environ.get(
+        "FINETUNED_LORA_MODEL_PATH", ""
+    )
+    if fileutils.is_gcs_path(self.finetuned_lora_model_path):
+      fileutils.download_gcs_dir_to_local(
+          self.finetuned_lora_model_path, constants.LOCAL_MODEL_DIR
+      )
+      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" {self.finetuned_lora_model_path}, and precision"
+        f" {self.precision_mode}."
+    )
+
+    self.pipeline = None
+    self.model = None
+    self.tokenizer = None
+    if self.task == TEXT_TO_IMAGE_LORA:
+      pipeline = StableDiffusionPipeline.from_pretrained(
+          self.base_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()
+      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)
+      logging.debug("Initialized the tokenizer for sequence classification.")
+      model = AutoModelForSequenceClassification.from_pretrained(
+          self.base_model_id, torch_dtype=torch.float16
+      )
+      logging.debug("Initialized the base model for sequence classification.")
+      if self.finetuned_lora_model_path:
+        model = PeftModel.from_pretrained(model, self.finetuned_lora_model_path)
+        logging.debug("Initialized the LoRA model for sequence classification.")
+      model.to(self.map_location)
+      self.model = model
+      self.tokenizer = tokenizer
+    elif (
+        self.task == CAUSAL_LANGUAGE_MODELING_LORA or self.task == INSTRUCT_LORA
+    ):
+      tokenizer = AutoTokenizer.from_pretrained(self.base_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,
+          )
+        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,
+          )
+      else:
+        model = AutoModelForCausalLM.from_pretrained(
+            self.base_model_id,
+            return_dict=True,
+            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)
+        logging.debug("Initialized the LoRA model.")
+      pipeline = transformers.pipeline(
+          "text-generation",
+          model=model,
+          tokenizer=tokenizer,
+      )
+      self.tokenizer = tokenizer
+      self.pipeline = pipeline
+    else:
+      raise ValueError(f"Invalid TASK: {self.task}")
+
+    self.initialized = True
+    logging.info("The PEFT handler was initialized.")
+
+  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.
+    max_length = _MAX_LENGTH_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 "top_k" in data[0]:
+      top_k = data[0]["top_k"]
+
+    return prompts, max_length, top_k
+
+  def inference(self, data: Any, *args, **kwargs) -> List[Image.Image]:
+    """Runs the inference."""
+    prompts, max_length, top_k = data
+    logging.debug(
+        f"Inference prompts={prompts}, max_length={max_length}, 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.to(self.map_location)
+      with torch.no_grad():
+        outputs = self.model(**encoded_input)
+      predictions = outputs.logits.argmax(dim=-1)
+      predicted_results = predictions.tolist()
+    elif (
+        self.task == CAUSAL_LANGUAGE_MODELING_LORA or self.task == INSTRUCT_LORA
+    ):
+      predicted_results = self.pipeline(
+          prompts,
+          max_length=max_length,
+          do_sample=True,
+          top_k=top_k,
+          num_return_sequences=1,
+          eos_token_id=self.tokenizer.eos_token_id,
+      )
+    else:
+      raise ValueError(f"Invalid TASK: {self.task}")
+    return predicted_results
+
+  def postprocess(self, data: Any) -> List[str]:
+    """Postprocesses output 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
+      ]
+    else:
+      outputs = data
+    return outputs
+
+
+# pylint: enable=logging-fstring-interpolation

community-content/vertex_vision_model_garden/model_oss/peft/hf_transformers_lazy_download.patch

@@ -0,0 +1,131 @@
+diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
+index 45459ed..32527f4 100644
+--- a/src/transformers/modeling_utils.py
++++ b/src/transformers/modeling_utils.py
+@@ -32,6 +32,8 @@ import torch
+ from packaging import version
+ from torch import Tensor, nn
+ from torch.nn import CrossEntropyLoss
++from huggingface_hub import hf_hub_download
++from google.cloud import storage
+
+ from .activations import get_activation
+ from .configuration_utils import PretrainedConfig
+@@ -442,6 +444,29 @@ def load_state_dict(checkpoint_file: Union[str, os.PathLike]):
+     """
+     Reads a PyTorch checkpoint file, returning properly formatted errors if they arise.
+     """
++    delete_download = False
++    tmp_dir = "/tmp/model"
++    os.makedirs(tmp_dir, exist_ok=True)
++    if isinstance(checkpoint_file, dict):
++        # Download model file from huggingface
++        print(f"==> Download model from HF: {checkpoint_file}")
++        checkpoint_file = hf_hub_download(
++            local_dir=tmp_dir, local_dir_use_symlinks=False, force_download=True, resume_download=True, **checkpoint_file)
++        delete_download = True
++    else:
++        with open(checkpoint_file, "rb") as f:
++            is_gcs_file = (f.read(2) == b"gs")
++        if is_gcs_file:
++            # Download model file from GCS
++            with open(checkpoint_file, "r") as f:
++                gcs_file = f.read()
++            checkpoint_file = os.path.join(tmp_dir, gcs_file.split("/")[-1])
++            print(f"==> Download model from GCS: {gcs_file} to: {checkpoint_file}")
++            client = storage.Client()
++            with open(checkpoint_file, 'wb') as f:
++                client.download_blob_to_file(gcs_file, f)
++            delete_download = True
++
+     if checkpoint_file.endswith(".safetensors") and is_safetensors_available():
+         # Check format of the archive
+         with safe_open(checkpoint_file, framework="pt") as f:
+@@ -455,9 +480,9 @@ def load_state_dict(checkpoint_file: Union[str, os.PathLike]):
+             raise NotImplementedError(
+                 f"Conversion from a {metadata['format']} safetensors archive to PyTorch is not implemented yet."
+             )
+-        return safe_load_file(checkpoint_file)
++        state_dict = safe_load_file(checkpoint_file)
+     try:
+-        return torch.load(checkpoint_file, map_location="cpu")
++        state_dict = torch.load(checkpoint_file, map_location="cpu")
+     except Exception as e:
+         try:
+             with open(checkpoint_file) as f:
+@@ -478,6 +503,10 @@ def load_state_dict(checkpoint_file: Union[str, os.PathLike]):
+                 f"at '{checkpoint_file}'. "
+                 "If you tried to load a PyTorch model from a TF 2.0 checkpoint, please set from_tf=True."
+             )
++    if delete_download:
++        print(f"==> Delete downloaded model: {checkpoint_file}")
++        os.remove(checkpoint_file)
++    return state_dict
+
+
+ def set_initialized_submodules(model, state_dict_keys):
+@@ -3179,7 +3208,10 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix
+             return mismatched_keys
+
+         if resolved_archive_file is not None:
+-            folder = os.path.sep.join(resolved_archive_file[0].split(os.path.sep)[:-1])
++            if isinstance(resolved_archive_file, str):
++                folder = os.path.sep.join(resolved_archive_file[0].split(os.path.sep)[:-1])
++            else:
++                folder = None
+         else:
+             folder = None
+         if device_map is not None and is_safetensors:
+diff --git a/src/transformers/utils/hub.py b/src/transformers/utils/hub.py
+index ffed743..4b15770 100644
+--- a/src/transformers/utils/hub.py
++++ b/src/transformers/utils/hub.py
+@@ -414,20 +414,34 @@ def cached_file(
+     user_agent = http_user_agent(user_agent)
+     try:
+         # Load from URL or cache if already cached
+-        resolved_file = hf_hub_download(
+-            path_or_repo_id,
+-            filename,
+-            subfolder=None if len(subfolder) == 0 else subfolder,
+-            repo_type=repo_type,
+-            revision=revision,
+-            cache_dir=cache_dir,
+-            user_agent=user_agent,
+-            force_download=force_download,
+-            proxies=proxies,
+-            resume_download=resume_download,
+-            use_auth_token=use_auth_token,
+-            local_files_only=local_files_only,
+-        )
++        if filename.endswith(".bin"):
++            # NOTE: To save disk we do not download bin file eagerly. Do not support safetensors.
++            resolved_file = dict(
++                repo_id=path_or_repo_id,
++                filename=filename,
++                subfolder=None if len(subfolder) == 0 else subfolder,
++                repo_type=repo_type,
++                revision=revision,
++                user_agent=user_agent,
++                proxies=proxies,
++                use_auth_token=use_auth_token,
++            )
++            print(f"--> Apply lazy download to bin file: {resolved_file}")
++        else:
++            resolved_file = hf_hub_download(
++                path_or_repo_id,
++                filename,
++                subfolder=None if len(subfolder) == 0 else subfolder,
++                repo_type=repo_type,
++                revision=revision,
++                cache_dir=cache_dir,
++                user_agent=user_agent,
++                force_download=force_download,
++                proxies=proxies,
++                resume_download=resume_download,
++                use_auth_token=use_auth_token,
++                local_files_only=local_files_only,
++            )
+
+     except RepositoryNotFoundError:
+         raise EnvironmentError(

community-content/vertex_vision_model_garden/model_oss/peft/instruct_lora.py

@@ -0,0 +1,97 @@
+"""Instruct/Chat with LoRA models."""
+
+# pylint: disable=g-importing-member
+from datasets import load_dataset
+from peft import LoraConfig
+import torch
+from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
+from transformers import BitsAndBytesConfig
+from transformers import TrainingArguments
+from trl import SFTTrainer
+
+
+def finetune_instruct(
+    pretrained_model_id: str,
+    dataset_name: str,
+    output_dir: str,
+    lora_rank: int = 64,
+    lora_alpha: int = 16,
+    lora_dropout: float = 0.1,
+    warmup_ratio: int = 0.03,
+    max_steps: int = 10,
+    max_seq_length: int = 512,
+    learning_rate: float = 2e-4,
+) -> None:
+  """Finetunes instruct."""
+  dataset = load_dataset(dataset_name, split="train")
+
+  bnb_config = BitsAndBytesConfig(
+      load_in_4bit=True,
+      bnb_4bit_quant_type="nf4",
+      bnb_4bit_compute_dtype=torch.float16,
+  )
+
+  model = AutoModelForCausalLM.from_pretrained(
+      pretrained_model_id,
+      quantization_config=bnb_config,
+      trust_remote_code=True,
+  )
+  model.config.use_cache = False
+
+  tokenizer = AutoTokenizer.from_pretrained(
+      pretrained_model_id, trust_remote_code=True
+  )
+  tokenizer.pad_token = tokenizer.eos_token
+
+  peft_config = LoraConfig(
+      lora_alpha=lora_alpha,
+      lora_dropout=lora_dropout,
+      r=lora_rank,
+      bias="none",
+      task_type="CAUSAL_LM",
+      target_modules=[
+          "query_key_value",
+          "dense",
+          "dense_h_to_4h",
+          "dense_4h_to_h",
+      ],
+  )
+
+  per_device_train_batch_size = 4
+  gradient_accumulation_steps = 4
+  optim = "paged_adamw_32bit"
+  save_steps = 10
+  logging_steps = 10
+  max_grad_norm = 0.3
+  lr_scheduler_type = "constant"
+
+  training_arguments = TrainingArguments(
+      output_dir=output_dir,
+      per_device_train_batch_size=per_device_train_batch_size,
+      gradient_accumulation_steps=gradient_accumulation_steps,
+      optim=optim,
+      save_steps=save_steps,
+      logging_steps=logging_steps,
+      learning_rate=learning_rate,
+      fp16=True,
+      max_grad_norm=max_grad_norm,
+      max_steps=max_steps,
+      warmup_ratio=warmup_ratio,
+      group_by_length=True,
+      lr_scheduler_type=lr_scheduler_type,
+  )
+
+  trainer = SFTTrainer(
+      model=model,
+      train_dataset=dataset,
+      peft_config=peft_config,
+      dataset_text_field="text",
+      max_seq_length=max_seq_length,
+      tokenizer=tokenizer,
+      args=training_arguments,
+  )
+  for name, module in trainer.model.named_modules():
+    if "norm" in name:
+      module = module.to(torch.float32)
+  trainer.train()

community-content/vertex_vision_model_garden/model_oss/peft/main.py

@@ -0,0 +1,177 @@
+"""Main function to start PEFT finetuning."""
+import subprocess
+
+from absl import app
+from absl import flags
+from absl import logging
+
+from peft import causal_language_modeling_lora
+from peft import instruct_lora
+from peft import sequence_classification_lora
+from util import constants
+from util import fileutils
+
+_TASK = flags.DEFINE_string(
+    'task',
+    constants.CAUSAL_LANGUAGE_MODELING_LORA,
+    'The supported PEFT tasks.',
+)
+
+_PRETRAINED_MODEL_ID = flags.DEFINE_string(
+    'pretrained_model_id',
+    None,
+    'The pretrained model id. Supported models can be causal language modeling'
+    ' models from https://github.com/huggingface/peft/tree/main.',
+    required=True,
+)
+
+_DATASET_NAME = flags.DEFINE_string(
+    'dataset_name',
+    None,
+    'The dataset name in huggingface.',
+    required=True,
+)
+
+_OUTPUT_DIR = flags.DEFINE_string(
+    'output_dir',
+    None,
+    'The output directory.',
+    required=True,
+)
+
+_PRECISION_MODE = flags.DEFINE_string(
+    'precision_mode',
+    constants.PRECISION_MODE_16,
+    'Supported finetuning precision_modes are `{}` and `{}`.'.format(
+        constants.PRECISION_MODE_8, constants.PRECISION_MODE_16
+    ),
+)
+
+_LORA_RANK = flags.DEFINE_integer(
+    'lora_rank',
+    16,
+    'The rank of the update matrices, expressed in int. Lower rank results in'
+    ' smaller update matrices with fewer trainable parameters, referring to'
+    ' https://huggingface.co/docs/peft/conceptual_guides/lora.',
+)
+
+_LORA_ALPHA = flags.DEFINE_integer(
+    'lora_alpha',
+    32,
+    'LoRA scaling factor, referring to'
+    ' https://huggingface.co/docs/peft/conceptual_guides/lora.',
+)
+
+_LORA_DROPOUT = flags.DEFINE_float(
+    'lora_dropout',
+    0.05,
+    'dropout probability of the LoRA layers, referring to'
+    ' https://huggingface.co/docs/peft/task_guides/token-classification-lora.',
+)
+
+_WARMUP_STEPS = flags.DEFINE_integer(
+    'warmup_steps',
+    10,
+    'Number of steps for the warmup in the learning rate scheduler.',
+)
+
+_WARMUP_RATIO = flags.DEFINE_float(
+    'warmup_ratio',
+    0.03,
+    'The warmup ratio in the learning rate scheduler.',
+)
+
+_MAX_STEPS = flags.DEFINE_integer(
+    'max_steps',
+    10,
+    'Total number of training steps.',
+)
+
+_MAX_SEQ_LENGTH = flags.DEFINE_integer(
+    'max_seq_length',
+    512,
+    'The maximum sequence length.',
+)
+
+_NUM_EPOCHS = flags.DEFINE_integer(
+    'num_epochs',
+    20,
+    'The number of training epochs.',
+)
+
+_BATCH_SIZE = flags.DEFINE_integer(
+    'batch_size',
+    32,
+    'The batch size.',
+)
+
+_LEARNING_RATE = flags.DEFINE_float(
+    'learning_rate',
+    2e-4,
+    'The learning rate after the potential warmup period.',
+)
+
+
+def main(_) -> None:
+  task = _TASK.value
+  pretrained_model_id = _PRETRAINED_MODEL_ID.value
+  local_pretrained_model_id = None
+  if pretrained_model_id.startswith(constants.GCS_URI_PREFIX):
+    logging.info(
+        'Start to copy pretrained models locally: %s.', pretrained_model_id
+    )
+    fileutils.download_gcs_dir_to_local(
+        pretrained_model_id, constants.LOCAL_BASE_MODEL_DIR
+    )
+    local_pretrained_model_id = constants.LOCAL_BASE_MODEL_DIR
+    logging.info(
+        'Finished copying pretrained models locally to: %s.',
+        local_pretrained_model_id,
+    )
+  if task == constants.TEXT_TO_IMAGE_LORA:
+    subprocess.run(['/bin/bash', 'train.sh'], check=True)
+  elif task == constants.SEQUENCE_CLASSIFICATION_LORA:
+    sequence_classification_lora.finetune_sequence_classification(
+        pretrained_model_id=pretrained_model_id,
+        dataset_name=_DATASET_NAME.value,
+        output_dir=_OUTPUT_DIR.value,
+        lora_rank=_LORA_RANK.value,
+        lora_alpha=_LORA_ALPHA.value,
+        lora_dropout=_LORA_DROPOUT.value,
+        num_epochs=_NUM_EPOCHS.value,
+        batch_size=_BATCH_SIZE.value,
+        learning_rate=_LEARNING_RATE.value,
+    )
+  elif task == constants.CAUSAL_LANGUAGE_MODELING_LORA:
+    causal_language_modeling_lora.finetune_causal_language_modeling(
+        pretrained_model_id=pretrained_model_id,
+        dataset_name=_DATASET_NAME.value,
+        output_dir=_OUTPUT_DIR.value,
+        precision_mode=_PRECISION_MODE.value,
+        lora_rank=_LORA_RANK.value,
+        lora_alpha=_LORA_ALPHA.value,
+        lora_dropout=_LORA_DROPOUT.value,
+        warmup_steps=_WARMUP_STEPS.value,
+        max_steps=_MAX_STEPS.value,
+        learning_rate=_LEARNING_RATE.value,
+        local_pretrained_model_id=local_pretrained_model_id,
+    )
+  elif task == constants.INSTRUCT_LORA:
+    instruct_lora.finetune_instruct(
+        pretrained_model_id=pretrained_model_id,
+        dataset_name=_DATASET_NAME.value,
+        output_dir=_OUTPUT_DIR.value,
+        lora_rank=_LORA_RANK.value,
+        lora_alpha=_LORA_ALPHA.value,
+        lora_dropout=_LORA_DROPOUT.value,
+        warmup_ratio=_WARMUP_RATIO.value,
+        max_steps=_MAX_STEPS.value,
+        max_seq_length=_MAX_SEQ_LENGTH.value,
+        learning_rate=_LEARNING_RATE.value,
+    )
+  else:
+    raise ValueError('The task {} is not supported.'.format(task))
+
+
+if __name__ == '__main__':
+  app.run(main)

community-content/vertex_vision_model_garden/model_oss/peft/sequence_classification_lora.py

@@ -0,0 +1,133 @@
+"""Sequence classification with LoRA models."""
+
+# pylint: disable=g-importing-member
+
+from datasets import load_dataset
+import evaluate
+from peft import get_peft_model
+from peft import LoraConfig
+import torch
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from transformers import AutoModelForSequenceClassification
+from transformers import AutoTokenizer
+from transformers import get_linear_schedule_with_warmup
+
+
+def finetune_sequence_classification(
+    pretrained_model_id: str,
+    dataset_name: str,
+    output_dir: str,
+    lora_rank: int = 8,
+    lora_alpha: int = 16,
+    lora_dropout: float = 0.1,
+    num_epochs: int = 20,
+    batch_size: int = 32,
+    learning_rate: float = 3e-4,
+) -> None:
+  """Finetunes sequence classification."""
+  task = "mrpc"
+  device = "cuda"
+
+  peft_config = LoraConfig(
+      task_type="SEQ_CLS",
+      inference_mode=False,
+      r=lora_rank,
+      lora_alpha=lora_alpha,
+      lora_dropout=lora_dropout,
+  )
+  if any(k in pretrained_model_id for k in ("gpt", "opt", "bloom")):
+    padding_side = "left"
+  else:
+    padding_side = "right"
+
+  tokenizer = AutoTokenizer.from_pretrained(
+      pretrained_model_id, padding_side=padding_side
+  )
+  if getattr(tokenizer, "pad_token_id") is None:
+    tokenizer.pad_token_id = tokenizer.eos_token_id
+
+  datasets = load_dataset(dataset_name, task)
+  metric = evaluate.load(dataset_name, task)
+
+  def tokenize_function(examples):
+    # max_length=None => use the model max length (it's actually the default)
+    outputs = tokenizer(
+        examples["sentence1"],
+        examples["sentence2"],
+        truncation=True,
+        max_length=None,
+    )
+    return outputs
+
+  tokenized_datasets = datasets.map(
+      tokenize_function,
+      batched=True,
+      remove_columns=["idx", "sentence1", "sentence2"],
+  )
+
+  # We also rename the 'label' column to 'labels' which is the expected name for
+  # labels by the models of the transformers library.
+  tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+  def collate_fn(examples):
+    return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+  # Instantiate dataloaders.
+  train_dataloader = DataLoader(
+      tokenized_datasets["train"],
+      shuffle=True,
+      collate_fn=collate_fn,
+      batch_size=batch_size,
+  )
+  eval_dataloader = DataLoader(
+      tokenized_datasets["validation"],
+      shuffle=False,
+      collate_fn=collate_fn,
+      batch_size=batch_size,
+  )
+
+  model = AutoModelForSequenceClassification.from_pretrained(
+      pretrained_model_id, return_dict=True
+  )
+  model = get_peft_model(model, peft_config)
+  model.print_trainable_parameters()
+
+  optimizer = AdamW(params=model.parameters(), lr=learning_rate)
+
+  # Instantiate scheduler
+  lr_scheduler = get_linear_schedule_with_warmup(
+      optimizer=optimizer,
+      num_warmup_steps=0.06 * (len(train_dataloader) * num_epochs),
+      num_training_steps=(len(train_dataloader) * num_epochs),
+  )
+
+  model.to(device)
+  for epoch in range(num_epochs):
+    model.train()
+    for _, batch in enumerate(tqdm(train_dataloader)):
+      batch.to(device)
+      outputs = model(**batch)
+      loss = outputs.loss
+      loss.backward()
+      optimizer.step()
+      lr_scheduler.step()
+      optimizer.zero_grad()
+
+    model.eval()
+    for _, batch in enumerate(tqdm(eval_dataloader)):
+      batch.to(device)
+      with torch.no_grad():
+        outputs = model(**batch)
+      predictions = outputs.logits.argmax(dim=-1)
+      references = batch["labels"]
+      metric.add_batch(
+          predictions=predictions,
+          references=references,
+      )
+
+    eval_metric = metric.compute()
+    print(f"epoch {epoch}:", eval_metric)
+
+  model.save_pretrained(output_dir)

community-content/vertex_vision_model_garden/model_oss/peft/train.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+
+# Setup accelerate config before running trainer.
+python -c "from accelerate.utils import write_basic_config; write_basic_config(mixed_precision='fp16')"
+
+accelerate launch "$@"

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

@@ -0,0 +1,115 @@
+FROM pytorch/torchserve:0.7.1-gpu
+
+USER root
+
+ENV infer_port=7080
+ENV mng_port=7081
+ENV model_name="pic2word"
+ENV PATH="/home/model-server/:${PATH}"
+
+# Copy license.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        wget
+
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Install dependencies.
+ENV PIP_ROOT_USER_ACTION=ignore
+RUN python3 -m pip install --upgrade pip
+RUN pip install google-cloud-storage==2.7.0
+RUN pip install open_clip_torch==2.20.0
+RUN pip install numpy==1.22.0
+RUN pip install scikit-image==0.21.0
+RUN pip install scikit-learn==1.0.2
+RUN pip install torch==2.0.0
+RUN pip install torchvision==0.15.2
+RUN pip install tensorboard==2.13.0
+RUN pip install ase==3.21.1
+RUN pip install braceexpand==0.1.7
+RUN pip install cached-property==1.5.2
+RUN pip install configparser==5.0.2
+RUN pip install cycler==0.10.0
+RUN pip install decorator==4.4.2
+RUN pip install docker-pycreds==0.4.0
+RUN pip install gitdb==4.0.7
+RUN pip install gitpython==3.1.30
+RUN pip install googledrivedownloader==0.4
+RUN pip install h5py==3.1.0
+RUN pip install isodate==0.6.0
+RUN pip install jinja2==3.0.1
+RUN pip install kiwisolver==1.3.1
+RUN pip install littleutils==0.2.2
+RUN pip install llvmlite==0.36.0
+RUN pip install markupsafe==2.0.1
+RUN pip install matplotlib==3.3.4
+RUN pip install networkx==2.5.1
+RUN pip install numba==0.53.1
+RUN pip install ogb==1.3.1
+RUN pip install outdated==0.2.1
+RUN pip install pathtools==0.1.2
+RUN pip install promise==2.3
+RUN pip install psutil==5.8.0
+RUN pip install pyarrow==4.0.0
+RUN pip install pyparsing==2.4.7
+RUN pip install python-louvain==0.15
+RUN pip install pyyaml==5.4.1
+RUN pip install rdflib==5.0.0
+RUN pip install sentry-sdk==1.14.0
+RUN pip install shortuuid==1.0.1
+RUN pip install sklearn==0.0
+RUN pip install smmap==4.0.0
+RUN pip install subprocess32==3.5.4
+RUN pip install torch-geometric==1.7.0
+RUN pip install wandb==0.10.30
+RUN pip install wilds==1.1.0
+RUN pip install ftfy==6.1.1
+RUN pip install regex==2023.6.3
+RUN pip install webdataset==0.2.48
+RUN pip install requests==2.31.0
+RUN pip install hydra-core==1.3.2
+RUN pip install omegaconf==2.3.0
+RUN pip install fairseq==0.10.0
+RUN pip install bitarray==2.7.6
+
+# Get 'composed_image_retrieval' repository from github.
+RUN git clone https://github.com/google-research/composed_image_retrieval
+# Set workdir to composed_image_retrieval.
+WORKDIR ./composed_image_retrieval
+# Using git reset command to pin it down to a specific version.
+RUN git reset --hard 8c053297c2fae9cd17ddcded48445a4f47208dbd
+
+# Fix issue introduced by installing composed_image_retrieval
+# https://github.com/huggingface/transformers/issues/8638#issuecomment-790772391
+RUN pip uninstall dataclasses -y
+
+# Copy model artifacts.
+COPY model_oss/pic2word/handler.py /home/model-server/handler.py
+
+# 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}" >> /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.
+CMD ["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_vision_model_garden/model_oss/pic2word/handler.py

@@ -0,0 +1,167 @@
+"""Custom handler for Pic2Word."""
+
+from argparse import Namespace  # pylint: disable=g-importing-member
+import os
+from typing import Any
+
+from absl import logging
+from data import CustomFolder
+from eval_utils import visualize_results
+from model.clip import load
+from model.model import convert_weights
+from model.model import IM2TEXT
+from params import get_project_root
+import torch
+from torch.utils.data import DataLoader
+from ts.torch_handler.base_handler import BaseHandler
+
+from util import fileutils
+
+# The COCO dataset is stored in a publicly accessible bucket.
+_COCO_STORAGE_DIR = "gs://pic2word-bucket/data/coco/"
+_COCO_LOCAL_DIR = "/home/model-server/composed_image_retrieval/data/coco/"
+_COCO_VAL2017_PATH = "coco/val2017"
+_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/"
+_DATA_DIR = "data"
+_CHECKPOINT_DIR = "checkpoint/pic2word_model.pt"
+_REQUEST_PROMPTS = "prompts"
+_REQUEST_OUTPUT_STORAGE_DIR = "output_storage_dir"
+_REQUEST_IMAGE_PATH = "image_path"
+_REQUEST_IMAGE_FILE_NAME = "image_file_name"
+_RESPONSE_MSG = "Successfully retrieved images."
+
+
+class ModelHandler(BaseHandler):
+  """A custom model handler implementation."""
+
+  def __init__(self):
+    self.initialized = False
+    self.gpu = 0
+    self.model = None
+    self.dataloader = None
+    self.prompt = None
+    self.output_storage_dir = None
+
+  def initialize(self, context: Any):
+    """Initialize."""
+    logging.info("Initializing pic2word.")
+
+    # Download COCO dataset. The model looks for this folder specifically
+    # during image retrieval to generate a response for each request.
+    # This is a publicly accessible bucket.
+    fileutils.download_gcs_dir_to_local(
+        _COCO_STORAGE_DIR,
+        _COCO_LOCAL_DIR,
+    )
+
+    # Load the model.
+
+    self.initialized = True
+
+    torch.cuda.set_device(self.gpu)
+    model, _, preprocess_val = load(_MODEL_NAME, jit=False)
+
+    img2text = IM2TEXT(
+        embed_dim=model.embed_dim,
+        output_dim=model.token_embedding.weight.shape[1],
+    )
+
+    model.cuda(self.gpu)
+    img2text.cuda(self.gpu)
+
+    convert_weights(model)
+    convert_weights(img2text)
+
+    self.model = model
+    self.img2text = img2text
+
+    # Load the dataset
+    logging.info("Loading dataset.")
+
+    root_project = os.path.join(get_project_root(), _DATA_DIR)
+    dataset = CustomFolder(
+        os.path.join(root_project, _COCO_VAL2017_PATH), transform=preprocess_val
+    )
+
+    # Initialize the dataloader. This is used to create the pickle file from
+    # the dataset.
+    dataloader = DataLoader(
+        dataset,
+        batch_size=64,
+        shuffle=False,
+        num_workers=1,
+        pin_memory=True,
+        drop_last=False,
+    )
+
+    self.dataloader = dataloader
+
+    logging.info("Finished initializing Pic2Word server.")
+
+  def preprocess(self, data: Any) -> str:
+    """Preprocess input data."""
+    logging.info("Preprocessing Pic2Word inference request.")
+    query = data[0]
+
+    self.output_storage_dir = query[_REQUEST_OUTPUT_STORAGE_DIR]
+    prompts = query[_REQUEST_PROMPTS]
+    prompts = prompts.split(",")
+    self.prompt = prompts
+
+    image_path = query[_REQUEST_IMAGE_PATH]
+    # The query image is only supported via GCS bucket upload.
+    fileutils.download_gcs_dir_to_local(image_path, _LOCAL_QUERY_PATH)
+    image_file_name = query[_REQUEST_IMAGE_FILE_NAME]
+
+    query_file = f"./query/{image_file_name}"
+
+    logging.info("Setting model args.")
+
+    args = {
+        "openai-pretrained": True,
+        "resume": _CHECKPOINT_DIR,
+        "retrieval_data": _COCO_DATASET_NAME,
+        "query_file": query_file,
+        "demo_out": _OUTPUT_LOCAL_DIR,
+        "prompts": prompts,
+        "distributed": False,
+        "dp": False,
+        "gpu": 0,
+        "model": _MODEL_NAME,
+        "world_size": 1,
+    }
+    model_input = Namespace(**args)
+
+    logging.info("Finished preprocessing Pic2Word inference request.")
+    return model_input
+
+  def inference(self, model_input: Any):
+    """Runs inference."""
+    logging.info("Running model-inference.")
+    visualize_results(
+        model=self.model,
+        img2text=self.img2text,
+        args=model_input,
+        prompt=self.prompt,
+        dataloader=self.dataloader,
+    )
+
+  def postprocess(self):
+    """Upload the output images to the bucket."""
+    logging.info("Running request postprocess.")
+    fileutils.upload_local_dir_to_gcs(
+        _IMAGE_OUTPUT_LOCAL_DIR, self.output_storage_dir
+    )
+
+  def handle(self, data: Any, context: Any) -> 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

community-content/vertex_vision_model_garden/model_oss/transformers/dockerfile/serve.Dockerfile

@@ -0,0 +1,79 @@
+# Dockerfile for serving dockers for transformers.
+#
+# To build:
+# docker build -f model_oss/transformers/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}
+# Switch to this base image for gpu serve.
+
+FROM pytorch/torchserve:0.7.0-gpu
+
+USER root
+
+ENV infer_port=7080
+ENV mng_port=7081
+ENV model_name="transformers_serving"
+ENV PATH="/home/model-server/:${PATH}"
+
+# Install libraries.
+ENV PIP_ROOT_USER_ACTION=ignore
+RUN python3 -m pip install --upgrade pip
+RUN pip install torch==1.13.1
+RUN pip install torchvision==0.14.1
+RUN pip install transformers==4.27.4
+RUN pip install datasets==2.9.0
+RUN pip install accelerate==0.17.0
+RUN pip install triton==2.0.0.dev20221120
+RUN pip install xformers==0.0.16
+RUN pip install google-cloud-storage==2.7.0
+RUN pip install absl-py==1.4.0
+
+# Install libraries for document-question-answering.
+RUN apt-get update
+RUN apt-get install -y --no-install-recommends tesseract-ocr
+RUN pip install tesseract==0.1.3
+RUN pip install pytesseract==0.3.10
+
+# Install tools.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Copy model artifacts.
+COPY model_oss/transformers/handler.py /home/model-server/handler.py
+COPY model_oss/util/ /home/model-server/util/
+ENV PYTHONPATH /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}" >> /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.
+CMD ["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_vision_model_garden/model_oss/transformers/handler.py

@@ -0,0 +1,233 @@
+"""Custom handler for huggingface/transformers models."""
+
+# pylint: disable=g-multiple-import
+# pylint: disable=g-importing-member
+
+import logging
+import os
+from typing import Any, List, Optional, Tuple
+
+from PIL import Image
+import torch
+from transformers import (
+    AutoProcessor,
+    AutoTokenizer,
+    Blip2ForConditionalGeneration,
+    Blip2Processor,
+    BlipForConditionalGeneration,
+    BlipForQuestionAnswering,
+    BlipProcessor,
+    CLIPModel,
+)
+from transformers import pipeline
+from ts.torch_handler.base_handler import BaseHandler
+
+from util import constants
+from util import fileutils
+from util import image_format_converter
+
+DEFAULT_MODEL_ID = "openai/clip-vit-base-patch32"
+SALESFORCE_BLIP = "Salesforce/blip"
+SALESFORCE_BLIP2 = "Salesforce/blip2"
+FLAN_T5 = "flan-t5"
+BART_LARGE_CNN = "facebook/bart-large-cnn"
+
+ZERO_CLASSIFICATION = "zero-shot-image-classification"
+FEATURE_EMBEDDING = "feature-embedding"
+ZERO_DETECTION = "zero-shot-object-detection"
+IMAGE_CAPTIONING = "image-to-text"
+VQA = "visual-question-answering"
+DQA = "document-question-answering"
+SUMMARIZATION = "summarization"
+SUMMARIZATION_TEMPLATE = (
+    "Summarize the following news article:\n{input}\nSummary:\n"
+)
+
+
+class TransformersHandler(BaseHandler):
+  """Custom handler for huggingface/transformers models."""
+
+  def initialize(self, context: Any):
+    """Custom initialize."""
+
+    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
+    # The model id is 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.
+    self.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 self.model_id.startswith(constants.GCS_URI_PREFIX):
+      gcs_path = self.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", self.model_id, local_model_dir)
+      fileutils.download_gcs_dir_to_local(self.model_id, local_model_dir)
+      self.model_id = local_model_dir
+
+    self.task = (
+        ZERO_CLASSIFICATION
+        if os.environ.get("TASK") is None
+        else os.environ["TASK"]
+    )
+    logging.info(
+        "Handler initializing task:%s, model:%s", self.task, self.model_id
+    )
+
+    if SALESFORCE_BLIP in self.model_id:
+      # pipeline() hasn't been ready for Salesforce/blip models.
+      self.salesforce_blip = True
+      self._create_blip_model()
+    else:
+      self.salesforce_blip = False
+      if self.task == FEATURE_EMBEDDING:
+        self.model = CLIPModel.from_pretrained(self.model_id).to(
+            self.map_location
+        )
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
+        self.processor = AutoProcessor.from_pretrained(self.model_id)
+      elif self.task == SUMMARIZATION and FLAN_T5 in self.model_id:
+        self.pipeline = pipeline(
+            task=self.task,
+            model=self.model_id,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+        )
+      else:
+        self.pipeline = pipeline(
+            task=self.task, model=self.model_id, device=self.device
+        )
+
+    self.initialized = True
+    logging.info("Handler initialization done.")
+
+  def _create_blip_model(self):
+    """A helper for creating BLIP and BLIP2 models."""
+    if SALESFORCE_BLIP2 in self.model_id:
+      self.torch_type = torch.float16
+      self.processor = Blip2Processor.from_pretrained(self.model_id)
+      self.model = Blip2ForConditionalGeneration.from_pretrained(
+          self.model_id, torch_dtype=self.torch_type
+      ).to(self.map_location)
+    else:
+      self.torch_type = torch.float32
+      self.processor = BlipProcessor.from_pretrained(self.model_id)
+      if self.task == IMAGE_CAPTIONING:
+        self.model = BlipForConditionalGeneration.from_pretrained(
+            self.model_id
+        ).to(self.map_location)
+      elif self.task == VQA:
+        self.model = BlipForQuestionAnswering.from_pretrained(self.model_id).to(
+            self.map_location
+        )
+
+  def _reformat_detection_result(self, data: List[Any]) -> List[Any]:
+    """Reformat zero-shot-object-detection output."""
+    if not data:
+      return [data]
+    boxes = {}
+    boxes["label"] = data[0]["label"]
+    boxes["boxes"] = []
+    for item in data:
+      box = {}
+      box["score"] = item["score"]
+      box.update(item["box"])
+      boxes["boxes"].append(box)
+    outputs = [boxes]
+    return outputs
+
+  def preprocess(
+      self, data: Any
+  ) -> Tuple[Optional[List[str]], Optional[List[Image.Image]]]:
+    """Preprocess input data."""
+    texts = None
+    images = None
+    if "text" in data[0]:
+      texts = [item["text"] for item in data]
+    if "image" in data[0]:
+      images = [
+          image_format_converter.base64_to_image(item["image"]) for item in data
+      ]
+    return texts, images
+
+  def inference(self, data: Any, *args, **kwargs) -> List[Any]:
+    """Run the inference."""
+    texts, images = data
+    preds = None
+    if self.task == ZERO_CLASSIFICATION:
+      preds = self.pipeline(images=images, candidate_labels=texts)
+    elif self.task == ZERO_DETECTION:
+      # The object detection pipeline doesn't support batch prediction.
+      preds = self.pipeline(image=images[0], candidate_labels=texts[0])
+    elif self.task == IMAGE_CAPTIONING:
+      if self.salesforce_blip:
+        inputs = self.processor(images[0], return_tensors="pt").to(
+            self.map_location, self.torch_type
+        )
+        preds = self.model.generate(**inputs)
+        preds = [
+            self.processor.decode(preds[0], skip_special_tokens=True).strip()
+        ]
+      else:
+        preds = self.pipeline(images=images)
+    elif self.task == VQA:
+      # The VQA pipelines doesn't support batch prediction.
+      if self.salesforce_blip:
+        inputs = self.processor(images[0], texts[0], return_tensors="pt").to(
+            self.map_location, self.torch_type
+        )
+        preds = self.model.generate(**inputs)
+        preds = [
+            self.processor.decode(preds[0], skip_special_tokens=True).strip()
+        ]
+      else:
+        preds = self.pipeline(image=images[0], question=texts[0])
+    elif self.task == DQA:
+      # The DQA pipelines doesn't support batch prediction.
+      preds = self.pipeline(image=images[0], question=texts[0])
+    elif self.task == FEATURE_EMBEDDING:
+      preds = {}
+      if texts:
+        inputs = self.tokenizer(
+            text=texts, padding=True, return_tensors="pt"
+        ).to(self.map_location)
+        text_features = self.model.get_text_features(**inputs)
+        preds["text_features"] = text_features.detach().cpu().numpy().tolist()
+      if images:
+        inputs = self.processor(images=images, return_tensors="pt").to(
+            self.map_location
+        )
+        image_features = self.model.get_image_features(**inputs)
+        preds["image_features"] = image_features.detach().cpu().numpy().tolist()
+      preds = [preds]
+    elif self.task == SUMMARIZATION and FLAN_T5 in self.model_id:
+      texts = [SUMMARIZATION_TEMPLATE.format(input=text) for text in texts]
+      preds = self.pipeline(texts, max_length=130)
+    elif self.task == SUMMARIZATION and self.model_id == BART_LARGE_CNN:
+      preds = self.pipeline(
+          texts[0], max_length=130, min_length=30, do_sample=False
+      )
+    else:
+      raise ValueError(f"Invalid TASK: {self.task}")
+    return preds
+
+  def postprocess(self, data: Any) -> List[Any]:
+    if self.task == ZERO_DETECTION:
+      data = self._reformat_detection_result(data)
+    return data

community-content/vertex_vision_model_garden/model_oss/util/commons.py

@@ -0,0 +1,79 @@
+"""Common utility lib for prediction on images."""
+
+from typing import Any, Dict, List
+
+import numpy as np
+from PIL import Image
+import tensorflow as tf
+import yaml
+
+from util import image_format_converter
+
+
+def get_prediction_instances(image: Image.Image) -> List[Dict[str, Any]]:
+  """Gets prediction instances.
+
+  Args:
+    image: Image instance.
+
+  Returns:
+    List[Dict[str, Any]]: List of prediction instances.
+  """
+  instances = [{
+      "encoded_image": {"b64": image_format_converter.image_to_base64(image)},
+  }]
+  return instances
+
+
+def get_label_map(label_map_yaml_filepath: str) -> Dict[str, Any]:
+  """Gets the label map from a YAML file.
+
+  Args:
+      label_map_yaml_filepath: Filepath to the label map YAML file.
+
+  Returns:
+      dict: Label map.
+  """
+  with tf.io.gfile.GFile(label_map_yaml_filepath, "rb") as input_file:
+    label_map = yaml.safe_load(input_file.read())
+  return label_map
+
+
+def get_object_detection_endpoint_predictions(
+    detection_endpoint: ...,
+    input_image: np.ndarray,
+    detection_thresh: float = 0.2,
+) -> np.ndarray:
+  """Gets endpoint predictions.
+
+  Args:
+      detection_endpoint: image object detection endpoint.
+      input_image: Input image.
+      detection_thresh: Detection threshold.
+
+  Returns:
+      Object detection predictions from endpoints.
+  """
+  height, width, _ = input_image.shape
+  predictions = detection_endpoint.predict(
+      get_prediction_instances(Image.fromarray(input_image))
+  ).predictions
+  detection_scores = np.array(predictions[0]["detection_scores"])
+  detection_classes = np.array(predictions[0]["detection_classes"])
+  detection_boxes = np.array(
+      [
+          [b[1] * width, b[0] * height, b[3] * width, b[2] * height]
+          for b in predictions[0]["detection_boxes"]
+      ]
+  )
+  thresh_indices = [
+      x for x, val in enumerate(detection_scores) if val > detection_thresh
+  ]
+  preds_merge_conf = np.column_stack((
+      detection_boxes[thresh_indices],
+      detection_scores[thresh_indices],
+  ))
+  preds_merge_cls = np.column_stack(
+      (preds_merge_conf, detection_classes[thresh_indices])
+  )
+  return preds_merge_cls

community-content/vertex_vision_model_garden/model_oss/util/constants.py

@@ -0,0 +1,85 @@
+"""Vertex vision model garden util constants."""
+
+# TfVision Objectives.
+OBJECTIVE_IMAGE_CLASSIFICATION = 'icn'
+OBJECTIVE_IMAGE_OBJECT_DETECTION = 'iod'
+OBJECTIVE_IMAGE_SEGMENTATION = 'isg'
+OBJECTIVE_VIDEO_CLASSIFICATION = 'vcn'
+OBJECTIVE_VIDEO_ACTION_RECOGNITION = 'var'
+
+# PyTorch Models.
+OBJECTIVE_TIMM = 'timm'
+
+# Input file types.
+INPUT_FILE_TYPE_CSV = 'csv'
+INPUT_FILE_TYPE_JSONL = 'jsonl'
+INPUT_FILE_TYPE_COCO_JSON = 'coco_json'
+
+# Output file types.
+OUTPUT_FILE_TYPE_TFRECORD = 'tfrecord'
+OUTPUT_FILE_TYPE_COCO_JSON = 'coco_json'
+
+# Best evaluation metrics.
+IMAGE_CLASSIFICATION_SINGLE_LABEL_BEST_EVAL_METRIC = 'accuracy'
+IMAGE_CLASSIFICATION_MULTI_LABEL_BEST_EVAL_METRIC = 'meanPR-AUC'
+
+IMAGE_OBJECT_DETECTION_BEST_EVAL_METRIC = 'AP50'
+IMAGE_SEGMENTATION_BEST_EVAL_METRIC = 'mean_iou'
+
+VIDEO_CLASSIFICATION_BEST_EVAL_METRIC = 'accuracy'
+
+# Best checkpoints.
+BEST_CKPT_DIRNAME = 'best_ckpt'
+BEST_CKPT_EVAL_FILENAME = 'info.json'
+BEST_CKPT_STEP_NAME = 'best_ckpt_global_step'
+BEST_CKPT_METRIC_COMP = 'higher'
+
+# Reported hyperparameter tuning metric tag.
+HP_METRIC_TAG = 'model_performance'
+
+# HPT trial prefix.
+TRIAL_PREFIX = 'trial_'
+
+# ML uses from user input.
+ML_USE_TRAINING = 'training'
+ML_USE_VALIDATION = 'validation'
+ML_USE_TEST = 'test'
+
+# COCO json keys
+COCO_JSON_ANNOTATIONS = 'annotations'
+COCO_JSON_ANNOTATION_IMAGE_ID = 'image_id'
+COCO_JSON_ANNOTATION_CATEGORY_ID = 'category_id'
+COCO_JSON_CATEGORIES = 'categories'
+COCO_JSON_CATEGORY_ID = 'id'
+COCO_JSON_CATEGORY_NAME = 'name'
+COCO_JSON_FILE_NAME = 'file_name'
+COCO_JSON_IMAGES = 'images'
+COCO_JSON_IMAGE_ID = 'id'
+COCO_JSON_IMAGE_WIDTH = 'width'
+COCO_JSON_IMAGE_HEIGHT = 'height'
+COCO_JSON_IMAGE_COCO_URL = 'coco_url'
+COCO_ANNOTATION_BBOX = 'bbox'
+
+# GCS prefixes
+GCS_URI_PREFIX = 'gs://'
+GCSFUSE_URI_PREFIX = '/gcs/'
+
+LOCAL_EVALUATION_RESULT_DIR = '/tmp/evaluation_result_dir'
+LOCAL_MODEL_DIR = '/tmp/model_dir'
+LOCAL_BASE_MODEL_DIR = '/tmp/base_model_dir'
+LOCAL_DATA_DIR = '/tmp/data'
+
+# Huggingface files.
+HF_MODEL_WEIGHTS_SUFFIX = '.bin'
+
+# PEFT finetuning constants.
+TEXT_TO_IMAGE_LORA = 'text-to-image-lora'
+SEQUENCE_CLASSIFICATION_LORA = 'sequence-classification-lora'
+CAUSAL_LANGUAGE_MODELING_LORA = 'causal-language-modeling-lora'
+INSTRUCT_LORA = 'instruct-lora'
+
+# Precision modes for loading model weights.
+PRECISION_MODE_4 = '4bit'
+PRECISION_MODE_8 = '8bit'
+PRECISION_MODE_16 = 'float16'
+PRECISION_MODE_32 = 'float32'

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

@@ -0,0 +1,257 @@
+"""Fileutil lib to copy files between gcs and local."""
+
+import glob
+import os
+import pathlib
+import shutil
+from typing import Tuple
+import uuid
+
+from absl import logging
+from google.cloud import storage
+
+from util import constants
+
+
+def generate_tmp_path(extension: str = '') -> str:
+  """Generates a temporary file path with UUID.
+
+  Args:
+    extension: File extension, e.g. '.jpg', '.avi'. If not given, no extension
+      will be appended to the filename.
+
+  Returns:
+    Generated file path.
+  """
+  return os.path.join(constants.LOCAL_DATA_DIR, uuid.uuid1().hex) + extension
+
+
+def force_gcs_fuse_path(gcs_uri: str) -> str:
+  """Converts gs:// uris to their /gcs/ equivalents. No-op for other uris."""
+  if is_gcs_path(gcs_uri):
+    return (
+        constants.GCSFUSE_URI_PREFIX + gcs_uri[len(constants.GCS_URI_PREFIX) :]
+    )
+  else:
+    return gcs_uri
+
+
+def download_gcs_file_to_local_dir(gcs_uri: str, local_dir: str):
+  """Download a gcs file to a local dir.
+
+  Args:
+    gcs_uri: A string of file path on GCS.
+    local_dir: A string of local directory.
+  """
+  if not is_gcs_path(gcs_uri):
+    raise ValueError(
+        f'{gcs_uri} is not a GCS path starting with {constants.GCS_URI_PREFIX}.'
+    )
+  filename = os.path.basename(gcs_uri)
+  download_gcs_file_to_local(gcs_uri, os.path.join(local_dir, filename))
+
+
+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 is_gcs_path(gcs_uri):
+    raise ValueError(
+        f'{gcs_uri} is not a GCS path starting with {constants.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_gcs_dir_to_local(
+    gcs_dir: str, local_dir: str, skip_hf_model_bin: bool = False
+):
+  """Downloads files in a GCS directory to a local directory.
+
+  For example:
+    download_gcs_dir_to_local(gs://bucket/foo, /tmp/bar)
+    gs://bucket/foo/a -> /tmp/bar/a
+    gs://bucket/foo/b/c -> /tmp/bar/b/c
+
+  Arguments:
+    gcs_dir: A string of directory path on GCS.
+    local_dir: A string of local directory path.
+    skip_hf_model_bin: True to skip downloading HF model bin files.
+  """
+  if not is_gcs_path(gcs_dir):
+    raise ValueError(f'{gcs_dir} is not a GCS path starting with gs://.')
+  bucket_name = gcs_dir.split('/')[2]
+  prefix = gcs_dir[len(constants.GCS_URI_PREFIX + bucket_name) :].strip('/')
+  client = storage.Client()
+  blobs = client.list_blobs(bucket_name, prefix=prefix)
+  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(constants.HF_MODEL_WEIGHTS_SUFFIX)
+        and skip_hf_model_bin
+    ):
+      logging.info('Skip downloading model bin %s', file_path)
+      with open(local_file_path, 'w') as f:
+        f.write(f'{constants.GCS_URI_PREFIX}{bucket_name}/{prefix}/{file_path}')
+    else:
+      logging.info('Downloading %s to %s', file_path, local_file_path)
+      blob.download_to_filename(local_file_path)
+
+
+def upload_local_dir_to_gcs(local_dir: str, gcs_dir: str):
+  """Uploads local dir to gcs.
+
+  For example:
+    upload_local_dir_to_gcs(/tmp/bar, gs://bucket/foo)
+    gs://bucket/foo/a -> /tmp/bar/a
+    gs://bucket/foo/b/c -> /tmp/bar/b/c
+
+  Arguments:
+    local_dir: A string of local directory path.
+    gcs_dir: A string of directory path on GCS.
+  """
+  bucket_name = gcs_dir.split('/')[2]
+  blob_dir = '/'.join(gcs_dir.split('/')[3:])
+  client = storage.Client()
+  bucket = client.bucket(bucket_name)
+  for local_file in glob.glob(local_dir + '/**'):
+    if os.path.isfile(local_file):
+      logging.info(
+          'Uploading %s to %s',
+          local_file,
+          os.path.join(constants.GCS_URI_PREFIX, bucket_name, blob_dir),
+      )
+      blob = bucket.blob(os.path.join(blob_dir, os.path.basename(local_file)))
+      blob.upload_from_filename(local_file)
+
+
+def upload_file_to_gcs_path(
+    source_path: str,
+    destination_uri: str,
+):
+  """Uploads local files to GCS uri.
+
+  After upload the destination_uri will contain the same data as the
+  source_path.
+
+  Args:
+      source_path: Required. Path of the local data to copy to GCS.
+      destination_uri: Required. GCS URI where the data should be uploaded.
+
+  Raises:
+      RuntimeError: When source_path does not exist.
+      GoogleCloudError: When the upload process fails.
+  """
+  source_path_obj = pathlib.Path(source_path)
+  if not source_path_obj.exists():
+    raise RuntimeError(f'Source path does not exist: {source_path}')
+
+  storage_client = storage.Client()
+  source_file_path = source_path
+  destination_file_uri = destination_uri
+  logging.info('Uploading "%s" to "%s"', source_file_path, destination_file_uri)
+  destination_blob = storage.Blob.from_string(
+      destination_file_uri, client=storage_client
+  )
+  destination_blob.upload_from_filename(filename=source_file_path)
+
+
+def is_gcs_path(input_path: str) -> bool:
+  """Checks if the input path is a Google Cloud Storage (GCS) path.
+
+  Args:
+      input_path: The input path to be checked.
+
+  Returns:
+      True if the input path is a GCS path, False otherwise.
+  """
+  return input_path.startswith(constants.GCS_URI_PREFIX)
+
+
+def release_text_assets(
+    output_bucket: str, local_text_file_name: str, remote_text_file_name: str
+) -> None:
+  """Releases text assets.
+
+  Args:
+    output_bucket: gcs output bucket.
+    local_text_file_name: Local text file name.
+    remote_text_file_name: Remote text file name.
+
+  Returns:
+    None
+  """
+  remote_file_path = '{}/{}'.format(output_bucket, remote_text_file_name)
+  logging.info('Uploading "%s" to "%s"', local_text_file_name, remote_file_path)
+  upload_file_to_gcs_path(local_text_file_name, remote_file_path)
+  os.remove(local_text_file_name)
+
+
+def upload_video_from_local_to_gcs(
+    output_bucket: str,
+    local_video_file_name: str,
+    remote_video_file_name: str,
+    temp_local_video_file_name: str,
+) -> None:
+  """Uploads video from local to gcs buckent and releases video assets.
+
+  Args:
+    output_bucket: GCS bucket address.
+    local_video_file_name: Local video file name.
+    remote_video_file_name: Remote video file name.
+    temp_local_video_file_name: Temporary local video file name.
+
+  Returns:
+    None
+  """
+  upload_file_to_gcs_path(
+      temp_local_video_file_name,
+      '{}/{}'.format(output_bucket, remote_video_file_name),
+  )
+  shutil.rmtree(local_video_file_name, ignore_errors=True)
+  shutil.rmtree(temp_local_video_file_name, ignore_errors=True)
+
+
+def download_video_from_gcs_to_local(video_file_path: str) -> Tuple[str, str]:
+  """Downloads video from gcs to local folders.
+
+  Args:
+    video_file_path: Path to the video file.
+
+  Returns:
+    Local and remote video file paths.
+  """
+  _, 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])
+  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
+
+
+def get_output_video_file(video_output_file_path: str) -> str:
+  """Gets the output video file name for writing video.
+
+  Args:
+    video_output_file_path: Path to the video output file.
+
+  Returns:
+    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
+  )
+  return out_local_video_file_name

community-content/vertex_vision_model_garden/model_oss/util/hypertune_utils.py

@@ -0,0 +1,22 @@
+"""Utility functions for Vertex Hyperparameter Tuning Jobs."""
+
+import os
+
+from absl import logging
+
+
+_ENVIRONMENT_VARIABLE_FOR_TRIAL_ID = 'CLOUD_ML_TRIAL_ID'
+
+
+def get_trial_id_from_environment() -> str:
+  """Gets the trial id from environment variable.
+
+  Returns:
+    The trial id from environement or '0' if not found.
+  """
+  if _ENVIRONMENT_VARIABLE_FOR_TRIAL_ID not in os.environ:
+    logging.warning(
+        'Environment variable %s not found, return 0 as default trial id.',
+        _ENVIRONMENT_VARIABLE_FOR_TRIAL_ID,
+    )
+  return os.environ.get(_ENVIRONMENT_VARIABLE_FOR_TRIAL_ID, '0')

community-content/vertex_vision_model_garden/model_oss/util/image_format_converter.py

@@ -0,0 +1,20 @@
+"""Image format converter util lib."""
+
+import base64
+import io
+
+from PIL import Image
+
+
+def image_to_base64(image: Image.Image) -> str:
+  """Convert a PIL image to a base64 string."""
+  buffer = io.BytesIO()
+  image.save(buffer, format="JPEG")
+  image_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
+  return image_str
+
+
+def base64_to_image(image_str: str) -> Image.Image:
+  """Convert a base64 string to a PIL image."""
+  image = Image.open(io.BytesIO(base64.b64decode(image_str)))
+  return image

community-content/vertex_vision_model_garden/model_oss/vot/dockerfile/serve.Dockerfile

@@ -0,0 +1,87 @@
+# Dockerfile for basic serving dockers for vot.
+#
+# To build:
+# docker build -f model_oss/vot/dockerfile/serving.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}
+
+# Switch to this base image for gpu serve.
+FROM pytorch/torchserve:0.7.0-gpu
+
+USER root
+
+ENV infer_port=7080
+ENV mng_port=7081
+ENV model_name="vot_serving"
+ENV PATH="/home/model-server/:${PATH}"
+
+ENV PYTHONPATH "${PYTHONPATH}:/automl_vision/vot"
+
+# Install libraries.
+ENV PIP_ROOT_USER_ACTION=ignore
+RUN python3 -m pip install --upgrade pip
+RUN pip install accelerate==0.17.0
+RUN pip install datasets==2.9.0
+RUN pip install bytetracker==0.3.2
+RUN pip install imageio[ffmpeg]==2.31.1
+RUN pip install google-cloud-aiplatform==1.25.0
+RUN pip install google-cloud-storage==2.9.0
+RUN pip install fastapi==0.96.0
+RUN pip install lap==0.4.0
+RUN pip install numpy==1.24.3
+RUN pip install opencv-python==4.7.0.72
+RUN pip install Pillow==9.5.0
+RUN pip install protobuf==3.19.6
+RUN pip install pandas==2.0.2
+RUN pip install pycocotools==2.0.6
+RUN pip install scipy==1.10.1
+RUN pip install tensorflow==2.11.1
+RUN pip install torch==2.0.1
+RUN pip install torchvision==0.15.2
+RUN pip install triton==2.0.0.dev20221120
+RUN pip install uvicorn==0.22.0
+
+# Install tools.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        curl \
+        wget \
+        vim
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Copy model artifacts.
+COPY model_oss/vot/handler.py /home/model-server/handler.py
+COPY model_oss/vot/visualization_utils.py /home/model-server/vot/
+COPY model_oss/util/ /home/model-server/util/
+ENV PYTHONPATH /home/model-server/
+
+# Create torchserve configuration file.
+RUN echo \
+    "default_response_timeout=3600\n" \
+    "service_envelope=json\n" \
+    "inference_address=http://0.0.0.0:${infer_port}\n"  \
+    "management_address=http://0.0.0.0:${mng_port}" >> /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.
+CMD ["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_vision_model_garden/model_oss/vot/handler.py

@@ -0,0 +1,196 @@
+"""Custom handler for video object tracking models."""
+
+import logging
+import os
+import tempfile
+from typing import Any, List, Optional, Tuple
+
+from bytetracker import BYTETracker
+import cv2
+from google.cloud import aiplatform
+import imageio.v2 as iio
+from PIL import Image
+import tensorflow as tf
+import torch
+from ts.torch_handler.base_handler import BaseHandler
+
+from util import commons
+from util import fileutils
+import visualization_utils
+
+_VIDEO_URI = "video_uri"
+_DATA = "data"
+_TRACK_THRESHOLD = 0.45
+_TRACK_BUFFER = 25
+_MATCH_THRESHOLD = 0.8
+
+
+class VideoObjectTrackingHandler(BaseHandler):
+  """Custom handler for video object tracking models."""
+
+  def initialize(self, context: Any) -> None:
+    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
+
+    detection_endpoint_id = os.environ.get("DETECTION_ENDPOINT", None)
+    if detection_endpoint_id:
+      self.detection_endpoint = aiplatform.Endpoint(detection_endpoint_id)
+      endpoint_label_map = os.environ.get("LABEL_MAP", None)
+      if endpoint_label_map:
+        endpoint_label_map_file = endpoint_label_map
+        self.label_map = commons.get_label_map(endpoint_label_map_file)
+      else:
+        raise ValueError(
+            "LABEL MAP must be provided with DETECTION ENDPOINT:"
+            f" {self.detection_endpoint}"
+        )
+
+    self.track_thresh = os.environ.get("TRACK_THRESHOLD", _TRACK_THRESHOLD)
+    self.track_buffer = os.environ.get("TRACK_BUFFER", _TRACK_BUFFER)
+    self.match_thresh = os.environ.get("MATCH_THRESHOLD", _MATCH_THRESHOLD)
+    self.save_video_results = bool(int(os.environ.get("SAVE_VIDEO_RESULTS", 0)))
+    self.output_bucket = os.environ.get("OUTPUT_BUCKET", None)
+    if not self.output_bucket:
+      raise ValueError("Empty Output Bucket.")
+    self.initialized = True
+    logging.info("Handler initialization done.")
+
+  def preprocess(
+      self, data: Any
+  ) -> Tuple[Optional[List[str]], Optional[List[Image.Image]]]:
+    """Preprocesses the input data.
+
+    Args:
+        data (Any): Input data.
+
+    Returns:
+        List of videos uris.
+    """
+    video_uris = None
+    if _VIDEO_URI in data[0]:
+      video_uris = [item[_VIDEO_URI] for item in data]
+    # TorchServe's default handlers expect each instance
+    # to be wrapped in a data field for batch prediction.
+    if _DATA in data[0]:
+      video_uris = [item[_DATA][_VIDEO_URI] for item in data]
+    return video_uris
+
+  def inference(self, data: Any, *args, **kwargs) -> List[Any]:
+    """Runs object detection and tracking inference on a video frame by frame.
+
+    If using yolo detection, the function uses the ultralytics yolo models for
+    IOD, otherwise it uses the provided IOD endpoint and associated the selected
+    tracking method to the detections.
+
+    Args:
+        data: List of video files.
+        *args: Additional arguments.
+        **kwargs: Additional keyword arguments.
+
+    Returns:
+        List of video frame annotations and/or output decorated video uris.
+    """
+    gcs_video_files = data
+    video_preds = []
+    for gcs_video_file in gcs_video_files:
+      results_info = {}
+      temp_text_file = tempfile.NamedTemporaryFile(delete=False, mode="w+t")
+      local_video_file_name, remote_video_file_name = (
+          fileutils.download_video_from_gcs_to_local(gcs_video_file)
+      )
+      remote_text_file_name = remote_video_file_name.replace(
+          "overlay.mp4", "annotations.txt"
+      )
+
+      cap = cv2.VideoCapture(local_video_file_name)
+      fps = cap.get(cv2.CAP_PROP_FPS)
+      temp_local_video_file_name = fileutils.get_output_video_file(
+          local_video_file_name
+      )
+      if self.save_video_results:
+        self.video_writer = iio.get_writer(
+            temp_local_video_file_name,
+            format="FFMPEG",
+            mode="I",
+            fps=float(fps),
+            codec="h264",
+        )
+      self.tracker = BYTETracker(
+          track_thresh=self.track_thresh,
+          track_buffer=self.track_buffer,
+          match_thresh=self.match_thresh,
+          frame_rate=fps,
+      )
+      frame_idx = 1
+      while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+          break
+
+        dets_np = commons.get_object_detection_endpoint_predictions(
+            self.detection_endpoint, frame
+        )
+        dets_tf = tf.convert_to_tensor(dets_np)
+        online_targets = self.tracker.update(dets_tf, None)
+        if online_targets.size > 0:
+          frame = visualization_utils.overlay_tracking_results(
+              frame_idx,
+              frame,
+              online_targets,
+              label_map=self.label_map,
+              temp_text_file_path=temp_text_file.name,
+          )
+
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        if self.save_video_results:
+          self.video_writer.append_data(frame)
+        logging.info(
+            "Finished processing frame %s for video %s.",
+            frame_idx,
+            gcs_video_file,
+        )
+        frame_idx += 1
+
+      self.video_writer.close()
+      cap.release()
+      if self.save_video_results:
+        fileutils.upload_video_from_local_to_gcs(
+            self.output_bucket,
+            local_video_file_name,
+            remote_video_file_name,
+            temp_local_video_file_name,
+        )
+        results_info["output_video"] = "{}/{}".format(
+            self.output_bucket, remote_video_file_name
+        )
+
+      fileutils.release_text_assets(
+          self.output_bucket,
+          temp_text_file.name,
+          remote_text_file_name,
+      )
+      results_info["annotations"] = "{}/{}".format(
+          self.output_bucket, remote_text_file_name
+      )
+      video_preds.append(results_info)
+
+    return video_preds
+
+  def handle(self, data: Any, context: Any) -> List[Any]:
+    model_input = self.preprocess(data)
+    model_out = self.inference(model_input)
+    output = self.postprocess(model_out)
+    return output
+
+  def postprocess(self, inference_result: List[Any]) -> List[Any]:
+    return inference_result

community-content/vertex_vision_model_garden/model_oss/vot/visualization_utils.py

@@ -0,0 +1,168 @@
+"""Image and bounding box visualization util lib."""
+
+from typing import Dict, List, Optional
+import cv2
+import numpy as np
+from PIL import ImageColor
+
+
+def draw_bounding_box_on_image(
+    image: np.ndarray,
+    ymin: float,
+    xmin: float,
+    ymax: float,
+    xmax: float,
+    color: str,
+    thickness: int = 4,
+    display_str_list: Optional[List[str]] = None,
+) -> np.ndarray:
+  """Draws a bounding box on an image.
+
+  Args:
+    image: The image to draw the bounding box on.
+    ymin: The minimum y-coordinate of the bounding box.
+    xmin: The minimum x-coordinate of the bounding box.
+    ymax: The maximum y-coordinate of the bounding box.
+    xmax: The maximum x-coordinate of the bounding box.
+    color: The color of the bounding box.
+    thickness: The thickness of the bounding box lines. Defaults to 4.
+    display_str_list: List of strings to display in new line inside the bounding
+      box.
+
+  Returns:
+    An image with a bounding box.
+  """
+  color = ImageColor.getrgb(color)
+  cv2.rectangle(
+      image, (int(xmin), int(ymin)), (int(xmax), int(ymax)), color, thickness
+  )
+  # Display the strings below the bounding box
+  for i, display_str in enumerate(display_str_list):
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    scale = 0.4
+    thickness = 1
+    text_width, text_height = cv2.getTextSize(
+        display_str, font, scale, thickness
+    )[0]
+    text_bottom = int(ymin - i * text_height)
+    text_left = int(xmin)
+    cv2.rectangle(
+        image,
+        (text_left, text_bottom - text_height),
+        (text_left + text_width, text_bottom),
+        color,
+        -1,
+    )
+    cv2.putText(
+        image,
+        display_str,
+        (text_left, text_bottom),
+        font,
+        scale,
+        (0, 0, 0),
+        thickness,
+    )
+
+  return image
+
+
+def draw_boxes(
+    image: np.ndarray,
+    boxes: List[List[float]],
+    track_ids: List[int],
+    class_names: List[str],
+    scores: List[float],
+    max_boxes: int = 40,
+    min_score: float = 0.05,
+) -> np.ndarray:
+  """Overlays labeled boxes on an image with formatted scores and label names.
+
+  Args:
+      image: The image to overlay the boxes on.
+      boxes: List of bounding box coordinates [xmin, ymin, xmax, ymax].
+      track_ids: List of track IDs corresponding to each box.
+      class_names: List of class names corresponding to each box.
+      scores: List of scores corresponding to each box.
+      max_boxes: Maximum number of boxes to draw. Defaults to 40.
+      min_score: Minimum score threshold for displaying a box. Defaults to 0.05.
+
+  Returns:
+      PIL.Image.Image: The image with the labeled boxes overlay.
+  """
+  colors = list(ImageColor.colormap.values())
+  for i in range(min(len(boxes), max_boxes)):
+    if scores[i] >= min_score:
+      xmin, ymin, xmax, ymax = boxes[i]
+      display_str = "{}-{}: {}%".format(
+          track_ids[i], class_names[i], int(100 * scores[i])
+      )
+      color = colors[hash(class_names[i]) % len(colors)]
+      image = draw_bounding_box_on_image(
+          image,
+          ymin,
+          xmin,
+          ymax,
+          xmax,
+          color,
+          display_str_list=[display_str],
+      )
+  return image
+
+
+def overlay_tracking_results(
+    frame_idx: int,
+    image_np: np.ndarray,
+    tracker_outputs: np.ndarray,
+    model_names: Optional[Dict[int, str]] = None,
+    label_map: Optional[Dict[str, Dict[int, str]]] = None,
+    temp_text_file_path: Optional[str] = None,
+) -> np.ndarray:
+  """Overlays the results on the image.
+
+  Args:
+      frame_idx: frame index.
+      image_np: Input image.
+      tracker_outputs: Tracker outputs.
+      model_names: label map for yolo models.
+      label_map: label map for IOD detector model.
+      temp_text_file_path: tempfile to save annotations.
+
+  Returns:
+      Decorated output frame.
+  """
+  dboxes = tracker_outputs[:, :4]
+  dtracks = tracker_outputs[:, 4]
+  dclasses = tracker_outputs[:, 5]
+  dscores = tracker_outputs[:, 6]
+  dclasses_as_text = []
+  for detection_class in dclasses:
+    if model_names:
+      dclasses_as_text.append(model_names[int(detection_class)])
+    elif label_map:
+      dclasses_as_text.append(label_map["label_map"][int(detection_class)])
+    else:
+      dclasses_as_text.append("")
+
+  plotted_img = np.array(
+      draw_boxes(
+          image=image_np,
+          boxes=dboxes,
+          track_ids=dtracks,
+          class_names=dclasses_as_text,
+          scores=dscores,
+      )
+  )
+  track_anno_list = []
+  for i, box in enumerate(dboxes):
+    result_list = [dtracks[i], dscores[i], dclasses[i]]
+    xyxy_anno = [np.round(item.item(), 2) for item in box]
+    tracks_anno = (
+        [frame_idx]
+        + [np.round(item.item(), 2) for item in result_list]
+        + xyxy_anno
+    )
+    track_anno_list.append(tracks_anno)
+    with open(temp_text_file_path, "a") as file:
+      file.write(", ".join([str(item) for item in tracks_anno]) + "\n")
+
+  return plotted_img

notebooks/community/CODEOWNERS

@@ -30,10 +30,62 @@
 /notebooks/community/neo4j/graph_paysim.ipynb @benofben @laeg
 /notebooks/community/ml_ops/stage1/get_started_with_visionapi_and_vertex_datasets.ipynb @mansari
 /notebooks/community/pipelines/google_cloud_pipeline_components_bqml_pipeline_demand_forecasting.ipynb @inardini
+/notebooks/community/cohere/cohere_embedding_with_matching_engine.ipynb @stewart-co
 /notebooks/community/ml_ops/stage2/get_started_vertex_hpt_r_kernel.ipynb @fhirschmann
 /notebooks/community/ml_ops/stage2/get_started_vertex_training_r_using_r_kernel.ipynb @fhirschmann
+/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/vizier/conversions_vertex_vizier_and_open_source_vizier.ipynb @halio-g
 /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_huggingface_local_inference.ipynb @dstnluong-google
+/notebooks/community/model_garden/model_garden_mediapipe_image_classification.ipynb @schmidt-sebastian
+/notebooks/community/model_garden/model_garden_mediapipe_image_generation.ipynb @schmidt-sebastian
+/notebooks/community/model_garden/model_garden_mediapipe_gesture_recognition.ipynb @schmidt-sebastian
+/notebooks/community/model_garden/model_garden_mediapipe_object_detection.ipynb @schmidt-sebastian
+/notebooks/community/model_garden/model_garden_mediapipe_text_classification.ipynb @schmidt-sebastian
+/notebooks/community/model_garden/model_garden_proprietary_image_classification.ipynb @weigary
+/notebooks/community/model_garden/model_garden_proprietary_image_object_detection.ipynb @weigary
+/notebooks/community/model_garden/model_garden_tfvision_image_classification.ipynb @genquan9
+/notebooks/community/model_garden/model_garden_tfvision_image_object_detection.ipynb @genquan9
+/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_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_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
+/notebooks/community/model_garden/model_garden_pytorch_blip_vqa.ipynb @xiangxu-google
+/notebooks/community/model_garden/model_garden_pytorch_vilt_vqa.ipynb @xiangxu-google
+/notebooks/community/model_garden/model_garden_pytorch_vit_gpt2_image_captioning.ipynb @xiangxu-google
+/notebooks/community/model_garden/model_garden_pytorch_clip.ipynb @xiangxu-google
+/notebooks/community/model_garden/model_garden_pytorch_owlvit.ipynb @xiangxu-google
+/notebooks/community/model_garden/model_garden_pytorch_layoutml_document_qa.ipynb @xiangxu-google
+/notebooks/community/model_garden/model_garden_pytorch_blip2.ipynb @xiangxu-google
+/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_starcoder.ipynb @xcchen1
+/notebooks/community/model_garden/model_garden_jax_vision_transformer.ipynb @lavraicse
+/notebooks/community/model_garden/model_garden_jax_fvlm.ipynb @lavraicse
+/notebooks/community/model_garden/model_garden_pytorch_text_to_video_zero_shot.ipynb @bingatgoogle
+/notebooks/community/model_garden/model_garden_pytorch_text_to_video.ipynb @KCFindstr
+/notebooks/community/generative_ai/text_embedding_api_semantic_search_with_scann.ipynb @henrytansetiawan
+/notebooks/community/bigquery_ml_inference/bq_ml_with_vision_translation_nlp.ipynb @deaconsmith
+/notebooks/community/model_garden/model_garden_keras_stable_diffusion.ipynb @genquan9
+/notebooks/community/model_garden/model_garden_keras_yolov8.ipynb @@dstnluong-google
+/notebooks/community/model_garden/model_garden_pytorch_sam.ipynb @huguensjean
+/notebooks/community/model_garden/model_garden_pytorch_pic2word.ipynb @jismailyan
+/notebooks/community/model_garden/model_garden_pytorch_peft.ipynb @genquan9
+/notebooks/community/model_garden/model_garden_pytorch_openllama_peft.ipynb @genquan9
+/notebooks/community/model_garden/model_garden_pytorch_falcon_instruct_peft.ipynb @genquan9
+/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/bigquery_ml/README.md

@@ -0,0 +1,19 @@
+
+[Unstructured data analytics with BigQuery ML and Vertex AI pre-trained models](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/bigquery_ml/bq_ml_with_vision_translation_nlp.ipynb)
+
+```
+Learn how to analyze unstructured data within BigQuery using BigQuery's inference engine. You will use BigQuery ML to connect to three pretrained Vertex AI APIs - Vision API, Translation API and Natural Language Processing API.
+
+The steps performed include:
+
+- Define pre-trained models for Vision AI, Translation AI and NLP AI in BigQuery ML
+- Call the Vision API (`ML.ANNOTATE_IMAGE`) to detect text in images stored in Cloud Storage
+  You will need to create an object table in BigQuery to do this
+- Call the Translation API (`ML.TRANSLATE`) to detect the language of text, and translate non-English movie titles to English
+- Call the Natural Language API (`ML.UNDERSTAND_TEXT`) to run sentiment analysis over movie reviews stored in BigQuery
+
+```
+
+   Check out the [blog for this notebook](https://cloud.google.com/blog/products/data-analytics/how-simplify-unstructured-data-analytics-using-bigquery-ml-and-vertex-ai).
+   Learn more about [BigQuery ML inference engine](https://cloud.google.com/bigquery/docs/reference/standard-sql/inference-overview).
+

notebooks/community/cohere/README.md

@@ -0,0 +1,3 @@
+# README
+
+These are notebooks [Cohere](https://cohere.ai/) built in collaboration with Google.  They demonstrate how to use Cohere's modeling API along with Vertex AI.

notebooks/community/generative_ai/text_embedding_api_semantic_search_with_scann.ipynb

@@ -0,0 +1,373 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "KSP1duKDeaDR",
+      "metadata": {
+        "id": "KSP1duKDeaDR"
+      },
+      "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",
+      "id": "67b2c5cc-8fc6-4082-9052-69fa0377d770",
+      "metadata": {
+        "id": "67b2c5cc-8fc6-4082-9052-69fa0377d770"
+      },
+      "source": [
+        "# Semantic Search using Embeddings\n",
+        "\n",
+        "Semantic search is a type of search that uses the meaning of words and phrases to find relevant results.\n",
+        "\n",
+        "In this tutorial, we will demonstrate how to do semantic search with embeddings generated from the news text and using [Google ScaNN: Efficient Vector Similarity Search](https://ai.googleblog.com/2020/07/announcing-scann-efficient-vector.html) to retrieve the most relevant news semantically.\n",
+        "\n",
+        "## Pre-requisites:\n",
+        "- Vertex LLM SDK\n",
+        "- ScaNN [github](https://github.com/google-research/google-research/tree/master/scann)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "FyyMdUeAJIVv",
+      "metadata": {
+        "id": "FyyMdUeAJIVv"
+      },
+      "source": [
+        "## Install Vertex LLM SDK\n",
+        "\n",
+        "DISCLAIMER: Text Embedding API is now in Experimental Preview. This release focuses on validating model prototypes and these models are not guaranteed to be released. Use of Text Embedding API is governed by the Google Cloud Terms of Service, the Pre-GA Offerings Terms of the GCP Service Specific Terms. The Acceptance Use Policy, and the Generative AI Prohibited Use Policy. Vertex Text Embedding API’s features may be unstable, change in backward-incompatible ways, and are not guaranteed to be released. There are no SLAs provided and no technical support obligations. GCP’s Cloud Data Processing Addendum does not apply to Pre-GA Offerings and customers should not use Text Embedding API to process personal data or other data subject to legal or regulatory compliance requirements. See description of launch stage for details.\n",
+        "\n",
+        "The information in this documentation is provided to the customer on an “as is” and “with all faults” basis without any warranty of any kind, either express or implied. Google does not warrant or guarantee the correctness, accuracy or reliability of the information in here. In no event will Google or its affiliates or licensors be liable for any damage or harm to customers from customer’s use of these materials."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "h6KaE3XRJdxc",
+      "metadata": {
+        "id": "h6KaE3XRJdxc"
+      },
+      "outputs": [],
+      "source": [
+        "from google.colab import auth as google_auth\n",
+        "\n",
+        "google_auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "snBUuUamoJPz",
+      "metadata": {
+        "id": "snBUuUamoJPz"
+      },
+      "outputs": [],
+      "source": [
+        "!pip3 install google-cloud-aiplatform>=1.25 \"shapely<2.0.0\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "zgVQcE0ewO8W",
+      "metadata": {
+        "id": "zgVQcE0ewO8W"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"cloud-nl-llm-embedding\"  # @param {type:\"string\"}\n",
+        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
+        "\n",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "4xFzXmPbY7FC",
+      "metadata": {
+        "id": "4xFzXmPbY7FC"
+      },
+      "source": [
+        "**Attention**: you would need to restart runtime so that the right package is installed."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "xnfG88OPZI18",
+      "metadata": {
+        "id": "xnfG88OPZI18"
+      },
+      "source": [
+        "## Import TextEmbeddingModel"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "1e6e1b98-a632-44a2-afb8-fc212018ef4f",
+      "metadata": {
+        "id": "1e6e1b98-a632-44a2-afb8-fc212018ef4f"
+      },
+      "outputs": [],
+      "source": [
+        "from vertexai.preview.language_models import TextEmbeddingModel\n",
+        "\n",
+        "model = TextEmbeddingModel.from_pretrained(\"textembedding-gecko@001\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "1suA-1HuaGj6",
+      "metadata": {
+        "id": "1suA-1HuaGj6"
+      },
+      "source": [
+        "## Install ScaNN Package"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "770255d3-54dd-48c8-bbdd-fbc0be41f085",
+      "metadata": {
+        "id": "770255d3-54dd-48c8-bbdd-fbc0be41f085"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install scann"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "dae340cb-0583-4e7e-a562-6817ee4d7f6d",
+      "metadata": {
+        "id": "dae340cb-0583-4e7e-a562-6817ee4d7f6d"
+      },
+      "source": [
+        "## Imports packages"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "412d00f1-08db-4880-8ced-52a9583757b8",
+      "metadata": {
+        "id": "412d00f1-08db-4880-8ced-52a9583757b8"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "import time\n",
+        "\n",
+        "import numpy as np\n",
+        "import pandas as pd\n",
+        "import scann"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "f50f22f3-ec85-463e-b6fe-5c8e6b80b07b",
+      "metadata": {
+        "id": "f50f22f3-ec85-463e-b6fe-5c8e6b80b07b"
+      },
+      "source": [
+        "## Create Embedding Dataset.\n",
+        "\n",
+        "The dataset is solely to demonstrate the use of the Text Embedding API with a vector database. It is not intended to be used for any other purpose, such as evaluating models. The dataset is small and does not represent a comprehensive sample of all possible text."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "2OUg-Qf8iFta",
+      "metadata": {
+        "id": "2OUg-Qf8iFta"
+      },
+      "outputs": [],
+      "source": [
+        "!gsutil cp gs://cloud-samples-data/vertex-ai/dataset-management/datasets/bert_finetuning/wide_and_deep_trainer_container_tests_input.jsonl ."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "BNPapKXviHlE",
+      "metadata": {
+        "id": "BNPapKXviHlE"
+      },
+      "outputs": [],
+      "source": [
+        "records = []\n",
+        "with open(\"wide_and_deep_trainer_container_tests_input.jsonl\") as f:\n",
+        "    for line in f:\n",
+        "        record = json.loads(line)\n",
+        "        records.append(record)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "Z1Y9Bx2miJba",
+      "metadata": {
+        "id": "Z1Y9Bx2miJba"
+      },
+      "outputs": [],
+      "source": [
+        "# Peek at the data.\n",
+        "df = pd.DataFrame(records)\n",
+        "df.head(50)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "1v7zUwoWiPl-",
+      "metadata": {
+        "id": "1v7zUwoWiPl-"
+      },
+      "outputs": [],
+      "source": [
+        "def get_embedding(text):\n",
+        "    get_embedding.counter += 1\n",
+        "    try:\n",
+        "        if get_embedding.counter % 100 == 0:\n",
+        "            time.sleep(3)\n",
+        "        return model.get_embeddings([text])[0].values\n",
+        "    except:\n",
+        "        return []\n",
+        "\n",
+        "\n",
+        "get_embedding.counter = 0\n",
+        "\n",
+        "# This may take several minutes to complete.\n",
+        "df[\"embedding\"] = df[\"textContent\"].apply(lambda x: get_embedding(x))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "ba4f49b6-65e1-49ea-988b-c4e195219deb",
+      "metadata": {
+        "id": "ba4f49b6-65e1-49ea-988b-c4e195219deb"
+      },
+      "source": [
+        "## Create an Index"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "245bc8cd-038b-484a-acb2-3a705d4cc4cf",
+      "metadata": {
+        "id": "245bc8cd-038b-484a-acb2-3a705d4cc4cf"
+      },
+      "outputs": [],
+      "source": [
+        "record_count = len(records)\n",
+        "dataset = np.empty((record_count, 768))\n",
+        "for i in range(record_count):\n",
+        "    dataset[i] = df.embedding[i]\n",
+        "\n",
+        "normalized_dataset = dataset / np.linalg.norm(dataset, axis=1)[:, np.newaxis]\n",
+        "# configure ScaNN as a tree - asymmetric hash hybrid with reordering\n",
+        "# anisotropic quantization as described in the paper; see README\n",
+        "\n",
+        "# use scann.scann_ops.build() to instead create a TensorFlow-compatible searcher\n",
+        "searcher = (\n",
+        "    scann.scann_ops_pybind.builder(normalized_dataset, 10, \"dot_product\")\n",
+        "    .tree(\n",
+        "        num_leaves=record_count,\n",
+        "        num_leaves_to_search=record_count,\n",
+        "        training_sample_size=record_count,\n",
+        "    )\n",
+        "    .score_ah(2, anisotropic_quantization_threshold=0.2)\n",
+        "    .reorder(100)\n",
+        "    .build()\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "9f1689eb-c27c-4566-afc7-7fbc55552aad",
+      "metadata": {
+        "id": "9f1689eb-c27c-4566-afc7-7fbc55552aad"
+      },
+      "source": [
+        "## Queries the Index"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "494079df-d8f6-4a6a-b26b-0b0477791adc",
+      "metadata": {
+        "id": "494079df-d8f6-4a6a-b26b-0b0477791adc"
+      },
+      "outputs": [],
+      "source": [
+        "def search(query):\n",
+        "    start = time.time()\n",
+        "    query = model.get_embeddings([query])[0].values\n",
+        "    neighbors, distances = searcher.search(query, final_num_neighbors=3)\n",
+        "    end = time.time()\n",
+        "\n",
+        "    for id, dist in zip(neighbors, distances):\n",
+        "        print(f\"[docid:{id}] [{dist}] -- {df.textContent[int(id)][:125]}...\")\n",
+        "    print(\"Latency (ms):\", 1000 * (end - start))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "50497681-6112-4147-a13e-afedb72b54f5",
+      "metadata": {
+        "id": "50497681-6112-4147-a13e-afedb72b54f5"
+      },
+      "outputs": [],
+      "source": [
+        "search(\"tell me about shark or animal\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "c814f835-3e2a-4366-a334-ccc636f00b83",
+      "metadata": {
+        "id": "c814f835-3e2a-4366-a334-ccc636f00b83"
+      },
+      "outputs": [],
+      "source": [
+        "search(\"tell me about an important moment or event in your life\")"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "text_embedding_api_semantic_search_with_scann.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/matching_engine/matching_engine_for_indexing.ipynb

@@ -8,7 +8,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2021 Google LLC\n",
+        "# Copyright 2023 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",
@@ -24,6 +24,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "JAPoU8Sm5E6e"
@@ -32,20 +33,28 @@
         "<table align=\"left\">\n",
         "\n",
         "  <td>\n",
-        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/matching_engine/matching_engine_for_indexing.ipynb\">\n",
-        "      Run in Google Cloud Notebooks\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/matching_engine/matching_engine_for_indexing.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\">\n",
+        "      Run in Colab\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
-        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/matching_engine/matching_engine_for_indexing.ipynb\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/matching_engine/matching_engine_for_indexing.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/community/matching_engine/matching_engine_for_indexing.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>"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "tvgnzT1CKxrO"
@@ -53,25 +62,49 @@
       "source": [
         "## Overview\n",
         "\n",
-        "This example demonstrates how to use the GCP ANN Service. It is a high scale, low latency solution, to find similar vectors (or more specifically \"embeddings\") for a large corpus. Moreover, it is a fully managed offering, further reducing operational overhead. It is built upon [Approximate Nearest Neighbor (ANN) technology](https://ai.googleblog.com/2020/07/announcing-scann-efficient-vector.html) developed by Google Research.\n",
-        "\n",
-        "### Dataset\n",
-        "\n",
-        "The dataset used for this tutorial is the [GloVe dataset](https://nlp.stanford.edu/projects/glove/).\n",
-        "\n",
+        "This example demonstrates how to use Vertex AI Matching Engine. It is a high scale, low latency solution, to find similar vectors (or more specifically \"embeddings\") for a large corpus. Moreover, it is a fully managed offering, further reducing operational overhead. It is built upon [Approximate Nearest Neighbor (ANN) technology](https://ai.googleblog.com/2020/07/announcing-scann-efficient-vector.html) developed by Google Research."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "56e5f9699c6c"
+      },
+      "source": [
         "### Objective\n",
         "\n",
         "In this notebook, you will learn how to create Approximate Nearest Neighbor (ANN) Index, query against indexes, and validate the performance of the index. \n",
         "\n",
         "The steps performed include:\n",
         "\n",
-        "* Create ANN Index and Brute Force Index\n",
+        "* Create a Vertex AI Matching Engine Index and Brute Force Index\n",
         "* Create an IndexEndpoint with VPC Network\n",
-        "* Deploy ANN Index and Brute Force Index\n",
-        "* Perform online query\n",
-        "* Compute recall\n",
-        "\n",
+        "* Deploy a Vertex AI Matching Engine Index and Brute Force Index\n",
+        "* Perform online queries\n",
+        "* Submit batch queries\n",
+        "* Compute recall metric"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0aaef374550b"
+      },
+      "source": [
+        "### Dataset\n",
         "\n",
+        "The dataset used for this tutorial is the [GloVe dataset](https://nlp.stanford.edu/projects/glove/)."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5e2eba58ad71"
+      },
+      "source": [
         "### Costs \n",
         "\n",
         "This tutorial uses billable components of Google Cloud:\n",
@@ -87,6 +120,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "S5zc4kbEiYCm"
@@ -94,79 +128,47 @@
       "source": [
         "## Before you begin\n",
         "\n",
-        "* **Prepare a VPC network**.  To reduce any network overhead that might lead to unnecessary increase in overhead latency, it is best to call the ANN endpoints from your VPC via a direct [VPC Peering](https://cloud.google.com/vertex-ai/docs/general/vpc-peering) connection. The following section describes how to setup a VPC Peering connection if you don't have one. This is a one-time initial setup task. You can also reuse existing VPC network and skip this section.\n",
-        "* **WARNING:** The match service gRPC API (to create online queries against your deployed index) has to be executed in a Google Cloud Notebook instance that is created with the following requirements:\n",
-        "  * **In the same region as where your ANN service is deployed** (for example, if you set `REGION = \"us-central1\"` as same as the tutorial, the notebook instance has to be in `us-central1`).\n",
-        "  * **Make sure you select the VPC network you created for ANN service** (instead of using the \"default\" one). That is, you will have to create the VPC network below and then create a new notebook instance that uses that VPC.  \n",
-        "  * If you run it in the colab or a Google Cloud Notebook instance in a different VPC network or region, the gRPC API will fail to peer the network (InactiveRPCError)."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "lW2LneA5mmmP"
-      },
-      "outputs": [],
-      "source": [
-        "PROJECT_ID = \"<your_project_id>\"  # @param {type:\"string\"}\n",
-        "NETWORK_NAME = \"ucaip-haystack-vpc-network\"  # @param {type:\"string\"}\n",
-        "PEERING_RANGE_NAME = \"ucaip-haystack-range\"\n",
+        "### Set up your Google Cloud project\n",
         "\n",
-        "# Create a VPC network\n",
-        "! gcloud compute networks create {NETWORK_NAME} --bgp-routing-mode=regional --subnet-mode=auto --project={PROJECT_ID}\n",
+        "**The following steps are required, regardless of your notebook environment.**\n",
         "\n",
-        "# Add necessary firewall rules\n",
-        "! gcloud compute firewall-rules create {NETWORK_NAME}-allow-icmp --network {NETWORK_NAME} --priority 65534 --project {PROJECT_ID} --allow icmp\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",
-        "! gcloud compute firewall-rules create {NETWORK_NAME}-allow-internal --network {NETWORK_NAME} --priority 65534 --project {PROJECT_ID} --allow all --source-ranges 10.128.0.0/9\n",
+        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
         "\n",
-        "! gcloud compute firewall-rules create {NETWORK_NAME}-allow-rdp --network {NETWORK_NAME} --priority 65534 --project {PROJECT_ID} --allow tcp:3389\n",
+        "3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
         "\n",
-        "! gcloud compute firewall-rules create {NETWORK_NAME}-allow-ssh --network {NETWORK_NAME} --priority 65534 --project {PROJECT_ID} --allow tcp:22\n",
-        "\n",
-        "# Reserve IP range\n",
-        "! gcloud compute addresses create {PEERING_RANGE_NAME} --global --prefix-length=16 --network={NETWORK_NAME} --purpose=VPC_PEERING --project={PROJECT_ID} --description=\"peering range for uCAIP Haystack.\"\n",
-        "\n",
-        "# Set up peering with service networking\n",
-        "! gcloud services vpc-peerings connect --service=servicenetworking.googleapis.com --network={NETWORK_NAME} --ranges={PEERING_RANGE_NAME} --project={PROJECT_ID}"
+        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
-        "id": "d3uj8x73nDX_"
-      },
-      "source": [
-        "* Authentication: `$ gcloud auth login` rerun this in Google Cloud Notebook terminal when you are logged out and need the credential again."
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "i7EUnXsZhAGF"
+        "id": "4700b0e39c5d"
       },
       "source": [
         "### Installation\n",
         "\n",
-        "Download and install the latest (preview) version of the Vertex SDK for Python."
+        "Download and install the latest version of the Vertex AI SDK for Python."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "wyy5Lbnzg5fi"
+        "id": "014470c6a8de"
       },
       "outputs": [],
       "source": [
-        "! pip install -U git+https://github.com/googleapis/python-aiplatform.git@main-test --user"
+        "! pip install -U git+https://github.com/googleapis/python-aiplatform.git@main --user"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
-        "id": "irSMQn6gZ19l"
+        "id": "cf00462144f7"
       },
       "source": [
         "Install the `h5py` to prepare sample dataset, and the `grpcio-tools` for querying against the index. "
@@ -176,11 +178,15 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "-h5sqwOEZ5Yq"
+        "id": "3f3e45e5a1d1"
       },
       "outputs": [],
       "source": [
-        "! pip install -U grpcio-tools --user\n",
+        "! pip install protobuf==3.20.*\n",
+        "! pip install -U google-api-python-client==1.8.0 --user\n",
+        "! pip install -U grpcio-tools==1.47.0 --user\n",
+        "! pip install -U grpcio==1.47.0 --user\n",
+        "! pip install -U grpcio-status==1.47.0 --user\n",
         "! pip install -U h5py --user"
       ]
     },
@@ -199,7 +205,7 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "EzrelQZ22IZj"
+        "id": "aa1d87bdc90b"
       },
       "outputs": [],
       "source": [
@@ -215,79 +221,216 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
-        "id": "BF1j6f9HApxa"
+        "id": "249da91c1011"
       },
       "source": [
-        "### Set up your Google Cloud project\n",
+        "### Set your project ID\n",
         "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
-        "\n",
-        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager).\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, and Service Networking API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component,servicenetworking.googleapis.com).\n",
-        "\n",
-        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
-        "Cloud SDK uses the right project for all the commands in this notebook.\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": "WReHDGG5g0XY"
-      },
-      "source": [
-        "#### Set your project ID\n",
-        "\n",
-        "**If you don't know your project ID**, you may be able to get your project ID using `gcloud`."
+        "**If you don't know your project ID**, try the following:\n",
+        "* Run `gcloud config list`.\n",
+        "* Run `gcloud projects list`.\n",
+        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "oM1iC_MfAts1"
+        "id": "10e0d2ee8c45"
       },
       "outputs": [],
       "source": [
-        "import os\n",
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
         "\n",
-        "PROJECT_ID = \"\"\n",
-        "\n",
-        "# Get your Google Cloud project ID from gcloud\n",
-        "if not os.getenv(\"IS_TESTING\"):\n",
-        "    shell_output=!gcloud config list --format 'value(core.project)' 2>/dev/null\n",
-        "    PROJECT_ID = shell_output[0]\n",
-        "    print(\"Project ID: \", PROJECT_ID)"
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
-        "id": "qJYoRfYng0XZ"
+        "id": "3fbfae3ff12a"
       },
       "source": [
-        "Otherwise, set your project ID here."
+        "### Set the region\n",
+        "\n",
+        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations).\n",
+        "* **WARNING:** \n",
+        "    *  **Make sure to [choose a region where Vertex AI services are available](https://cloud.google.com/vertex-ai/docs/general/locations#available_regions).**\n",
+        "    *  **If you use Vertex Workbench, the Notebook instance needs to be in the same region where your Vertex AI Matching Engine is deployed.** (for example, if you set `REGION = \"us-central1\"` as same as the tutorial, the notebook instance has to be in `us-central1`)."
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "riG_qUokg0XZ"
+        "id": "71c3fd82024e"
       },
       "outputs": [],
       "source": [
-        "if PROJECT_ID == \"\" or PROJECT_ID is None:\n",
-        "    PROJECT_ID = \"<your_project_id>\"  # @param {type:\"string\"}"
+        "REGION = \"us-central1\"  # @param {type: \"string\"}\n",
+        "\n",
+        "# Set the regions\n",
+        "! gcloud config set ai_platform/region {REGION}"
       ]
     },
     {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "60c5a0f69ad8"
+      },
+      "source": [
+        "### Authenticate your Google Cloud account\n",
+        "\n",
+        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d118c95af93f"
+      },
+      "source": [
+        "**1. Vertex AI Workbench**\n",
+        "* Do nothing as you are already authenticated."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3035286fcdda"
+      },
+      "source": [
+        "**2. Local JupyterLab instance, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "455882ec0f11"
+      },
+      "outputs": [],
+      "source": [
+        "# ! gcloud auth login"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5097f3233d53"
+      },
+      "source": [
+        "**3. Colab, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2b88e46ac2c8"
+      },
+      "outputs": [],
+      "source": [
+        "# from google.colab import auth\n",
+        "# auth.authenticate_user()"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "fcdbb8929927"
+      },
+      "source": [
+        "**4. Service account or other**\n",
+        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "7c6eef70dfdb"
+      },
+      "source": [
+        "### Prepare a VPC network\n",
+        "\n",
+        "To reduce any network overhead that might lead to unnecessary increase in overhead latency, it is best to call the Vertex AI Matching Engine endpoints from your VPC via a direct [VPC Peering](https://cloud.google.com/vertex-ai/docs/general/vpc-peering) connection. The following section describes how to setup a VPC Peering connection if you don't have one. This is a one-time initial setup task. You can also reuse existing VPC network and skip this section.\n",
+        "\n",
+        "* **WARNING:** The match service gRPC API (to create online queries against your deployed index) has to be executed in a Google Cloud Notebook instance that is created with the following requirements:\n",
+        "    * **Make sure you select the VPC network you created for Vertex AI Matching Engine service** (instead of using the \"default\" one). That is, you will have to create the VPC network below and then create a new notebook instance that uses that VPC.  \n",
+        "    * If you run it in the colab or a Google Cloud Notebook instance in a different VPC network or region, the gRPC API will fail to peer the network (InactiveRPCError)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ab38a8cc634c"
+      },
+      "outputs": [],
+      "source": [
+        "NETWORK_NAME = \"ucaip-haystack-vpc-network\"  # @param {type:\"string\"}\n",
+        "PEERING_RANGE_NAME = \"ucaip-haystack-range\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ec6bf3199835"
+      },
+      "outputs": [],
+      "source": [
+        "# Create a VPC network\n",
+        "! gcloud compute networks create {NETWORK_NAME} --bgp-routing-mode=regional --subnet-mode=auto --project={PROJECT_ID}\n",
+        "\n",
+        "# Add necessary firewall rules\n",
+        "! gcloud compute firewall-rules create {NETWORK_NAME}-allow-icmp --network {NETWORK_NAME} --priority 65534 --project {PROJECT_ID} --allow icmp\n",
+        "\n",
+        "! gcloud compute firewall-rules create {NETWORK_NAME}-allow-internal --network {NETWORK_NAME} --priority 65534 --project {PROJECT_ID} --allow all --source-ranges 10.128.0.0/9\n",
+        "\n",
+        "! gcloud compute firewall-rules create {NETWORK_NAME}-allow-rdp --network {NETWORK_NAME} --priority 65534 --project {PROJECT_ID} --allow tcp:3389\n",
+        "\n",
+        "! gcloud compute firewall-rules create {NETWORK_NAME}-allow-ssh --network {NETWORK_NAME} --priority 65534 --project {PROJECT_ID} --allow tcp:22\n",
+        "\n",
+        "# Reserve IP range\n",
+        "! gcloud compute addresses create {PEERING_RANGE_NAME} --global --prefix-length=16 --network={NETWORK_NAME} --purpose=VPC_PEERING --project={PROJECT_ID} --description=\"peering range for uCAIP Haystack.\""
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ddbace09fe81"
+      },
+      "source": [
+        "Create the VPC Peering. If you are running this from Vertex AI Workbench it is possible you might need your notebook's instance service or user account to have the Service Networking Admin Role"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "d329aa3c54d3"
+      },
+      "outputs": [],
+      "source": [
+        "# Set up peering with service networking\n",
+        "! gcloud services vpc-peerings connect --service=servicenetworking.googleapis.com --network={NETWORK_NAME} --ranges={PEERING_RANGE_NAME} --project={PROJECT_ID}"
+      ]
+    },
+    {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "zgPO1eR3CYjk"
@@ -297,13 +440,11 @@
         "\n",
         "**The following steps are required, regardless of your notebook environment.**\n",
         "\n",
-        "Set the name of your Cloud Storage bucket below. It must be unique across all\n",
+        "Create a storage bucket to store intermediate artifacts such as datasets. Set the name of your Cloud Storage bucket below. It must be unique across all\n",
         "Cloud Storage buckets.\n",
         "\n",
-        "You may also change the `REGION` variable, which is used for operations\n",
-        "throughout the rest of this notebook. Make sure to [choose a region where Vertex AI services are\n",
-        "available](https://cloud.google.com/vertex-ai/docs/general/locations#available_regions). You may\n",
-        "not use a Multi-Regional Storage bucket for training with Vertex AI."
+        "* **WARNING:** \n",
+        "    * **You may not use a Multi-Regional Storage bucket for training with Vertex AI.**"
       ]
     },
     {
@@ -314,8 +455,7 @@
       },
       "outputs": [],
       "source": [
-        "BUCKET_NAME = \"gs://[your-bucket-name]\"  # @param {type:\"string\"}\n",
-        "REGION = \"us-central1\"  # @param {type:\"string\"}"
+        "BUCKET_NAME = \"gs://[your-bucket-name-unique]\"  # @param {type:\"string\"}"
       ]
     },
     {
@@ -328,10 +468,14 @@
       "source": [
         "from datetime import datetime\n",
         "\n",
-        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "UUID = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
         "\n",
-        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"gs://[your-bucket-name]\":\n",
-        "    BUCKET_NAME = \"gs://\" + PROJECT_ID + \"aip-\" + TIMESTAMP"
+        "if (\n",
+        "    BUCKET_NAME == \"\"\n",
+        "    or BUCKET_NAME is None\n",
+        "    or BUCKET_NAME == \"gs://[your-bucket-name-unique]\"\n",
+        "):\n",
+        "    BUCKET_NAME = \"gs://\" + PROJECT_ID + \"aip-\" + UUID"
       ]
     },
     {
@@ -351,7 +495,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil mb -l $REGION $BUCKET_NAME"
+        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_NAME"
       ]
     },
     {
@@ -416,10 +560,7 @@
       },
       "outputs": [],
       "source": [
-        "REGION = \"us-central1\"\n",
         "ENDPOINT = \"{}-aiplatform.googleapis.com\".format(REGION)\n",
-        "NETWORK_NAME = \"ucaip-haystack-vpc-network\"  # @param {type:\"string\"}\n",
-        "\n",
         "\n",
         "AUTH_TOKEN = !gcloud auth print-access-token\n",
         "PROJECT_NUMBER = !gcloud projects list --filter=\"PROJECT_ID:'{PROJECT_ID}'\" --format='value(PROJECT_NUMBER)'\n",
@@ -429,10 +570,7 @@
         "\n",
         "print(\"ENDPOINT: {}\".format(ENDPOINT))\n",
         "print(\"PROJECT_ID: {}\".format(PROJECT_ID))\n",
-        "print(\"REGION: {}\".format(REGION))\n",
-        "\n",
-        "!gcloud config set project {PROJECT_ID}\n",
-        "!gcloud config set ai_platform/region {REGION}"
+        "print(\"REGION: {}\".format(REGION))"
       ]
     },
     {
@@ -523,12 +661,13 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "QuVl8DrWG8NS"
       },
       "source": [
-        "Upload the training data to GCS."
+        "Upload the training data to Google Cloud Storage"
       ]
     },
     {
@@ -539,9 +678,9 @@
       },
       "outputs": [],
       "source": [
-        "# NOTE: Everything in this GCS DIR will be DELETED before uploading the data.\n",
+        "# NOTE: Everything in this Google Cloud Storage directory will be DELETED before uploading the data\n",
         "\n",
-        "! gsutil rm -rf {BUCKET_NAME}/*"
+        "! gsutil rm -raf {BUCKET_NAME}/** 2> /dev/null || true"
       ]
     },
     {
@@ -567,21 +706,23 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "mglUPwHpJH98"
       },
       "source": [
-        "## Create Indexes\n"
+        "## Create the indexes\n"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "qhIBCQ7dDSbW"
       },
       "source": [
-        "### Create ANN Index (for Production Usage)"
+        "### Create Vertex AI Matching Engine index (for production usage)"
       ]
     },
     {
@@ -597,6 +738,16 @@
         ")"
       ]
     },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "14e1ed031d66"
+      },
+      "source": [
+        "Set constants"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -611,14 +762,15 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "svLYiDf0OD2G"
       },
       "source": [
-        "Create the ANN index configuration:\n",
+        "#### Create the Vertex AI Matching Engine index configuration\n",
         "\n",
-        "Please read the documentation to understand the various configuration parameters that can be used to tune the index\n"
+        "Please read the [documentation](https://cloud.google.com/vertex-ai/docs/matching-engine/configuring-indexes) to understand the various configuration parameters that can be used to tune the index"
       ]
     },
     {
@@ -656,9 +808,9 @@
         "    }\n",
         ")\n",
         "\n",
-        "ann_index = {\n",
+        "matching_engine_index = {\n",
         "    \"display_name\": DISPLAY_NAME,\n",
-        "    \"description\": \"Glove 100 ANN index\",\n",
+        "    \"description\": \"Glove 100 Vertex AI Matching Engine Index\",\n",
         "    \"metadata\": struct_pb2.Value(struct_value=metadata),\n",
         "}"
       ]
@@ -671,7 +823,9 @@
       },
       "outputs": [],
       "source": [
-        "ann_index = index_client.create_index(parent=PARENT, index=ann_index)"
+        "matching_engine_index = index_client.create_index(\n",
+        "    parent=PARENT, index=matching_engine_index\n",
+        ")"
       ]
     },
     {
@@ -686,7 +840,7 @@
         "# This will take ~45 min.\n",
         "\n",
         "while True:\n",
-        "    if ann_index.done():\n",
+        "    if matching_engine_index.done():\n",
         "        break\n",
         "    print(\"Poll the operation to create index...\")\n",
         "    time.sleep(60)"
@@ -700,17 +854,18 @@
       },
       "outputs": [],
       "source": [
-        "INDEX_RESOURCE_NAME = ann_index.result().name\n",
+        "INDEX_RESOURCE_NAME = matching_engine_index.result().name\n",
         "INDEX_RESOURCE_NAME"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "kSsqZuyoA1SG"
       },
       "source": [
-        "### Create Brute Force Index (for Ground Truth)\n",
+        "### Create brute force index (for ground truth)\n",
         "\n",
         "The brute force index uses a naive brute force method to find the nearest neighbors. This method is not fast or efficient. Hence brute force indices are not recommended for production usage. They are to be used to find the \"ground truth\" set of neighbors, so that the \"ground truth\" set can be used to measure recall of the indices being tuned for production usage. To ensure an apples to apples comparison, the `distanceMeasureType` and `featureNormType`, `dimensions` of the brute force index should match those of the production indices being tuned.\n",
         "\n",
@@ -725,8 +880,6 @@
       },
       "outputs": [],
       "source": [
-        "from google.protobuf import *\n",
-        "\n",
         "algorithmConfig = struct_pb2.Struct(\n",
         "    fields={\"bruteForceConfig\": struct_pb2.Value(struct_value=struct_pb2.Struct())}\n",
         ")\n",
@@ -796,12 +949,13 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "mglUPwHpJH98"
       },
       "source": [
-        "## Update Indexes\n",
+        "## Update the indexes\n",
         "\n",
         "Create incremental data file.\n"
       ]
@@ -863,10 +1017,10 @@
         "    }\n",
         ")\n",
         "\n",
-        "ann_index = {\n",
+        "matching_engine_index = {\n",
         "    \"name\": INDEX_RESOURCE_NAME,\n",
         "    \"display_name\": DISPLAY_NAME,\n",
-        "    \"description\": \"Glove 100 ANN index\",\n",
+        "    \"description\": \"Glove 100 Vertex AI Matching Engine Index\",\n",
         "    \"metadata\": struct_pb2.Value(struct_value=metadata),\n",
         "}"
       ]
@@ -879,7 +1033,7 @@
       },
       "outputs": [],
       "source": [
-        "ann_index = index_client.update_index(index=ann_index)"
+        "matching_engine_index = index_client.update_index(index=matching_engine_index)"
       ]
     },
     {
@@ -894,7 +1048,7 @@
         "# This will take ~45 min.\n",
         "\n",
         "while True:\n",
-        "    if ann_index.done():\n",
+        "    if matching_engine_index.done():\n",
         "        break\n",
         "    print(\"Poll the operation to update index...\")\n",
         "    time.sleep(60)"
@@ -908,17 +1062,18 @@
       },
       "outputs": [],
       "source": [
-        "INDEX_RESOURCE_NAME = ann_index.result().name\n",
+        "INDEX_RESOURCE_NAME = matching_engine_index.result().name\n",
         "INDEX_RESOURCE_NAME"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "qV2xjAnDDObD"
       },
       "source": [
-        "## Create an IndexEndpoint with VPC Network"
+        "## Create an index endpoint with VPC network"
       ]
     },
     {
@@ -997,21 +1152,23 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "np2cgVuuIe9k"
       },
       "source": [
-        "## Deploy Indexes"
+        "## Deploy the indexes"
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "8Ew1UgcIIiJG"
       },
       "source": [
-        "### Deploy ANN Index"
+        "### Deploy a Vertex AI Matching Engine index"
       ]
     },
     {
@@ -1022,7 +1179,7 @@
       },
       "outputs": [],
       "source": [
-        "DEPLOYED_INDEX_ID = \"ann_glove_deployed\""
+        "DEPLOYED_INDEX_ID = \"matching_engine_glove_deployed\""
       ]
     },
     {
@@ -1033,13 +1190,23 @@
       },
       "outputs": [],
       "source": [
-        "deploy_ann_index = {\n",
+        "deploy_matching_engine_index = {\n",
         "    \"id\": DEPLOYED_INDEX_ID,\n",
         "    \"display_name\": DEPLOYED_INDEX_ID,\n",
         "    \"index\": INDEX_RESOURCE_NAME,\n",
         "}"
       ]
     },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cb6d956d7419"
+      },
+      "source": [
+        "If errors occur with the next command wait some minutes for the index endpoint to be created and retry."
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -1049,7 +1216,7 @@
       "outputs": [],
       "source": [
         "r = index_endpoint_client.deploy_index(\n",
-        "    index_endpoint=INDEX_ENDPOINT_NAME, deployed_index=deploy_ann_index\n",
+        "    index_endpoint=INDEX_ENDPOINT_NAME, deployed_index=deploy_matching_engine_index\n",
         ")"
       ]
     },
@@ -1082,12 +1249,13 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "RNZnXmO5AhDO"
       },
       "source": [
-        "### Deploy Brute Force Index"
+        "### Deploy brute force index"
       ]
     },
     {
@@ -1158,12 +1326,13 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "6LCGvBNvBd8D"
       },
       "source": [
-        "## Create Online Queries\n",
+        "## Create online queries\n",
         "\n",
         "After you built your indexes, you may query against the deployed index through the online querying gRPC API (Match service) within the virtual machine instances from the same region (for example 'us-central1' in this tutorial).  \n",
         "\n",
@@ -1178,7 +1347,15 @@
         "\n",
         "* Compile the protocal buffer (see below)\n",
         "* Obtain the index endpoint\n",
-        "* Use a code-generated stub to make the call, passing the parameter values"
+        "* Use a code-generated stub to make the call, passing the parameter values\n",
+        "\n",
+        "### Troubleshooting connectivity issues\n",
+        "\n",
+        "In case you have connectivity errors please perform the following:\n",
+        "\n",
+        "* Verify that the index endpoint, index, and VPC are all in the same Google Cloud project\n",
+        "* Verify that the index endpoint, index, and VPC are all in the same region and it is a valid (e.g. us-central1)\n",
+        "* Verify the Network does not have a firewall rule which denies all egress connections. Else, disable this rule or overwrite it with another rule that allows connection to the index endpoint IP"
       ]
     },
     {
@@ -1351,12 +1528,13 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "8wXTSgz1Bl0x"
       },
       "source": [
-        "Obtain the Private Endpoint: "
+        "Obtain the private endpoint: "
       ]
     },
     {
@@ -1521,12 +1699,13 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "_mNwdU9_B_Ez"
       },
       "source": [
-        "### Batch Query\n",
+        "## Submit a batch query\n",
         "\n",
         "You can run multiple queries in a single RPC call using the BatchMatch API:"
       ]
@@ -1764,18 +1943,20 @@
         "]\n",
         "\n",
         "batch_request = match_service_pb2.BatchMatchRequest()\n",
-        "batch_request_ann = match_service_pb2.BatchMatchRequest.BatchMatchRequestPerIndex()\n",
+        "batch_request_matching_engine = (\n",
+        "    match_service_pb2.BatchMatchRequest.BatchMatchRequestPerIndex()\n",
+        ")\n",
         "batch_request_brute_force = (\n",
         "    match_service_pb2.BatchMatchRequest.BatchMatchRequestPerIndex()\n",
         ")\n",
-        "batch_request_ann.deployed_index_id = DEPLOYED_INDEX_ID\n",
+        "batch_request_matching_engine.deployed_index_id = DEPLOYED_INDEX_ID\n",
         "batch_request_brute_force.deployed_index_id = DEPLOYED_BRUTE_FORCE_INDEX_ID\n",
         "for query in queries:\n",
-        "    batch_request_ann.requests.append(get_request(query, DEPLOYED_INDEX_ID))\n",
+        "    batch_request_matching_engine.requests.append(get_request(query, DEPLOYED_INDEX_ID))\n",
         "    batch_request_brute_force.requests.append(\n",
         "        get_request(query, DEPLOYED_BRUTE_FORCE_INDEX_ID)\n",
         "    )\n",
-        "batch_request.requests.append(batch_request_ann)\n",
+        "batch_request.requests.append(batch_request_matching_engine)\n",
         "batch_request.requests.append(batch_request_brute_force)\n",
         "\n",
         "response = stub.BatchMatch(batch_request)\n",
@@ -1783,14 +1964,15 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "_mNwdU9_B_Ez"
       },
       "source": [
-        "### Compute Recall\n",
+        "### Compute the recall metric\n",
         "\n",
-        "Use deployed brute force Index as the ground truth to calculate the recall of ANN Index:"
+        "Use the deployed brute force index as the ground truth to calculate the recall of the Vertex AI Matching Engine index:"
       ]
     },
     {
@@ -1835,6 +2017,7 @@
       ]
     },
     {
+      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "TpV-iwP9qw9c"
@@ -1844,7 +2027,18 @@
         "\n",
         "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
         "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
-        "You can also manually delete resources that you created by running the following code."
+        "\n",
+        "Otherwise, you can delete the individual resources you created in this tutorial:"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "390c331dc7d9"
+      },
+      "source": [
+        "### Delete the Vertex AI Matching Engine resources"
       ]
     },
     {
@@ -1869,6 +2063,31 @@
       "source": [
         "index_endpoint_client.delete_index_endpoint(name=INDEX_ENDPOINT_NAME)"
       ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ff14a85c85fb"
+      },
+      "source": [
+        "### Delete the Google Cloud Storage bucket"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "68d4781faac4"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "\n",
+        "delete_bucket = False\n",
+        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "    ! gsutil -m rm -r $BUCKET_NAME"
+      ]
     }
   ],
   "metadata": {

notebooks/community/ml_ops/stage3/get_started_vertex_pipelines_sklearn_with_prediction.ipynb

@@ -560,6 +560,30 @@
         "    print(\"Service Account:\", SERVICE_ACCOUNT)"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "set_service_account:pipelines"
+      },
+      "source": [
+        "#### Set service account access for Vertex AI Pipelines\n",
+        "\n",
+        "Run the following commands to grant your service account access to read and write pipeline artifacts in the bucket that you created in the previous step -- you only need to run these once per service account."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "set_service_account:pipelines"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectCreator $BUCKET_URI\n",
+        "\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectViewer $BUCKET_URI"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {

notebooks/community/ml_ops/stage3/get_started_vertex_pipelines_xgboost_with_experiments.ipynb

@@ -557,6 +557,30 @@
         "    print(\"Service Account:\", SERVICE_ACCOUNT)"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "set_service_account:pipelines"
+      },
+      "source": [
+        "#### Set service account access for Vertex AI Pipelines\n",
+        "\n",
+        "Run the following commands to grant your service account access to read and write pipeline artifacts in the bucket that you created in the previous step -- you only need to run these once per service account."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 13,
+      "metadata": {
+        "id": "set_service_account:pipelines"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectCreator $BUCKET_URI\n",
+        "\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectViewer $BUCKET_URI"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {

notebooks/community/ml_ops/stage3/get_started_vertex_pipelines_xgboost_with_prediction.ipynb

@@ -196,7 +196,7 @@
         "    USER_FLAG = \"--user\"\n",
         "\n",
         "! pip3 install {USER_FLAG} --upgrade --quiet google-cloud-aiplatform \\\n",
-        "                                             google-cloud-pipeline-components \\\n",
+        "                                             google-cloud-pipeline-components==1.0.25 \\\n",
         "                                             kfp "
       ]
     },
@@ -557,6 +557,30 @@
         "    print(\"Service Account:\", SERVICE_ACCOUNT)"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "set_service_account:pipelines"
+      },
+      "source": [
+        "#### Set service account access for Vertex AI Pipelines\n",
+        "\n",
+        "Run the following commands to grant your service account access to read and write pipeline artifacts in the bucket that you created in the previous step -- you only need to run these once per service account."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 13,
+      "metadata": {
+        "id": "set_service_account:pipelines"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectCreator $BUCKET_URI\n",
+        "\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectViewer $BUCKET_URI"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {

notebooks/community/ml_ops/stage3/get_started_with_dataflow_flex_template_component.ipynb

@@ -0,0 +1,832 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "copyright"
+      },
+      "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": "title:generic,gcp"
+      },
+      "source": [
+        "# E2E ML on GCP: MLOps stage 3 : formalization: get started with the Dataflow Flex Template component\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage3/get_started_with_dataflow_pipeline_components.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://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage3/get_started_with_dataflow_pipeline_components.ipynb\">\n",
+        "<img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> \n",
+        "        Run in Colab\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/community/ml_ops/stage3/get_started_with_dataflow_pipeline_components.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",
+        "    \n",
+        "</table>\n",
+        "<br/><br/><br/>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "overview:mlops"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "\n",
+        "This tutorial demonstrates how to use Vertex AI for E2E MLOps on Google Cloud in production. This tutorial covers stage 3 : formalization: get started with the Dataflow Flex Template component."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "objective:mlops,stage3,get_started_dataflow_pipeline_components"
+      },
+      "source": [
+        "### Objective\n",
+        "\n",
+        "In this tutorial, you learn how to use prebuilt `Google Cloud Pipeline Components` to execute `Dataflow` [Flex Template](https://cloud.google.com/dataflow/docs/guides/templates/using-flex-templates) jobs.\n",
+        "\n",
+        "This tutorial uses the following Google Cloud ML services:\n",
+        "\n",
+        "- `Vertex AI Pipelines`\n",
+        "- `Google Cloud Pipeline Components`\n",
+        "- `Dataflow`\n",
+        "\n",
+        "The steps performed include:\n",
+        "\n",
+        "- Defining a pipeline step to execute a Dataflow Flex Template job within a Vertex AI pipeline.\n",
+        "- Execute a Vertex AI pipeline."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dataset:gsod,lrg"
+      },
+      "source": [
+        "### Dataset\n",
+        "\n",
+        "The dataset used for this tutorial is the sample data from the [Apache Beam Mobile Gaming Pipeline example](https://beam.apache.org/get-started/mobile-gaming-example/)."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0c997d8d92ce"
+      },
+      "source": [
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "* Dataflow\n",
+        "\n",
+        "Learn about [Vertex AI\n",
+        "pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage\n",
+        "pricing](https://cloud.google.com/storage/pricing), and [Dataflow pricing](https://cloud.google.com/dataflow/pricing)\n",
+        "and use the [Pricing\n",
+        "Calculator](https://cloud.google.com/products/calculator/)\n",
+        "to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "install_mlops"
+      },
+      "source": [
+        "## Installations\n",
+        "\n",
+        "Install the required packages for executing the notebook."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "install_mlops"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "\n",
+        "# The Vertex AI Workbench Notebook product has specific requirements\n",
+        "IS_WORKBENCH_NOTEBOOK = os.getenv(\"DL_ANACONDA_HOME\")\n",
+        "IS_USER_MANAGED_WORKBENCH_NOTEBOOK = os.path.exists(\n",
+        "    \"/opt/deeplearning/metadata/env_version\"\n",
+        ")\n",
+        "\n",
+        "# Vertex AI Notebook requires dependencies to be installed with '--user'\n",
+        "USER_FLAG = \"\"\n",
+        "if IS_WORKBENCH_NOTEBOOK:\n",
+        "    USER_FLAG = \"--user\"\n",
+        "\n",
+        "! pip3 install --pre --upgrade google-cloud-pipeline-components $USER_FLAG -q"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "restart"
+      },
+      "source": [
+        "### Restart the kernel\n",
+        "\n",
+        "Once you've installed the additional packages, you need to restart the notebook kernel so it can find the packages."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "restart"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "\n",
+        "if not os.getenv(\"IS_TESTING\"):\n",
+        "    # Automatically restart kernel after installs\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BF1j6f9HApxa"
+      },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "### Set up your Google Cloud project\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\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 Dataflow API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,dataflow.googleapis.com).\n",
+        "\n",
+        "1. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "\n",
+        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
+        "Cloud SDK uses the right project for all the commands in this notebook.\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": "project_id"
+      },
+      "source": [
+        "#### Set your project ID\n",
+        "\n",
+        "**If you don't know your project ID**, you may be able to get your project ID using `gcloud`."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "set_project_id"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "autoset_project_id"
+      },
+      "outputs": [],
+      "source": [
+        "if PROJECT_ID == \"\" or PROJECT_ID is None or PROJECT_ID == \"[your-project-id]\":\n",
+        "    # Get your GCP project id from gcloud\n",
+        "    shell_output = ! gcloud config list --format 'value(core.project)' 2>/dev/null\n",
+        "    PROJECT_ID = shell_output[0]\n",
+        "    print(\"Project ID:\", PROJECT_ID)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "250cb8c648d5"
+      },
+      "outputs": [],
+      "source": [
+        "! gcloud config set project $PROJECT_ID"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "region"
+      },
+      "source": [
+        "#### Region\n",
+        "\n",
+        "You can also change the `REGION` variable, which is used for operations\n",
+        "throughout the rest of this notebook.  Below are regions supported for Vertex AI. We recommend that you choose the region closest to you.\n",
+        "\n",
+        "- Americas: `us-central1`\n",
+        "- Europe: `europe-west4`\n",
+        "- Asia Pacific: `asia-east1`\n",
+        "\n",
+        "You may not use a multi-regional bucket for training with Vertex AI. Not all regions provide support for all Vertex AI services.\n",
+        "\n",
+        "Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "region"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"[your-region]\"  # @param {type: \"string\"}\n",
+        "\n",
+        "if REGION == \"[your-region]\":\n",
+        "    REGION = \"us-central1\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "timestamp"
+      },
+      "source": [
+        "#### Timestamp\n",
+        "\n",
+        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a timestamp for each instance session, and append the timestamp onto the name of resources you create in this tutorial."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "timestamp"
+      },
+      "outputs": [],
+      "source": [
+        "from datetime import datetime\n",
+        "\n",
+        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "927085b84a07"
+      },
+      "source": [
+        "### Authenticate your Google Cloud account\n",
+        "\n",
+        "**If you are using Vertex AI Workbench Notebooks**, your environment is already authenticated. Skip this step.\n",
+        "\n",
+        "**If you are using Colab**, run the cell below and follow the instructions when prompted to authenticate your account via oAuth.\n",
+        "\n",
+        "**Otherwise**, follow these steps:\n",
+        "\n",
+        "In the Cloud Console, go to the [Create service account key](https://console.cloud.google.com/apis/credentials/serviceaccountkey) page.\n",
+        "\n",
+        "**Click Create service account**.\n",
+        "\n",
+        "In the **Service account name** field, enter a name, and click **Create**.\n",
+        "\n",
+        "In the **Grant this service account access to project** section, click the Role drop-down list. Type \"Vertex\" into the filter box, and select **Vertex Administrator**. Type \"Storage Object Admin\" into the filter box, and select **Storage Object Admin**.\n",
+        "\n",
+        "Click Create. A JSON file that contains your key downloads to your local environment.\n",
+        "\n",
+        "Enter the path to your service account key as the GOOGLE_APPLICATION_CREDENTIALS variable in the cell below and run the cell."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "89788a802687"
+      },
+      "outputs": [],
+      "source": [
+        "# If you are running this notebook in Colab, run this cell and follow the\n",
+        "# instructions to authenticate your GCP account. This provides access to your\n",
+        "# Cloud Storage bucket and lets you submit training jobs and prediction\n",
+        "# requests.\n",
+        "\n",
+        "import os\n",
+        "import sys\n",
+        "\n",
+        "# If on Vertex AI Workbench, then don't execute this code\n",
+        "IS_COLAB = False\n",
+        "if not os.path.exists(\"/opt/deeplearning/metadata/env_version\") and not os.getenv(\n",
+        "    \"DL_ANACONDA_HOME\"\n",
+        "):\n",
+        "    if \"google.colab\" in sys.modules:\n",
+        "        IS_COLAB = True\n",
+        "        from google.colab import auth as google_auth\n",
+        "\n",
+        "        google_auth.authenticate_user()\n",
+        "\n",
+        "    # If you are running this notebook locally, replace the string below with the\n",
+        "    # path to your service account key and run this cell to authenticate your GCP\n",
+        "    # account.\n",
+        "    elif not os.getenv(\"IS_TESTING\"):\n",
+        "        %env GOOGLE_APPLICATION_CREDENTIALS ''"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "40ed98f5cc48"
+      },
+      "source": [
+        "#### If you are using Colab Notebooks, set the project using gcloud config."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "fde1a355f1e9"
+      },
+      "outputs": [],
+      "source": [
+        "if not os.path.exists(\"/opt/deeplearning/metadata/env_version\"):\n",
+        "    if \"google.colab\" in sys.modules:\n",
+        "        ! gcloud config set project $PROJECT_ID"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "bucket:mbsdk"
+      },
+      "source": [
+        "### Create a Cloud Storage bucket\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\n",
+        "\n",
+        "When you initialize the Vertex SDK for Python, you specify a Cloud Storage staging bucket. The staging bucket is where all the data associated with your dataset and model resources are retained across sessions.\n",
+        "\n",
+        "Set the name of your Cloud Storage bucket below. Bucket names must be globally unique across all Google Cloud projects, including those outside of your organization."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bucket"
+      },
+      "outputs": [],
+      "source": [
+        "BUCKET_URI = \"gs://[your-bucket-name]\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "autoset_bucket"
+      },
+      "outputs": [],
+      "source": [
+        "if BUCKET_URI == \"\" or BUCKET_URI is None or BUCKET_URI == \"gs://[your-bucket-name]\":\n",
+        "    BUCKET_URI = \"gs://\" + PROJECT_ID + \"aip-\" + TIMESTAMP"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "create_bucket"
+      },
+      "source": [
+        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "create_bucket"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil mb -l $REGION $BUCKET_URI"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "validate_bucket"
+      },
+      "source": [
+        "Finally, validate access to your Cloud Storage bucket by examining its contents:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "validate_bucket"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil ls -al $BUCKET_URI"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "set_service_account"
+      },
+      "source": [
+        "#### Service Account\n",
+        "\n",
+        "**If you don't know your service account**, try to get your service account using `gcloud` command by executing the second cell below."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "set_service_account"
+      },
+      "outputs": [],
+      "source": [
+        "SERVICE_ACCOUNT = \"[your-service-account]\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "autoset_service_account"
+      },
+      "outputs": [],
+      "source": [
+        "if (\n",
+        "    SERVICE_ACCOUNT == \"\"\n",
+        "    or SERVICE_ACCOUNT is None\n",
+        "    or SERVICE_ACCOUNT == \"[your-service-account]\"\n",
+        "):\n",
+        "    # Get your service account from gcloud\n",
+        "    if not IS_COLAB:\n",
+        "        shell_output = !gcloud auth list 2>/dev/null\n",
+        "        SERVICE_ACCOUNT = shell_output[2].replace(\"*\", \"\").strip()\n",
+        "\n",
+        "    if IS_COLAB:\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(\"Service Account:\", SERVICE_ACCOUNT)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "set_service_account:pipelines"
+      },
+      "source": [
+        "#### Set service account access for Vertex AI Pipelines\n",
+        "\n",
+        "Run the following commands to grant your service account access to read and write pipeline artifacts in the bucket that you created in the previous step -- you only need to run these once per service account."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "set_service_account:pipelines"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectCreator $BUCKET_URI\n",
+        "\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectViewer $BUCKET_URI"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "setup_vars"
+      },
+      "source": [
+        "### Set up variables\n",
+        "\n",
+        "Next, set up some variables used throughout the tutorial.\n",
+        "### Import libraries and define constants"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "import_aip:mbsdk"
+      },
+      "outputs": [],
+      "source": [
+        "import google.cloud.aiplatform as aip\n",
+        "from google_cloud_pipeline_components.experimental.dataflow import \\\n",
+        "    DataflowFlexTemplateJobOp\n",
+        "from google_cloud_pipeline_components.v1.wait_gcp_resources import \\\n",
+        "    WaitGcpResourcesOp\n",
+        "from kfp import dsl\n",
+        "from kfp.v2 import compiler"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "init_aip:mbsdk"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "init_aip:mbsdk"
+      },
+      "outputs": [],
+      "source": [
+        "aip.init(project=PROJECT_ID, staging_bucket=BUCKET_URI)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "writefile:wc.py"
+      },
+      "source": [
+        "### Prepare a Flex Template job to convert a CSV file to Parquet\n",
+        "\n",
+        "In this tutorial, you use the Google-provided [File Format Conversion template](https://cloud.google.com/dataflow/docs/guides/templates/provided/file-format-conversion) to convert a CSV file to Parquet format."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "writefile:requirements,wc"
+      },
+      "source": [
+        "#### Create the Avro schema file\n",
+        "\n",
+        "First, create an Avro schema file that describes the example data."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "writefile:requirements,wc"
+      },
+      "outputs": [],
+      "source": [
+        "%%writefile gaming_schema.avsc\n",
+        "\n",
+        "{\n",
+        "   \"type\" : \"record\",\n",
+        "   \"name\" : \"user_score\",\n",
+        "   \"fields\" : [\n",
+        "      { \"name\" : \"user\" , \"type\" : \"string\" },\n",
+        "      { \"name\" : \"team\" , \"type\" : \"string\" },\n",
+        "      { \"name\" : \"score\" , \"type\" : \"int\" },\n",
+        "      { \"name\" : \"ts_epoch\" , \"type\" : \"long\" },\n",
+        "      { \"name\" : \"ts_str\" , \"type\" : \"string\" }\n",
+        "   ]\n",
+        "}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "copy_to_gcs:wc"
+      },
+      "source": [
+        "#### Copy the Avro schema file to Cloud Storage\n",
+        "\n",
+        "Next, you copy the Avro schema file to your Cloud Storage bucket.\n",
+        "\n",
+        "Additional, you set the following:\n",
+        "\n",
+        "- The Cloud Storage location of the Flex Template definition file. In this tutorial, you use the location of the File Conversion template.\n",
+        "- A `Dict` containing the user parameters required by the Flex Template. These parameters include:\n",
+        "\n",
+        "    - `inputFileFormat`: The file format of the input files.\n",
+        "    - `outputFileFormat`: The file format of the output files.\n",
+        "    - `inputFileSpec`: The input filepattern to read from.\n",
+        "    - `outputBucket`: The Cloud Storage path to write the output files.\n",
+        "    - `schema`: The Cloud Storage path top the avro schema file used for the conversion.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "copy_to_gcs:wc"
+      },
+      "outputs": [],
+      "source": [
+        "GCS_AVRO_SCHEMA = BUCKET_URI + \"/gaming_schema.avsc\"\n",
+        "! gsutil cp gaming_schema.avsc $GCS_AVRO_SCHEMA\n",
+        "\n",
+        "GCS_FLEX_TEMPLATE_PATH = \"gs://dataflow-templates/latest/flex/File_Format_Conversion\"\n",
+        "GCS_CONVERT_IN = \"gs://dataflow-samples/game/5000_gaming_data.csv\"\n",
+        "GCS_CONVERT_OUT = BUCKET_URI + \"/parquet_out/\"\n",
+        "\n",
+        "TEMPLATE_PARAMETERS = {\n",
+        "    \"inputFileFormat\": \"csv\",\n",
+        "    \"outputFileFormat\": \"parquet\",\n",
+        "    \"inputFileSpec\": GCS_CONVERT_IN,\n",
+        "    \"outputBucket\": GCS_CONVERT_OUT,\n",
+        "    \"schema\": GCS_AVRO_SCHEMA,\n",
+        "}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "create_dataflow_pipeline:wc"
+      },
+      "source": [
+        "### Create and execute the pipeline job\n",
+        "\n",
+        "In this example, the `DataflowFlexTemplateJobOp` component takes the following parameters:\n",
+        "\n",
+        "- `project_id`: The project ID.\n",
+        "- `location`: The region.\n",
+        "- `container_spec_gcs_path`: The Cloud Storage path to a file that contains the Flex Template definition. This file contains a json serialized `ContainerSpec` as content.\n",
+        "- `temp_location`: The Cloud Storage path to use for temporary files.\n",
+        "- `parameters`: The parameters for the Flex Template.\n",
+        "\n",
+        "Learn more about [Google Cloud Pipeline Components for Dataflow](https://google-cloud-pipeline-components.readthedocs.io/en/google-cloud-pipeline-components-2.0.0b2/google_cloud_pipeline_components.experimental.dataflow.html)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "create_dataflow_pipeline:wc"
+      },
+      "outputs": [],
+      "source": [
+        "PIPELINE_ROOT = \"{}/pipeline_root/dataflow_file_conversion\".format(BUCKET_URI)\n",
+        "\n",
+        "\n",
+        "@dsl.pipeline(\n",
+        "    name=\"dataflow-file-conversion\", description=\"Dataflow file format conversion\"\n",
+        ")\n",
+        "def pipeline(\n",
+        "    project_id: str = PROJECT_ID,\n",
+        "    location: str = REGION,\n",
+        "    container_spec_gcs_path: str = GCS_FLEX_TEMPLATE_PATH,\n",
+        "    temp_location: str = PIPELINE_ROOT,\n",
+        "    parameters: dict = TEMPLATE_PARAMETERS,\n",
+        "):\n",
+        "    flex_template_op = DataflowFlexTemplateJobOp(\n",
+        "        project=project_id,\n",
+        "        location=location,\n",
+        "        container_spec_gcs_path=container_spec_gcs_path,\n",
+        "        temp_location=temp_location,\n",
+        "        parameters=parameters,\n",
+        "    )\n",
+        "\n",
+        "    _ = WaitGcpResourcesOp(gcp_resources=flex_template_op.outputs[\"gcp_resources\"])\n",
+        "\n",
+        "\n",
+        "compiler.Compiler().compile(\n",
+        "    pipeline_func=pipeline, package_path=\"dataflow_file_conversion.yaml\"\n",
+        ")\n",
+        "\n",
+        "pipeline = aip.PipelineJob(\n",
+        "    display_name=\"dataflow_file_conversion\",\n",
+        "    template_path=\"dataflow_file_conversion.yaml\",\n",
+        "    pipeline_root=PIPELINE_ROOT,\n",
+        "    enable_caching=False,\n",
+        ")\n",
+        "\n",
+        "pipeline.run()\n",
+        "\n",
+        "! gsutil ls $GCS_CONVERT_OUT\n",
+        "\n",
+        "! rm -f dataflow_file_conversion.yaml gaming_schema.avsc"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "delete_pipeline"
+      },
+      "source": [
+        "### Delete a pipeline job\n",
+        "\n",
+        "After a pipeline job is completed, you can delete the pipeline job with the method `delete()`.  Prior to completion, a pipeline job can be canceled with the method `cancel()`."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "delete_pipeline"
+      },
+      "outputs": [],
+      "source": [
+        "pipeline.delete()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cleanup:mbsdk"
+      },
+      "source": [
+        "# Cleaning up\n",
+        "\n",
+        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
+        "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",
+        "- Dataset\n",
+        "- Cloud Storage Bucket"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "cleanup:mbsdk"
+      },
+      "outputs": [],
+      "source": [
+        "# Warning: Setting this to true will delete everything in your bucket\n",
+        "delete_bucket = False\n",
+        "\n",
+        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "    ! gsutil rm -r $BUCKET_URI"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "get_started_with_dataflow_flex_template_component.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/ml_ops/stage3/get_started_with_hpt_pipeline_components.ipynb

@@ -1333,7 +1333,6 @@
         "Next, you compile the pipeline and then exeute it. The pipeline takes the following parameters, which are passed as the dictionary `parameter_values`:\n",
         "\n",
         "- `display_name`: A human readable name for the pipeline job.\n",
-        "- `import_file`: The Cloud Storage location to the dataset.\n",
         "- `worker_pool_specs`: The the machine and container, and auto-scaling requirements, as well as command line arguments.\n",
         "- `study_spec_metrics`: The metrics to optimize in the study trials.\n",
         "- `study_spec_parameters`: The parameters to tune."

notebooks/community/ml_ops/stage4/get_started_with_vertex_ml_metadata.ipynb

@@ -29,22 +29,22 @@
         "id": "title:generic,gcp"
       },
       "source": [
-        "# E2E ML on GCP: MLOps stage 4 : formalization: get started with Vertex ML Metadata\n",
+        "# Get started with Vertex ML Metadata\n",
         "\n",
         "<table align=\"left\">\n",
         "  <td>\n",
-        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage4/get_started_with_vertex_ml_metadata.ipynb\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/ml_metadata/get_started_with_vertex_ml_metadata.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/ml_ops/stage4/get_started_with_vertex_ml_metadata.ipynb\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/ml_metadata/get_started_with_vertex_ml_metadata.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/community/ml_ops/stage4/get_started_with_vertex_ml_metadata.ipynb\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/ml_metadata/get_started_with_vertex_ml_metadata.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",
@@ -62,7 +62,9 @@
         "## Overview\n",
         "\n",
         "\n",
-        "This tutorial demonstrates how to use Vertex AI for E2E MLOps on Google Cloud in production. This tutorial covers stage 4 : formalization: get started with Vertex ML Metadata."
+        "This tutorial demonstrates how to use Vertex ML Metadata.\n",
+        "\n",
+        "Learn more about [Vertex ML Metadata](https://cloud.google.com/vertex-ai/docs/ml-metadata)."
       ]
     },
     {
@@ -144,19 +146,32 @@
       "source": [
         "import os\n",
         "\n",
-        "# The Vertex AI Workbench Notebook product has specific requirements\n",
-        "IS_WORKBENCH_NOTEBOOK = os.getenv(\"DL_ANACONDA_HOME\")\n",
-        "IS_USER_MANAGED_WORKBENCH_NOTEBOOK = os.path.exists(\n",
-        "    \"/opt/deeplearning/metadata/env_version\"\n",
-        ")\n",
+        "! pip3 install --upgrade google-cloud-aiplatform[tensorboard] \\\n",
+        "                         google-cloud-pipeline-components  --quiet"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "restart"
+      },
+      "source": [
+        "### Colab only: Uncomment the following cell to restart the kernel"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "D-ZBOjErv5mM"
+      },
+      "outputs": [],
+      "source": [
+        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
+        "# import IPython\n",
         "\n",
-        "# Vertex AI Notebook requires dependencies to be installed with '--user'\n",
-        "USER_FLAG = \"\"\n",
-        "if IS_WORKBENCH_NOTEBOOK:\n",
-        "    USER_FLAG = \"--user\"\n",
-        "\n",
-        "! pip3 install --upgrade google-cloud-aiplatform[tensorboard] $USER_FLAG -q\n",
-        "! pip3 install --upgrade google-cloud-pipeline-components $USER_FLAG -q"
+        "# app = IPython.Application.instance()\n",
+        "# app.kernel.do_shutdown(True)"
       ]
     },
     {
@@ -170,63 +185,27 @@
         "Once you've installed the additional packages, you need to restart the notebook kernel so it can find the packages."
       ]
     },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "restart"
-      },
-      "outputs": [],
-      "source": [
-        "import os\n",
-        "\n",
-        "if not os.getenv(\"IS_TESTING\"):\n",
-        "    # Automatically restart kernel after installs\n",
-        "    import IPython\n",
-        "\n",
-        "    app = IPython.Application.instance()\n",
-        "    app.kernel.do_shutdown(True)"
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
         "id": "before_you_begin"
       },
       "source": [
-        "## Before you begin\n",
-        "\n",
-        "### GPU runtime\n",
-        "\n",
-        "*Make sure you're running this notebook in a GPU runtime if you have that option. In Colab, select* **Runtime > Change Runtime Type > GPU**\n",
-        "\n",
-        "### Set up your Google Cloud project\n",
-        "\n",
-        "**The following steps are required, regardless of your notebook environment.**\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",
-        "2. [Make sure that billing is enabled for your project.](https://cloud.google.com/billing/docs/how-to/modify-project)\n",
-        "\n",
-        "3. [Enable the following APIs: Vertex AI APIs, Compute Engine APIs, and Cloud Storage.](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component,storage-component.googleapis.com)\n",
-        "\n",
-        "4. If you are running this notebook locally, you need to install the [Cloud SDK]((https://cloud.google.com/sdk)).\n",
-        "\n",
-        "5. Enter your project ID in the cell below. Then run the  cell to make sure the\n",
-        "Cloud SDK uses the right project for all the commands in this notebook.\n",
-        "\n",
-        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$`."
+        "## Before you begin"
       ]
     },
     {
       "cell_type": "markdown",
       "metadata": {
-        "id": "project_id"
+        "id": "before_you_begin:nogpu"
       },
       "source": [
-        "#### Set your project ID\n",
+        "### Set your project ID\n",
         "\n",
-        "**If you don't know your project ID**, you may be able to get your project ID using `gcloud`."
+        "**If you don't know your project ID**, try the following:\n",
+        "* Run `gcloud config list`.\n",
+        "* Run `gcloud projects list`.\n",
+        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
       ]
     },
     {
@@ -237,33 +216,10 @@
       },
       "outputs": [],
       "source": [
-        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "autoset_project_id"
-      },
-      "outputs": [],
-      "source": [
-        "if PROJECT_ID == \"\" or PROJECT_ID is None or PROJECT_ID == \"[your-project-id]\":\n",
-        "    # Get your GCP project id from gcloud\n",
-        "    shell_output = ! gcloud config list --format 'value(core.project)' 2>/dev/null\n",
-        "    PROJECT_ID = shell_output[0]\n",
-        "    print(\"Project ID:\", PROJECT_ID)"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "set_gcloud_project_id"
-      },
-      "outputs": [],
-      "source": [
-        "! gcloud config set project $PROJECT_ID"
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
       ]
     },
     {
@@ -274,16 +230,7 @@
       "source": [
         "#### Region\n",
         "\n",
-        "You can also change the `REGION` variable, which is used for operations\n",
-        "throughout the rest of this notebook.  Below are regions supported for Vertex AI. We recommend that you choose the region closest to you.\n",
-        "\n",
-        "- Americas: `us-central1`\n",
-        "- Europe: `europe-west4`\n",
-        "- Asia Pacific: `asia-east1`\n",
-        "\n",
-        "You may not use a multi-regional bucket for training with Vertex AI. Not all regions provide support for all Vertex AI services.\n",
-        "\n",
-        "Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
       ]
     },
     {
@@ -294,34 +241,7 @@
       },
       "outputs": [],
       "source": [
-        "REGION = \"[your-region]\"  # @param {type: \"string\"}\n",
-        "\n",
-        "if REGION == \"[your-region]\":\n",
-        "    REGION = \"us-central1\""
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "timestamp"
-      },
-      "source": [
-        "#### Timestamp\n",
-        "\n",
-        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a timestamp for each instance session, and append the timestamp onto the name of resources you create in this tutorial."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "timestamp"
-      },
-      "outputs": [],
-      "source": [
-        "from datetime import datetime\n",
-        "\n",
-        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")"
+        "REGION = \"us-central1\"  # @param {type: \"string\"}"
       ]
     },
     {
@@ -332,57 +252,70 @@
       "source": [
         "### Authenticate your Google Cloud account\n",
         "\n",
-        "**If you are using Vertex AI Workbench Notebooks**, your environment is already authenticated. Skip this step.\n",
-        "\n",
-        "**If you are using Colab**, run the cell below and follow the instructions when prompted to authenticate your account via oAuth.\n",
-        "\n",
-        "**Otherwise**, follow these steps:\n",
-        "\n",
-        "In the Cloud Console, go to the [Create service account key](https://console.cloud.google.com/apis/credentials/serviceaccountkey) page.\n",
-        "\n",
-        "**Click Create service account**.\n",
-        "\n",
-        "In the **Service account name** field, enter a name, and click **Create**.\n",
-        "\n",
-        "In the **Grant this service account access to project** section, click the Role drop-down list. Type \"Vertex\" into the filter box, and select **Vertex Administrator**. Type \"Storage Object Admin\" into the filter box, and select **Storage Object Admin**.\n",
-        "\n",
-        "Click Create. A JSON file that contains your key downloads to your local environment.\n",
-        "\n",
-        "Enter the path to your service account key as the GOOGLE_APPLICATION_CREDENTIALS variable in the cell below and run the cell."
+        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "FvQeFm3Gv5mR"
+      },
+      "source": [
+        "**1. Vertex AI Workbench**\n",
+        "* Do nothing as you are already authenticated."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ad1138a125ea"
+      },
+      "source": [
+        "**2. Local JupyterLab instance, uncomment and run:**"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "gcp_authenticate"
+        "id": "ce6043da7b33"
+      },
+      "outputs": [],
+      "source": [
+        "# ! gcloud auth login"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0367eac06a10"
+      },
+      "source": [
+        "**3. Colab, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "21ad4dbb4a61"
       },
       "outputs": [],
       "source": [
-        "# If you are running this notebook in Colab, run this cell and follow the\n",
-        "# instructions to authenticate your GCP account. This provides access to your\n",
-        "# Cloud Storage bucket and lets you submit training jobs and prediction\n",
-        "# requests.\n",
-        "\n",
-        "import os\n",
-        "import sys\n",
-        "\n",
-        "# If on Vertex AI Workbench, then don't execute this code\n",
         "IS_COLAB = False\n",
-        "if not os.path.exists(\"/opt/deeplearning/metadata/env_version\") and not os.getenv(\n",
-        "    \"DL_ANACONDA_HOME\"\n",
-        "):\n",
-        "    if \"google.colab\" in sys.modules:\n",
-        "        IS_COLAB = True\n",
-        "        from google.colab import auth as google_auth\n",
-        "\n",
-        "        google_auth.authenticate_user()\n",
-        "\n",
-        "    # If you are running this notebook locally, replace the string below with the\n",
-        "    # path to your service account key and run this cell to authenticate your GCP\n",
-        "    # account.\n",
-        "    elif not os.getenv(\"IS_TESTING\"):\n",
-        "        %env GOOGLE_APPLICATION_CREDENTIALS ''"
+        "# from google.colab import auth\n",
+        "# auth.authenticate_user()\n",
+        "# IS_COLAB = True"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c13224697bfb"
+      },
+      "source": [
+        "**4. Service account or other**\n",
+        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
       ]
     },
     {
@@ -393,11 +326,7 @@
       "source": [
         "### Create a Cloud Storage bucket\n",
         "\n",
-        "**The following steps are required, regardless of your notebook environment.**\n",
-        "\n",
-        "When you initialize the Vertex SDK for Python, you specify a Cloud Storage staging bucket. The staging bucket is where all the data associated with your dataset and model resources are retained across sessions.\n",
-        "\n",
-        "Set the name of your Cloud Storage bucket below. Bucket names must be globally unique across all Google Cloud projects, including those outside of your organization."
+        "Create a storage bucket to store intermediate artifacts such as datasets."
       ]
     },
     {
@@ -408,21 +337,7 @@
       },
       "outputs": [],
       "source": [
-        "BUCKET_NAME = \"[your-bucket-name]\"  # @param {type:\"string\"}\n",
-        "BUCKET_URI = f\"gs://{BUCKET_NAME}\""
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "autoset_bucket"
-      },
-      "outputs": [],
-      "source": [
-        "if BUCKET_NAME == \"\" or BUCKET_NAME is None or BUCKET_NAME == \"[your-bucket-name]\":\n",
-        "    BUCKET_NAME = PROJECT_ID + \"aip-\" + TIMESTAMP\n",
-        "    BUCKET_URI = \"gs://\" + BUCKET_NAME"
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
       ]
     },
     {
@@ -445,26 +360,6 @@
         "! gsutil mb -l $REGION $BUCKET_URI"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "validate_bucket"
-      },
-      "source": [
-        "Finally, validate access to your Cloud Storage bucket by examining its contents:"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "validate_bucket"
-      },
-      "outputs": [],
-      "source": [
-        "! gsutil ls -al $BUCKET_URI"
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -545,19 +440,7 @@
       "source": [
         "### Set up variables\n",
         "\n",
-        "Next, set up some variables used throughout the tutorial.\n",
-        "### Import libraries and define constants"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "import_aip:mbsdk"
-      },
-      "outputs": [],
-      "source": [
-        "import google.cloud.aiplatform as aip"
+        "Next, set up some variables used throughout the tutorial."
       ]
     },
     {
@@ -579,7 +462,10 @@
       },
       "outputs": [],
       "source": [
-        "import google.cloud.aiplatform_v1beta1 as aip_beta"
+        "from datetime import datetime\n",
+        "\n",
+        "import google.cloud.aiplatform_v1beta1 as aip_beta\n",
+        "from google.cloud import aiplatform"
       ]
     },
     {
@@ -707,7 +593,7 @@
       "outputs": [],
       "source": [
         "metadata_store = clients[\"metadata\"].create_metadata_store(\n",
-        "    parent=PARENT, metadata_store_id=\"my-metadata-store\"\n",
+        "    parent=PARENT, metadata_store_id=\"my-metadata-store-unique\"\n",
         ")\n",
         "\n",
         "metadata_store_id = str(metadata_store.result())[7:-2]\n",
@@ -1144,7 +1030,7 @@
       "source": [
         "from kfp.v2 import compiler, dsl\n",
         "from kfp.v2.dsl import (Artifact, Dataset, Input, Metrics, Model, Output,\n",
-        "                        OutputPath, component, pipeline)"
+        "                        OutputPath, component)"
       ]
     },
     {
@@ -1196,7 +1082,7 @@
       "outputs": [],
       "source": [
         "@component(\n",
-        "    packages_to_install=[\"google-cloud-bigquery\", \"pandas\", \"pyarrow\"],\n",
+        "    packages_to_install=[\"google-cloud-bigquery\", \"pandas\", \"pyarrow\", \"db-dtypes\"],\n",
         "    base_image=\"python:3.9\",\n",
         "    output_component_file=\"create_dataset.yaml\",\n",
         ")\n",
@@ -1214,7 +1100,7 @@
         "\n",
         "\n",
         "@component(\n",
-        "    packages_to_install=[\"sklearn\", \"pandas\", \"joblib\"],\n",
+        "    packages_to_install=[\"scikit-learn\", \"pandas\", \"joblib\"],\n",
         "    base_image=\"python:3.9\",\n",
         "    output_component_file=\"beans_model_component.yaml\",\n",
         ")\n",
@@ -1298,7 +1184,7 @@
         "    # A name for the pipeline.\n",
         "    name=\"mlmd-pipeline\",\n",
         ")\n",
-        "def pipeline(\n",
+        "def my_pipeline(\n",
         "    bq_table: str = \"\",\n",
         "    output_data_path: str = \"data.csv\",\n",
         "    project: str = PROJECT_ID,\n",
@@ -1332,20 +1218,22 @@
       "source": [
         "NOW = datetime.now().isoformat().replace(\".\", \":\")[:-7]\n",
         "\n",
-        "compiler.Compiler().compile(pipeline_func=pipeline, package_path=\"mlmd_pipeline.json\")\n",
+        "compiler.Compiler().compile(\n",
+        "    pipeline_func=my_pipeline, package_path=\"mlmd_pipeline.json\"\n",
+        ")\n",
         "\n",
-        "run1 = aip.PipelineJob(\n",
+        "run1 = aiplatform.PipelineJob(\n",
         "    display_name=\"mlmd-pipeline\",\n",
         "    template_path=\"mlmd_pipeline.json\",\n",
-        "    job_id=\"mlmd-pipeline-small-{}\".format(TIMESTAMP),\n",
+        "    job_id=\"mlmd-pipeline-small-unique\",\n",
         "    parameter_values={\"bq_table\": \"sara-vertex-demos.beans_demo.small_dataset\"},\n",
         "    enable_caching=True,\n",
         ")\n",
         "\n",
-        "run2 = aip.PipelineJob(\n",
+        "run2 = aiplatform.PipelineJob(\n",
         "    display_name=\"mlmd-pipeline\",\n",
         "    template_path=\"mlmd_pipeline.json\",\n",
-        "    job_id=\"mlmd-pipeline-large-{}\".format(TIMESTAMP),\n",
+        "    job_id=\"mlmd-pipeline-large-unique\",\n",
         "    parameter_values={\"bq_table\": \"sara-vertex-demos.beans_demo.large_dataset\"},\n",
         "    enable_caching=True,\n",
         ")\n",
@@ -1382,7 +1270,7 @@
       },
       "outputs": [],
       "source": [
-        "df = aip.get_pipeline_df(pipeline=\"mlmd-pipeline\")\n",
+        "df = aiplatform.get_pipeline_df(pipeline=\"mlmd-pipeline\")\n",
         "print(df)"
       ]
     },
@@ -1466,6 +1354,10 @@
       },
       "outputs": [],
       "source": [
+        "metadata_store_id = (\n",
+        "    f\"projects/{PROJECT_ID}/locations/{REGION}/metadataStores/my-metadata-store-unique\"\n",
+        ")\n",
+        "\n",
         "clients[\"metadata\"].delete_metadata_store(name=metadata_store_id)"
       ]
     },

notebooks/community/ml_ops/stage5/get_started_with_vertex_endpoint_and_shared_vm.ipynb

@@ -1028,6 +1028,9 @@
         "deployment_resource_pool.dedicated_resources.min_replica_count = MIN_NODES\n",
         "deployment_resource_pool.dedicated_resources.max_replica_count = MAX_NODES\n",
         "deployment_resource_pool.dedicated_resources.machine_spec.machine_type = DEPLOY_COMPUTE\n",
+        "if DEPLOY_NGPU:\n",
+        "    deployment_resource_pool.dedicated_resources.machine_spec.accelerator_type = DEPLOY_GPU\n",
+        "    deployment_resource_pool.dedicated_resources.machine_spec.accelerator_count = DEPLOY_NGPU\n",
         "\n",
         "request = aip_beta.CreateDeploymentResourcePoolRequest(\n",
         "    parent=f\"projects/{PROJECT_ID}/locations/{REGION}\",\n",

notebooks/community/model_garden/model_garden_huggingface_local_inference.ipynb

@@ -0,0 +1,358 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7d9bbf86da5e"
+      },
+      "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."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4dc4391f6be7"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Hugging Face Local Inference\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_huggingface_local_inference.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_huggingface_local_inference.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_huggingface_local_inference.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>"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4e8a0fdd6f44"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to run local inference with various Hugging Face models by using [Colab](https://colab.research.google.com/) and installing the necessary libraries or by deploying a [Vertex AI Workbench Instance](https://cloud.google.com/vertex-ai-workbench) with preinstalled transformer and diffuser libraries.\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "* Run local inference with various transformer or diffusion models.\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "69453bf7230e"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "68990d91bc5f"
+      },
+      "source": [
+        "### Colab only"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3d342b32fb08"
+      },
+      "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",
+        "    ! pip3 install --upgrade pip\n",
+        "    ! pip3 install torchvision==0.14.1\n",
+        "    ! pip3 install transformers==4.27.1\n",
+        "    ! pip3 install diffusers==0.15.1\n",
+        "    ! apt-get update\n",
+        "    ! apt-get install -y --no-install-recommends tesseract-ocr\n",
+        "    ! pip3 install tesseract==0.1.3\n",
+        "    ! pip3 install pytesseract==0.3.10\n",
+        "    ! pip3 install datasets==2.9.0\n",
+        "    ! pip3 install accelerate==0.18.0\n",
+        "    ! pip3 install triton==2.0.0.dev20221120\n",
+        "    ! pip3 install xformers==0.0.16\n",
+        "    ! pip3 install modelscope==1.4.2\n",
+        "    ! pip3 install open_clip_torch==2.17.1\n",
+        "    ! pip3 install pytorch-lightning==1.9.5\n",
+        "    ! pip3 install opencv-python-headless==4.7.0.72\n",
+        "    # Install gdown for downloading example training images.\n",
+        "    ! pip3 install gdown\n",
+        "    # Remove wrong cublas version.\n",
+        "    ! pip3 uninstall nvidia_cublas_cu11 --yes\n",
+        "\n",
+        "    # Restart the notebook kernel after installs.\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "05e23144b125"
+      },
+      "source": [
+        "### Workbench only\n",
+        "\n",
+        "1.   Follow [this link](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_huggingfacE_local_inference.ipynb) to deploy the notebook to a Vertex AI Workbench Instance.\n",
+        "2.   Select `Create a new Notebook`.\n",
+        "3.   Click `Advanced Options`.\n",
+        "4.   In the **Environment** tab, select `Debian 10` for **Operating System** and select `Custom Container` for **Environment**.\n",
+        "5.   Set the **Docker container image** field to `us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/transformers-notebook`.\n",
+        "6.   In the **Machine Type** tab, select a 1 `T4` GPU and select `Install NVIDIA GPU driver automatically for me`.\n",
+        "7.   Click `Create` to create the Vertex AI Workbench instance.\n"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ad1a690839d5"
+      },
+      "source": [
+        "## Sample code"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0a4008240483"
+      },
+      "source": [
+        "#### [runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5) (Text-to-image)\n",
+        "Generate photo-realistic images given any text input."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "5ec6b474d1be"
+      },
+      "outputs": [],
+      "source": [
+        "import torch\n",
+        "from diffusers import StableDiffusionPipeline\n",
+        "\n",
+        "model_id = \"runwayml/stable-diffusion-v1-5\"\n",
+        "pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)\n",
+        "pipe = pipe.to(\"cuda\")\n",
+        "\n",
+        "prompt = \"a photo of an astronaut riding a horse on mars\"\n",
+        "image = pipe(prompt).images[0]\n",
+        "\n",
+        "display(image)"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ae94b9b23a52"
+      },
+      "source": [
+        "#### [runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5) (Text guided image-to-image)\n",
+        "Generate an image based on an initial image and a text prompt."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0acd70f4d08a"
+      },
+      "outputs": [],
+      "source": [
+        "from io import BytesIO\n",
+        "\n",
+        "import requests\n",
+        "import torch\n",
+        "from diffusers import StableDiffusionImg2ImgPipeline\n",
+        "from PIL import Image\n",
+        "\n",
+        "device = \"cuda\"\n",
+        "model_id_or_path = \"runwayml/stable-diffusion-v1-5\"\n",
+        "pipe = StableDiffusionImg2ImgPipeline.from_pretrained(\n",
+        "    model_id_or_path, torch_dtype=torch.float16\n",
+        ")\n",
+        "pipe = pipe.to(device)\n",
+        "\n",
+        "url = \"https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg\"\n",
+        "\n",
+        "response = requests.get(url)\n",
+        "init_image = Image.open(BytesIO(response.content)).convert(\"RGB\")\n",
+        "init_image = init_image.resize((768, 512))\n",
+        "\n",
+        "prompt = \"A fantasy landscape, trending on artstation\"\n",
+        "\n",
+        "images = pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5).images\n",
+        "display(images[0])"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e76b3fe8d10c"
+      },
+      "source": [
+        "#### [runwayml/stable-diffusion-inpainting](https://huggingface.co/runwayml/stable-diffusion-inpainting) (Image-inpainting)\n",
+        "Generate an image based on an original image and prompt, only editing the areas denoted by a mask image."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "8bc3238be4e7"
+      },
+      "outputs": [],
+      "source": [
+        "from io import BytesIO\n",
+        "\n",
+        "import requests\n",
+        "import torch\n",
+        "from diffusers import StableDiffusionInpaintPipeline\n",
+        "from PIL import Image\n",
+        "\n",
+        "image_url = \"https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png\"\n",
+        "image_response = requests.get(image_url)\n",
+        "init_image = Image.open(BytesIO(image_response.content)).convert(\"RGB\")\n",
+        "display(init_image)\n",
+        "\n",
+        "mask_url = \"https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png\"\n",
+        "mask_response = requests.get(mask_url)\n",
+        "mask_image = Image.open(BytesIO(mask_response.content)).convert(\"RGB\")\n",
+        "\n",
+        "pipe = StableDiffusionInpaintPipeline.from_pretrained(\n",
+        "    \"runwayml/stable-diffusion-inpainting\",\n",
+        "    revision=\"fp16\",\n",
+        "    torch_dtype=torch.float16,\n",
+        ")\n",
+        "pipe.to(\"cuda\")\n",
+        "\n",
+        "prompt = \"Face of a yellow cat, high resolution, sitting on a park bench\"\n",
+        "images = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images\n",
+        "display(images[0])"
+      ]
+    },
+    {
+      "attachments": {},
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "1ade95a9b20e"
+      },
+      "source": [
+        "#### [impira/layoutlm-document-qa](https://huggingface.co/impira/layoutlm-document-qa) (Document question answering)\n",
+        "Answer questions about a given document."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "492d9f1de3f2"
+      },
+      "outputs": [],
+      "source": [
+        "from transformers import pipeline\n",
+        "\n",
+        "nlp = pipeline(\n",
+        "    \"document-question-answering\",\n",
+        "    model=\"impira/layoutlm-document-qa\",\n",
+        ")\n",
+        "\n",
+        "print(\n",
+        "    nlp(\n",
+        "        \"https://templates.invoicehome.com/invoice-template-us-neat-750px.png\",\n",
+        "        \"What is the invoice number?\",\n",
+        "    )\n",
+        ")\n",
+        "# [{'score': 0.9943977, 'answer': 'us-001', 'start': 15, 'end': 15}]\n",
+        "\n",
+        "print(\n",
+        "    nlp(\n",
+        "        \"https://miro.medium.com/max/787/1*iECQRIiOGTmEFLdWkVIH2g.jpeg\",\n",
+        "        \"What is the purchase amount?\",\n",
+        "    )\n",
+        ")\n",
+        "# [{'score': 0.9912159, 'answer': '$1,000,000,000', 'start': 97, 'end': 97}]\n",
+        "\n",
+        "print(\n",
+        "    nlp(\n",
+        "        \"https://www.accountingcoach.com/wp-content/uploads/2013/10/income-statement-example@2x.png\",\n",
+        "        \"What are the 2020 net sales?\",\n",
+        "    )\n",
+        ")\n",
+        "# [{'score': 0.978011429309845, 'answer': '$ 3,980', 'start': 15, 'end': 16}]"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_huggingface_local_inference.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_jax_fvlm.ipynb

@@ -0,0 +1,924 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ur8xi4C7S06n"
+      },
+      "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": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden - JAX F-VLM\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_f_vlm.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/community/model_garden/model_garden_jax_f_vlm.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>                                                                                               <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_f_vlm.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.9"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates serving a [JAX F-VLM model](https://github.com/google-research/google-research/tree/master/fvlm) for [open-vocabulary object detection](https://arxiv.org/abs/2209.15639) task and deploying them on Vertex AI for online prediction."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d975e698c9a4"
+      },
+      "source": [
+        "### Objective\n",
+        "\n",
+        "In this tutorial, you learn how to:\n",
+        "\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 image classification.\n",
+        "\n",
+        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "\n",
+        "- Vertex AI Model Registry\n",
+        "- Vertex AI Online Prediction"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "08d289fa873f"
+      },
+      "source": [
+        "### Dataset\n",
+        "\n",
+        "This notebook uses the following prediction image as an example:\n",
+        "\n",
+        "Image: https://pixabay.com/nl/photos/het-fruit-eten-citroen-limoen-3134631/\n",
+        "\n",
+        "Creative Commons License: https://pixabay.com/nl/service/terms/\n",
+        "\n",
+        "You can use your own custom prediction image as well as by modifying the `DEMO_IMAGE_PATH` variable in this notebook below."
+      ]
+    },
+    {
+      "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": "i7EUnXsZhAGF"
+      },
+      "source": [
+        "## Installation\n",
+        "\n",
+        "Install the following packages required to execute this notebook."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2b4ef9b72d43"
+      },
+      "outputs": [],
+      "source": [
+        "# Install the packages.\n",
+        "! pip3 install --upgrade google-cloud-aiplatform\n",
+        "# Get F-VLM repository by using svn to avoid downloading entire google-research repository.\n",
+        "! apt install subversion\n",
+        "! rm -rf ./fvlm\n",
+        "! svn export -r 59152 https://github.com/google-research/google-research/trunk/fvlm\n",
+        "# Note: The following libraries are pinned down versions of:\n",
+        "# https://github.com/google-research/google-research/blob/master/fvlm/requirements.txt\n",
+        "! pip3 install tensorflow==2.12.0\n",
+        "! pip3 install numpy==1.23.5\n",
+        "! pip3 install jax==0.4.14\n",
+        "! pip3 install jaxlib==0.4.14+cuda11.cudnn86\n",
+        "! pip3 install flax==0.7.1\n",
+        "! pip3 install torch==2.0.1+cu118\n",
+        "! pip3 install torchvision==0.15.2+cu118\n",
+        "! pip3 install opencv-python==4.7.0.72\n",
+        "! pip3 install tqdm==4.65.0\n",
+        "! pip3 install git+https://github.com/openai/CLIP.git@a1d071733d7111c9c014f024669f959182114e33\n",
+        "! pip3 install Pillow==9.5.0"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "_9q83As4G2Yn"
+      },
+      "source": [
+        "Download the F-VLM checkpoints into the `fvlm/checkpoints` folder."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "eXVI6s57FPyg"
+      },
+      "outputs": [],
+      "source": [
+        "%cd fvlm/checkpoints\n",
+        "! ./download.sh\n",
+        "%cd ../../"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "58707a750154"
+      },
+      "source": [
+        "### Colab only"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "f200f10a1da3"
+      },
+      "outputs": [],
+      "source": [
+        "# Automatically restart kernel after installs so that your environment can access the new packages.\n",
+        "import IPython\n",
+        "\n",
+        "app = IPython.Application.instance()\n",
+        "app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BF1j6f9HApxa"
+      },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "### Set up your Google Cloud project\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\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. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk).\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WReHDGG5g0XY"
+      },
+      "source": [
+        "#### Set your project ID\n",
+        "\n",
+        "**If you don't know your project ID**, try the following:\n",
+        "* Run `gcloud config list`.\n",
+        "* Run `gcloud projects list`.\n",
+        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "oM1iC_MfAts1"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "region"
+      },
+      "source": [
+        "#### Region\n",
+        "\n",
+        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "twgKk-LsLmX3"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"us-central1\"  # @param {type: \"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "sBCra4QMA2wR"
+      },
+      "source": [
+        "### Authenticate your Google Cloud account\n",
+        "\n",
+        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "74ccc9e52986"
+      },
+      "source": [
+        "**1. Vertex AI Workbench**\n",
+        "* Do nothing as you are already authenticated."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "de775a3773ba"
+      },
+      "source": [
+        "**2. Local JupyterLab instance, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "254614fa0c46"
+      },
+      "outputs": [],
+      "source": [
+        "# ! gcloud auth login"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ef21552ccea8"
+      },
+      "source": [
+        "**3. Colab, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "603adbbf0532"
+      },
+      "outputs": [],
+      "source": [
+        "# from google.colab import auth\n",
+        "# auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f6b2ccc891ed"
+      },
+      "source": [
+        "**4. Service account or other**\n",
+        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Create a Cloud Storage bucket\n",
+        "\n",
+        "Create a storage bucket to store intermediate artifacts such as datasets."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MzGDU7TWdts_"
+      },
+      "outputs": [],
+      "source": [
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-EcIXiGsCePi"
+      },
+      "source": [
+        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NIq7R4HZCfIc"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "960505627ddf"
+      },
+      "source": [
+        "### Import libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "PyQmSRbKA8r-"
+      },
+      "outputs": [],
+      "source": [
+        "import base64\n",
+        "import functools\n",
+        "import os\n",
+        "import sys\n",
+        "from io import BytesIO\n",
+        "\n",
+        "import jax\n",
+        "import numpy as np\n",
+        "import tensorflow as tf\n",
+        "import tqdm\n",
+        "from PIL import Image\n",
+        "\n",
+        "sys.path.append(\"./fvlm\")\n",
+        "import inputs\n",
+        "import jax_clip\n",
+        "import utils\n",
+        "from google.cloud import aiplatform\n",
+        "from google.protobuf import json_format\n",
+        "from google.protobuf.struct_pb2 import Value"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "init_aip:mbsdk,all"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "Initialize the Vertex AI SDK for Python for your project."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "vS1hQiGuLmX4"
+      },
+      "outputs": [],
+      "source": [
+        "staging_bucket = os.path.join(BUCKET_URI, \"jax_fvlm_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 prediction docker image.\n",
+        "OPTIMIZED_TF_RUNTIME_IMAGE_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-cpu.nightly:latest\"\n",
+        ")\n",
+        "# The local path to the F-VLM folder.\n",
+        "F_VLM_FOLDER = \"./fvlm\"\n",
+        "# The F-VLM model to use. Choose between 'resnet_50', 'resnet_50x4', or 'resnet_50x16'.\n",
+        "MODEL = \"resnet_50\"\n",
+        "# The list of object categories to detect. For example: \"person, car, oven\".\n",
+        "CATEGORIES = [\n",
+        "    \"kiwi\",\n",
+        "    \"orange\",\n",
+        "    \"lemon\",\n",
+        "    \"blackberry\",\n",
+        "    \"pine cone\",\n",
+        "    \"red orange\",\n",
+        "    \"table\",\n",
+        "    \"spoon\",\n",
+        "    \"pine needles\",\n",
+        "    \"seed\",\n",
+        "]\n",
+        "# An upper bound on the number of classes.\n",
+        "MAX_NUM_CLS = 91\n",
+        "# The max number of boxes to draw on the output image.\n",
+        "MAX_BOXES_TO_DRAW = 25\n",
+        "# The minimum score required to draw a detected object.\n",
+        "MIN_SCORE_THRESH = 0.2  # @param {type:\"slider\", min:0, max:0.9, step:0.05}\n",
+        "# The local path to the output image.\n",
+        "OUTPUT_IMAGE_PATH = \"./output.jpg\"\n",
+        "# The original F-VLM SavedModel folder which takes image and text embeddings as inputs.\n",
+        "SAVED_MODEL_DIR = f'{F_VLM_FOLDER}/checkpoints/{MODEL.replace(\"resnet_\",\"r\")}'\n",
+        "# The converted SavedModel folder which takes jpeg bytes and text-embeddings bytes as inputs.\n",
+        "CONVERTED_SAVED_MODEL_DIR = \"./converted_saved_model\"\n",
+        "# The Cloud Storage location for the converted SavedModel.\n",
+        "GCS_CONVERTED_SAVED_MODEL_DIR = f\"{BUCKET_URI}/fvlm_saved_model\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0c250872074f"
+      },
+      "source": [
+        "### Define common functions\n",
+        "\n",
+        "This section defines functions for:\n",
+        "\n",
+        "- Loading and converting input image into the required prediction format.\n",
+        "- Visualization of detection outputs."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "XcYUGwr-AJGY"
+      },
+      "outputs": [],
+      "source": [
+        "def convert_numpy_array_to_byte_string_via_tf_tensor(np_array):\n",
+        "    \"\"\"Serializes a numpy array to tensor bytes.\"\"\"\n",
+        "    tensor_array = tf.convert_to_tensor(np_array)\n",
+        "    tensor_byte_string = tf.io.serialize_tensor(tensor_array)\n",
+        "    return tensor_byte_string.numpy()\n",
+        "\n",
+        "\n",
+        "def generate_text_embeddings(categories):\n",
+        "    \"\"\"Generates text embeddings in numpy format from object categories.\"\"\"\n",
+        "    clip_text_fn = jax_clip.get_clip_text_fn(MODEL)\n",
+        "    class_clip_features = []\n",
+        "    print(\"Computing custom category text embeddings.\")\n",
+        "    for cls_name in tqdm.tqdm(categories, total=len(categories)):\n",
+        "        cls_feat = clip_text_fn(cls_name)\n",
+        "        class_clip_features.append(cls_feat)\n",
+        "    text_embeddings = np.concatenate(class_clip_features, axis=0)\n",
+        "    embed_path = (\n",
+        "        f'{F_VLM_FOLDER}/data/{MODEL.replace(\"resnet_\", \"r\")}_bg_empty_embed.npy'\n",
+        "    )\n",
+        "    background_embedding, empty_embeddings = np.load(embed_path)\n",
+        "    background_embedding = background_embedding[np.newaxis, Ellipsis]\n",
+        "    empty_embeddings = empty_embeddings[np.newaxis, Ellipsis]\n",
+        "    tile_empty_embeddings = np.tile(\n",
+        "        empty_embeddings, (MAX_NUM_CLS - len(categories) - 1, 1)\n",
+        "    )\n",
+        "    # Concatenate 'background' and 'empty' embeddings.\n",
+        "    text_embeddings = np.concatenate(\n",
+        "        (background_embedding, text_embeddings, tile_empty_embeddings), axis=0\n",
+        "    )\n",
+        "    return text_embeddings\n",
+        "\n",
+        "\n",
+        "def get_jpeg_bytes(local_image_path, new_width=-1):\n",
+        "    \"\"\"Returns jpeg bytes given an image path and resizes if required.\"\"\"\n",
+        "    image = Image.open(local_image_path)\n",
+        "    if new_width <= 0:\n",
+        "        new_image = image\n",
+        "    else:\n",
+        "        width, height = image.size\n",
+        "        print(\"original input image size: \", width, \" , \", height)\n",
+        "        new_height = int(height * new_width / width)\n",
+        "        print(\"new input image size: \", new_width, \" , \", new_height)\n",
+        "        new_image = image.resize((new_width, new_height))\n",
+        "    buffered = BytesIO()\n",
+        "    new_image.save(buffered, format=\"JPEG\")\n",
+        "    return buffered.getvalue()\n",
+        "\n",
+        "\n",
+        "def generate_prediction_output_image(\n",
+        "    input_image_path, prediction_output, output_image_path\n",
+        "):\n",
+        "    \"\"\"Generates prediction output image with detected objects and bounding boxes.\"\"\"\n",
+        "    # Generate tensors from prediction outputs.\n",
+        "    prediction_output_tensor = {}\n",
+        "    for key, val in prediction_output.items():\n",
+        "        prediction_output_tensor[key] = tf.expand_dims(\n",
+        "            tf.convert_to_tensor(val), axis=0\n",
+        "        )\n",
+        "    prediction_output_tensor[\"num_detections\"] = tf.cast(\n",
+        "        prediction_output_tensor[\"num_detections\"], tf.int32\n",
+        "    )\n",
+        "    # Generate image embeddings for the input image.\n",
+        "    with open(input_image_path, \"rb\") as f:\n",
+        "        np_image = np.array(Image.open(f))\n",
+        "    parser_fn = inputs.get_maskrcnn_parser()\n",
+        "    data = parser_fn({\"image\": np_image, \"source_id\": np.array([0])})\n",
+        "    np_data = jax.tree_map(lambda x: x.numpy()[np.newaxis, Ellipsis], data)\n",
+        "    image_embeddings = np_data.pop(\"images\")\n",
+        "    labels = np_data.pop(\"labels\")\n",
+        "    # Generate visualization.\n",
+        "    print(\"Preparing visualization.\")\n",
+        "    categories = CATEGORIES\n",
+        "    id_mapping = {(i + 1): c for i, c in enumerate(categories)}\n",
+        "    id_mapping[0] = \"background\"\n",
+        "    for k in range(len(categories) + 2, MAX_NUM_CLS):\n",
+        "        id_mapping[k] = \"empty\"\n",
+        "    category_index = inputs.get_category_index(id_mapping)\n",
+        "    maskrcnn_visualizer_fn = functools.partial(\n",
+        "        utils.visualize_boxes_and_labels_on_image_array,\n",
+        "        category_index=category_index,\n",
+        "        use_normalized_coordinates=False,\n",
+        "        max_boxes_to_draw=MAX_BOXES_TO_DRAW,\n",
+        "        min_score_thresh=MIN_SCORE_THRESH,\n",
+        "        skip_labels=False,\n",
+        "    )\n",
+        "    vis_image = utils.visualize_instance_segmentations(\n",
+        "        prediction_output_tensor,\n",
+        "        image_embeddings,\n",
+        "        labels[\"image_info\"],\n",
+        "        maskrcnn_visualizer_fn,\n",
+        "    )\n",
+        "    pil_vis_image = Image.fromarray(vis_image, mode=\"RGB\")\n",
+        "    pil_vis_image.save(output_image_path)\n",
+        "    print(\"Completed saving the output image at: \", output_image_path)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ayNrua2txk0B"
+      },
+      "source": [
+        "# Convert F-VLM SavedModel to support smaller input size\n",
+        "\n",
+        "The F-VLM SavedModel takes image embeddings and text embeddings as input. But you can not send these inputs directly for Vertex AI Online Prediction because there is a limit of 1.5 MB on the prediction request size. So you will first convert the SavedModel format to take jpeg bytes and text-embeddings bytes as an input instead. This modified input format will meet the 1.5 MB limit requirement."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ef-svu2Ix1OA"
+      },
+      "outputs": [],
+      "source": [
+        "def preprocess_jpeg_byte_string(tensor_byte_string):\n",
+        "    \"\"\"Converts jpeg bytes to image embeddings as an input for the original F-VLM SavedModel.\"\"\"\n",
+        "    decoded_image_tensor = tf.io.decode_jpeg(tensor_byte_string, channels=3)\n",
+        "    parser_fn = inputs.get_maskrcnn_parser()\n",
+        "    parser_output = parser_fn({\"image\": decoded_image_tensor})\n",
+        "    image_embeddings_tensor = parser_output[\"images\"]\n",
+        "    return image_embeddings_tensor\n",
+        "\n",
+        "\n",
+        "def preprocess_text_embeddings_byte_string(tensor_byte_string):\n",
+        "    \"\"\"Converts text-embeddings bytes to text-embeddings as an input for the original F-VLM SavedModel.\"\"\"\n",
+        "    return tf.io.parse_tensor(tensor_byte_string, tf.float32)\n",
+        "\n",
+        "\n",
+        "def get_serve_fn(model):\n",
+        "    \"\"\"Creates a serving function for the modified SavedModel which takes jpeg bytes and text-embeddings bytes as an input.\"\"\"\n",
+        "\n",
+        "    @tf.function(\n",
+        "        input_signature=[\n",
+        "            tf.TensorSpec([None], tf.string),\n",
+        "            tf.TensorSpec([None], tf.string),\n",
+        "        ]\n",
+        "    )\n",
+        "    def serve_fn(image_jpeg_bytes_inputs, text_embeddings_bytes_inputs):\n",
+        "        image_embeddings_tensor = tf.map_fn(\n",
+        "            preprocess_jpeg_byte_string, image_jpeg_bytes_inputs, dtype=tf.bfloat16\n",
+        "        )\n",
+        "        text_embeddings_tensor = tf.map_fn(\n",
+        "            preprocess_text_embeddings_byte_string,\n",
+        "            text_embeddings_bytes_inputs,\n",
+        "            dtype=tf.float32,\n",
+        "        )\n",
+        "        return model({\"image\": image_embeddings_tensor, \"text\": text_embeddings_tensor})\n",
+        "\n",
+        "    return serve_fn\n",
+        "\n",
+        "\n",
+        "! rm -rf {CONVERTED_SAVED_MODEL_DIR}\n",
+        "model = tf.saved_model.load(SAVED_MODEL_DIR)\n",
+        "signatures = {\n",
+        "    \"serving_default\": get_serve_fn(model=model).get_concrete_function(\n",
+        "        tf.TensorSpec(shape=[None], dtype=tf.string), tf.TensorSpec([None], tf.string)\n",
+        "    )\n",
+        "}\n",
+        "tf.saved_model.save(model, CONVERTED_SAVED_MODEL_DIR, signatures=signatures)\n",
+        "print(\"Saved the converted SavedModel to directory: \", CONVERTED_SAVED_MODEL_DIR)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "1cJQEETi1jsg"
+      },
+      "source": [
+        "Copy the local converted TF SavedModel to Cloud Storage."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "6hlTWKxh11dF"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil -m rm -R -f {GCS_CONVERTED_SAVED_MODEL_DIR}\n",
+        "! gsutil -m cp -R {CONVERTED_SAVED_MODEL_DIR} {GCS_CONVERTED_SAVED_MODEL_DIR}\n",
+        "! gsutil ls {GCS_CONVERTED_SAVED_MODEL_DIR}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "iILhhP3TfO8B"
+      },
+      "source": [
+        "## Run online prediction\n",
+        "Run online prediction with the converted TF SavedModel."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ExIyCnKf3a94"
+      },
+      "source": [
+        "Upload TF SavedModel and deploy it to an endpoint for prediction. This step can take up to 15 minutes to finish."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "t0xYDT0BxP0W"
+      },
+      "outputs": [],
+      "source": [
+        "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",
+        ")\n",
+        "\n",
+        "jax_fvlm_endpoint = jax_fvlm_model.deploy(\n",
+        "    deployed_model_display_name=\"jax_vlm_deployed\",\n",
+        "    traffic_split={\"0\": 100},\n",
+        "    machine_type=\"n1-highmem-16\",\n",
+        "    min_replica_count=1,\n",
+        "    max_replica_count=1,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "w99wNhz_3ruV"
+      },
+      "source": [
+        "Prepare input prediction image.\n",
+        "\n",
+        "Note: You can modify the input image as required."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0Itg0k1s30t3"
+      },
+      "outputs": [],
+      "source": [
+        "# Local path to the prediction image.\n",
+        "DEMO_IMAGE_PATH = \"./prediction_image.jpg\"\n",
+        "# Download the prediction image.\n",
+        "! wget -O {DEMO_IMAGE_PATH} https://cdn.pixabay.com/photo/2018/02/06/12/37/fruit-3134631_1280.jpg"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "B1Q7AbmJ4QxZ"
+      },
+      "source": [
+        "Prepare jpeg bytes and text-embeddings bytes inputs."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "qxj4Xv_DhHXj"
+      },
+      "outputs": [],
+      "source": [
+        "image_jpeg_bytes_inputs = get_jpeg_bytes(\n",
+        "    local_image_path=DEMO_IMAGE_PATH, new_width=1024\n",
+        ")\n",
+        "text_embeddings = generate_text_embeddings(categories=CATEGORIES)\n",
+        "text_embeddings_bytes_inputs = convert_numpy_array_to_byte_string_via_tf_tensor(\n",
+        "    text_embeddings\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Ys88lEkK4XDp"
+      },
+      "source": [
+        "Use base-64 encoding followed by UTF-8 decoding to package the bytes inputs and then send them to the endpoint for prediction. The Vertex AI Prediction service will automatically convert these input strings back to bytes based on the `b64` keyword.\n",
+        "\n",
+        "**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",
+        "The subsequent predictions take 4 seconds to finish."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Mj4sqTAG4sU5"
+      },
+      "outputs": [],
+      "source": [
+        "instances_list = [\n",
+        "    {\n",
+        "        \"image_jpeg_bytes_inputs\": {\n",
+        "            \"b64\": base64.b64encode(image_jpeg_bytes_inputs).decode(\"utf-8\")\n",
+        "        },\n",
+        "        \"text_embeddings_bytes_inputs\": {\n",
+        "            \"b64\": base64.b64encode(text_embeddings_bytes_inputs).decode(\"utf-8\")\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "instances = [json_format.ParseDict(s, Value()) for s in instances_list]\n",
+        "prediction_output = jax_fvlm_endpoint.predict(instances=instances).predictions[0]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "S3qC-MrN6SEs"
+      },
+      "source": [
+        "Generate output image with predicted bounding boxes, labels, and probabilities. The output image will be saved to `OUTPUT_IMAGE_PATH`."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "wnPY2MFN6fL6"
+      },
+      "outputs": [],
+      "source": [
+        "generate_prediction_output_image(\n",
+        "    input_image_path=DEMO_IMAGE_PATH,\n",
+        "    prediction_output=prediction_output,\n",
+        "    output_image_path=OUTPUT_IMAGE_PATH,\n",
+        ")"
+      ]
+    },
+    {
+      "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\n",
+        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
+        "\n",
+        "Otherwise, you can delete the individual resources you created in this tutorial:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "sx_vKniMq9ZX"
+      },
+      "outputs": [],
+      "source": [
+        "# Delete endpoint resource.\n",
+        "jax_fvlm_endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete model resource.\n",
+        "jax_fvlm_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"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_jax_fvlm.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

@@ -0,0 +1,879 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ur8xi4C7S06n"
+      },
+      "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": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden - JAX Vision Transformer\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_vision_transformer.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/community/model_garden/model_garden_jax_vision_transformer.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_vision_transformer.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.9"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates finetuning a [JAX ViT-B16 model](https://github.com/google-research/vision_transformer#available-vit-models) for image classification task on GPU and deploying them on Vertex AI for online prediction.\n",
+        "\n",
+        "Learn more about [Generative AI Support in Vertex AI](https://cloud.google.com/blog/products/ai-machine-learning/vertex-ai-model-garden-and-generative-ai-studio)."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d975e698c9a4"
+      },
+      "source": [
+        "### Objective\n",
+        "\n",
+        "In this tutorial, you learn how fine-tune, deploy and predict with a Vertex AI pretrained JAX Vision Transformer based model.\n",
+        "\n",
+        "This tutorial uses the following Google Cloud ML services and resources:\n",
+        "\n",
+        "- Vertex AI Model Garden\n",
+        "- Vertex AI Training\n",
+        "- Vertex AI Model Registry\n",
+        "- Vertex AI Online Prediction\n",
+        "\n",
+        "The steps performed are:\n",
+        "\n",
+        "- Finetune a JAX Vision Transformer based 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 image classification.\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "08d289fa873f"
+      },
+      "source": [
+        "### Dataset\n",
+        "\n",
+        "This notebook uses the [tf_flowers dataset](https://www.tensorflow.org/datasets/catalog/tf_flowers) and has a section which shows how to download and prepare it. You can follow similar process to use your own custom dataset too."
+      ]
+    },
+    {
+      "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": "i7EUnXsZhAGF"
+      },
+      "source": [
+        "## Installation\n",
+        "\n",
+        "Install the following packages required to execute this notebook."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2b4ef9b72d43"
+      },
+      "outputs": [],
+      "source": [
+        "# Install the packages.\n",
+        "! pip3 install --upgrade google-cloud-aiplatform"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "58707a750154"
+      },
+      "source": [
+        "### Colab only"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "f200f10a1da3"
+      },
+      "outputs": [],
+      "source": [
+        "# Automatically restart kernel after installs so that your environment can access the new packages.\n",
+        "# import IPython\n",
+        "\n",
+        "# app = IPython.Application.instance()\n",
+        "# app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "BF1j6f9HApxa"
+      },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "### Set up your Google Cloud project\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\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. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk).\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WReHDGG5g0XY"
+      },
+      "source": [
+        "#### Set your project ID\n",
+        "\n",
+        "**If you don't know your project ID**, try the following:\n",
+        "* Run `gcloud config list`.\n",
+        "* Run `gcloud projects list`.\n",
+        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "oM1iC_MfAts1"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Set the project id\n",
+        "! gcloud config set project {PROJECT_ID}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "region"
+      },
+      "source": [
+        "#### Region\n",
+        "\n",
+        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "twgKk-LsLmX3"
+      },
+      "outputs": [],
+      "source": [
+        "REGION = \"us-central1\"  # @param {type: \"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "sBCra4QMA2wR"
+      },
+      "source": [
+        "### Authenticate your Google Cloud account\n",
+        "\n",
+        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "74ccc9e52986"
+      },
+      "source": [
+        "**1. Vertex AI Workbench**\n",
+        "* Do nothing as you are already authenticated."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "de775a3773ba"
+      },
+      "source": [
+        "**2. Local JupyterLab instance, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "254614fa0c46"
+      },
+      "outputs": [],
+      "source": [
+        "# ! gcloud auth login"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ef21552ccea8"
+      },
+      "source": [
+        "**3. Colab, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "603adbbf0532"
+      },
+      "outputs": [],
+      "source": [
+        "# from google.colab import auth\n",
+        "# auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f6b2ccc891ed"
+      },
+      "source": [
+        "**4. Service account or other**\n",
+        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "zgPO1eR3CYjk"
+      },
+      "source": [
+        "### Create a Cloud Storage bucket\n",
+        "\n",
+        "Create a storage bucket to store intermediate artifacts such as datasets."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MzGDU7TWdts_"
+      },
+      "outputs": [],
+      "source": [
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-EcIXiGsCePi"
+      },
+      "source": [
+        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "NIq7R4HZCfIc"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "960505627ddf"
+      },
+      "source": [
+        "### Import libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "PyQmSRbKA8r-"
+      },
+      "outputs": [],
+      "source": [
+        "import base64\n",
+        "import glob\n",
+        "import os\n",
+        "import random\n",
+        "import shutil\n",
+        "from datetime import datetime\n",
+        "from io import BytesIO\n",
+        "\n",
+        "import numpy as np\n",
+        "from google.cloud import aiplatform\n",
+        "from google.protobuf import json_format\n",
+        "from google.protobuf.struct_pb2 import Value\n",
+        "from PIL import Image"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "init_aip:mbsdk,all"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "Initialize the Vertex AI SDK for Python for your project."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "vS1hQiGuLmX4"
+      },
+      "outputs": [],
+      "source": [
+        "staging_bucket = os.path.join(BUCKET_URI, \"jax_vit_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-restricted/vertex-vision-model-garden-dockers/jax-vit-train-gpu\"\n",
+        "# The pre-built TF SavedModel conversion docker image.\n",
+        "MODEL_CONVERSION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/jax-vit-model-conversion\"\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",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0c250872074f"
+      },
+      "source": [
+        "### Define common functions\n",
+        "\n",
+        "This section defines functions for:\n",
+        "\n",
+        "- Splitting the [tf_flowers dataset](https://www.tensorflow.org/datasets/catalog/tf_flowers) images into `train` and `test` folders.\n",
+        "- Converting a Cloud Storage path such as `gs://bucket-name` to GCSFuse path format such as `/gcsfuse/bucket-name`.\n",
+        "- Encoding a local image file to a string for prediction input."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "XcYUGwr-AJGY"
+      },
+      "outputs": [],
+      "source": [
+        "def split(base_dir, test_ratio=0.1):\n",
+        "    \"\"\"Splits images and moves them to train and test folders.\"\"\"\n",
+        "    paths = glob.glob(f\"{base_dir}/*/*.jpg\")\n",
+        "    random.shuffle(paths)\n",
+        "    counts = dict(test=0, train=0)\n",
+        "    for i, path in enumerate(paths):\n",
+        "        split = \"test\" if i < test_ratio * len(paths) else \"train\"\n",
+        "        *_, class_name, basename = path.split(\"/\")\n",
+        "        dst = f\"{base_dir}/{split}/{class_name}/{basename}\"\n",
+        "        if not os.path.isdir(os.path.dirname(dst)):\n",
+        "            os.makedirs(os.path.dirname(dst))\n",
+        "        shutil.move(path, dst)\n",
+        "        counts[split] += 1\n",
+        "    print(f'Moved {counts[\"train\"]:,} train and {counts[\"test\"]:,} test images.')\n",
+        "\n",
+        "\n",
+        "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 load_bytes_from_local_image(local_image_path, new_width=-1):\n",
+        "    \"\"\"Returns encoded image string for prediction input.\"\"\"\n",
+        "    image = Image.open(local_image_path)\n",
+        "    if new_width <= 0:\n",
+        "        new_image = image\n",
+        "    else:\n",
+        "        width, height = image.size\n",
+        "        print(\"original input image size: \", width, \" , \", height)\n",
+        "        new_height = int(height * new_width / width)\n",
+        "        print(\"new input image size: \", new_width, \" , \", new_height)\n",
+        "        new_image = image.resize((new_width, new_height))\n",
+        "    buffered = BytesIO()\n",
+        "    new_image.save(buffered, format=\"JPEG\")\n",
+        "    encoded_string = base64.b64encode(buffered.getvalue()).decode(\"utf-8\")\n",
+        "    return encoded_string"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "X0wWrfhDC8ni"
+      },
+      "source": [
+        "### Prepare dataset\n",
+        "\n",
+        "If you are not using [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview#all_datasets), then you need to prepare your dataset and store it on Cloud Storage. The following example shows\n",
+        "how to do this for the [tf_flowers dataset](https://www.tensorflow.org/datasets/catalog/tf_flowers). If using TensorFlow Datasets, you pass\n",
+        "the dataset name such as `tf_flowers` to the `--config.dataset` flag and bypass this section."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "LW31Ws1RN9AC"
+      },
+      "outputs": [],
+      "source": [
+        "local_flower_data_directory = \"./flower_photos\"  # @param {type:\"string\"}\n",
+        "FLOWER_DATA_GCS_PATH = os.path.join(BUCKET_URI, \"flower_dataset\")\n",
+        "# The flower dataset has 5 classes.\n",
+        "NUM_CLASSES = 5\n",
+        "# NOTE: For custom dataset, the training code picks the class names\n",
+        "# from the folder structure and then sorts them to create a mapping\n",
+        "# from class-index to class-name. This is why the mapping below\n",
+        "# looks different from default `tf_flowers` documentation.\n",
+        "LABEL_IDX_TO_STR = {\n",
+        "    0: \"daisy\",\n",
+        "    1: \"dandelion\",\n",
+        "    2: \"roses\",\n",
+        "    3: \"sunflowers\",\n",
+        "    4: \"tulips\",\n",
+        "}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "heMhYO-DD4II"
+      },
+      "outputs": [],
+      "source": [
+        "# Download flower data to a local directory.\n",
+        "! rm -rf $local_flower_data_directory;\n",
+        "! (cd \"./\" && curl https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz | tar xz)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YtWxe2y8Gqzl"
+      },
+      "outputs": [],
+      "source": [
+        "# Since the default file format of above \"tf_flowers\" dataset is\n",
+        "# flower_photos/{class_name}/{filename}.jpg\n",
+        "# we first need to split it into a \"train\" (90%) and a \"test\" (10%) set:\n",
+        "# flower_photos/train/{class_name}/{filename}.jpg\n",
+        "# flower_photos/test/{class_name}/{filename}.jpg\n",
+        "\n",
+        "split(local_flower_data_directory)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "g043ydQ_wlpk"
+      },
+      "outputs": [],
+      "source": [
+        "# Move Flower data from local directory to Cloud Storage.\n",
+        "# This step takes around 2 mins to finish.\n",
+        "! gsutil -m cp -R $local_flower_data_directory/train/* $FLOWER_DATA_GCS_PATH/train/\n",
+        "! gsutil -m cp -R $local_flower_data_directory/test/* $FLOWER_DATA_GCS_PATH/test/"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aCpLmWPMpJQ8"
+      },
+      "source": [
+        "## Finetune with JAX Vision Transformer\n",
+        "\n",
+        "Create and run the training job with the model-garden JAX vision transformer training docker using the Vertex AI SDK. The training uses one V100 GPU and runs for around 10 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 = \"jax_vision_transformer\" + TIMESTAMP\n",
+        "\n",
+        "finetuning_workdir = os.path.join(BUCKET_URI, JOB_NAME)\n",
+        "pre_trained_dir = \"gs://vit_models/imagenet21k\"\n",
+        "docker_args_list = [\n",
+        "    \"--config\",\n",
+        "    \"vit_jax/configs/vit.py:b16\",\n",
+        "    \"--config.dataset\",\n",
+        "    f\"{gcs_fuse_path(FLOWER_DATA_GCS_PATH)}\",\n",
+        "    \"--config.pp.train\",\n",
+        "    \"train\",\n",
+        "    \"--config.pp.test\",\n",
+        "    \"test\",\n",
+        "    \"--config.pretrained_dir\",\n",
+        "    f\"{gcs_fuse_path(pre_trained_dir)}\",\n",
+        "    \"--config.batch\",\n",
+        "    \"128\",\n",
+        "    \"--config.batch_eval\",\n",
+        "    \"128\",\n",
+        "    \"--config.base_lr\",\n",
+        "    \"0.01\",\n",
+        "    \"--config.shuffle_buffer\",\n",
+        "    \"1000\",\n",
+        "    \"--config.total_steps\",\n",
+        "    \"100\",\n",
+        "    \"--config.warmup_steps\",\n",
+        "    \"10\",\n",
+        "    \"--config.pp.crop\",\n",
+        "    \"224\",\n",
+        "    \"--workdir\",\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",
+        "NUM_GPU = 1\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-standard-4\",\n",
+        "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
+        "    accelerator_count=NUM_GPU,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-2qiROKIONnI"
+      },
+      "source": [
+        "## Convert JAX Vision Transformer model to TF SavedModel\n",
+        "\n",
+        "Convert the previously fine-tuned JAX model to a TF SavedModel for online prediction."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "6Y7slAFLOwlV"
+      },
+      "outputs": [],
+      "source": [
+        "# Set up model conversion docker arguments.\n",
+        "# Note: Many of the arguments below are similar to the training job\n",
+        "# such as the model name and train and test data related parameters.\n",
+        "\n",
+        "jax_checkpoint_dir = finetuning_workdir\n",
+        "\n",
+        "TIMESTAMP = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
+        "JOB_NAME = \"jax_model_conversion\" + TIMESTAMP\n",
+        "saved_model_dir = os.path.join(BUCKET_URI, \"jax2tf_\" + TIMESTAMP)\n",
+        "\n",
+        "docker_args_list = [\n",
+        "    \"--config\",\n",
+        "    \"vit_jax/configs/vit.py:b16\",\n",
+        "    \"--num_classes\",\n",
+        "    f\"{NUM_CLASSES}\",\n",
+        "    \"--saved_model_dir\",\n",
+        "    f\"{saved_model_dir}\",\n",
+        "    \"--jax_checkpoint_dir\",\n",
+        "    f\"{jax_checkpoint_dir}\",\n",
+        "    \"--config.pretrained_dir\",\n",
+        "    f\"{pre_trained_dir}\",\n",
+        "    \"--config.dataset\",\n",
+        "    f\"{gcs_fuse_path(FLOWER_DATA_GCS_PATH)}\",\n",
+        "    \"--config.pp.train\",\n",
+        "    \"train\",\n",
+        "    \"--config.pp.test\",\n",
+        "    \"test\",\n",
+        "    \"--config.pp.crop\",\n",
+        "    \"224\",\n",
+        "]\n",
+        "print(docker_args_list)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0Acfh1VWUsTL"
+      },
+      "outputs": [],
+      "source": [
+        "# Create and run the model conversion 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 = MODEL_CONVERSION_DOCKER_URI\n",
+        "job = aiplatform.CustomContainerTrainingJob(\n",
+        "    display_name=JOB_NAME,\n",
+        "    container_uri=container_uri,\n",
+        ")\n",
+        "model_conversion_workdir = os.path.join(BUCKET_URI, JOB_NAME)\n",
+        "model = job.run(\n",
+        "    args=docker_args_list,\n",
+        "    base_output_dir=f\"{model_conversion_workdir}\",\n",
+        "    replica_count=1,\n",
+        "    machine_type=\"n1-standard-4\",\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "iILhhP3TfO8B"
+      },
+      "source": [
+        "## Run online prediction\n",
+        "\n",
+        "Run online prediction with the converted TF SavedModel."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "XswgX6JqRwFK"
+      },
+      "source": [
+        "Upload TF SavedModel and deploy it to an endpoint for prediction. This step takes around 15 minutes to finish."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "74yqis5ufO8B"
+      },
+      "outputs": [],
+      "source": [
+        "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",
+        ")\n",
+        "\n",
+        "jax_vit_endpoint = jax_vit_model.deploy(\n",
+        "    deployed_model_display_name=\"jax_vit_deployed\",\n",
+        "    traffic_split={\"0\": 100},\n",
+        "    machine_type=\"n1-standard-4\",\n",
+        "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
+        "    accelerator_count=1,\n",
+        "    min_replica_count=1,\n",
+        "    max_replica_count=1,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "iiozz1aVR7Pe"
+      },
+      "source": [
+        "Load a local test image file, encode it into a string, send it to the endpoint for prediction, and then generate the final class label from the predicted class probabilities."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "qxj4Xv_DhHXj"
+      },
+      "outputs": [],
+      "source": [
+        "test_directory = os.path.join(local_flower_data_directory, \"test/tulips\")\n",
+        "local_test_image_path = os.path.join(test_directory, os.listdir(test_directory)[0])\n",
+        "print(local_test_image_path)\n",
+        "instances_list = [\n",
+        "    {\n",
+        "        \"bytes_inputs\": {\n",
+        "            \"b64\": load_bytes_from_local_image(local_test_image_path, new_width=240)\n",
+        "        }\n",
+        "    }\n",
+        "]\n",
+        "instances = [json_format.ParseDict(s, Value()) for s in instances_list]\n",
+        "results = jax_vit_endpoint.predict(instances=instances)\n",
+        "logits = results.predictions[0]\n",
+        "predicted_label = LABEL_IDX_TO_STR[int(np.argmax(logits))]\n",
+        "print(\"predicted_label: \", predicted_label)"
+      ]
+    },
+    {
+      "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\n",
+        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
+        "\n",
+        "Otherwise, you can delete the individual resources you created in this tutorial:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "sx_vKniMq9ZX"
+      },
+      "outputs": [],
+      "source": [
+        "# Delete endpoint resource.\n",
+        "jax_vit_endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete model resource.\n",
+        "jax_vit_model.delete()\n",
+        "\n",
+        "# Delete Cloud Storage objects that were created.\n",
+        "delete_bucket = True\n",
+        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "    ! gsutil -m rm -r $BUCKET_URI"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_jax_vision_transformer.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_keras_stable_diffusion.ipynb

@@ -0,0 +1,745 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ur8xi4C7S06n"
+      },
+      "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": "TirJ-SGQseby"
+      },
+      "source": [
+        "# Vertex AI Model Garden Keras Stable Diffusion\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_keras_stable_diffusion.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/community/model_garden/model_garden_keras_stable_diffusion.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_keras_stable_diffusion.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": "dwGLvtIeECLK"
+      },
+      "source": [
+        "**_NOTE_**: This notebook has been tested in the following environment:\n",
+        "\n",
+        "* Python version = 3.9\n",
+        "\n",
+        "You can open this notebook directly in colab, or create [google managed](https://cloud.google.com/vertex-ai/docs/workbench/managed/create-instance) or [user managed](https://cloud.google.com/vertex-ai/docs/workbench/user-managed/create-new) workbench instances."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to use [Keras Stable Diffusion](https://keras.io/api/keras_cv/models/stable_diffusion) in Vertex AI Model Garden."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0z9r_mBmDeYh"
+      },
+      "source": [
+        "### Objective\n",
+        "\n",
+        "* Run local inferences for pretrained or customized models\n",
+        "\n",
+        "* Deploy pretrained or customized models in Google Cloud Vertex AI\n",
+        "\n",
+        "* Finetune models in Google Cloud Vertex AI"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "xxo28lDtDxn-"
+      },
+      "source": [
+        "### Dataset\n",
+        "\n",
+        "We use the dataset\n",
+        "[Pokémon BLIP captions](https://huggingface.co/datasets/lambdalabs/pokemon-blip-captions) to show how to finetune the stable diffusion models.\n",
+        "However, we'll use a slightly different version which was derived from the original\n",
+        "dataset to fit better with `tf.data`. Refer to\n",
+        "[the documentation](https://huggingface.co/datasets/sayakpaul/pokemon-blip-original-version)\n",
+        "for more details."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "AEnkHABrDijz"
+      },
+      "source": [
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI\n",
+        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
+        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
+        "Calculator](https://cloud.google.com/products/calculator/)\n",
+        "to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "z__i0w0lCAsW"
+      },
+      "source": [
+        "## Installation\n",
+        "\n",
+        "Install the following packages required to execute this notebook."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Jvqs-ehKlaYh"
+      },
+      "outputs": [],
+      "source": [
+        "if \"google.colab\" in str(get_ipython()):\n",
+        "    # Configs for colab notebooks.\n",
+        "    ! pip3 install --upgrade --quiet google-cloud-aiplatform\n",
+        "\n",
+        "    # Automatically restart kernel after installs\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)\n",
+        "\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()\n",
+        "\n",
+        "# Configs for all notebooks.\n",
+        "! pip3 install --quiet keras-cv==0.4.1"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KEukV6uRk_S3"
+      },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "### Set up your Google Cloud project\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\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",
+        "1. If you are running this notebook locally, you will 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",
+        "\n",
+        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
+        "Cloud SDK uses the right project for all the commands in this notebook.\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": "BF1j6f9HApxa"
+      },
+      "source": [
+        "### Set your project, region and buckets\n",
+        "\n",
+        "**If you don't know your project ID**, try the following:\n",
+        "* Run `gcloud config list`.\n",
+        "* Run `gcloud projects list`.\n",
+        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)\n",
+        "\n",
+        "You can change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations).\n",
+        "\n",
+        "You can create a storage bucket to store intermediate artifacts such as datasets, trained models etc."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YjNCFxq0JxlA"
+      },
+      "outputs": [],
+      "source": [
+        "# The project and bucket are for experiments below.\n",
+        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "# The form for BUCKET_URI is gs://<bucket-name>.\n",
+        "BUCKET_URI = \"\"  # @param {type:\"string\"}\n",
+        "REGION = \"us-central1\"  # @param {type: \"string\"}\n",
+        "\n",
+        "import os\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "EXPERIMENT_BUCKET = os.path.join(BUCKET_URI, \"keras\")\n",
+        "DATA_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"data\")\n",
+        "MODEL_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"model\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "uDjp76aaLZY9"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "Initialize the Vertex AI SDK for Python for your project."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "5uv7-iDKLbO0"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ZZFPe_GezXg8"
+      },
+      "source": [
+        "### Define constants and common functions"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "XcYUGwr-AJGY"
+      },
+      "outputs": [],
+      "source": [
+        "import base64\n",
+        "import os\n",
+        "from datetime import datetime\n",
+        "from io import BytesIO\n",
+        "\n",
+        "import matplotlib.pyplot as plt\n",
+        "from google.cloud import storage\n",
+        "from PIL import Image\n",
+        "\n",
+        "GCS_URI_PREFIX = \"gs://\"\n",
+        "\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_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_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
+        "PREDICTION_MACHINE_TYPE = \"n1-standard-8\"\n",
+        "DEPLOY_JOB_PREFIX = \"deploy\"\n",
+        "\n",
+        "\n",
+        "def get_job_name_with_datetime(prefix: str):\n",
+        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
+        "\n",
+        "\n",
+        "def download_data_to_gcs(tar_filepath, gcs_bucket):\n",
+        "    filename_with_ext = os.path.basename(tar_filepath)\n",
+        "    filename_without_ext = filename_with_ext.replace(\".tar.gz\", \"\")\n",
+        "    print(\"Download files from: \", tar_filepath)\n",
+        "    ! wget $tar_filepath -O $filename_with_ext\n",
+        "    ! mkdir -p $filename_without_ext\n",
+        "    ! tar -xvf $filename_with_ext -C .\n",
+        "\n",
+        "    ! gsutil -m cp -r $filename_without_ext $gcs_bucket/\n",
+        "    gcs_path = os.path.join(gcs_bucket, filename_without_ext)\n",
+        "    print(\"Upload files to: \", gcs_path)\n",
+        "    return gcs_path\n",
+        "\n",
+        "\n",
+        "def download_gcs_file_to_local(gcs_uri: str, local_path: str):\n",
+        "    \"\"\"Download a gcs file to a local path.\n",
+        "\n",
+        "    Args:\n",
+        "      gcs_uri: A string of file path on GCS.\n",
+        "      local_path: A string of local file path.\n",
+        "    \"\"\"\n",
+        "    if not gcs_uri.startswith(GCS_URI_PREFIX):\n",
+        "        raise ValueError(f\"{gcs_uri} is not a GCS path starting with {GCS_URI_PREFIX}.\")\n",
+        "    client = storage.Client()\n",
+        "    os.makedirs(os.path.dirname(local_path), exist_ok=True)\n",
+        "    with open(local_path, \"wb\") as f:\n",
+        "        client.download_blob_to_file(gcs_uri, f)\n",
+        "\n",
+        "\n",
+        "def deploy_model(model_path, service_account):\n",
+        "\n",
+        "    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_PATH\": f\"{model_path}\",\n",
+        "        \"IMAGE_WIDTH\": f\"{RESOLUTION}\",\n",
+        "        \"IMAGE_HEIGHT\": f\"{RESOLUTION}\",\n",
+        "    }\n",
+        "\n",
+        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{deploy_model_name}-endpoint\")\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_predict_route=\"/predict\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=serving_env,\n",
+        "    )\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=PREDICTION_MACHINE_TYPE,\n",
+        "        accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
+        "        accelerator_count=1,\n",
+        "        min_replica_count=1,\n",
+        "        max_replica_count=1,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "    )\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "def base64_to_image(image_str):\n",
+        "    image = Image.open(BytesIO(base64.b64decode(image_str)))\n",
+        "    return image\n",
+        "\n",
+        "\n",
+        "def display_image(image):\n",
+        "    _ = plt.figure(figsize=(20, 15))\n",
+        "    plt.grid(False)\n",
+        "    plt.imshow(image)\n",
+        "\n",
+        "\n",
+        "def display_image_grid(imgs, rows=2, cols=2):\n",
+        "    w, h = imgs[0].size\n",
+        "    grid = Image.new(\"RGB\", size=(cols * w, rows * h))\n",
+        "    for i, img in enumerate(imgs):\n",
+        "        grid.paste(img, box=(i % cols * w, i // cols * h))\n",
+        "    return grid"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "epo-RHXzcBBT"
+      },
+      "source": [
+        "## Run inferences\n",
+        "\n",
+        "This section shows how to run inferences with Keras Stable Diffusion models.\n",
+        "\n",
+        "1. Run inferences locally\n",
+        "2. Run inferences with serving dockers\n",
+        "\n",
+        "You can run inferences with pre-trained models from Keras team, or your own finetuned models.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "6zsa9vnBHhvO"
+      },
+      "outputs": [],
+      "source": [
+        "# Sets the model_path to empty to load the pre-trained model from Keras team.\n",
+        "# Sets the model_path to a gcs uri to load the finetuned models.\n",
+        "model_path = \"\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ld39hkcIceE2"
+      },
+      "source": [
+        "### Run inferences locally\n",
+        "Local inferences can finish in seconds with GPUs.\n",
+        "\n",
+        "Load models first."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "G1nCKVSac3Y5"
+      },
+      "outputs": [],
+      "source": [
+        "from keras_cv.models import StableDiffusion\n",
+        "\n",
+        "model = StableDiffusion(img_height=RESOLUTION, img_width=RESOLUTION, jit_compile=True)\n",
+        "if model_path.startswith(GCS_URI_PREFIX):\n",
+        "    local_model_path = \"/tmp/saved_model.h5\"\n",
+        "    download_gcs_file_to_local(model_path, local_model_path)\n",
+        "    model.diffusion_model.load_weights(local_model_path)\n",
+        "elif model_path:\n",
+        "    model.diffusion_model.load_weights(model_path)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ABaCSIWuP-_G"
+      },
+      "source": [
+        "Then run inferences."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "pnyeVsh8RNI5"
+      },
+      "outputs": [],
+      "source": [
+        "batch_size = 1\n",
+        "img = model.text_to_image(\n",
+        "    prompt=\"a squirrel in Picasso style\",\n",
+        "    batch_size=batch_size,  # How many images to generate at once\n",
+        "    num_steps=25,  # Number of iterations (controls image quality)\n",
+        "    seed=123,  # A fixed seed guarantees the same prompt always generates the same image\n",
+        ")\n",
+        "for i in range(batch_size):\n",
+        "    display_image(img[i])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kY87SU9Adq4o"
+      },
+      "source": [
+        "### Serve models with dockers\n",
+        "When serve models with dockers, we will deploy models in Google Cloud Vertex AI. The default setting will use 1 V100 GPU for deployment.\n",
+        "\n",
+        "Please create a Service Account for serving with dockers if you do not have one yet.\n",
+        "\n",
+        "The model deployment will take ~10 minutes to finish."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "yCB9vu7RenY6"
+      },
+      "outputs": [],
+      "source": [
+        "# The service account looks like:\n",
+        "# '<account_name>@<project>.iam.gserviceaccount.com'\n",
+        "# Please go to https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console\n",
+        "# and create service account with `Vertex AI User` and `Storage Object Admin` roles.\n",
+        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "model, endpoint = deploy_model(\n",
+        "    model_path=model_path,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        ")\n",
+        "\n",
+        "endpoint_id = endpoint.name\n",
+        "print(\"endpoint id is: \", endpoint_id)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "72_BW_BgfvYT"
+      },
+      "source": [
+        "Once deployed, you can send a batch of text prompts to the endpoint to generate images.\n",
+        "\n",
+        "Note, the inference time for the first request for a fresh deployment will need more time to process and take ~45 seconds on one V100 GPU. The inferences for further request is ~12 seconds on one V100 GPU per image."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "U_jrNcZ5eVbH"
+      },
+      "outputs": [],
+      "source": [
+        "# # Loads an existing endpoint as below.\n",
+        "# endpoint_id = <An Existing 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",
+        "\n",
+        "instances = [\n",
+        "    {\"prompt\": \"a squirrel in Picasso style\"},\n",
+        "    {\"prompt\": \"a dog in Picasso style\"},\n",
+        "    {\"prompt\": \"a cat in Picasso style\"},\n",
+        "    {\"prompt\": \"a deer in Picasso style\"},\n",
+        "]\n",
+        "\n",
+        "parameters = {\n",
+        "    \"batch_size\": 1,  # How many images to generate at once\n",
+        "    \"num_steps\": 25,  # Number of iterations (controls image quality)\n",
+        "    \"seed\": 123,  # A fixed seed guarantees the same prompt always generates the same image\n",
+        "}\n",
+        "response = endpoint.predict(instances=instances, parameters=parameters)\n",
+        "# prediction['predicted_image'] will contains the prediction images in a batch.\n",
+        "# The batch size in this example is 1, and the visualization only parses the\n",
+        "# first predicted image.\n",
+        "images = [\n",
+        "    base64_to_image(prediction[\"predicted_image\"][0])\n",
+        "    for prediction in response.predictions\n",
+        "]\n",
+        "display_image_grid(images, rows=2, cols=2)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "LiQF7fm6f842"
+      },
+      "source": [
+        "### Clean up"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "eqJyypt-f9K6"
+      },
+      "outputs": [],
+      "source": [
+        "# Undeploys models and deletes endpoints.\n",
+        "endpoint.delete(force=True)\n",
+        "# Deletes models.\n",
+        "model.delete()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "RB_xY9ipr7ZU"
+      },
+      "source": [
+        "## Finetune models\n",
+        "This section shows how to finetune Keras Stable diffusion models with training dockers.\n",
+        "\n",
+        "If you would like to use finetuned models, please go to the section `Run inferences`."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "OD3TtaWs5b4v"
+      },
+      "source": [
+        "### Download data\n",
+        " We download the data to GCS storage for the experiments with training dockers."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2TVB8MU-5i-q"
+      },
+      "outputs": [],
+      "source": [
+        "# Skips this step if you have already downloaded the dataset.\n",
+        "download_data_to_gcs(\n",
+        "    \"https://huggingface.co/datasets/sayakpaul/pokemon-blip-original-version/resolve/main/pokemon_dataset.tar.gz\",\n",
+        "    DATA_BUCKET,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Ee7Hzq8O5jgF"
+      },
+      "source": [
+        "### Start training jobs\n",
+        "We finetune 512*512 stable diffusion models with 1 epoch and it takes ~15 minutes to finish using 1 A100 GPU with default settings."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "riG_qUokg0XZ"
+      },
+      "outputs": [],
+      "source": [
+        "data_csv = os.path.join(DATA_BUCKET, \"pokemon_dataset/data.csv\")\n",
+        "epochs = 1\n",
+        "\n",
+        "train_job_name = get_job_name_with_datetime(TRAINING_JOB_PREFIX)\n",
+        "model_dir = os.path.join(MODEL_BUCKET, train_job_name)\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": TRAIN_MACHINE_TYPE,\n",
+        "            \"accelerator_type\": TRAIN_ACCELERATOR_TYPE,\n",
+        "            \"accelerator_count\": TRAIN_NUM_GPU,\n",
+        "        },\n",
+        "        \"replica_count\": 1,\n",
+        "        \"disk_spec\": {\n",
+        "            \"boot_disk_type\": \"pd-ssd\",\n",
+        "            \"boot_disk_size_gb\": 500,\n",
+        "        },\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": TRAIN_CONTAINER_URI,\n",
+        "            \"command\": [],\n",
+        "            \"env\": [\n",
+        "                {\n",
+        "                    \"name\": \"RESOLUTION\",\n",
+        "                    \"value\": f\"{RESOLUTION}\",\n",
+        "                },\n",
+        "            ],\n",
+        "            \"args\": [\n",
+        "                f\"--epochs={epochs}\",\n",
+        "                f\"--input_csv_path={data_csv}\",\n",
+        "                f\"--output_model_dir={model_dir}\",\n",
+        "            ],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "train_job = aiplatform.CustomJob(\n",
+        "    display_name=train_job_name,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=STAGING_BUCKET,\n",
+        ")\n",
+        "\n",
+        "train_job.run()\n",
+        "\n",
+        "model_path = os.path.join(model_dir, \"saved_model.h5\")\n",
+        "print(\"The trained model is saved as: \", model_path)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "wBlQ6FQlJhBi"
+      },
+      "source": [
+        "After the training finishes, you can use `model_path` and then go to the `Run inferences` section above to run predictions."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kkH2nrpdp4sp"
+      },
+      "source": [
+        "### Clean up"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Ax6vQVZhp9pR"
+      },
+      "outputs": [],
+      "source": [
+        "train_job.delete()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "1dijQDiZWegt"
+      },
+      "source": [
+        "## References\n",
+        "\n",
+        "- [Fine-tuning Stable Diffusion](https://keras.io/examples/generative/finetune_stable_diffusion/)\n",
+        "- [StableDiffusion image-generation model](https://keras.io/api/keras_cv/models/stable_diffusion/)\n",
+        "- [High-performance image generation using Stable Diffusion in KerasCV](https://keras.io/guides/keras_cv/generate_images_with_stable_diffusion/)"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "name": "model_garden_keras_stable_diffusion.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_keras_yolov8.ipynb

@@ -0,0 +1,771 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ur8xi4C7S06n"
+      },
+      "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": "TirJ-SGQseby"
+      },
+      "source": [
+        "# Vertex AI Model Garden Keras YOLOv8\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_keras_yolov8.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/community/model_garden/model_garden_keras_yolov8.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_keras_yolov8.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": "dwGLvtIeECLK"
+      },
+      "source": [
+        "**_NOTE_**: This notebook has been tested in the following environment:\n",
+        "\n",
+        "* Python version = 3.9\n",
+        "\n",
+        "You can open this notebook directly in Colab, or create [google managed](https://cloud.google.com/vertex-ai/docs/workbench/managed/create-instance) or [user managed](https://cloud.google.com/vertex-ai/docs/workbench/user-managed/create-new) Workbench instances."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tvgnzT1CKxrO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to use [Keras YOLOv8](https://keras.io/api/keras_cv/models/tasks/yolo_v8_detector/) in Vertex AI Model Garden."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0z9r_mBmDeYh"
+      },
+      "source": [
+        "### Objective\n",
+        "\n",
+        "* Run local inferences for pretrained or customized models\n",
+        "\n",
+        "* Deploy pretrained or customized models in Google Cloud Vertex AI\n",
+        "\n",
+        "* Finetune models in Google Cloud Vertex AI"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "AEnkHABrDijz"
+      },
+      "source": [
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI\n",
+        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
+        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
+        "Calculator](https://cloud.google.com/products/calculator/)\n",
+        "to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "af989c0e437d"
+      },
+      "source": [
+        "### Dataset\n",
+        "\n",
+        "\n",
+        "The dataset used for this tutorial is the Salads category of the [OpenImages dataset](https://www.tensorflow.org/datasets/catalog/open_images_v4) from [TensorFlow Datasets](https://www.tensorflow.org/datasets/catalog/overview). This dataset does not require any feature engineering. The version of the dataset you will use in this tutorial is stored in a public Cloud Storage bucket. The trained model predicts the bounding box locations and corresponding type of salad items in an image from a class of five items: Salad, Seafood, Tomato, Baked Goods, or Cheese."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "z__i0w0lCAsW"
+      },
+      "source": [
+        "## Installation\n",
+        "\n",
+        "Install the following packages required to execute this notebook."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Jvqs-ehKlaYh"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "    # Configs for Colab notebooks.\n",
+        "    ! pip3 install --upgrade --quiet google-cloud-aiplatform\n",
+        "\n",
+        "    # Automatically restart kernel after installs\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)\n",
+        "\n",
+        "    from google.colab import auth as google_auth\n",
+        "\n",
+        "    google_auth.authenticate_user()\n",
+        "\n",
+        "# Configs for all notebooks.\n",
+        "! pip3 install --quiet keras-cv==0.6.1\n",
+        "! pip3 install --quiet keras-core==0.1.0"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "KEukV6uRk_S3"
+      },
+      "source": [
+        "## Before you begin\n",
+        "\n",
+        "### Set up your Google Cloud project\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\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",
+        "1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
+        "\n",
+        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
+        "Cloud SDK uses the right project for all the commands in this notebook.\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": "BF1j6f9HApxa"
+      },
+      "source": [
+        "### Set your project, region and buckets\n",
+        "\n",
+        "**If you don't know your project ID**, try the following:\n",
+        "* Run `gcloud config list`.\n",
+        "* Run `gcloud projects list`.\n",
+        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)\n",
+        "\n",
+        "You can change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations).\n",
+        "\n",
+        "You can create a storage bucket to store intermediate artifacts such as datasets, trained models etc."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YjNCFxq0JxlA"
+      },
+      "outputs": [],
+      "source": [
+        "# The project and bucket are for experiments below.\n",
+        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "# The form for BUCKET_URI is gs://<bucket-name>.\n",
+        "BUCKET_URI = \"\"  # @param {type:\"string\"}\n",
+        "REGION = \"us-central1\"  # @param {type: \"string\"}\n",
+        "\n",
+        "import os\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "MODEL_BUCKET = os.path.join(STAGING_BUCKET, \"keras_yolov8\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "uDjp76aaLZY9"
+      },
+      "source": [
+        "### Initialize Vertex AI SDK for Python\n",
+        "\n",
+        "Initialize the Vertex AI SDK for Python for your project."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "5uv7-iDKLbO0"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ZZFPe_GezXg8"
+      },
+      "source": [
+        "### Define constants and common functions"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "XcYUGwr-AJGY"
+      },
+      "outputs": [],
+      "source": [
+        "import base64\n",
+        "import io\n",
+        "from datetime import datetime\n",
+        "from typing import Dict, List, Union\n",
+        "\n",
+        "import keras_cv\n",
+        "import numpy as np\n",
+        "import tensorflow as tf\n",
+        "import yaml\n",
+        "from google.protobuf import json_format\n",
+        "from google.protobuf.struct_pb2 import Value\n",
+        "from keras_cv import visualization\n",
+        "from PIL import Image\n",
+        "\n",
+        "TRAIN_MACHINE_TYPE = \"n1-highmem-16\"\n",
+        "TRAIN_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
+        "TRAIN_NUM_GPU = 2\n",
+        "TRAIN_CONTAINER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/keras-yolov8-train\"\n",
+        ")\n",
+        "TRAINING_JOB_PREFIX = \"train_yolov8\"\n",
+        "\n",
+        "UPLOAD_JOB_PREFIX = \"upload_yolov8\"\n",
+        "DEPLOY_JOB_PREFIX = \"deploy_yolov8\"\n",
+        "SERVING_CONTAINER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-12:latest\"\n",
+        ")\n",
+        "SERVING_ACCELERATOR_TYPE = \"NVIDIA_TESLA_T4\"\n",
+        "SERVING_MACHINE_TYPE = \"n1-standard-4\"\n",
+        "SERVING_CONTAINER_ARGS = [\"--allow_precompilation\", \"--allow_compression\"]\n",
+        "\n",
+        "RESOLUTION = 512\n",
+        "\n",
+        "def get_job_name_with_datetime(prefix: str):\n",
+        "  \"\"\"Generates a 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 load_img(path):\n",
+        "    \"\"\"Reads image from path and return PIL.Image instance.\"\"\"\"\n",
+        "    img = tf.io.read_file(path)\n",
+        "    img = tf.image.decode_jpeg(img, channels=3)\n",
+        "    return Image.fromarray(np.uint8(img)).convert(\"RGB\")\n",
+        "\n",
+        "\n",
+        "def decode_image(image_str_tensor: tf.string) -> tf.float32:\n",
+        "    \"\"\"Converts and resizes image bytes to image tensor.\"\"\"\"\n",
+        "    image = tf.io.decode_image(image_str_tensor, 3, expand_animations=False)\n",
+        "    image = tf.image.resize(image, (RESOLUTION, RESOLUTION))\n",
+        "    return image\n",
+        "\n",
+        "\n",
+        "def get_label_map(label_map_yaml_filepath):\n",
+        "    \"\"\"Returns class id to label mapping given a filepath to the label map.\"\"\"\n",
+        "    with tf.io.gfile.GFile(label_map_yaml_filepath, \"rb\") as input_file:\n",
+        "        label_map = yaml.safe_load(input_file.read())[\"label_map\"]\n",
+        "    return label_map\n",
+        "\n",
+        "\n",
+        "def get_prediction_instances(test_filepath, new_width=-1):\n",
+        "    \"\"\"Generate instance from image path to pass to Vertex AI Endpoint for prediction.\"\"\"\n",
+        "    if new_width <= 0:\n",
+        "        with tf.io.gfile.GFile(test_filepath, \"rb\") as input_file:\n",
+        "            encoded_string = base64.b64encode(input_file.read()).decode(\"utf-8\")\n",
+        "    else:\n",
+        "        img = load_img(test_filepath)\n",
+        "        width, height = img.size\n",
+        "        print(\"original input image size: \", width, \" , \", height)\n",
+        "        new_height = int(height * new_width / width)\n",
+        "        new_img = img.resize((new_width, new_height))\n",
+        "        print(\"resized input image size: \", new_width, \" , \", new_height)\n",
+        "        buffered = io.BytesIO()\n",
+        "        new_img.save(buffered, format=\"JPEG\")\n",
+        "        encoded_string = base64.b64encode(buffered.getvalue()).decode(\"utf-8\")\n",
+        "\n",
+        "    instances = [\n",
+        "        {\n",
+        "            \"encoded_image\": {\"b64\": encoded_string},\n",
+        "        }\n",
+        "    ]\n",
+        "    return instances\n",
+        "\n",
+        "\n",
+        "def predict_custom_trained_model(\n",
+        "    project: str,\n",
+        "    endpoint_id: str,\n",
+        "    instances: Union[Dict, List[Dict]],\n",
+        "    location: str = \"us-central1\",\n",
+        "):\n",
+        "    # The AI Platform services require regional API endpoints.\n",
+        "    client_options = {\"api_endpoint\": f\"{location}-aiplatform.googleapis.com\"}\n",
+        "    # Initialize client that will be used to create and send requests.\n",
+        "    # This client only needs to be created once, and can be reused for multiple requests.\n",
+        "    client = aiplatform.gapic.PredictionServiceClient(client_options=client_options)\n",
+        "    parameters_dict = {}\n",
+        "    parameters = json_format.ParseDict(parameters_dict, Value())\n",
+        "    endpoint = client.endpoint_path(\n",
+        "        project=project, location=location, endpoint=endpoint_id\n",
+        "    )\n",
+        "    response = client.predict(\n",
+        "        endpoint=endpoint, instances=instances, parameters=parameters\n",
+        "    )\n",
+        "    return response.predictions, response.deployed_model_id"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "epo-RHXzcBBT"
+      },
+      "source": [
+        "## Run local inferences with pretrained model\n",
+        "\n",
+        "This section shows how to run inferences locally with YOLOv8-M pretrained on PascalVOC 2012 object detection task, which consists of 20 classes.\n",
+        "\n",
+        "Load image from Cloud Storage and decode as Tensor."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "6zsa9vnBHhvO"
+      },
+      "outputs": [],
+      "source": [
+        "test_filepath = \"\"  # @param {type:\"string\"}\n",
+        "img_bytes = tf.io.read_file(test_filepath)\n",
+        "image = tf.expand_dims(decode_image(img_bytes), axis=0)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2wC-pSYR0jjU"
+      },
+      "source": [
+        "Load model pretrained on PascalVOC 2012."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7nvPEly_4Vm6"
+      },
+      "outputs": [],
+      "source": [
+        "model = keras_cv.models.YOLOV8Detector.from_preset(\n",
+        "    \"yolo_v8_m_pascalvoc\",\n",
+        "    bounding_box_format=\"xywh\",\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ZrijGrxT0lvC"
+      },
+      "source": [
+        "Then run inferences and visualize results."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "65yEa4N0xcTS"
+      },
+      "outputs": [],
+      "source": [
+        "decoded = model.predict(image)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "n8-X3gA5xV_l"
+      },
+      "outputs": [],
+      "source": [
+        "# Classes in PascalVOC 2012 dataset.\n",
+        "class_ids = [\n",
+        "    \"Aeroplane\",\n",
+        "    \"Bicycle\",\n",
+        "    \"Bird\",\n",
+        "    \"Boat\",\n",
+        "    \"Bottle\",\n",
+        "    \"Bus\",\n",
+        "    \"Car\",\n",
+        "    \"Cat\",\n",
+        "    \"Chair\",\n",
+        "    \"Cow\",\n",
+        "    \"Dining Table\",\n",
+        "    \"Dog\",\n",
+        "    \"Horse\",\n",
+        "    \"Motorbike\",\n",
+        "    \"Person\",\n",
+        "    \"Potted Plant\",\n",
+        "    \"Sheep\",\n",
+        "    \"Sofa\",\n",
+        "    \"Train\",\n",
+        "    \"Tvmonitor\",\n",
+        "    \"Total\",\n",
+        "]\n",
+        "class_mapping = dict(zip(range(len(class_ids)), class_ids))\n",
+        "\n",
+        "visualization.plot_bounding_box_gallery(\n",
+        "    image,\n",
+        "    value_range=(0, 255),\n",
+        "    rows=1,\n",
+        "    cols=1,\n",
+        "    y_pred=decoded,\n",
+        "    scale=5,\n",
+        "    font_scale=0.7,\n",
+        "    bounding_box_format=\"xywh\",\n",
+        "    class_mapping=class_mapping,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "RB_xY9ipr7ZU"
+      },
+      "source": [
+        "## Finetune models\n",
+        "This section shows how to finetune the Keras YOLOv8 model with training dockers and then deploy to Vertex AI Endpoint resource. The accepted dataset format is a CSV formatted as it would for [AutoML Image Object Detection](https://cloud.google.com/vertex-ai/docs/image-data/object-detection/prepare-data#input-files), without an `ML_USE` column."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "pkNc7jyq1js1"
+      },
+      "outputs": [],
+      "source": [
+        "input_csv_path = \"gs://cloud-samples-data/vision/salads.csv\"  # @param {type:\"string\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Ee7Hzq8O5jgF"
+      },
+      "source": [
+        "### Start training jobs\n",
+        "The following code block shows some of the possible hyperparameters that can be set. The settings are for demonstration purposes only. Parameters such as `batch_size`, `learning_rate`, and `epochs` be overridden when used. `backbone` must be one of the following:\n",
+        "* `yolo_v8_xs_backbone`\n",
+        "* `yolo_v8_s_backbone`\n",
+        "* `yolo_v8_m_backbone`\n",
+        "* `yolo_v8_l_backbone`\n",
+        "* `yolo_v8_xl_backbone`\n",
+        "* `yolo_v8_xs_backbone_coco`\n",
+        "* `yolo_v8_s_backbone_coco`\n",
+        "* `yolo_v8_m_backbone_coco`\n",
+        "* `yolo_v8_l_backbone_coco`\n",
+        "* `yolo_v8_xl_backbone_coco`\n",
+        "\n",
+        "If looking for a preset with pretrained weights, choose one of `yolo_v8_xs_backbone_coco`, `yolo_v8_s_backbone_coco`, `yolo_v8_m_backbone_coco`, `yolo_v8_l_backbone_coco`, `yolo_v8_xl_backbone_coco`."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "riG_qUokg0XZ"
+      },
+      "outputs": [],
+      "source": [
+        "# Hyperparameters\n",
+        "epochs = 10\n",
+        "learning_rate = 0.0005\n",
+        "fpn_depth = 3\n",
+        "confidence_threshold = 0.02\n",
+        "iou_threshold = 0.3\n",
+        "backbone = \"yolo_v8_xl_backbone_coco\"\n",
+        "\n",
+        "train_job_name = get_job_name_with_datetime(TRAINING_JOB_PREFIX)\n",
+        "model_dir = os.path.join(MODEL_BUCKET, train_job_name)\n",
+        "worker_pool_specs = [\n",
+        "    {\n",
+        "        \"machine_spec\": {\n",
+        "            \"machine_type\": TRAIN_MACHINE_TYPE,\n",
+        "            \"accelerator_type\": TRAIN_ACCELERATOR_TYPE,\n",
+        "            \"accelerator_count\": TRAIN_NUM_GPU,\n",
+        "        },\n",
+        "        \"replica_count\": 1,\n",
+        "        \"disk_spec\": {\n",
+        "            \"boot_disk_type\": \"pd-ssd\",\n",
+        "            \"boot_disk_size_gb\": 500,\n",
+        "        },\n",
+        "        \"container_spec\": {\n",
+        "            \"image_uri\": TRAIN_CONTAINER_URI,\n",
+        "            \"command\": [],\n",
+        "            \"env\": [\n",
+        "                {\n",
+        "                    \"name\": \"RESOLUTION\",\n",
+        "                    \"value\": f\"{RESOLUTION}\",\n",
+        "                },\n",
+        "            ],\n",
+        "            \"args\": [\n",
+        "                f\"--input_csv_path={input_csv_path}\",\n",
+        "                f\"--output_model_dir={model_dir}\",\n",
+        "                f\"--epochs={epochs}\",\n",
+        "                f\"--pretrained_backbone={backbone}\",\n",
+        "                f\"--fpn_depth={fpn_depth}\",\n",
+        "                f\"--learning_rate={learning_rate}\",\n",
+        "                f\"--confidence_threshold={confidence_threshold}\",\n",
+        "                f\"--iou_threshold={iou_threshold}\",\n",
+        "            ],\n",
+        "        },\n",
+        "    }\n",
+        "]\n",
+        "\n",
+        "train_job = aiplatform.CustomJob(\n",
+        "    display_name=train_job_name,\n",
+        "    project=PROJECT_ID,\n",
+        "    worker_pool_specs=worker_pool_specs,\n",
+        "    staging_bucket=STAGING_BUCKET,\n",
+        ")\n",
+        "\n",
+        "train_job.run()\n",
+        "\n",
+        "print(\"The trained model is saved in: \", model_dir)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9KBJ0ySVYX47"
+      },
+      "source": [
+        "### Prediction\n",
+        "This section shows how to deploy and make online predictions with the model.\n",
+        "\n",
+        "1. Upload and deploy models\n",
+        "2. Run predictions"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "K6rUSSKmYZJ6"
+      },
+      "outputs": [],
+      "source": [
+        "upload_job_name = get_job_name_with_datetime(UPLOAD_JOB_PREFIX)\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",
+        ")\n",
+        "\n",
+        "print(\"The uploaded model name is: \", upload_job_name)\n",
+        "\n",
+        "deploy_model_name = get_job_name_with_datetime(DEPLOY_JOB_PREFIX)\n",
+        "\n",
+        "endpoint = model.deploy(\n",
+        "    deployed_model_display_name=deploy_model_name,\n",
+        "    machine_type=SERVING_MACHINE_TYPE,\n",
+        "    traffic_split={\"0\": 100},\n",
+        "    accelerator_type=SERVING_ACCELERATOR_TYPE,\n",
+        "    accelerator_count=1,\n",
+        "    min_replica_count=1,\n",
+        "    max_replica_count=1,\n",
+        ")\n",
+        "print(\"The deployed job name is: \", deploy_model_name)\n",
+        "\n",
+        "endpoint_id = endpoint.name\n",
+        "print(\"endpoint id is: \", endpoint_id)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "a879effaf402"
+      },
+      "source": [
+        "Load image from Cloud Storage, resize, and encode."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "VDznWEMmbwj4"
+      },
+      "outputs": [],
+      "source": [
+        "test_filepath = \"gs://cloud-ml-data/img/openimage/1302/4677521502_6f2767039c_o.jpg\"  # @param {type:\"string\"}\n",
+        "image_bytes = tf.io.read_file(test_filepath)\n",
+        "image_resized = tf.expand_dims(decode_image(image_bytes), axis=0)\n",
+        "\n",
+        "instances = get_prediction_instances(test_filepath, new_width=640)\n",
+        "\n",
+        "predictions, _ = predict_custom_trained_model(\n",
+        "    project=PROJECT_ID, location=REGION, endpoint_id=endpoint_id, instances=instances\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "14e889492871"
+      },
+      "source": [
+        "Run online predictions using the endpoint and visualize the result."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2bx1cW0IdXqp"
+      },
+      "outputs": [],
+      "source": [
+        "predictions_dict = {\n",
+        "    \"boxes\": tf.expand_dims(predictions[0][\"boxes\"], axis=0),\n",
+        "    \"classes\": tf.expand_dims(predictions[0][\"classes\"], axis=0),\n",
+        "    \"confidence\": tf.expand_dims(predictions[0][\"confidence\"], axis=0),\n",
+        "    \"num_detections\": predictions[0][\"num_detections\"],\n",
+        "}\n",
+        "\n",
+        "label_map = get_label_map(os.path.join(model_dir, \"label_map.yaml\"))\n",
+        "\n",
+        "visualization.plot_bounding_box_gallery(\n",
+        "    image_resized,\n",
+        "    value_range=(0, 255),\n",
+        "    rows=1,\n",
+        "    cols=1,\n",
+        "    y_pred=predictions_dict,\n",
+        "    scale=5,\n",
+        "    font_scale=0.7,\n",
+        "    bounding_box_format=\"xywh\",\n",
+        "    class_mapping=label_map,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "kkH2nrpdp4sp"
+      },
+      "source": [
+        "### Clean up"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Ax6vQVZhp9pR"
+      },
+      "outputs": [],
+      "source": [
+        "# Deletes custom train jobs.\n",
+        "train_job.delete()\n",
+        "# Undeploys models and deletes endpoints.\n",
+        "endpoint.delete(force=True)\n",
+        "# Deletes models.\n",
+        "model.delete()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "1dijQDiZWegt"
+      },
+      "source": [
+        "## References\n",
+        "\n",
+        "- [Efficient Object Detection with YOLOV8 and KerasCV](https://keras.io/examples/vision/yolov8/)\n",
+        "- [Keras YOLOv8 API Documentation](https://keras.io/api/keras_cv/models/tasks/yolo_v8_detector/)\n",
+        "- [Keras YOLOv8 Backbones](https://keras.io/api/keras_cv/models/backbones/yolo_v8/)"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_keras_yolov8.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}