{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2025 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Import, Deploy, and Serve custom open models on Vertex AI using Vertex AI Model Garden SDK.\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/open-models/get_started_with_model_garden_sdk_custom_import.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fopen-models%2Fget_started_with_model_garden_sdk_custom_import.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/open-models/get_started_with_model_garden_sdk_custom_import.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/get_started_with_model_garden_sdk_custom_import.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/get_started_with_model_garden_sdk_custom_import.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/get_started_with_model_garden_sdk_custom_import.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/get_started_with_model_garden_sdk_custom_import.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/get_started_with_model_garden_sdk_custom_import.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/open-models/get_started_with_model_garden_sdk_custom_import.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "84f0f73a0f76"
      },
      "source": [
        "| Authors |\n",
        "| --- |\n",
        "| [Ivan Nardini](https://github.com/inardini) |\n",
        "| [Eliza Huang](https://github.com/lizzij) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "The Vertex AI Model Garden allows you to not only use Google's state-of-the-art models but also to bring your own. \n",
        "\n",
        "This tutorial shows how to take an open-source model from Hugging Face, importing it into Vertex AI, and serving it on a scalable, production-ready endpoint using Vertex AI Model Garden SDK Import Custom Model Weights feature. \n",
        "\n",
        "This feature gives you the flexibility to use a vast ecosystem of open models while leveraging the Google Cloud's enterprise-grade infrastructure on Vertex AI for serving models.\n",
        "\n",
        "### What You'll Learn\n",
        "\n",
        "  * How to efficiently transfer model artifacts from Hugging Face Hub directly to Google Cloud Storage (GCS).\n",
        "  * How to register your model with the Vertex AI Model Registry using a GCS URI.\n",
        "  * How to deploy the custom model to a dedicated Vertex AI Endpoint with a single command.\n",
        "  * How to run inference on your newly deployed model.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install required packages\n",
        "\n",
        "Install the necessary Python packages for this tutorial.\n",
        "\n",
        "*Note: The hf_transfer extra enables faster downloads using Rust-based transfer. This can speed up large model downloads by 2-5x*\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --force-reinstall --quiet 'google-cloud-aiplatform>=1.105.0' 'openai' 'google-auth' 'requests' 'huggingface_hub[hf_transfer]'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "# import sys\n",
        "\n",
        "# if \"google.colab\" in sys.modules:\n",
        "#     from google.colab import auth\n",
        "\n",
        "#     auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "# Use the environment variable if the user doesn't provide Project ID.\n",
        "import os\n",
        "\n",
        "import vertexai\n",
        "\n",
        "# fmt: off\n",
        "PROJECT_ID = \"[your-project-id]\"  # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
        "\n",
        "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
        "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
        "\n",
        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
        "\n",
        "BUCKET_NAME = \"[your-bucket-name]\"  # @param {type: \"string\", placeholder: \"[your-bucket-name]\", isTemplate: true}\n",
        "# fmt: on\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}\"\n",
        "\n",
        "! gcloud storage buckets create {BUCKET_URI} --project={PROJECT_ID} --location={LOCATION}\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4oWLaELWrRPz"
      },
      "source": [
        "### Authenticate your HuggingFace account\n",
        "\n",
        "To download models from Hugging Face, especially gated ones, you need to authenticate.\n",
        "\n",
        "The `interpreter_login` function provides an easy way to do this within a notebook environment.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BQcw4H0srWNE"
      },
      "outputs": [],
      "source": [
        "from huggingface_hub import interpreter_login\n",
        "\n",
        "interpreter_login()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5303c05f7aa6"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6fc324893334"
      },
      "outputs": [],
      "source": [
        "import shutil\n",
        "from pathlib import Path\n",
        "\n",
        "import google.auth\n",
        "import openai\n",
        "from google.cloud import storage\n",
        "from google.cloud.storage import transfer_manager\n",
        "from huggingface_hub import snapshot_download\n",
        "from tqdm.auto import tqdm\n",
        "from vertexai.preview import model_garden"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w0Lwlu7DM3Bw"
      },
      "source": [
        "### Helpers\n",
        "\n",
        "Before Vertex AI can deploy a custom model, its artifacts must be accessible in a Google Cloud Storage bucket. This helper function, `transfer_model`, automates the entire process.\n",
        "\n",
        "It handles:\n",
        "\n",
        "1.  **Fast Downloads**: Enables `hf_transfer` for accelerated downloads from Hugging Face.\n",
        "2.  **Local Staging**: Downloads the model files to a temporary local directory.\n",
        "3.  **Efficient Uploads**: Uses the `transfer_manager` to upload files to GCS in parallel chunks, which is much faster for large model weights.\n",
        "4.  **Automatic Cleanup**: Removes the temporary local files after the upload is complete."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Fi0qLHMxrgfc"
      },
      "outputs": [],
      "source": [
        "def transfer_model(\n",
        "    model_id: str,\n",
        "    bucket_name: str,\n",
        "    location: str,\n",
        "    artifact_path: str = \"\",\n",
        "    exclude_patterns: list | None = None,\n",
        "    chunk_size: int = 32 * 1024 * 1024,  # 32 MiB\n",
        "    workers: int = 8,\n",
        ") -> str:\n",
        "    \"\"\"Download model from HuggingFace and upload to GCS with parallel transfers.\"\"\"\n",
        "    # Defaults\n",
        "    exclude_patterns = exclude_patterns or [\n",
        "        \"*.bin\",\n",
        "        \"*.pth\",\n",
        "        \"*.gguf\",\n",
        "        \".gitattributes\",\n",
        "    ]\n",
        "    artifact_path = artifact_path or model_id.split(\"/\")[-1]\n",
        "    local_dir = Path(f\"tmp/{model_id.replace('/', '--')}\")\n",
        "\n",
        "    # Enable fast transfers\n",
        "    os.environ[\"HF_HUB_ENABLE_HF_TRANSFER\"] = \"1\"\n",
        "\n",
        "    try:\n",
        "        # Download from HuggingFace\n",
        "        print(f\"📥 Downloading {model_id}...\")\n",
        "        snapshot_download(\n",
        "            repo_id=model_id,\n",
        "            local_dir=str(local_dir),\n",
        "            ignore_patterns=exclude_patterns,\n",
        "            resume_download=True,\n",
        "        )\n",
        "\n",
        "        # Setup GCS\n",
        "        client = storage.Client()\n",
        "        bucket = client.bucket(bucket_name)\n",
        "\n",
        "        # Create bucket if needed\n",
        "        if not bucket.exists():\n",
        "            print(f\"📦 Creating bucket {bucket_name}...\")\n",
        "            bucket = client.create_bucket(bucket_name, location=location)\n",
        "\n",
        "        # Get files to upload\n",
        "        files = [f for f in local_dir.rglob(\"*\") if f.is_file()]\n",
        "        total_size = sum(f.stat().st_size for f in files) / (1024**3)\n",
        "\n",
        "        print(f\"☁️  Uploading {len(files)} files ({total_size:.2f} GB) to GCS...\")\n",
        "\n",
        "        # Upload files with parallel chunks\n",
        "        for file_path in tqdm(files, desc=\"Files\", unit=\"file\"):\n",
        "            blob_name = f\"{artifact_path}/{file_path.relative_to(local_dir)}\"\n",
        "            blob = bucket.blob(blob_name)\n",
        "\n",
        "            # Use parallel upload for large files (>32MB), regular for small files\n",
        "            if file_path.stat().st_size > chunk_size:\n",
        "                transfer_manager.upload_chunks_concurrently(\n",
        "                    str(file_path), blob, chunk_size=chunk_size, max_workers=workers\n",
        "                )\n",
        "            else:\n",
        "                blob.upload_from_filename(str(file_path))\n",
        "\n",
        "        gcs_path = f\"gs://{bucket_name}/{artifact_path}\"\n",
        "        print(f\"✅ Done! Model available at {gcs_path}\")\n",
        "        return gcs_path\n",
        "\n",
        "    finally:\n",
        "        # Cleanup\n",
        "        if local_dir.exists():\n",
        "            shutil.rmtree(local_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HX7UhfWdsN3w"
      },
      "source": [
        "## Transfer Your Model from Hugging Face to GCS\n",
        "\n",
        "Now, let's use the function to transfer a fine-tuned Gemma model from Hugging Face to our GCS bucket.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uqf_RmUwv7z5"
      },
      "outputs": [],
      "source": [
        "# The Hugging Face model we want to import.\n",
        "hf_model_id = \"xsanskarx/thinkygemma-4b\"\n",
        "\n",
        "# This command will download the model and upload it to your GCS bucket.\n",
        "# This may take several minutes depending on the model size and your connection.\n",
        "imported_custom_model_uri = transfer_model(hf_model_id, BUCKET_NAME, LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hUhWp6uHwa7s"
      },
      "source": [
        "## Import and Deploy the Model with Model Garden\n",
        "\n",
        "With the model artifacts in GCS, we can now use the `model_garden` SDK to register and deploy it.\n",
        "\n",
        "First, we create a `CustomModel` object, pointing it to the GCS URI where our model is stored."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TAKZG1sBwfRs"
      },
      "outputs": [],
      "source": [
        "model = model_garden.CustomModel(\n",
        "    gcs_uri=imported_custom_model_uri,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "215f72a4a900"
      },
      "source": [
        "After you initiate the model, use `list_deploy_options()` method to discover the verified deployment configurations supported by a specific model.\n",
        "\n",
        "This is important to verify if you have enough resources to deploy the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0a490035154e"
      },
      "outputs": [],
      "source": [
        "deploy_options = model.list_deploy_options()\n",
        "print(deploy_options)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6f2815215036"
      },
      "source": [
        "Next, we use the suggested configuration to set the serving infrastructure for our model's endpoint and deploy the model.\n",
        "\n",
        "The real magic happens with the **`model.deploy()`** command. This single line of code abstracts away a complex series of operations:\n",
        "\n",
        "  * It creates a new **Vertex AI Model** resource from your GCS artifacts.\n",
        "  * It provisions the specified compute resources.\n",
        "  * It creates a **Vertex AI Endpoint**.\n",
        "  * It deploys the model to that endpoint, making it ready to serve inference requests."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "18b2o2qx0jIh"
      },
      "outputs": [],
      "source": [
        "# Define the machine resources for the serving endpoint.\n",
        "# g2-standard-24 and 2 L4 GPUs is a robust configuration for many ~7B models.\n",
        "MACHINE_TYPE = \"g2-standard-24\"\n",
        "ACCELERATOR_TYPE = \"NVIDIA_L4\"\n",
        "ACCELERATOR_COUNT = 2\n",
        "\n",
        "# Deploy the model. This will create a Vertex AI Endpoint and deploy the model to it.\n",
        "# This step can take 15-20 minutes as it provisions the necessary hardware.\n",
        "endpoint = model.deploy(\n",
        "    machine_type=MACHINE_TYPE,\n",
        "    accelerator_type=ACCELERATOR_TYPE,\n",
        "    accelerator_count=ACCELERATOR_COUNT,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oynBjezZvWIq"
      },
      "source": [
        "## Run Inference on Your Deployed Model\n",
        "\n",
        "Once the deployment is complete, your model is live and ready to serve predictions. You can interact with the endpoint using the `predict()` method.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0LHF_UR_vd39"
      },
      "outputs": [],
      "source": [
        "response = endpoint.predict(\n",
        "    instances=[{\"prompt\": \"how many r does strawberry have?\"}],\n",
        "    use_dedicated_endpoint=True,\n",
        ")\n",
        "print(response.predictions)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tmuwCPShxso1"
      },
      "source": [
        "You can also get your prediction using ChatCompletion API from the OpenAI SDK."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "acGaSDaH9YAp"
      },
      "outputs": [],
      "source": [
        "creds, project = google.auth.default()\n",
        "auth_req = google.auth.transport.requests.Request()\n",
        "creds.refresh(auth_req)\n",
        "\n",
        "endpoint_url = f\"https://{endpoint.gca_resource.dedicated_endpoint_dns}/v1beta1/{endpoint.resource_name}\"\n",
        "\n",
        "client = openai.OpenAI(base_url=endpoint_url, api_key=creds.token)\n",
        "\n",
        "prediction = client.chat.completions.create(\n",
        "    model=\"\", messages=[{\"role\": \"user\", \"content\": \"Tell me a joke\"}], temperature=0.7\n",
        ")\n",
        "\n",
        "print(prediction.choices[0].message.content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LE7mv18CvvZs"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To avoid incurring ongoing charges to your Google Cloud account, it's important to clean up the resources you've created. The following commands will undeploy the model from the endpoint and then delete the endpoint itself.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r-LKGlR2v32X"
      },
      "outputs": [],
      "source": [
        "delete_endpoint = True\n",
        "delete_bucket = True\n",
        "\n",
        "if delete_endpoint:\n",
        "    endpoint.delete(force=True)\n",
        "\n",
        "if delete_bucket:\n",
        "    !gcloud storage rm --recursive {BUCKET_URI}"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "get_started_with_model_garden_sdk_custom_import.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
