notebooks/community/model_garden/model_garden_mammut.ipynb

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