hugging_face_to_model_registry.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyO+Zl2ZJw23uX477M3Qooz6",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/rastringer/vertex-ai-examples/blob/main/hugging_face_to_model_registry.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vdo0DdnlAh8Y"
      },
      "outputs": [],
      "source": [
        "!pip install transformers torchserve torch-model-archiver torch-workflow-archiver  google-cloud-aiplatform captum"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Uploading and deploying Hugging Face models on Vertex AI Model Registry\n",
        "\n",
        "In this notebook, we will outline the process of bringing Hugging Face models to Vertex AI.\n",
        "\n",
        "To get started, we need a free Hugging Face account and [token](https://huggingface.co/docs/hub/en/security-tokens)."
      ],
      "metadata": {
        "id": "RM0yeORHhjOE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from huggingface_hub import notebook_login\n",
        "\n",
        "notebook_login()"
      ],
      "metadata": {
        "id": "iYay70VRBjlP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Download model and tokenizer"
      ],
      "metadata": {
        "id": "n2zUCQ2Ohpuc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import AutoModel, AutoTokenizer\n",
        "\n",
        "model_name = \"nlpaueb/legal-bert-base-uncased\"\n",
        "model = AutoModel.from_pretrained(model_name)\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
        "\n",
        "# Save the model and tokenizer\n",
        "model.save_pretrained(\"./legal-bert-model\")\n",
        "tokenizer.save_pretrained(\"./legal-bert-model\")"
      ],
      "metadata": {
        "id": "gJxUgL-VBj9A"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Handler\n",
        "\n",
        "The handler file denotes how to process input, run inference, and and postprocessing of outputs. Writing the handler file is the most complex step in the process.\n",
        "\n",
        "We will experiment with a generalized [handler file](https://github.com/pytorch/serve/blob/master/examples/Huggingface_Transformers/Transformer_handler_generalized.py) from Hugging Face, which handles a variety of classes that can be used with their Transformers library.\n"
      ],
      "metadata": {
        "id": "G1-mRQG-Wqxg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%writefile legal_bert_handler.py\n",
        "\n",
        "import ast\n",
        "import json\n",
        "import logging\n",
        "import os\n",
        "\n",
        "import torch\n",
        "import transformers\n",
        "from captum.attr import LayerIntegratedGradients\n",
        "from transformers import (\n",
        "    AutoModelForCausalLM,\n",
        "    AutoModelForQuestionAnswering,\n",
        "    AutoModelForSequenceClassification,\n",
        "    AutoModelForTokenClassification,\n",
        "    AutoTokenizer,\n",
        "    GPT2TokenizerFast,\n",
        ")\n",
        "\n",
        "from ts.torch_handler.base_handler import BaseHandler\n",
        "\n",
        "logger = logging.getLogger(__name__)\n",
        "logger.info(\"Transformers version %s\", transformers.__version__)\n",
        "\n",
        "\n",
        "class TransformersSeqClassifierHandler(BaseHandler):\n",
        "    \"\"\"\n",
        "    Transformers handler class for sequence, token classification and question answering.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self):\n",
        "        super(TransformersSeqClassifierHandler, self).__init__()\n",
        "        self.setup_config = None\n",
        "        self.initialized = False\n",
        "\n",
        "    def initialize(self, ctx):\n",
        "        \"\"\"In this initialize function, the BERT model is loaded and\n",
        "        the Layer Integrated Gradients Algorithm for Captum Explanations\n",
        "        is initialized here.\n",
        "        Args:\n",
        "            ctx (context): It is a JSON Object containing information\n",
        "            pertaining to the model artifacts parameters.\n",
        "        \"\"\"\n",
        "        self.manifest = ctx.manifest\n",
        "        self.model_yaml_config = (\n",
        "            ctx.model_yaml_config\n",
        "            if ctx is not None and hasattr(ctx, \"model_yaml_config\")\n",
        "            else {}\n",
        "        )\n",
        "        properties = ctx.system_properties\n",
        "        model_dir = properties.get(\"model_dir\")\n",
        "        serialized_file = self.manifest[\"model\"][\"serializedFile\"]\n",
        "        model_pt_path = os.path.join(model_dir, serialized_file)\n",
        "\n",
        "        self.device = torch.device(\n",
        "            \"cuda:\" + str(properties.get(\"gpu_id\"))\n",
        "            if torch.cuda.is_available() and properties.get(\"gpu_id\") is not None\n",
        "            else \"cpu\"\n",
        "        )\n",
        "\n",
        "        # read configs for the mode, model_name, etc. from the handler config\n",
        "        self.setup_config = self.model_yaml_config.get(\"handler\", {})\n",
        "        if not self.setup_config:\n",
        "            logger.warning(\"Missing the handler config\")\n",
        "\n",
        "        # Loading the model and tokenizer from checkpoint and config files based on the user's choice of mode\n",
        "        # further setup config can be added.\n",
        "        if self.setup_config[\"save_mode\"] == \"torchscript\":\n",
        "            self.model = torch.jit.load(model_pt_path, map_location=self.device)\n",
        "        elif self.setup_config[\"save_mode\"] == \"pretrained\":\n",
        "            if self.setup_config[\"mode\"] == \"sequence_classification\":\n",
        "                self.model = AutoModelForSequenceClassification.from_pretrained(\n",
        "                    model_dir\n",
        "                )\n",
        "            elif self.setup_config[\"mode\"] == \"question_answering\":\n",
        "                self.model = AutoModelForQuestionAnswering.from_pretrained(model_dir)\n",
        "            elif self.setup_config[\"mode\"] == \"token_classification\":\n",
        "                self.model = AutoModelForTokenClassification.from_pretrained(model_dir)\n",
        "            elif self.setup_config[\"mode\"] == \"text_generation\":\n",
        "                self.model = AutoModelForCausalLM.from_pretrained(model_dir)\n",
        "            else:\n",
        "                logger.warning(\"Missing the operation mode.\")\n",
        "            # Using the Better Transformer integration to speedup the inference\n",
        "            if self.setup_config[\"BetterTransformer\"]:\n",
        "                from optimum.bettertransformer import BetterTransformer\n",
        "\n",
        "                try:\n",
        "                    self.model = BetterTransformer.transform(self.model)\n",
        "                except RuntimeError as error:\n",
        "                    logger.warning(\n",
        "                        \"HuggingFace Optimum is not supporting this model,for the list of supported models, please refer to this doc,https://huggingface.co/docs/optimum/bettertransformer/overview\"\n",
        "                    )\n",
        "            # HF GPT2 models options can be gpt2, gpt2-medium, gpt2-large, gpt2-xl\n",
        "            # this basically place different model blocks on different devices,\n",
        "            # https://github.com/huggingface/transformers/blob/v4.17.0/src/transformers/models/gpt2/modeling_gpt2.py#L962\n",
        "            if (\n",
        "                self.setup_config[\"model_parallel\"]\n",
        "                and \"gpt2\" in self.setup_config[\"model_name\"]\n",
        "            ):\n",
        "                self.model.parallelize()\n",
        "            else:\n",
        "                self.model.to(self.device)\n",
        "\n",
        "        else:\n",
        "            logger.warning(\"Missing the checkpoint or state_dict.\")\n",
        "\n",
        "        if \"gpt2\" in self.setup_config[\"model_name\"]:\n",
        "            self.tokenizer = GPT2TokenizerFast.from_pretrained(\n",
        "                \"gpt2\", pad_token=\"<|endoftext|>\"\n",
        "            )\n",
        "\n",
        "        elif any(\n",
        "            fname\n",
        "            for fname in os.listdir(model_dir)\n",
        "            if fname.startswith(\"vocab.\") and os.path.isfile(fname)\n",
        "        ):\n",
        "            self.tokenizer = AutoTokenizer.from_pretrained(\n",
        "                model_dir, do_lower_case=self.setup_config[\"do_lower_case\"]\n",
        "            )\n",
        "        else:\n",
        "            self.tokenizer = AutoTokenizer.from_pretrained(\n",
        "                self.setup_config[\"model_name\"],\n",
        "                do_lower_case=self.setup_config[\"do_lower_case\"],\n",
        "            )\n",
        "\n",
        "        self.model.eval()\n",
        "\n",
        "        pt2_value = self.model_yaml_config.get(\"pt2\", {})\n",
        "        if \"compile\" in pt2_value:\n",
        "            compile_options = pt2_value[\"compile\"]\n",
        "            if compile_options[\"enable\"] == True:\n",
        "                del compile_options[\"enable\"]\n",
        "\n",
        "                compile_options_str = \", \".join(\n",
        "                    [f\"{k} {v}\" for k, v in compile_options.items()]\n",
        "                )\n",
        "                self.model = torch.compile(\n",
        "                    self.model,\n",
        "                    **compile_options,\n",
        "                )\n",
        "                logger.info(f\"Compiled model with {compile_options_str}\")\n",
        "        logger.info(\"Transformer model from path %s loaded successfully\", model_dir)\n",
        "\n",
        "        # Read the mapping file, index to object name\n",
        "        mapping_file_path = os.path.join(model_dir, \"index_to_name.json\")\n",
        "        # Question answering does not need the index_to_name.json file.\n",
        "        if not (\n",
        "            self.setup_config[\"mode\"] == \"question_answering\"\n",
        "            or self.setup_config[\"mode\"] == \"text_generation\"\n",
        "        ):\n",
        "            if os.path.isfile(mapping_file_path):\n",
        "                with open(mapping_file_path) as f:\n",
        "                    self.mapping = json.load(f)\n",
        "            else:\n",
        "                logger.warning(\"Missing the index_to_name.json file.\")\n",
        "        self.initialized = True\n",
        "\n",
        "    def preprocess(self, requests):\n",
        "        \"\"\"Basic text preprocessing, based on the user's chocie of application mode.\n",
        "        Args:\n",
        "            requests (str): The Input data in the form of text is passed on to the preprocess\n",
        "            function.\n",
        "        Returns:\n",
        "            list : The preprocess function returns a list of Tensor for the size of the word tokens.\n",
        "        \"\"\"\n",
        "        input_ids_batch = None\n",
        "        attention_mask_batch = None\n",
        "        for idx, data in enumerate(requests):\n",
        "            input_text = data.get(\"data\")\n",
        "            if input_text is None:\n",
        "                input_text = data.get(\"body\")\n",
        "            if isinstance(input_text, (bytes, bytearray)):\n",
        "                input_text = input_text.decode(\"utf-8\")\n",
        "            if (\n",
        "                self.setup_config[\"captum_explanation\"]\n",
        "                and not self.setup_config[\"mode\"] == \"question_answering\"\n",
        "            ):\n",
        "                input_text_target = ast.literal_eval(input_text)\n",
        "                input_text = input_text_target[\"text\"]\n",
        "            max_length = self.setup_config[\"max_length\"]\n",
        "            logger.info(\"Received text: '%s'\", input_text)\n",
        "            # preprocessing text for sequence_classification, token_classification or text_generation\n",
        "            if self.setup_config[\"mode\"] in {\n",
        "                \"sequence_classification\",\n",
        "                \"token_classification\",\n",
        "                \"text_generation\",\n",
        "            }:\n",
        "                inputs = self.tokenizer.encode_plus(\n",
        "                    input_text,\n",
        "                    max_length=int(max_length),\n",
        "                    pad_to_max_length=True,\n",
        "                    add_special_tokens=True,\n",
        "                    return_tensors=\"pt\",\n",
        "                )\n",
        "\n",
        "            # preprocessing text for question_answering.\n",
        "            elif self.setup_config[\"mode\"] == \"question_answering\":\n",
        "                # TODO Reading the context from a pickeled file or other fromats that\n",
        "                # fits the requirements of the task in hand. If this is done then need to\n",
        "                # modify the following preprocessing accordingly.\n",
        "\n",
        "                # the sample text for question_answering in the current version\n",
        "                # should be formated as dictionary with question and text as keys\n",
        "                # and related text as values.\n",
        "                # we use this format here seperate question and text for encoding.\n",
        "\n",
        "                question_context = ast.literal_eval(input_text)\n",
        "                question = question_context[\"question\"]\n",
        "                context = question_context[\"context\"]\n",
        "                inputs = self.tokenizer.encode_plus(\n",
        "                    question,\n",
        "                    context,\n",
        "                    max_length=int(max_length),\n",
        "                    pad_to_max_length=True,\n",
        "                    add_special_tokens=True,\n",
        "                    return_tensors=\"pt\",\n",
        "                )\n",
        "            input_ids = inputs[\"input_ids\"].to(self.device)\n",
        "            attention_mask = inputs[\"attention_mask\"].to(self.device)\n",
        "            # making a batch out of the recieved requests\n",
        "            # attention masks are passed for cases where input tokens are padded.\n",
        "            if input_ids.shape is not None:\n",
        "                if input_ids_batch is None:\n",
        "                    input_ids_batch = input_ids\n",
        "                    attention_mask_batch = attention_mask\n",
        "                else:\n",
        "                    input_ids_batch = torch.cat((input_ids_batch, input_ids), 0)\n",
        "                    attention_mask_batch = torch.cat(\n",
        "                        (attention_mask_batch, attention_mask), 0\n",
        "                    )\n",
        "        return (input_ids_batch, attention_mask_batch)\n",
        "\n",
        "    @torch.inference_mode\n",
        "    def inference(self, input_batch):\n",
        "        \"\"\"Predict the class (or classes) of the received text using the\n",
        "        serialized transformers checkpoint.\n",
        "        Args:\n",
        "            input_batch (list): List of Text Tensors from the pre-process function is passed here\n",
        "        Returns:\n",
        "            list : It returns a list of the predicted value for the input text\n",
        "        \"\"\"\n",
        "        input_ids_batch, attention_mask_batch = input_batch\n",
        "        inferences = []\n",
        "        # Handling inference for sequence_classification.\n",
        "        if self.setup_config[\"mode\"] == \"sequence_classification\":\n",
        "            predictions = self.model(input_ids_batch, attention_mask_batch)\n",
        "            print(\n",
        "                \"This the output size from the Seq classification model\",\n",
        "                predictions[0].size(),\n",
        "            )\n",
        "            print(\"This the output from the Seq classification model\", predictions)\n",
        "\n",
        "            num_rows, num_cols = predictions[0].shape\n",
        "            for i in range(num_rows):\n",
        "                out = predictions[0][i].unsqueeze(0)\n",
        "                y_hat = out.argmax(1).item()\n",
        "                predicted_idx = str(y_hat)\n",
        "                inferences.append(self.mapping[predicted_idx])\n",
        "        # Handling inference for question_answering.\n",
        "        elif self.setup_config[\"mode\"] == \"question_answering\":\n",
        "            # the output should be only answer_start and answer_end\n",
        "            # we are outputing the words just for demonstration.\n",
        "            if self.setup_config[\"save_mode\"] == \"pretrained\":\n",
        "                outputs = self.model(input_ids_batch, attention_mask_batch)\n",
        "                answer_start_scores = outputs.start_logits\n",
        "                answer_end_scores = outputs.end_logits\n",
        "            else:\n",
        "                answer_start_scores, answer_end_scores = self.model(\n",
        "                    input_ids_batch, attention_mask_batch\n",
        "                )\n",
        "            print(\n",
        "                \"This the output size for answer start scores from the question answering model\",\n",
        "                answer_start_scores.size(),\n",
        "            )\n",
        "            print(\n",
        "                \"This the output for answer start scores from the question answering model\",\n",
        "                answer_start_scores,\n",
        "            )\n",
        "            print(\n",
        "                \"This the output size for answer end scores from the question answering model\",\n",
        "                answer_end_scores.size(),\n",
        "            )\n",
        "            print(\n",
        "                \"This the output for answer end scores from the question answering model\",\n",
        "                answer_end_scores,\n",
        "            )\n",
        "\n",
        "            num_rows, num_cols = answer_start_scores.shape\n",
        "            # inferences = []\n",
        "            for i in range(num_rows):\n",
        "                answer_start_scores_one_seq = answer_start_scores[i].unsqueeze(0)\n",
        "                answer_start = torch.argmax(answer_start_scores_one_seq)\n",
        "                answer_end_scores_one_seq = answer_end_scores[i].unsqueeze(0)\n",
        "                answer_end = torch.argmax(answer_end_scores_one_seq) + 1\n",
        "                prediction = self.tokenizer.convert_tokens_to_string(\n",
        "                    self.tokenizer.convert_ids_to_tokens(\n",
        "                        input_ids_batch[i].tolist()[answer_start:answer_end]\n",
        "                    )\n",
        "                )\n",
        "                inferences.append(prediction)\n",
        "            logger.info(\"Model predicted: '%s'\", prediction)\n",
        "        # Handling inference for token_classification.\n",
        "        elif self.setup_config[\"mode\"] == \"token_classification\":\n",
        "            outputs = self.model(input_ids_batch, attention_mask_batch)[0]\n",
        "            print(\n",
        "                \"This the output size from the token classification model\",\n",
        "                outputs.size(),\n",
        "            )\n",
        "            print(\"This the output from the token classification model\", outputs)\n",
        "            num_rows = outputs.shape[0]\n",
        "            for i in range(num_rows):\n",
        "                output = outputs[i].unsqueeze(0)\n",
        "                predictions = torch.argmax(output, dim=2)\n",
        "                tokens = self.tokenizer.tokenize(\n",
        "                    self.tokenizer.decode(input_ids_batch[i])\n",
        "                )\n",
        "                if self.mapping:\n",
        "                    label_list = self.mapping[\"label_list\"]\n",
        "                label_list = label_list.strip(\"][\").split(\", \")\n",
        "                prediction = [\n",
        "                    (token, label_list[prediction])\n",
        "                    for token, prediction in zip(tokens, predictions[0].tolist())\n",
        "                ]\n",
        "                inferences.append(prediction)\n",
        "            logger.info(\"Model predicted: '%s'\", prediction)\n",
        "\n",
        "        # Handling inference for text_generation.\n",
        "        if self.setup_config[\"mode\"] == \"text_generation\":\n",
        "            if self.setup_config[\"model_parallel\"]:\n",
        "                # Need to move the first device, as the trasnformer model has been placed there\n",
        "                # https://github.com/huggingface/transformers/blob/v4.17.0/src/transformers/models/gpt2/modeling_gpt2.py#L970\n",
        "                input_ids_batch = input_ids_batch.to(\"cuda:0\")\n",
        "            outputs = self.model.generate(\n",
        "                input_ids_batch,\n",
        "                max_new_tokens=self.setup_config[\"max_length\"],\n",
        "                do_sample=True,\n",
        "                top_p=0.95,\n",
        "                top_k=60,\n",
        "            )\n",
        "            for i, x in enumerate(outputs):\n",
        "                inferences.append(\n",
        "                    self.tokenizer.decode(outputs[i], skip_special_tokens=True)\n",
        "                )\n",
        "\n",
        "            logger.info(\"Generated text: '%s'\", inferences)\n",
        "\n",
        "        print(\"Generated text\", inferences)\n",
        "        return inferences\n",
        "\n",
        "    def postprocess(self, inference_output):\n",
        "        \"\"\"Post Process Function converts the predicted response into Torchserve readable format.\n",
        "        Args:\n",
        "            inference_output (list): It contains the predicted response of the input text.\n",
        "        Returns:\n",
        "            (list): Returns a list of the Predictions and Explanations.\n",
        "        \"\"\"\n",
        "        return inference_output\n",
        "\n",
        "    def get_insights(self, input_batch, text, target):\n",
        "        \"\"\"This function initialize and calls the layer integrated gradient to get word importance\n",
        "        of the input text if captum explanation has been selected through setup_config\n",
        "        Args:\n",
        "            input_batch (int): Batches of tokens IDs of text\n",
        "            text (str): The Text specified in the input request\n",
        "            target (int): The Target can be set to any acceptable label under the user's discretion.\n",
        "        Returns:\n",
        "            (list): Returns a list of importances and words.\n",
        "        \"\"\"\n",
        "\n",
        "        if self.setup_config[\"captum_explanation\"]:\n",
        "            embedding_layer = getattr(self.model, self.setup_config[\"embedding_name\"])\n",
        "            embeddings = embedding_layer.embeddings\n",
        "            self.lig = LayerIntegratedGradients(captum_sequence_forward, embeddings)\n",
        "        else:\n",
        "            logger.warning(\"Captum Explanation is not chosen and will not be available\")\n",
        "\n",
        "        if isinstance(text, (bytes, bytearray)):\n",
        "            text = text.decode(\"utf-8\")\n",
        "        text_target = ast.literal_eval(text)\n",
        "\n",
        "        if not self.setup_config[\"mode\"] == \"question_answering\":\n",
        "            text = text_target[\"text\"]\n",
        "        self.target = text_target[\"target\"]\n",
        "\n",
        "        input_ids, ref_input_ids, attention_mask = construct_input_ref(\n",
        "            text, self.tokenizer, self.device, self.setup_config[\"mode\"]\n",
        "        )\n",
        "        all_tokens = get_word_token(input_ids, self.tokenizer)\n",
        "        response = {}\n",
        "        response[\"words\"] = all_tokens\n",
        "        if (\n",
        "            self.setup_config[\"mode\"] == \"sequence_classification\"\n",
        "            or self.setup_config[\"mode\"] == \"token_classification\"\n",
        "        ):\n",
        "            attributions, delta = self.lig.attribute(\n",
        "                inputs=input_ids,\n",
        "                baselines=ref_input_ids,\n",
        "                target=self.target,\n",
        "                additional_forward_args=(attention_mask, 0, self.model),\n",
        "                return_convergence_delta=True,\n",
        "            )\n",
        "\n",
        "            attributions_sum = summarize_attributions(attributions)\n",
        "            response[\"importances\"] = attributions_sum.tolist()\n",
        "            response[\"delta\"] = delta[0].tolist()\n",
        "\n",
        "        elif self.setup_config[\"mode\"] == \"question_answering\":\n",
        "            attributions_start, delta_start = self.lig.attribute(\n",
        "                inputs=input_ids,\n",
        "                baselines=ref_input_ids,\n",
        "                target=self.target,\n",
        "                additional_forward_args=(attention_mask, 0, self.model),\n",
        "                return_convergence_delta=True,\n",
        "            )\n",
        "            attributions_end, delta_end = self.lig.attribute(\n",
        "                inputs=input_ids,\n",
        "                baselines=ref_input_ids,\n",
        "                target=self.target,\n",
        "                additional_forward_args=(attention_mask, 1, self.model),\n",
        "                return_convergence_delta=True,\n",
        "            )\n",
        "            attributions_sum_start = summarize_attributions(attributions_start)\n",
        "            attributions_sum_end = summarize_attributions(attributions_end)\n",
        "            response[\"importances_answer_start\"] = attributions_sum_start.tolist()\n",
        "            response[\"importances_answer_end\"] = attributions_sum_end.tolist()\n",
        "            response[\"delta_start\"] = delta_start[0].tolist()\n",
        "            response[\"delta_end\"] = delta_end[0].tolist()\n",
        "\n",
        "        return [response]\n",
        "\n",
        "\n",
        "def construct_input_ref(text, tokenizer, device, mode):\n",
        "    \"\"\"For a given text, this function creates token id, reference id and\n",
        "    attention mask based on encode which is faster for captum insights\n",
        "    Args:\n",
        "        text (str): The text specified in the input request\n",
        "        tokenizer (AutoTokenizer Class Object): To word tokenize the input text\n",
        "        device (cpu or gpu): Type of the Environment the server runs on.\n",
        "    Returns:\n",
        "        input_id(Tensor): It attributes to the tensor of the input tokenized words\n",
        "        ref_input_ids(Tensor): Ref Input IDs are used as baseline for the attributions\n",
        "        attention mask() :  The attention mask is a binary tensor indicating the position\n",
        "         of the padded indices so that the model does not attend to them.\n",
        "    \"\"\"\n",
        "    if mode == \"question_answering\":\n",
        "        question_context = ast.literal_eval(text)\n",
        "        question = question_context[\"question\"]\n",
        "        context = question_context[\"context\"]\n",
        "        text_ids = tokenizer.encode(question, context, add_special_tokens=False)\n",
        "\n",
        "    text_ids = tokenizer.encode(text, add_special_tokens=False)\n",
        "    # construct input token ids\n",
        "    logger.info(\"text_ids %s\", text_ids)\n",
        "    logger.info(\"[tokenizer.cls_token_id] %s\", [tokenizer.cls_token_id])\n",
        "    input_ids = [tokenizer.cls_token_id] + text_ids + [tokenizer.sep_token_id]\n",
        "    logger.info(\"input_ids %s\", input_ids)\n",
        "\n",
        "    input_ids = torch.tensor([input_ids], device=device)\n",
        "    # construct reference token ids\n",
        "    ref_input_ids = (\n",
        "        [tokenizer.cls_token_id]\n",
        "        + [tokenizer.pad_token_id] * len(text_ids)\n",
        "        + [tokenizer.sep_token_id]\n",
        "    )\n",
        "    ref_input_ids = torch.tensor([ref_input_ids], device=device)\n",
        "    # construct attention mask\n",
        "    attention_mask = torch.ones_like(input_ids)\n",
        "    return input_ids, ref_input_ids, attention_mask\n",
        "\n",
        "\n",
        "def captum_sequence_forward(inputs, attention_mask=None, position=0, model=None):\n",
        "    \"\"\"This function is used to get the predictions from the model and this function\n",
        "    can be used independent of the type of the BERT Task.\n",
        "    Args:\n",
        "        inputs (list): Input for Predictions\n",
        "        attention_mask (list, optional): The attention mask is a binary tensor indicating the position\n",
        "         of the padded indices so that the model does not attend to them, it defaults to None.\n",
        "        position (int, optional): Position depends on the BERT Task.\n",
        "        model ([type], optional): Name of the model, it defaults to None.\n",
        "    Returns:\n",
        "        list: Prediction Outcome\n",
        "    \"\"\"\n",
        "    model.eval()\n",
        "    model.zero_grad()\n",
        "    pred = model(inputs, attention_mask=attention_mask)\n",
        "    pred = pred[position]\n",
        "    return pred\n",
        "\n",
        "\n",
        "def summarize_attributions(attributions):\n",
        "    \"\"\"Summarises the attribution across multiple runs\n",
        "    Args:\n",
        "        attributions ([list): attributions from the Layer Integrated Gradients\n",
        "    Returns:\n",
        "        list : Returns the attributions after normalizing them.\n",
        "    \"\"\"\n",
        "    attributions = attributions.sum(dim=-1).squeeze(0)\n",
        "    attributions = attributions / torch.norm(attributions)\n",
        "    return attributions\n",
        "\n",
        "\n",
        "def get_word_token(input_ids, tokenizer):\n",
        "    \"\"\"constructs word tokens from token id using the BERT's\n",
        "    Auto Tokenizer\n",
        "    Args:\n",
        "        input_ids (list): Input IDs from construct_input_ref method\n",
        "        tokenizer (class): The Auto Tokenizer Pre-Trained model object\n",
        "    Returns:\n",
        "        (list): Returns the word tokens\n",
        "    \"\"\"\n",
        "    indices = input_ids[0].detach().tolist()\n",
        "    tokens = tokenizer.convert_ids_to_tokens(indices)\n",
        "    # Remove unicode space character from BPE Tokeniser\n",
        "    tokens = [token.replace(\"Ġ\", \"\") for token in tokens]\n",
        "    return tokens"
      ],
      "metadata": {
        "id": "6eZeX817BvhY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!mkdir model-store"
      ],
      "metadata": {
        "id": "6BCemgYQ2LB3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### PyTorch models on Model Registry require a .mar file. This file is the packaging of all model artifacts into a TorchServe model archive file. We require a .json config, the model, and a handler file."
      ],
      "metadata": {
        "id": "fFrpfnwxGPlc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!torch-model-archiver --model-name model \\\n",
        "                     --version 1.0 \\\n",
        "                     --model-file \"./legal-bert-model/model.safetensors\" \\\n",
        "                     --handler \"./legal_bert_handler.py\" \\\n",
        "                     --extra-files \"./legal-bert-model/config.json,./legal-bert-model/vocab.txt\" \\\n",
        "                     --export-path ./model-store \\\n",
        "                     --force"
      ],
      "metadata": {
        "id": "QPM6G52YCKKa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Authenticate to GCP"
      ],
      "metadata": {
        "id": "Je088EYyGpi3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import auth\n",
        "auth.authenticate_user()"
      ],
      "metadata": {
        "id": "W3n3XbnWFwO0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "PROJECT_ID = \"your-project-id\"\n",
        "REGION = \"us-central1\"\n",
        "BUCKET_URI = \"gs://your-bucket-name\"\n",
        "MODEL_FOLDER = \"your-model-folder-name\""
      ],
      "metadata": {
        "id": "QhkwWoDzCQCr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!gsutil cp model-store/model.mar {BUCKET_URI}/{MODEL_FOLDER}/"
      ],
      "metadata": {
        "id": "sDrWuPkUFYQO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "PREBUILT_PREDICTION_CONTAINER = \"us-docker.pkg.dev/vertex-ai/prediction/pytorch-gpu.2-1:latest\"\n",
        "DISPLAY_NAME = \"legal-bert\"\n",
        "ARTIFACT_URI = BUCKET_URI"
      ],
      "metadata": {
        "id": "ptVb1LQCE1yV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Upload to Model Registry"
      ],
      "metadata": {
        "id": "cCM-EtWmcTf3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.cloud import aiplatform\n",
        "from typing import Optional, Sequence, Dict\n",
        "from google.cloud.aiplatform import explain\n",
        "\n",
        "def upload_model_sample(\n",
        "    project: str,\n",
        "    location: str,\n",
        "    display_name: str,\n",
        "    serving_container_image_uri: str,\n",
        "    artifact_uri: Optional[str] = None,\n",
        "    serving_container_predict_route: Optional[str] = None,\n",
        "    serving_container_health_route: Optional[str] = None,\n",
        "    description: Optional[str] = None,\n",
        "    serving_container_command: Optional[Sequence[str]] = None,\n",
        "    serving_container_args: Optional[Sequence[str]] = None,\n",
        "    serving_container_environment_variables: Optional[Dict[str, str]] = None,\n",
        "    serving_container_ports: Optional[Sequence[int]] = None,\n",
        "    instance_schema_uri: Optional[str] = None,\n",
        "    parameters_schema_uri: Optional[str] = None,\n",
        "    prediction_schema_uri: Optional[str] = None,\n",
        "    explanation_metadata: Optional[explain.ExplanationMetadata] = None,\n",
        "    explanation_parameters: Optional[explain.ExplanationParameters] = None,\n",
        "    sync: bool = True,\n",
        "):\n",
        "\n",
        "    aiplatform.init(project=project, location=location)\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=display_name,\n",
        "        artifact_uri=artifact_uri,\n",
        "        serving_container_image_uri=serving_container_image_uri,\n",
        "        serving_container_predict_route=serving_container_predict_route,\n",
        "        serving_container_health_route=serving_container_health_route,\n",
        "        instance_schema_uri=instance_schema_uri,\n",
        "        parameters_schema_uri=parameters_schema_uri,\n",
        "        prediction_schema_uri=prediction_schema_uri,\n",
        "        description=description,\n",
        "        serving_container_command=serving_container_command,\n",
        "        serving_container_args=serving_container_args,\n",
        "        serving_container_environment_variables=serving_container_environment_variables,\n",
        "        serving_container_ports=serving_container_ports,\n",
        "        explanation_metadata=explanation_metadata,\n",
        "        explanation_parameters=explanation_parameters,\n",
        "        sync=sync,\n",
        "    )\n",
        "\n",
        "    model.wait()\n",
        "\n",
        "    print(model.display_name)\n",
        "    print(model.resource_name)\n",
        "    return model"
      ],
      "metadata": {
        "id": "wZJfxgnTEna7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "upload_model_sample(project=PROJECT_ID,\n",
        "                    location=REGION,\n",
        "                    display_name=DISPLAY_NAME,\n",
        "                    artifact_uri=ARTIFACT_URI,\n",
        "                    serving_container_image_uri=PREBUILT_PREDICTION_CONTAINER,\n",
        ")"
      ],
      "metadata": {
        "id": "pPz98IE1GXi8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### List all models to check upload"
      ],
      "metadata": {
        "id": "AuTo8ao3cZS2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!gcloud ai models list --project={PROJECT_ID} --region={REGION}"
      ],
      "metadata": {
        "id": "6_FCLHrpHgYp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Create an endpoint"
      ],
      "metadata": {
        "id": "uEbP18rpceLe"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def create_endpoint(\n",
        "    project: str,\n",
        "    display_name: str,\n",
        "    location: str,\n",
        "):\n",
        "    aiplatform.init(project=project, location=location)\n",
        "\n",
        "    endpoint = aiplatform.Endpoint.create(\n",
        "        display_name=display_name,\n",
        "        project=project,\n",
        "        location=location,\n",
        "    )\n",
        "\n",
        "    print(endpoint.display_name)\n",
        "    print(endpoint.resource_name)\n",
        "    return endpoint"
      ],
      "metadata": {
        "id": "7XELZEaFIvcL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "create_endpoint(project=PROJECT_ID,\n",
        "                display_name=DISPLAY_NAME,\n",
        "                location=REGION,\n",
        ")"
      ],
      "metadata": {
        "id": "ZTy7JFCcQxg1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "ENDPOINT_NAME = \"legal-bert\""
      ],
      "metadata": {
        "id": "dHfxqzX9RD85"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!gcloud ai endpoints list \\\n",
        "  --region={REGION} \\\n",
        "  --project={PROJECT_ID} \\\n",
        "  #--filter=display_name=ENDPOINT_NAME"
      ],
      "metadata": {
        "id": "faJzFo1VQ0LN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Key gotcha here is that we need to pass an *endpoint object*, not just the name or id.\n",
        "\n",
        "Here is how we get the Vertex endpoint object for the most recent endpoint we created."
      ],
      "metadata": {
        "id": "x2QupYc7VQZY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "aiplatform.init(project=PROJECT_ID, location=REGION)"
      ],
      "metadata": {
        "id": "G5OeDiBxURBB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# If you have set an ENDPOINT_NAME variable, uncomment and try the first option.\n",
        "# endpoint = aiplatform.Endpoint.list(filter=f'display_name={ENDPOINT_NAME}')[0] # Retrieve existing endpoint\n",
        "endpoint = aiplatform.Endpoint.list()[0] # Retrieve latest endpoint\n",
        "endpoint"
      ],
      "metadata": {
        "id": "nGpqusjWVjKC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from typing import Tuple\n",
        "\n",
        "def deploy_model_with_dedicated_resources(\n",
        "    project,\n",
        "    location,\n",
        "    model_name: str,\n",
        "    machine_type: str,\n",
        "    endpoint: Optional[aiplatform.Endpoint] = None,\n",
        "    deployed_model_display_name: Optional[str] = None,\n",
        "    traffic_percentage: Optional[int] = 0,\n",
        "    traffic_split: Optional[Dict[str, int]] = None,\n",
        "    min_replica_count: int = 1,\n",
        "    max_replica_count: int = 1,\n",
        "    accelerator_type: Optional[str] = None,\n",
        "    accelerator_count: Optional[int] = None,\n",
        "    explanation_metadata: Optional[explain.ExplanationMetadata] = None,\n",
        "    explanation_parameters: Optional[explain.ExplanationParameters] = None,\n",
        "    metadata: Optional[Sequence[Tuple[str, str]]] = (),\n",
        "    sync: bool = True,\n",
        "):\n",
        "    \"\"\"\n",
        "    model_name: A fully-qualified model resource name or model ID.\n",
        "          Example: \"projects/123/locations/us-central1/models/456\" or\n",
        "          \"456\" when project and location are initialized or passed.\n",
        "    \"\"\"\n",
        "    model_name=\"56533589365358592\"\n",
        "    model = aiplatform.Model(model_name=model_name)\n",
        "\n",
        "    # The explanation_metadata and explanation_parameters should only be\n",
        "    # provided for a custom trained model and not an AutoML model.\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        deployed_model_display_name=deployed_model_display_name,\n",
        "        traffic_percentage=traffic_percentage,\n",
        "        traffic_split=traffic_split,\n",
        "        machine_type=machine_type,\n",
        "        min_replica_count=min_replica_count,\n",
        "        max_replica_count=max_replica_count,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "        explanation_metadata=explanation_metadata,\n",
        "        explanation_parameters=explanation_parameters,\n",
        "        metadata=metadata,\n",
        "        sync=sync,\n",
        "    )\n",
        "\n",
        "    model.wait()\n",
        "\n",
        "    print(model.display_name)\n",
        "    #print(model.resource_name)\n",
        "    return model"
      ],
      "metadata": {
        "id": "-aaKeyObRIBQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Deployment\n",
        "\n",
        "This step will take about 10-20 minutes"
      ],
      "metadata": {
        "id": "cZpe5WkBWCsc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "deploy_model_with_dedicated_resources(\n",
        "    project=PROJECT_ID,\n",
        "    location=REGION,\n",
        "    model_name=\"projects/{PROJECT_ID}/locations/us-central1/models/{model_name}}\",\n",
        "    machine_type=\"n1-standard-4\",\n",
        "    accelerator_type=\"NVIDIA_TESLA_T4\",\n",
        "    accelerator_count=1,\n",
        "    endpoint=endpoint,\n",
        "    deployed_model_display_name=\"legal-bert\"\n",
        ")\n",
        "\n"
      ],
      "metadata": {
        "id": "iNRk92kBSH_y"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "56533589365358592"
      ],
      "metadata": {
        "id": "qtJW7JHtEdfE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Inference"
      ],
      "metadata": {
        "id": "eQnfXkGmgouJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# If you're starting cold without running the previous cells, get the endpoint first:\n",
        "# endpoint = aiplatform.Endpoint.list(filter=f'display_name={ENDPOINT_NAME}')[0] # Retrieve existing endpoint\n"
      ],
      "metadata": {
        "id": "DMsb-0tvgq-Q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "endpoint"
      ],
      "metadata": {
        "id": "LWOKHCagiDm1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "instances = { \"data\":  \"tell me some highlights about UK trade union law\"}\n",
        "\n",
        "def endpoint_predict_sample(\n",
        "    project: str, location: str, instances: list, endpoint: str\n",
        "):\n",
        "    aiplatform.init(project=project, location=location)\n",
        "\n",
        "    # If running seperatetly from the workflow above, we will\n",
        "    # need to create an endpoint object.\n",
        "    # endpoint = aiplatform.Endpoint(endpoint)\n",
        "\n",
        "    # If you've run the previous cells in this notebook, we already\n",
        "    # have an endpoint object, so the following suffices:\n",
        "    endpoint = endpoint\n",
        "\n",
        "    prediction = endpoint.predict(instances=instances)\n",
        "    print(prediction)\n",
        "    return prediction"
      ],
      "metadata": {
        "id": "qM0VbPGcZdne"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "endpoint_predict_sample(\n",
        "    project=PROJECT_ID,\n",
        "    location=REGION,\n",
        "    instances=[instances],\n",
        "    endpoint=endpoint\n",
        ")"
      ],
      "metadata": {
        "id": "gecR2GO-U6Ol"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "instances = [{\"data\": \"What is the capital of France?\"}]  # Adjust the format based on your model's input\n",
        "prediction = endpoint.predict(instances=instances)\n",
        "print(prediction)"
      ],
      "metadata": {
        "id": "5-d9ZynJZzgy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import AutoModel, AutoTokenizer\n",
        "\n",
        "model_name = \"nlpaueb/legal-bert-base-uncased\"\n",
        "model = AutoModel.from_pretrained(model_name)\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
        "\n",
        "# Save the model and tokenizer\n",
        "model.save_pretrained(\"./legal-bert-model\")\n",
        "tokenizer.save_pretrained(\"./legal-bert-model\")"
      ],
      "metadata": {
        "id": "_SHfubnaWdNc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import tensorflow as tf\n",
        "from transformers import TFAutoModel\n",
        "\n",
        "# Load the PyTorch model\n",
        "pt_model = AutoModel.from_pretrained(\"./legal-bert-model\")\n",
        "\n",
        "# Convert to TensorFlow\n",
        "tf_model = TFAutoModel.from_pretrained(\"./legal-bert-model\", from_pt=True)\n",
        "\n",
        "# Save in SavedModel format\n",
        "tf_model.save(\"./legal-bert-tf\", save_format=\"tf\")"
      ],
      "metadata": {
        "id": "aUjlwTdkmpEl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!gsutil -m cp -r ./legal-bert-tf gs://genai-experiments/tf-legal-bert-model/"
      ],
      "metadata": {
        "id": "-O4S2uqUWZzR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
        "\n",
        "model = aiplatform.Model.upload(\n",
        "    display_name=\"legal-bert-base-uncased\",\n",
        "    artifact_uri=\"gs://genai-experiments/tf-legal-bert-model/legal-bert-tf/\",\n",
        "    serving_container_image_uri=\"us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-12:latest\"\n",
        ")\n",
        "\n",
        "print(f\"Model uploaded. Resource name: {model.resource_name}\")"
      ],
      "metadata": {
        "id": "KcbO2-JdW9wv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!gsutil ls gs://genai-experiments/tf-legal-bert-model/*"
      ],
      "metadata": {
        "id": "JtOxXy_lXNro"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "8A-zc0p1W9zE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "r0UlQ8u0W91L"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "PM-Lzj_lWoOB"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}