notebooks/community/model_garden/model_garden_gradio_streaming_chat_completions.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "L7tH6GT2M9F9"
      },
      "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": "BG8HkBaxM9F9"
      },
      "source": [
        "# Vertex AI Model Garden - Chat Completions With Streaming Playground\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_gradio_streaming_chat_completions.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_gradio_streaming_chat_completions.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": "bmPrfBTFM9F9"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates starting a playground based on [Gradio UI](https://www.gradio.app/) that allows users to interact with the instruction-tuned text generation models via a chatbot UI more easily.\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Chat with instruction-tuned text generation models deployed on the [Vertex Online Prediction](https://cloud.google.com/vertex-ai/docs/predictions/get-online-predictions) endpoints.\n",
        "- (Optional) One-click deploy demo models to [Vertex Online Prediction](https://cloud.google.com/vertex-ai/docs/predictions/get-online-predictions) endpoints.\n",
        "\n",
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B4ppASahFB9b"
      },
      "source": [
        "## Run the notebook"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "62VgpTrAGx9JQPwjG5RYFCJT"
      },
      "outputs": [],
      "source": [
        "# @title Setup Google Cloud project and install dependencies\n",
        "# Upgrade Vertex AI SDK.\n",
        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
        "! pip3 install --upgrade gradio~=4.40.0\n",
        "\n",
        "import os\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "# Get the default cloud project id.\n",
        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
        "\n",
        "# Get the default region for endpoints.\n",
        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
        "# Dedicated endpoint not supported yet\n",
        "use_dedicated_endpoint = False"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-zFBGiLWUVNd"
      },
      "outputs": [],
      "source": [
        "# @title Start the playground\n",
        "\n",
        "# @markdown This is a chatbot playground for instruction-tuned text generation models.\n",
        "# @markdown After the cell runs, this playground is available in a separate browser tab if you click the public URL,\n",
        "# @markdown i.e. [\"https://####.gradio.live\"](#) in the output of the cell.\n",
        "\n",
        "# @markdown **How to use:**\n",
        "# @markdown 1. **Important**: Notebook cell reruns create new public URLs. Previous URLs will stop working.\n",
        "# @markdown 1. Before you start, you need to select a Vertex prediction endpoint with a matching model\n",
        "# @markdown from the endpoint dropdown list in the same project and region where you run this notebook.\n",
        "# @markdown 1. This playground only supports new deployments with\n",
        "# @markdown text-generation-inference (`us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-hf-tgi-serve`),\n",
        "# @markdown vLLM (`us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve`),\n",
        "# @markdown or HexLLM (`us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve`).\n",
        "# @markdown\n",
        "# @markdown    **Endpoints deployed with older serving containers or before August 20, 2024 might not work**. We recommend deploying a new endpoint from the listed demo models inside the Gradio app.\n",
        "# @markdown 1. After experiments, do not forget to undeploy the models from [Vertex Online Prediction](https://console.cloud.google.com/vertex-ai/online-prediction/endpoints) to avoid continuous charges to the project.\n",
        "\n",
        "import dataclasses\n",
        "import json\n",
        "from typing import Callable, Tuple\n",
        "\n",
        "import gradio as gr\n",
        "import requests\n",
        "\n",
        "MAX_TOKENS = 512\n",
        "HF_TOKEN = \"\"\n",
        "\n",
        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240819_0916_RC00\"\n",
        "TGI_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-hf-tgi-serve:20240820_0936_RC01\"\n",
        "\n",
        "SERVER_TYPE_VLLM = \"vllm\"\n",
        "SERVER_TYPE_HEXLLM = \"hex-llm\"\n",
        "SERVER_TYPE_TGI = \"tgi\"\n",
        "SERVER_TYPES = [\n",
        "    SERVER_TYPE_VLLM,\n",
        "    SERVER_TYPE_HEXLLM,\n",
        "    SERVER_TYPE_TGI,\n",
        "]\n",
        "\n",
        "\n",
        "@dataclasses.dataclass\n",
        "class Endpoint:\n",
        "    display_name: str\n",
        "    location: str\n",
        "    resource_name: str\n",
        "    server_type: str\n",
        "\n",
        "\n",
        "PUBLIC_PLAYGROUND_ENDPOINTS = [\n",
        "    Endpoint(\n",
        "        display_name=\"Gemma-2-2b-it (Public playground)\",\n",
        "        location=\"us-west1\",\n",
        "        resource_name=\"playground:google/796\",\n",
        "        server_type=SERVER_TYPE_HEXLLM,\n",
        "    ),\n",
        "]\n",
        "\n",
        "\n",
        "@dataclasses.dataclass\n",
        "class DeployConfig:\n",
        "    display_name: str\n",
        "    model_name: str\n",
        "    func: Callable[[str], tuple[aiplatform.Model, aiplatform.Endpoint]]\n",
        "\n",
        "\n",
        "def deploy_model_vllm(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    service_account: str,\n",
        "    base_model_id: str = None,\n",
        "    machine_type: str = \"g2-standard-8\",\n",
        "    accelerator_type: str = \"NVIDIA_L4\",\n",
        "    accelerator_count: int = 1,\n",
        "    gpu_memory_utilization: float = 0.9,\n",
        "    max_model_len: int = 4096,\n",
        "    dtype: str = \"auto\",\n",
        "    enable_trust_remote_code: bool = False,\n",
        "    enforce_eager: bool = False,\n",
        "    enable_lora: bool = False,\n",
        "    max_loras: int = 1,\n",
        "    max_cpu_loras: int = 8,\n",
        "    use_dedicated_endpoint: bool = False,\n",
        "    max_num_seqs: int = 256,\n",
        "    model_type: str = None,\n",
        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
        "    endpoint = aiplatform.Endpoint.create(\n",
        "        display_name=f\"{model_name}-endpoint\",\n",
        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
        "    )\n",
        "\n",
        "    if not base_model_id:\n",
        "        base_model_id = model_id\n",
        "\n",
        "    # See https://docs.vllm.ai/en/latest/models/engine_args.html for a list of possible arguments with descriptions.\n",
        "    vllm_args = [\n",
        "        \"python\",\n",
        "        \"-m\",\n",
        "        \"vllm.entrypoints.api_server\",\n",
        "        \"--host=0.0.0.0\",\n",
        "        \"--port=8080\",\n",
        "        f\"--model={model_id}\",\n",
        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
        "        \"--swap-space=16\",\n",
        "        f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
        "        f\"--max-model-len={max_model_len}\",\n",
        "        f\"--dtype={dtype}\",\n",
        "        f\"--max-loras={max_loras}\",\n",
        "        f\"--max-cpu-loras={max_cpu_loras}\",\n",
        "        f\"--max-num-seqs={max_num_seqs}\",\n",
        "        \"--disable-log-stats\",\n",
        "    ]\n",
        "\n",
        "    if enable_trust_remote_code:\n",
        "        vllm_args.append(\"--trust-remote-code\")\n",
        "\n",
        "    if enforce_eager:\n",
        "        vllm_args.append(\"--enforce-eager\")\n",
        "\n",
        "    if enable_lora:\n",
        "        vllm_args.append(\"--enable-lora\")\n",
        "\n",
        "    if model_type:\n",
        "        vllm_args.append(f\"--model-type={model_type}\")\n",
        "\n",
        "    env_vars = {\n",
        "        \"MODEL_ID\": base_model_id,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "\n",
        "    # HF_TOKEN is not a compulsory field and may not be defined.\n",
        "    try:\n",
        "        if HF_TOKEN:\n",
        "            env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
        "    except NameError:\n",
        "        pass\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
        "        serving_container_args=vllm_args,\n",
        "        serving_container_ports=[8080],\n",
        "        serving_container_predict_route=\"/generate\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=env_vars,\n",
        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
        "        serving_container_deployment_timeout=7200,\n",
        "    )\n",
        "    print(\n",
        "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
        "    )\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
        "    )\n",
        "    print(\"endpoint_name:\", endpoint.name)\n",
        "\n",
        "    return model, endpoint\n",
        "\n",
        "\n",
        "def deploy_model_tgi(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    service_account: str,\n",
        "    machine_type: str = \"g2-standard-8\",\n",
        "    accelerator_type: str = \"NVIDIA_L4\",\n",
        "    accelerator_count: int = 1,\n",
        "    max_input_length: int = 2047,\n",
        "    max_total_tokens: int = 2048,\n",
        "    max_batch_prefill_tokens: int = 2048,\n",
        "    use_dedicated_endpoint: bool = False,\n",
        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    \"\"\"Deploys models with TGI on GPU in Vertex AI.\"\"\"\n",
        "    endpoint = aiplatform.Endpoint.create(\n",
        "        display_name=f\"{model_name}-endpoint\",\n",
        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
        "    )\n",
        "\n",
        "    env_vars = {\n",
        "        \"MODEL_ID\": model_id,\n",
        "        \"NUM_SHARD\": f\"{accelerator_count}\",\n",
        "        \"MAX_INPUT_LENGTH\": f\"{max_input_length}\",\n",
        "        \"MAX_TOTAL_TOKENS\": f\"{max_total_tokens}\",\n",
        "        \"MAX_BATCH_PREFILL_TOKENS\": f\"{max_batch_prefill_tokens}\",\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "\n",
        "    # HF_TOKEN is not a compulsory field and may not be defined.\n",
        "    try:\n",
        "        if HF_TOKEN:\n",
        "            env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
        "    except NameError:\n",
        "        pass\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=TGI_DOCKER_URI,\n",
        "        serving_container_ports=[8080],\n",
        "        serving_container_environment_variables=env_vars,\n",
        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
        "    )\n",
        "\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
        "    )\n",
        "    return model, endpoint\n",
        "\n",
        "\n",
        "DEPLOY_CONFIGS = [\n",
        "    DeployConfig(\n",
        "        display_name=\"microsoft/Phi-3-mini-4k-instruct (vLLM)\",\n",
        "        model_name=\"vllm-Phi-3-mini-4k-instruct\",\n",
        "        func=lambda x: deploy_model_vllm(x, \"microsoft/Phi-3-mini-4k-instruct\", None),\n",
        "    ),\n",
        "    DeployConfig(\n",
        "        display_name=\"Qwen/Qwen2-7B-Instruct (TGI)\",\n",
        "        model_name=\"tgi-Qwen2-7B-Instruct\",\n",
        "        func=lambda x: deploy_model_tgi(x, \"Qwen/Qwen2-7B-Instruct\", None),\n",
        "    ),\n",
        "]\n",
        "\n",
        "\n",
        "def get_server_type(endpoint: aiplatform.Endpoint) -> str | None:\n",
        "    \"\"\"Returns the model server type or None if not recognizable.\"\"\"\n",
        "    models = endpoint.list_models()\n",
        "    models: list[aiplatform.Model] = [aiplatform.Model(m.model) for m in models]\n",
        "    for server_type in SERVER_TYPES:\n",
        "        if any(server_type in model.container_spec.image_uri for model in models):\n",
        "            return server_type\n",
        "    return None\n",
        "\n",
        "\n",
        "def format_payload(messages: list[dict[str, str]]) -> dict[str, str]:\n",
        "    return {\n",
        "        \"messages\": messages,\n",
        "        \"max_tokens\": MAX_TOKENS,\n",
        "        \"stream\": True,\n",
        "    }\n",
        "\n",
        "\n",
        "def list_endpoints() -> list[tuple[str, str]]:\n",
        "    \"\"\"Returns all valid prediction endpoints for in the project and region.\"\"\"\n",
        "    endpoints = [\n",
        "        endpoint\n",
        "        for endpoint in aiplatform.Endpoint.list(order_by=\"create_time desc\")\n",
        "        if endpoint.traffic_split and get_server_type(endpoint)\n",
        "    ]\n",
        "    endpoints = [(e.display_name, e.resource_name) for e in endpoints]\n",
        "    endpoints.extend(\n",
        "        (e.display_name, e.resource_name) for e in PUBLIC_PLAYGROUND_ENDPOINTS\n",
        "    )\n",
        "    return endpoints\n",
        "\n",
        "\n",
        "class StreamingClient:\n",
        "    \"\"\"A wrapper for a streaming client.\"\"\"\n",
        "\n",
        "    endpoint: Endpoint | None = None\n",
        "\n",
        "    def set_endpoint(self, endpoint: str):\n",
        "        \"\"\"Sets the prediction endpoint.\"\"\"\n",
        "        playground_endpoint = [\n",
        "            e for e in PUBLIC_PLAYGROUND_ENDPOINTS if e.resource_name == endpoint\n",
        "        ]\n",
        "        if playground_endpoint:\n",
        "            self.endpoint = playground_endpoint[0]\n",
        "        else:\n",
        "            vertex_endpoint = aiplatform.Endpoint(endpoint)\n",
        "            server_type = get_server_type(vertex_endpoint)\n",
        "            self.endpoint = Endpoint(\n",
        "                display_name=vertex_endpoint.display_name,\n",
        "                location=vertex_endpoint.location,\n",
        "                resource_name=endpoint,\n",
        "                server_type=server_type,\n",
        "            )\n",
        "        print(\n",
        "            \"Selected endpoint:\",\n",
        "            self.endpoint.resource_name,\n",
        "            \"Server:\",\n",
        "            self.endpoint.server_type,\n",
        "        )\n",
        "\n",
        "    def predict(self, message: str, chat_history: list[tuple[str, str]]):\n",
        "        if not self.endpoint:\n",
        "            raise gr.Error(\"Select an endpoint first.\")\n",
        "\n",
        "        messages = []\n",
        "        for u, a in chat_history:\n",
        "            messages.append({\"role\": \"user\", \"content\": u})\n",
        "            messages.append({\"role\": \"assistant\", \"content\": a})\n",
        "        messages.append({\"role\": \"user\", \"content\": message})\n",
        "        payload = format_payload(messages)\n",
        "\n",
        "        is_playground_endpoint = self.endpoint.resource_name.startswith(\"playground:\")\n",
        "        if is_playground_endpoint:\n",
        "            url = f\"https://{self.endpoint.location}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{self.endpoint.location}/endpoints/openapi/chat/completions\"\n",
        "            payload[\"model\"] = self.endpoint.resource_name.removeprefix(\"playground:\")\n",
        "        else:\n",
        "            url = f\"https://{self.endpoint.location}-aiplatform.googleapis.com/v1beta1/{self.endpoint.resource_name}/chat/completions\"\n",
        "\n",
        "        access_token = ! gcloud auth print-access-token\n",
        "        access_token = access_token[0]\n",
        "        response = requests.post(\n",
        "            url,\n",
        "            headers={\"Authorization\": f\"Bearer {access_token}\"},\n",
        "            json=payload,\n",
        "            stream=True,\n",
        "        )\n",
        "        if not response.ok:\n",
        "            raise gr.Error(response)\n",
        "        prediction = \"\"\n",
        "        for chunk in response.iter_lines(chunk_size=8192, decode_unicode=False):\n",
        "            if chunk:\n",
        "                chunk = chunk.decode(\"utf-8\").removeprefix(\"data:\").strip()\n",
        "                if chunk == \"[DONE]\":\n",
        "                    break\n",
        "                data = json.loads(chunk)\n",
        "                if type(data) is not dict or \"error\" in data:\n",
        "                    raise gr.Error(data)\n",
        "                delta = data[\"choices\"][0][\"delta\"].get(\"content\")\n",
        "                if delta:\n",
        "                    prediction += delta\n",
        "                    yield prediction\n",
        "\n",
        "\n",
        "streaming_client = StreamingClient()\n",
        "\n",
        "\n",
        "def create_endpoint_selector():\n",
        "    \"\"\"Creates a dropdown list of prediction endpoints.\"\"\"\n",
        "\n",
        "    with gr.Row():\n",
        "        endpoints_dropdown = gr.Dropdown(\n",
        "            list_endpoints(),\n",
        "            label=\"Endpoint\",\n",
        "            scale=1,\n",
        "            info=\"Only TGI, vLLM, and HexLLM endpoints deployed after August 20, 2024 with a new container image support chat completions and streaming features. \"\n",
        "            + \"If you are not sure, you can deploy a demo endpoint directly from below. \",\n",
        "        )\n",
        "        endpoints_dropdown.input(\n",
        "            streaming_client.set_endpoint, inputs=[endpoints_dropdown], outputs=[]\n",
        "        )\n",
        "        refresh_btn = gr.Button(\"Refresh\", scale=0)\n",
        "        refresh_btn.click(\n",
        "            lambda: gr.Dropdown(choices=list_endpoints()),\n",
        "            inputs=[],\n",
        "            outputs=[endpoints_dropdown],\n",
        "        )\n",
        "\n",
        "\n",
        "def create_deploy_selector():\n",
        "    \"\"\"Creates a dropdown list of model deploy configs.\"\"\"\n",
        "\n",
        "    def find_deploy_config(display_name: str) -> DeployConfig:\n",
        "        \"\"\"Finds the deploy config from display name.\"\"\"\n",
        "        matches = [c for c in DEPLOY_CONFIGS if c.display_name == display_name]\n",
        "        if not matches:\n",
        "            raise gr.Error(\"Select a model to deploy first.\")\n",
        "        return matches[0]\n",
        "\n",
        "    def deploy(endpoint_name: str, display_name: str):\n",
        "        \"\"\"Deploys the model.\"\"\"\n",
        "        config = find_deploy_config(display_name)\n",
        "        gr.Info(f\"Deploying to {endpoint_name}...\")\n",
        "        config.func(endpoint_name)\n",
        "        gr.Info(f\"Deployed to {endpoint_name}. Refresh the endpoints to see it.\")\n",
        "\n",
        "    with gr.Row():\n",
        "        deploy_dropdown = gr.Dropdown(\n",
        "            [x.display_name for x in DEPLOY_CONFIGS],\n",
        "            label=\"Deploy Model\",\n",
        "            scale=1,\n",
        "            info=\"Model deployment will take ~20 minutes. After you finish your experiments, \"\n",
        "            + \"undeploy the endpoint from Vertex Online Prediction to avoid continuous charges to the project.\",\n",
        "        )\n",
        "        model_name = gr.Textbox(\n",
        "            label=\"Model Name\",\n",
        "            placeholder=\"Enter a custom model name for endpoint creation\",\n",
        "            interactive=True,\n",
        "        )\n",
        "        deploy_dropdown.change(\n",
        "            lambda x: find_deploy_config(x).model_name,\n",
        "            inputs=[deploy_dropdown],\n",
        "            outputs=[model_name],\n",
        "        )\n",
        "\n",
        "        deploy_btn = gr.Button(\"Deploy\", scale=0)\n",
        "        deploy_btn.click(\n",
        "            lambda: gr.Button(\"Deploying...\", interactive=False),\n",
        "            inputs=[],\n",
        "            outputs=[deploy_btn],\n",
        "        ).then(deploy, inputs=[model_name, deploy_dropdown], outputs=[]).then(\n",
        "            lambda: gr.Button(\"Deploy\", interactive=True), [], [deploy_btn]\n",
        "        )\n",
        "\n",
        "\n",
        "with gr.Blocks(title=\"Vertex Model Garden Chat\", fill_height=True) as demo:\n",
        "    create_endpoint_selector()\n",
        "    create_deploy_selector()\n",
        "    gr.ChatInterface(streaming_client.predict)\n",
        "\n",
        "\n",
        "show_debug_logs = True  # @param {type: \"boolean\"}\n",
        "demo.queue()\n",
        "demo.launch(share=True, inline=False, debug=show_debug_logs, show_error=True)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_gradio_streaming_chat_completions.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}