deploy_to_vertex_ai.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cc6e5f2b",
   "metadata": {},
   "source": [
    "# Deploying a Hugging Face model to Google Vertex AI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1dbbda7",
   "metadata": {},
   "source": [
    "Inspired by the [GCP tutorial]( https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/community-content/pytorch_text_classification_using_vertex_sdk_and_gcloud/pytorch-text-classification-vertex-ai-train-tune-deploy.ipynb) we will deploy a `sentence-transformers` model on a [Vertex AI](https://cloud.google.com/vertex-ai/docs/predictions/deploy-model-api) endpoint. We will use [TorchServe](https://pytorch.org/serve/) to serve a Hugging Face model available on the [Hub](hf.co). To accelerate inference we will also use features from the `optimum` [library](https://github.com/huggingface/optimum) to apply graph optimization and/or quantization to the model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "854097a7",
   "metadata": {},
   "source": [
    "### Set up your local development environment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df0c5ae2",
   "metadata": {},
   "source": [
    "1. Follow the Google Cloud guide to [setting up a Python development environment](https://cloud.google.com/python/docs/setup) \n",
    "2. [Install and initialize the Cloud SDK.](https://cloud.google.com/sdk/docs/) \n",
    "3. Create a virtual environment (virtualenv, pyenv) with Python 3 (<3.9) and activate the environment\n",
    "4. Launch jupyter notebook from this environment\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7626c33b",
   "metadata": {},
   "outputs": [],
   "source": [
    "https://cloud.google.com/products/calculator/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41229c4d",
   "metadata": {},
   "source": [
    "### Install packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e0237e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip -q install --upgrade google-cloud-aiplatform #Vertex AI sdk\n",
    "!pip -q install --upgrade transformers\n",
    "!pip -q install --upgrade datasets\n",
    "!pip -q install --upgrade 'optimum[onnxruntime]'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4b7143d",
   "metadata": {},
   "source": [
    "### Set up your Google Cloud project"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "445b5bee",
   "metadata": {},
   "source": [
    "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager)\n",
    "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project)\n",
    "1. Enable following APIs in your project required for running the tutorial\n",
    "    - [Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com)\n",
    "    - [Cloud Storage API](https://console.cloud.google.com/flows/enableapi?apiid=storage.googleapis.com)\n",
    "    - [Container Registry API](https://console.cloud.google.com/flows/enableapi?apiid=containerregistry.googleapis.com)\n",
    "    - [Cloud Build API](https://console.cloud.google.com/flows/enableapi?apiid=cloudbuild.googleapis.com)\n",
    "   \n",
    "1. Enter your project ID in the cell below. Then run the cell to make sure the Cloud SDK uses the right project for all the commands in this notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11bc4202",
   "metadata": {},
   "source": [
    "### Authenticate to gcloud"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef0bcad2",
   "metadata": {},
   "source": [
    " 1. In the Cloud Console, go to the [**Create service account key** page](https://console.cloud.google.com/apis/credentials/serviceaccountkey).,\n",
    " 2. Click **Create service account**.,\n",
    " 3. In the **Service account name** field, enter a name, and click **Create**,\n",
    " 4. In the **Grant this service account access to project** section, click the **Role** drop-down list. Type \\\"Vertex AI\\\" into the filter box, and select **Vertex AI Administrator**. Type \\\"Storage Object Admin\\\" into the filter box, and select **Storage Object Admin**.\n",
    " 5. Click *Create*. A JSON file that contains your key downloads to your local environment.\n",
    " 6. 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,
   "id": "ee9cbff5",
   "metadata": {},
   "outputs": [],
   "source": [
    "%env GOOGLE_APPLICATION_CREDENTIALS ./keys/huggingface-ml-e974975230cc.json #change to your service account key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a4d2a8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get your Google Cloud project ID using google.auth\n",
    "import google.auth\n",
    "\n",
    "_, PROJECT_ID = google.auth.default()\n",
    "print(\"Project ID: \", PROJECT_ID)\n",
    "\n",
    "#Or set it yourself manually\n",
    "PROJECT_ID = \"huggingface-ml\" "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01857fb5",
   "metadata": {},
   "source": [
    "### Create a cloud storage bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cb4c6ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "BUCKET_NAME = \"gs://florent-bucket\"  # <---CHANGE THIS TO YOUR BUCKET\n",
    "REGION = \"us-central1\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b366365",
   "metadata": {},
   "source": [
    "**If the bucket doesn't exist, run the following:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9c1fb26",
   "metadata": {},
   "outputs": [],
   "source": [
    "#! gsutil mb -l $REGION $BUCKET_NAME"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecd343f3",
   "metadata": {},
   "source": [
    "Access the content of the bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20adc9c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "! gsutil ls -al $BUCKET_NAME"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "827daccf",
   "metadata": {},
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d25c6e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import base64\n",
    "import json\n",
    "import os\n",
    "import random\n",
    "import sys\n",
    "import transformers\n",
    "\n",
    "import google.auth\n",
    "from google.cloud import aiplatform\n",
    "from google.cloud.aiplatform import gapic as aip\n",
    "from google.protobuf.json_format import MessageToDict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa5e61db",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Notebook runtime: {'GPU' if torch.cuda.is_available() else 'CPU'}\")\n",
    "print(f\"PyTorch version : {torch.__version__}\")\n",
    "print(f\"Transformers version : {transformers.__version__}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4818dd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "APP_NAME = \"test_sbert_embedder_optimum\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed45919e",
   "metadata": {},
   "source": [
    "## Deployment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "841c6e37",
   "metadata": {},
   "source": [
    "#### *Overview*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e05d7bb",
   "metadata": {},
   "source": [
    "Deploying a PyTorch model on [Vertex AI Predictions](https://cloud.google.com/vertex-ai/docs/predictions/getting-predictions) requires to use a custom container that serves online predictions. You will deploy a container running [PyTorch's TorchServe](https://pytorch.org/serve/) tool in order to serve predictions from a fine-tuned sentence transformer model `msmarco-distilbert-base-tas-b` available in [Hugging Face Transformers](https://huggingface.co/sentence-transformers/msmarco-distilbert-base-tas-b). "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7b69d42",
   "metadata": {},
   "source": [
    "Essentially, to deploy a PyTorch model on Vertex AI Predictions following are the steps:\n",
    "1. Package the trained model artifacts including [default](https://pytorch.org/serve/#default-handlers) or [custom](https://pytorch.org/serve/custom_service.html) handlers by creating an archive file using [Torch model archiver](https://github.com/pytorch/serve/tree/master/model-archiver),\n",
    "2. Build a [custom container](https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements) compatible with Vertex AI Predictions to serve the model using Torchserve\n",
    "3. Upload the model with custom container image to serve predictions as a Vertex AI Model resource,\n",
    "4. Create a Vertex AI Endpoint and [deploy the model](https://cloud.google.com/vertex-ai/docs/predictions/deploy-model-api) resource"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c1149d0",
   "metadata": {},
   "source": [
    "#### *How to improve latency*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cbf4e1c",
   "metadata": {},
   "source": [
    "Deployment of the model will be made here on CPU. To improve latency of the model we will use the [Hugging Face Optimum](https://github.com/huggingface/optimum) library to convert the model to the [ONNX (Open Neural Network eXchange)](http://onnx.ai/) format and apply graph optimization and/or quantization to improve inference time. To learn more about these techniques consult:\n",
    "- [Hugging Face Optimum documentation](https://huggingface.co/docs/optimum/quickstart)\n",
    "- [Convert Transformers to ONNX with Hugging Face Optimum](https://huggingface.co/blog/convert-transformers-to-onnx#2-what-is-hugging-face-optimum)\n",
    "- [Graph Optimizations in ONNX Runtime](https://onnxruntime.ai/docs/performance/graph-optimizations.html)\n",
    "- [Quantize ONNX Models](https://onnxruntime.ai/docs/performance/quantization.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03e98ae4",
   "metadata": {},
   "source": [
    "Those operations need to be performed before using the Torch model archiver. The ONNX exported model will then be loaded in the custom handler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "035c60c5",
   "metadata": {},
   "source": [
    "### Save model locally"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbf0aa39",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir ./predictor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0825e6ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModel\n",
    "\n",
    "model_name = \"sentence-transformers/msmarco-distilbert-base-tas-b\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "model = AutoModel.from_pretrained(model_name)\n",
    "\n",
    "\n",
    "pt_save_directory = \"./predictor/model/\"\n",
    "\n",
    "tokenizer.save_pretrained(pt_save_directory)\n",
    "model.save_pretrained(pt_save_directory)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78b713d2",
   "metadata": {},
   "source": [
    "### Apply optimum optimizations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2993f6da",
   "metadata": {},
   "source": [
    "Optimization is enough here for the latency we need but you can also apply quantization with `ORTQuantizer` if you need faster predictions. However this may affect the performance of the model. See the [documentation](https://huggingface.co/docs/optimum/main/en/pipelines#quantizing-with-ortquantizer)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebdfbf2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from optimum.onnxruntime import ORTModelForFeatureExtraction, ORTOptimizer\n",
    "from optimum.onnxruntime.configuration import OptimizationConfig\n",
    "from optimum.pipelines import pipeline\n",
    "\n",
    "\n",
    "pt_save_directory_optimum = \"./predictor/optimum/\"\n",
    "\n",
    "save_path = Path(\"optimum_model\")\n",
    "save_path.mkdir(exist_ok=True)\n",
    "\n",
    "#use ORTOptimizer to export the model and define quantization configuration\n",
    "optimizer = ORTOptimizer(model=model, tokenizer=tokenizer)\n",
    "optimization_config = OptimizationConfig(optimization_level=2)\n",
    "\n",
    "\n",
    "# apply the optimization configuration to the model\n",
    "optimizer.export(\n",
    "    onnx_model_path=save_path / \"model.onnx\",\n",
    "    onnx_optimized_model_output_path=save_path / \"model-optimized.onnx\",\n",
    "    optimization_config=optimization_config,\n",
    ")\n",
    "\n",
    "optimizer.model.config.save_pretrained(save_path) # saves config.json \n",
    "\n",
    "model = ORTModelForFeatureExtraction.from_pretrained(save_path, file_name=\"model-optimized.onnx\")\n",
    "\n",
    "tokenizer.save_pretrained(pt_save_directory_optimum)\n",
    "model.save_pretrained(pt_save_directory_optimum)\n",
    "\n",
    "#You can also push the model to the HF hub\n",
    "#model.push_to_hub(pt_save_directory_optimum,\n",
    "#                  repository_id=\"onnx-msmarco-distilbert-base-tas-b\",\n",
    "#                  use_auth_token=True\n",
    "#                  )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44d1b219",
   "metadata": {},
   "source": [
    "### Create a custom model handler "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd2eb813",
   "metadata": {},
   "source": [
    "Please refer to the [TorchServe documentation](https://pytorch.org/serve/custom_service.html) for defining a custom handler."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9de29252",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile predictor/custom_handler.py\n",
    "\n",
    "import os\n",
    "import json\n",
    "import logging\n",
    "\n",
    "import torch\n",
    "from transformers import AutoModel, AutoTokenizer\n",
    "from optimum.onnxruntime import ORTModelForFeatureExtraction, ORTOptimizer\n",
    "from optimum.onnxruntime.configuration import OptimizationConfig\n",
    "from optimum.pipelines import pipeline\n",
    "\n",
    "from ts.torch_handler.base_handler import BaseHandler\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "torch.set_num_threads(1)\n",
    "\n",
    "class SentenceTransformersHandler(BaseHandler):\n",
    "    \"\"\"\n",
    "    The handler takes an input string and returns the embedding \n",
    "    based on the serialized transformers checkpoint.\n",
    "    \"\"\"\n",
    "    def __init__(self):\n",
    "        super(SentenceTransformersHandler, self).__init__()\n",
    "        self.initialized = False\n",
    "\n",
    "    def initialize(self, ctx):\n",
    "        \"\"\" Loads the model.onnx file and initialized the model object.\n",
    "        Instantiates Tokenizer for preprocessor to use and a feature extraction pipeline\n",
    "        \"\"\"\n",
    "        self.manifest = ctx.manifest\n",
    "\n",
    "        properties = ctx.system_properties\n",
    "        model_dir = properties.get(\"model_dir\")\n",
    "        #self.device = torch.device(\"cuda:\" + str(properties.get(\"gpu_id\")) if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "        # Read model serialize/pt file\n",
    "        serialized_file = self.manifest[\"model\"][\"serializedFile\"]\n",
    "        model_pt_path = os.path.join(model_dir, serialized_file)\n",
    "        if not os.path.isfile(model_pt_path):\n",
    "            raise RuntimeError(\"Missing the model.onnx or pytorch_model.bin file\")\n",
    "        \n",
    "        # Load model\n",
    "        self.model = ORTModelForFeatureExtraction.from_pretrained(model_dir)\n",
    "        logger.debug('Transformer model from path {0} loaded successfully'.format(model_dir))\n",
    "        \n",
    "        # Ensure to use the same tokenizer used during training\n",
    "        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, model_max_length=128)\n",
    "        \n",
    "        # Create an optimum pipeline\n",
    "        self.pipeline = pipeline(\"feature-extraction\", model=self.model, tokenizer=self.tokenizer)\n",
    "\n",
    "        self.initialized = True\n",
    "\n",
    "    def preprocess(self, data):\n",
    "        \"\"\" Preprocessing input request by tokenizing\n",
    "            Extend with your own preprocessing steps as needed\n",
    "        \"\"\"\n",
    "        text = data[0].get(\"data\")\n",
    "        if text is None:\n",
    "            text = data[0].get(\"body\")\n",
    "        sentences = text.decode('utf-8')\n",
    "        logger.info(\"Received text: '%s'\", sentences)\n",
    "        return sentences\n",
    "\n",
    "    def inference(self, sentences):\n",
    "        \"\"\" Predict the class of a text using a trained transformer model.\n",
    "        \"\"\"\n",
    "        \n",
    "        def cls_pooling(pipeline_output):\n",
    "            \"\"\"\n",
    "            Return the [CLS] token embedding\n",
    "            \"\"\"\n",
    "            return [_h[0] for _h in pipeline_output]\n",
    "        \n",
    "        embeddings = cls_pooling(self.pipeline(sentences))\n",
    "\n",
    "        logger.info(f\"Model embedded: {len(embeddings)}\" )\n",
    "        return embeddings\n",
    "\n",
    "    def postprocess(self, inference_output):\n",
    "        return inference_output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f44f516",
   "metadata": {},
   "source": [
    "### Create custom container image"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5fb4a9a",
   "metadata": {},
   "source": [
    "**Create a Dockerfile with TorchServe as base image**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cbb394c",
   "metadata": {},
   "source": [
    "**NB**: to define the right Torchserve parameters such as `workers` please consult (https://github.com/pytorch/serve/blob/master/docs/performance_guide.md) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca79194d",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash -s $APP_NAME\n",
    "\n",
    "APP_NAME=$1\n",
    "\n",
    "cat << EOF > ./predictor/Dockerfile\n",
    "\n",
    "FROM pytorch/torchserve:latest-cpu\n",
    "\n",
    "# install dependencies\n",
    "RUN python3 -m pip install --upgrade pip\n",
    "RUN pip3 install transformers\n",
    "RUN pip3 install 'optimum[onnxruntime]'\n",
    "\n",
    "\n",
    "USER model-server\n",
    "\n",
    "# copy model artifacts, custom handler and other dependencies\n",
    "COPY custom_handler.py /home/model-server/\n",
    "COPY ./optimum/ / /home/model-server/\n",
    "\n",
    "# create torchserve configuration file\n",
    "USER root\n",
    "RUN printf \"\\nservice_envelope=json\" >> /home/model-server/config.properties\n",
    "RUN printf \"\\ninference_address=http://0.0.0.0:7080\" >> /home/model-server/config.properties\n",
    "RUN printf \"\\nmanagement_address=http://0.0.0.0:7081\" >> /home/model-server/config.properties\n",
    "RUN printf \"\\nmanagement_address=http://0.0.0.0:7081\" >> /home/model-server/config.properties\n",
    "\n",
    "# Consult https://github.com/pytorch/serve/blob/master/docs/performance_guide.md to define the right parameters\n",
    "RUN printf \"\\nworkers=4\" >> /home/model-server/config.properties\n",
    "\n",
    "# expose health and prediction listener ports from the image\n",
    "EXPOSE 7080\n",
    "EXPOSE 7081\n",
    "\n",
    "# create model archive file packaging model artifacts and dependencies\n",
    "RUN torch-model-archiver -f \\\n",
    "  --model-name=$APP_NAME \\\n",
    "  --version=1.0 \\\n",
    "  --serialized-file=/home/model-server/model.onnx \\\n",
    "  --handler=/home/model-server/custom_handler.py \\\n",
    "  --extra-files \"/home/model-server/config.json,/home/model-server/tokenizer.json,/home/model-server/tokenizer_config.json,/home/model-server/special_tokens_map.json,/home/model-server/vocab.txt\" \\\n",
    "  --export-path=/home/model-server/model-store\n",
    "\n",
    "# run Torchserve HTTP serve to respond to prediction requests\n",
    "CMD [\"torchserve\", \\\n",
    "     \"--start\", \\\n",
    "     \"--ts-config=/home/model-server/config.properties\", \\\n",
    "     \"--models\", \\\n",
    "     \"$APP_NAME=$APP_NAME.mar\", \\\n",
    "     \"--model-store\", \\\n",
    "     \"/home/model-server/model-store\"]\n",
    "EOF\n",
    "\n",
    "echo \"Writing ./predictor/Dockerfile\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca594d5b",
   "metadata": {},
   "source": [
    "**Build container**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2b160e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "CUSTOM_PREDICTOR_IMAGE_URI = f\"gcr.io/{PROJECT_ID}/pytorch_predict_{APP_NAME}\"\n",
    "print(f\"CUSTOM_PREDICTOR_IMAGE_URI = {CUSTOM_PREDICTOR_IMAGE_URI}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d09ce47",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!docker build \\\n",
    "  --tag=$CUSTOM_PREDICTOR_IMAGE_URI \\\n",
    "  ./predictor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d3bdd6f",
   "metadata": {},
   "source": [
    "**Run container locally**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d23d9ad8",
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker stop local_sbert_embedder_optimum\n",
    "!docker run -t -d --rm -p 7080:7080 --name=local_sbert_embedder_optimum $CUSTOM_PREDICTOR_IMAGE_URI\n",
    "!sleep 20"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97df6271",
   "metadata": {},
   "source": [
    "**Test API locally**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4816b243",
   "metadata": {},
   "source": [
    "1. Health check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f66cc11b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!curl http://localhost:7080/ping"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c175997",
   "metadata": {},
   "source": [
    "2. Send request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dcbddc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash -s $APP_NAME\n",
    "\n",
    "APP_NAME=$1\n",
    "\n",
    "cat > ./predictor/instances.json <<END\n",
    "{ \n",
    "   \"instances\": [\n",
    "     { \n",
    "       \"data\": {\n",
    "         \"b64\": \"$(echo 'I am creating an endpoint using TorchServe and HF transformers' | base64 --wrap=0)\"\n",
    "       }\n",
    "     }\n",
    "   ]\n",
    "}\n",
    "END\n",
    "\n",
    "curl -s -X POST \\\n",
    "  -H \"Content-Type: application/json; charset=utf-8\" \\\n",
    "  -d @./predictor/instances.json \\\n",
    "  http://localhost:7080/predictions/$APP_NAME/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f906fc1",
   "metadata": {},
   "source": [
    "3. Stop the container"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "726bd349",
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker stop local_sbert_embedder_optimum"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c24c5d0c",
   "metadata": {},
   "source": [
    "### Push image Container Registry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2611ed58",
   "metadata": {},
   "outputs": [],
   "source": [
    "!docker push $CUSTOM_PREDICTOR_IMAGE_URI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14e9565b",
   "metadata": {},
   "source": [
    "### Create model and endpoint to VertexAI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba10b3e0",
   "metadata": {},
   "source": [
    "We create a model resource on Vertex AI and deploy the model to a Vertex AI Endpoints. You must deploy a model to an endpoint before using the model. The deployed model runs the custom container image to serve predictions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02b99644",
   "metadata": {},
   "source": [
    "**Initialize the Vertex AI SDK for Python**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5f00242",
   "metadata": {},
   "outputs": [],
   "source": [
    "aiplatform.init(project=PROJECT_ID, staging_bucket=BUCKET_NAME)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e418380",
   "metadata": {},
   "source": [
    "**Create a Model resource with custom serving container**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "170d5ba9",
   "metadata": {},
   "outputs": [],
   "source": [
    "VERSION = 1\n",
    "model_display_name = f\"{APP_NAME}-v{VERSION}\"\n",
    "model_description = \"PyTorch based sentence transformers embedder with custom container\"\n",
    "\n",
    "MODEL_NAME = APP_NAME\n",
    "health_route = \"/ping\"\n",
    "predict_route = f\"/predictions/{MODEL_NAME}\"\n",
    "serving_container_ports = [7080]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6eab0c84",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = aiplatform.Model.upload(\n",
    "    display_name=model_display_name,\n",
    "    description=model_description,\n",
    "    serving_container_image_uri=CUSTOM_PREDICTOR_IMAGE_URI,\n",
    "    serving_container_predict_route=predict_route,\n",
    "    serving_container_health_route=health_route,\n",
    "    serving_container_ports=serving_container_ports,\n",
    ")\n",
    "\n",
    "model.wait()\n",
    "\n",
    "print(model.display_name)\n",
    "print(model.resource_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7343cea4",
   "metadata": {},
   "source": [
    "**Create an Endpoint for Model with Custom Container**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3863a48",
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint_display_name = f\"{APP_NAME}-endpoint\"\n",
    "endpoint = aiplatform.Endpoint.create(display_name=endpoint_display_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea432646",
   "metadata": {},
   "source": [
    "**Deploy the Model to Endpoint**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b772534",
   "metadata": {},
   "source": [
    "See more on the [documentation](https://cloud.google.com/vertex-ai/docs/predictions/deploy-model-api)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d52b599d",
   "metadata": {},
   "source": [
    "To select the right machine type according to your budget select go to [Google Cloud Pricing Calculator](https://cloud.google.com/products/calculator) and [Finding the ideal machine type](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute#finding_the_ideal_machine_type)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecad0cb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "traffic_percentage = 100\n",
    "machine_type = \"n1-standard-8\"\n",
    "deployed_model_display_name = model_display_name\n",
    "min_replica_count = 1\n",
    "max_replica_count = 3\n",
    "sync = True\n",
    "\n",
    "model.deploy(\n",
    "    endpoint=endpoint,\n",
    "    deployed_model_display_name=deployed_model_display_name,\n",
    "    machine_type=machine_type,\n",
    "    traffic_percentage=traffic_percentage,\n",
    "    min_replica_count=min_replica_count,\n",
    "    max_replica_count=max_replica_count,\n",
    "    sync=sync,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1add143",
   "metadata": {},
   "source": [
    "### Invoking the Endpoint with deployed Model using Vertex AI SDK to make predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49af35f0",
   "metadata": {},
   "source": [
    "**Get the endpoint id**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a405988",
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint_display_name = f\"{APP_NAME}-endpoint\"\n",
    "filter = f'display_name=\"{endpoint_display_name}\"'\n",
    "\n",
    "for endpoint_info in aiplatform.Endpoint.list(filter=filter):\n",
    "    print(\n",
    "        f\"Endpoint display name = {endpoint_info.display_name} resource id ={endpoint_info.resource_name} \"\n",
    "    )\n",
    "\n",
    "endpoint = aiplatform.Endpoint(endpoint_info.resource_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c937f2ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint.list_models()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a59ee57",
   "metadata": {},
   "source": [
    "**Formatting input for online prediction**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16cf7c7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_instances = [\n",
    "    b\"This is an example of model deployment using a sentence transformers model and optimum\",\n",
    "]*100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f260939",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(tokenizer(test_instances[0])[\"input_ids\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0268a7d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#test_instances"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a19a344",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "print(\"=\" * 100)\n",
    "for instance in test_instances:\n",
    "    print(f\"Input text: \\n\\t{instance.decode('utf-8')}\\n\")\n",
    "    b64_encoded = base64.b64encode(instance)\n",
    "    test_instance = [{\"data\": {\"b64\": f\"{str(b64_encoded.decode('utf-8'))}\"}}]\n",
    "    print(f\"Formatted input: \\n{json.dumps(test_instance, indent=4)}\\n\")\n",
    "    prediction = endpoint.predict(instances=test_instance)\n",
    "    #print(f\"Prediction response: \\n\\t{prediction}\")\n",
    "    print(\"=\" * 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db1a6d0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "prediction = endpoint.predict(instances=test_instance)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.10 64-bit ('venv_hf_3.9.10')",
   "language": "python",
   "name": "python3910jvsc74a57bd041d616a00e55b30810a056d2ec88dde9c7f0c29a440bd149fc2105d405ad956d"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}