Update owlvit notebook to support Vertex Workbench.

PiperOrigin-RevId: 731527931
GoogleCloudPlatform · Feb 27, 2025 · 3f2da94 · 3f2da94
1 parent 6dcd8ca
commit 3f2da94
Showing 1 changed file with 13 additions and 1 deletion.
diff --git a/notebooks/community/model_garden/model_garden_pytorch_owlvit.ipynb b/notebooks/community/model_garden/model_garden_pytorch_owlvit.ipynb
@@ -34,6 +34,11 @@
         "\n",
         "<table><tbody><tr>\n",
         "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_owlvit.ipynb\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_owlvit.ipynb\">\n",
         "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
         "    </a>\n",
@@ -122,6 +127,7 @@
         "import matplotlib.pyplot as plt\n",
         "from google.cloud import aiplatform\n",
         "\n",
+        "! pip install --upgrade tensorflow\n",
         "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
         "\n",
         "common_util = importlib.import_module(\n",
@@ -134,6 +140,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -227,7 +239,7 @@
         "    serving_env = {\n",
         "        \"MODEL_ID\": model_id,\n",
         "        \"TASK\": task,\n",
-        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "        \"DEPLOY_SOURCE\": common_util.get_deploy_source(),\n",
         "    }\n",
         "    # If the model_id is a GCS path, use artifact_uri to pass it to serving docker.\n",
         "    artifact_uri = model_id if model_id.startswith(\"gs://\") else None\n",