HuggingFace_ONNX_in_Spark_NLP_T5.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AcVmXaYCWVb7"
      },
      "source": [
        "## Import ONNX T5 models from HuggingFace 🤗 into Spark NLP 🚀\n",
        "\n",
        "Let's keep in mind a few things before we start 😊\n",
        "\n",
        "- ONNX support was introduced in  `Spark NLP 5.0.0`, enabling high performance inference for models.\n",
        "- ONNX support for the `T5Transformer` is only available since in `Spark NLP 5.2.0` and after. So please make sure you have upgraded to the latest Spark NLP release\n",
        "- You can import T5 models via `T5Model`. These models are usually under `Text2Text Generation` category and have `T5` in their labels\n",
        "- This is a very computationally expensive module especially on larger sequence. The use of an accelerator such as GPU is recommended.\n",
        "- Reference: [T5Model](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Model)\n",
        "- Some [example models](https://huggingface.co/models?other=T5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NDHYfM5_WVb7"
      },
      "source": [
        "## Export and Save HuggingFace model\n",
        "\n",
        "- Let's install `transformers` package with the `onnx` extension and it's dependencies. You don't need `onnx` to be installed for Spark NLP, however, we need it to load and save models from HuggingFace.\n",
        "- We lock `transformers` on version `4.35.2`. This doesn't mean it won't work with the future releases\n",
        "- We will also need `sentencepiece` for tokenization."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "gPTqlJuOWVb8",
        "outputId": "9774fb65-6fb2-43bd-caba-2b169f8feef5"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "tensorflow 2.13.1 requires protobuf!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.20.3, but you have protobuf 3.20.2 which is incompatible.\n",
            "tensorflow 2.13.1 requires typing-extensions<4.6.0,>=3.6.6, but you have typing-extensions 4.11.0 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
            "\u001b[0m"
          ]
        }
      ],
      "source": [
        "!pip install -q --upgrade transformers[onnx]==4.35.2 optimum sentencepiece"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
            "\u001b[0m"
          ]
        }
      ],
      "source": [
        "!pip install -q --upgrade --force-reinstall transformers[onnx]==4.35.2 optimum sentencepiece tensorflow"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZRqvjOehWVb8"
      },
      "source": [
        "- HuggingFace has an extension called Optimum which offers specialized model inference, including ONNX. We can use this to import and export ONNX models with `from_pretrained` and `save_pretrained`.\n",
        "- We'll use [google/flan-t5-base](https://huggingface.co/google/flan-t5-base) model from HuggingFace as an example\n",
        "- In addition to `T5Model` we also need to save the tokenizer. This is the same for every model, these are assets needed for tokenization inside Spark NLP.\n",
        "- If we want to optimize the model, a GPU will be needed. Make sure to select the correct runtime.\n",
        "0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "37l409GuWVb5"
      },
      "source": [
        "![JohnSnowLabs](https://sparknlp.org/assets/images/logo.png)\n",
        "\n",
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/JohnSnowLabs/spark-nlp/blob/master/examples/python/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_T5.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "WW1DaaCNWVb8"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/root/anaconda3/envs/sparknlp/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
            "  from .autonotebook import tqdm as notebook_tqdm\n"
          ]
        }
      ],
      "source": [
        "import transformers\n",
        "# Model name, either HF (e.g. \"google/flan-t5-base\") or a local path\n",
        "MODEL_NAME = \"/data/HW/proj2/best_model\"\n",
        "\n",
        "\n",
        "# Path to store the exported models\n",
        "EXPORT_PATH = \"/data/HW/proj2/exported_model\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "3zvy4XOLWVb9",
        "outputId": "8444306f-15e5-4e60-fbdb-d6158ebff309"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "2024-05-22 13:10:15.390121: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
            "2024-05-22 13:10:15.392066: I tensorflow/tsl/cuda/cudart_stub.cc:28] Could not find cuda drivers on your machine, GPU will not be used.\n",
            "2024-05-22 13:10:15.432109: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
            "To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
            "2024-05-22 13:10:16.190892: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n",
            "Framework not specified. Using pt to export the model.\n",
            "Using the export variant default. Available variants are:\n",
            "    - default: The default ONNX variant.\n",
            "\n",
            "***** Exporting submodel 1/3: T5Stack *****\n",
            "Using framework PyTorch: 2.1.2+cu121\n",
            "Overriding 1 configuration item(s)\n",
            "\t- use_cache -> False\n",
            "\n",
            "***** Exporting submodel 2/3: T5ForConditionalGeneration *****\n",
            "Using framework PyTorch: 2.1.2+cu121\n",
            "Overriding 1 configuration item(s)\n",
            "\t- use_cache -> True\n",
            "/root/anaconda3/envs/sparknlp/lib/python3.8/site-packages/transformers/modeling_utils.py:873: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
            "  if causal_mask.shape[1] < attention_mask.shape[1]:\n",
            "\n",
            "***** Exporting submodel 3/3: T5ForConditionalGeneration *****\n",
            "Using framework PyTorch: 2.1.2+cu121\n",
            "Overriding 1 configuration item(s)\n",
            "\t- use_cache -> True\n",
            "/root/anaconda3/envs/sparknlp/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py:508: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
            "  elif past_key_value.shape[2] != key_value_states.shape[1]:\n",
            "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
            "In-place op on output of tensor.shape. See https://pytorch.org/docs/master/onnx.html#avoid-inplace-operations-when-using-tensor-shape-in-tracing-mode\n",
            "Post-processing the exported models...\n",
            "Weight deduplication check in the ONNX export requires accelerate. Please install accelerate to run it.\n",
            "The two models proto have different outputs (33 and 17 outputs). Constant outputs will be added to unify the two models outputs. This is expected for encoder-decoder models where cached cross-attention key/values are constant outputs, omitted in the model with KV cache.\n",
            "Adding a constant output for present.0.encoder.key of shape [0, 6, 1, 64] in model2.\n",
            "Adding a constant output for present.0.encoder.value of shape [0, 6, 1, 64] in model2.\n",
            "Adding a constant output for present.1.encoder.key of shape [0, 6, 1, 64] in model2.\n",
            "Adding a constant output for present.1.encoder.value of shape [0, 6, 1, 64] in model2.\n",
            "Adding a constant output for present.2.encoder.key of shape [0, 6, 1, 64] in model2.\n",
            "Adding a constant output for present.2.encoder.value of shape [0, 6, 1, 64] in model2.\n",
            "Adding a constant output for present.3.encoder.key of shape [0, 6, 1, 64] in model2.\n",
            "Adding a constant output for present.3.encoder.value of shape [0, 6, 1, 64] in model2.\n",
            "Adding a constant output for present.4.encoder.key of shape [0, 6, 1, 64] in model2.\n",
            "Adding a constant output for present.4.encoder.value of shape [0, 6, 1, 64] in model2.\n",
            "Adding a constant output for present.5.encoder.key of shape [0, 6, 1, 64] in model2.\n",
            "Adding a constant output for present.5.encoder.value of shape [0, 6, 1, 64] in model2.\n",
            "Adding a constant output for present.6.encoder.key of shape [0, 6, 1, 64] in model2.\n",
            "Adding a constant output for present.6.encoder.value of shape [0, 6, 1, 64] in model2.\n",
            "Adding a constant output for present.7.encoder.key of shape [0, 6, 1, 64] in model2.\n",
            "Adding a constant output for present.7.encoder.value of shape [0, 6, 1, 64] in model2.\n",
            "\n",
            "Validating ONNX model /data/HW/proj2/exported_model/encoder_model.onnx...\n",
            "\t-[✓] ONNX model output names match reference model (last_hidden_state)\n",
            "\t- Validating ONNX Model output \"last_hidden_state\":\n",
            "\t\t-[✓] (2, 16, 512) matches (2, 16, 512)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\n",
            "Validating ONNX model /data/HW/proj2/exported_model/decoder_model_merged.onnx...\n",
            "\t-[✓] ONNX model output names match reference model (present.4.encoder.value, present.3.encoder.value, present.6.encoder.key, present.1.encoder.key, present.2.encoder.key, present.1.encoder.value, present.3.decoder.key, present.7.decoder.value, present.0.encoder.key, present.2.decoder.key, present.0.encoder.value, present.6.decoder.value, present.1.decoder.value, present.3.decoder.value, present.2.decoder.value, present.5.encoder.key, present.5.encoder.value, present.4.decoder.value, present.6.encoder.value, present.4.encoder.key, present.5.decoder.value, present.0.decoder.key, logits, present.7.decoder.key, present.6.decoder.key, present.7.encoder.key, present.3.encoder.key, present.0.decoder.value, present.2.encoder.value, present.7.encoder.value, present.1.decoder.key, present.5.decoder.key, present.4.decoder.key)\n",
            "\t- Validating ONNX Model output \"logits\":\n",
            "\t\t-[✓] (2, 16, 32128) matches (2, 16, 32128)\n",
            "\t\t-[x] values not close enough, max diff: 6.103515625e-05 (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.0.decoder.key\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.0.decoder.value\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.0.encoder.key\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.0.encoder.value\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.1.decoder.key\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.1.decoder.value\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.1.encoder.key\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.1.encoder.value\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.2.decoder.key\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.2.decoder.value\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.2.encoder.key\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.2.encoder.value\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.3.decoder.key\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.3.decoder.value\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[x] values not close enough, max diff: 1.239776611328125e-05 (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.3.encoder.key\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.3.encoder.value\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.4.decoder.key\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.4.decoder.value\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[x] values not close enough, max diff: 1.811981201171875e-05 (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.4.encoder.key\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.4.encoder.value\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.5.decoder.key\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.5.decoder.value\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[x] values not close enough, max diff: 4.315376281738281e-05 (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.5.encoder.key\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.5.encoder.value\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.6.decoder.key\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.6.decoder.value\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[x] values not close enough, max diff: 3.24249267578125e-05 (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.6.encoder.key\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.6.encoder.value\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.7.decoder.key\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.7.decoder.value\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[x] values not close enough, max diff: 2.574920654296875e-05 (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.7.encoder.key\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.7.encoder.value\":\n",
            "\t\t-[✓] (2, 6, 16, 64) matches (2, 6, 16, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\n",
            "Validating ONNX model /data/HW/proj2/exported_model/decoder_model_merged.onnx...\n",
            "\t-[✓] ONNX model output names match reference model (present.6.decoder.value, present.1.decoder.value, present.4.decoder.key, present.5.decoder.value, present.0.decoder.key, logits, present.1.decoder.key, present.3.decoder.key, present.3.decoder.value, present.2.decoder.value, present.6.decoder.key, present.7.decoder.key, present.7.decoder.value, present.5.decoder.key, present.2.decoder.key, present.4.decoder.value, present.0.decoder.value)\n",
            "\t- Validating ONNX Model output \"logits\":\n",
            "\t\t-[✓] (2, 1, 32128) matches (2, 1, 32128)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.0.decoder.key\":\n",
            "\t\t-[✓] (2, 6, 17, 64) matches (2, 6, 17, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.0.decoder.value\":\n",
            "\t\t-[✓] (2, 6, 17, 64) matches (2, 6, 17, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.1.decoder.key\":\n",
            "\t\t-[✓] (2, 6, 17, 64) matches (2, 6, 17, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.1.decoder.value\":\n",
            "\t\t-[✓] (2, 6, 17, 64) matches (2, 6, 17, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.2.decoder.key\":\n",
            "\t\t-[✓] (2, 6, 17, 64) matches (2, 6, 17, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.2.decoder.value\":\n",
            "\t\t-[✓] (2, 6, 17, 64) matches (2, 6, 17, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.3.decoder.key\":\n",
            "\t\t-[✓] (2, 6, 17, 64) matches (2, 6, 17, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.3.decoder.value\":\n",
            "\t\t-[✓] (2, 6, 17, 64) matches (2, 6, 17, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.4.decoder.key\":\n",
            "\t\t-[✓] (2, 6, 17, 64) matches (2, 6, 17, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.4.decoder.value\":\n",
            "\t\t-[✓] (2, 6, 17, 64) matches (2, 6, 17, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.5.decoder.key\":\n",
            "\t\t-[✓] (2, 6, 17, 64) matches (2, 6, 17, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.5.decoder.value\":\n",
            "\t\t-[✓] (2, 6, 17, 64) matches (2, 6, 17, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.6.decoder.key\":\n",
            "\t\t-[✓] (2, 6, 17, 64) matches (2, 6, 17, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.6.decoder.value\":\n",
            "\t\t-[✓] (2, 6, 17, 64) matches (2, 6, 17, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.7.decoder.key\":\n",
            "\t\t-[✓] (2, 6, 17, 64) matches (2, 6, 17, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "\t- Validating ONNX Model output \"present.7.decoder.value\":\n",
            "\t\t-[✓] (2, 6, 17, 64) matches (2, 6, 17, 64)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "The ONNX export succeeded with the warning: The maximum absolute difference between the output of the reference model and the ONNX exported model is not within the set tolerance 1e-05:\n",
            "- logits: max diff = 6.103515625e-05\n",
            "- present.3.decoder.value: max diff = 1.239776611328125e-05\n",
            "- present.4.decoder.value: max diff = 1.811981201171875e-05\n",
            "- present.5.decoder.value: max diff = 4.315376281738281e-05\n",
            "- present.6.decoder.value: max diff = 3.24249267578125e-05\n",
            "- present.7.decoder.value: max diff = 2.574920654296875e-05.\n",
            " The exported model was saved at: /data/HW/proj2/exported_model\n"
          ]
        }
      ],
      "source": [
        "# Export the model to ONNX using optimum\n",
        "\n",
        "# Export with optimizations (uncomment next line)\n",
        "# !optimum-cli export onnx --task text2text-generation-with-past --model {MODEL_NAME} --optimize O2 {EXPORT_PATH}\n",
        "# IMPORTANT - there is a bug in onnxruntime which crashes it when trying to optimize a T5 small model (or any derivative of it)\n",
        "# There are two ways to addess the problem:\n",
        "# 1. Go to onnx_model_bert.py in the onnxruntime module (the full path depends on the module version),\n",
        "#    find the BertOnnxModel class and comment the following line in the constructor:\n",
        "#    assert (num_heads == 0 and hidden_size == 0) or (num_heads > 0 and hidden_size % num_heads == 0)\n",
        "# 2. Disable optimization by removing '--optimize O2' (use line below).\n",
        "\n",
        "# Export without optimizations\n",
        "!optimum-cli export onnx --task text2text-generation-with-past --model {MODEL_NAME} {EXPORT_PATH}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wslRw6y3WVb9"
      },
      "source": [
        "Let's have a look inside these two directories and see what we are dealing with:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "BYTKHJ8bWVb9",
        "outputId": "55fc594f-518e-4159-ae52-c0ac537b958d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "total 810572\n",
            "-rw-r--r-- 1 root root      1532 May 22 13:10 config.json\n",
            "-rw-r--r-- 1 root root 232553641 May 22 13:10 decoder_model.onnx\n",
            "-rw-r--r-- 1 root root 232784326 May 22 13:10 decoder_model_merged.onnx\n",
            "-rw-r--r-- 1 root root 219953955 May 22 13:10 decoder_with_past_model.onnx\n",
            "-rw-r--r-- 1 root root 141456353 May 22 13:10 encoder_model.onnx\n",
            "-rw-r--r-- 1 root root       142 May 22 13:10 generation_config.json\n",
            "-rw-r--r-- 1 root root      2201 May 22 13:10 special_tokens_map.json\n",
            "-rw-r--r-- 1 root root    791656 May 22 13:10 spiece.model\n",
            "-rw-r--r-- 1 root root   2422256 May 22 13:10 tokenizer.json\n",
            "-rw-r--r-- 1 root root     20771 May 22 13:10 tokenizer_config.json\n"
          ]
        }
      ],
      "source": [
        "!ls -l {EXPORT_PATH}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3Gn6rot7WVb9"
      },
      "source": [
        "- As you can see, we need to move the sentence piece models `spiece.model` from the tokenizer to assets folder which Spark NLP will look for"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "ELeIg_ixWVb9"
      },
      "outputs": [],
      "source": [
        "! mkdir -p {EXPORT_PATH}/assets\n",
        "! mv -t {EXPORT_PATH}/assets {EXPORT_PATH}/spiece.model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "gSO-sFikWVb9",
        "outputId": "48ce8d04-2226-488c-c489-d6a175a0d1f6"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "total 776\n",
            "-rw-r--r-- 1 root root 791656 May 22 13:10 spiece.model\n"
          ]
        }
      ],
      "source": [
        "!ls -l {EXPORT_PATH}/assets"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XjPcKrH0WVb-"
      },
      "source": [
        "## Import and Save T5 in Spark NLP\n",
        "\n",
        "- Let's install and setup Spark NLP in Google Colab\n",
        "- This part is pretty easy via our simple script"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "GntK4k8sWVb-",
        "outputId": "696fc2e9-a99e-4dfd-aa02-7d377be290ba"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Installing PySpark 3.2.3 and Spark NLP 5.3.3\n",
            "setup Colab for PySpark 3.2.3 and Spark NLP 5.3.3\n",
            "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
            "\u001b[0m"
          ]
        }
      ],
      "source": [
        "! wget -q http://setup.johnsnowlabs.com/colab.sh -O - | bash"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c_VVcdp_WVb-"
      },
      "source": [
        "Let's start Spark with Spark NLP included via our simple `start()` function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "mVrq1ZlWWVb-"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            ":: loading settings :: url = jar:file:/opt/module/spark-3.5.0-bin-hadoop3/jars/ivy-2.5.1.jar!/org/apache/ivy/core/settings/ivysettings.xml\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Ivy Default Cache set to: /root/.ivy2/cache\n",
            "The jars for the packages stored in: /root/.ivy2/jars\n",
            "com.johnsnowlabs.nlp#spark-nlp_2.12 added as a dependency\n",
            ":: resolving dependencies :: org.apache.spark#spark-submit-parent-9880cb02-82c9-474a-8f38-96695ac04bf4;1.0\n",
            "\tconfs: [default]\n",
            "\tfound com.johnsnowlabs.nlp#spark-nlp_2.12;5.3.3 in central\n",
            "\tfound com.typesafe#config;1.4.2 in central\n",
            "\tfound org.rocksdb#rocksdbjni;6.29.5 in central\n",
            "\tfound com.amazonaws#aws-java-sdk-s3;1.12.500 in central\n",
            "\tfound com.amazonaws#aws-java-sdk-kms;1.12.500 in central\n",
            "\tfound com.amazonaws#aws-java-sdk-core;1.12.500 in central\n",
            "\tfound commons-logging#commons-logging;1.1.3 in central\n",
            "\tfound commons-codec#commons-codec;1.15 in central\n",
            "\tfound org.apache.httpcomponents#httpclient;4.5.13 in central\n",
            "\tfound org.apache.httpcomponents#httpcore;4.4.13 in central\n",
            "\tfound software.amazon.ion#ion-java;1.0.2 in central\n",
            "\tfound joda-time#joda-time;2.8.1 in central\n",
            "\tfound com.amazonaws#jmespath-java;1.12.500 in central\n",
            "\tfound com.github.universal-automata#liblevenshtein;3.0.0 in central\n",
            "\tfound com.google.protobuf#protobuf-java-util;3.0.0-beta-3 in central\n",
            "\tfound com.google.protobuf#protobuf-java;3.0.0-beta-3 in central\n",
            "\tfound com.google.code.gson#gson;2.3 in central\n",
            "\tfound it.unimi.dsi#fastutil;7.0.12 in central\n",
            "\tfound org.projectlombok#lombok;1.16.8 in central\n",
            "\tfound com.google.cloud#google-cloud-storage;2.20.1 in central\n",
            "\tfound com.google.guava#guava;31.1-jre in central\n",
            "\tfound com.google.guava#failureaccess;1.0.1 in central\n",
            "\tfound com.google.guava#listenablefuture;9999.0-empty-to-avoid-conflict-with-guava in central\n",
            "\tfound com.google.errorprone#error_prone_annotations;2.18.0 in central\n",
            "\tfound com.google.j2objc#j2objc-annotations;1.3 in central\n",
            "\tfound com.google.http-client#google-http-client;1.43.0 in central\n",
            "\tfound io.opencensus#opencensus-contrib-http-util;0.31.1 in central\n",
            "\tfound com.google.http-client#google-http-client-jackson2;1.43.0 in central\n",
            "\tfound com.google.http-client#google-http-client-gson;1.43.0 in central\n",
            "\tfound com.google.api-client#google-api-client;2.2.0 in central\n",
            "\tfound com.google.oauth-client#google-oauth-client;1.34.1 in central\n",
            "\tfound com.google.http-client#google-http-client-apache-v2;1.43.0 in central\n",
            "\tfound com.google.apis#google-api-services-storage;v1-rev20220705-2.0.0 in central\n",
            "\tfound com.google.code.gson#gson;2.10.1 in central\n",
            "\tfound com.google.cloud#google-cloud-core;2.12.0 in central\n",
            "\tfound io.grpc#grpc-context;1.53.0 in central\n",
            "\tfound com.google.auto.value#auto-value-annotations;1.10.1 in central\n",
            "\tfound com.google.auto.value#auto-value;1.10.1 in central\n",
            "\tfound javax.annotation#javax.annotation-api;1.3.2 in central\n",
            "\tfound com.google.cloud#google-cloud-core-http;2.12.0 in central\n",
            "\tfound com.google.http-client#google-http-client-appengine;1.43.0 in central\n",
            "\tfound com.google.api#gax-httpjson;0.108.2 in central\n",
            "\tfound com.google.cloud#google-cloud-core-grpc;2.12.0 in central\n",
            "\tfound io.grpc#grpc-alts;1.53.0 in central\n",
            "\tfound io.grpc#grpc-grpclb;1.53.0 in central\n",
            "\tfound org.conscrypt#conscrypt-openjdk-uber;2.5.2 in central\n",
            "\tfound io.grpc#grpc-auth;1.53.0 in central\n",
            "\tfound io.grpc#grpc-protobuf;1.53.0 in central\n",
            "\tfound io.grpc#grpc-protobuf-lite;1.53.0 in central\n",
            "\tfound io.grpc#grpc-core;1.53.0 in central\n",
            "\tfound com.google.api#gax;2.23.2 in central\n",
            "\tfound com.google.api#gax-grpc;2.23.2 in central\n",
            "\tfound com.google.auth#google-auth-library-credentials;1.16.0 in central\n",
            "\tfound com.google.auth#google-auth-library-oauth2-http;1.16.0 in central\n",
            "\tfound com.google.api#api-common;2.6.2 in central\n",
            "\tfound io.opencensus#opencensus-api;0.31.1 in central\n",
            "\tfound com.google.api.grpc#proto-google-iam-v1;1.9.2 in central\n",
            "\tfound com.google.protobuf#protobuf-java;3.21.12 in central\n",
            "\tfound com.google.protobuf#protobuf-java-util;3.21.12 in central\n",
            "\tfound com.google.api.grpc#proto-google-common-protos;2.14.2 in central\n",
            "\tfound org.threeten#threetenbp;1.6.5 in central\n",
            "\tfound com.google.api.grpc#proto-google-cloud-storage-v2;2.20.1-alpha in central\n",
            "\tfound com.google.api.grpc#grpc-google-cloud-storage-v2;2.20.1-alpha in central\n",
            "\tfound com.google.api.grpc#gapic-google-cloud-storage-v2;2.20.1-alpha in central\n",
            "\tfound com.google.code.findbugs#jsr305;3.0.2 in central\n",
            "\tfound io.grpc#grpc-api;1.53.0 in central\n",
            "\tfound io.grpc#grpc-stub;1.53.0 in central\n",
            "\tfound org.checkerframework#checker-qual;3.31.0 in central\n",
            "\tfound io.perfmark#perfmark-api;0.26.0 in central\n",
            "\tfound com.google.android#annotations;4.1.1.4 in central\n",
            "\tfound org.codehaus.mojo#animal-sniffer-annotations;1.22 in central\n",
            "\tfound io.opencensus#opencensus-proto;0.2.0 in central\n",
            "\tfound io.grpc#grpc-services;1.53.0 in central\n",
            "\tfound com.google.re2j#re2j;1.6 in central\n",
            "\tfound io.grpc#grpc-netty-shaded;1.53.0 in central\n",
            "\tfound io.grpc#grpc-googleapis;1.53.0 in central\n",
            "\tfound io.grpc#grpc-xds;1.53.0 in central\n",
            "\tfound com.navigamez#greex;1.0 in central\n",
            "\tfound dk.brics.automaton#automaton;1.11-8 in central\n",
            "\tfound com.johnsnowlabs.nlp#tensorflow-cpu_2.12;0.4.4 in central\n",
            "\tfound com.microsoft.onnxruntime#onnxruntime;1.17.0 in central\n",
            ":: resolution report :: resolve 791ms :: artifacts dl 29ms\n",
            "\t:: modules in use:\n",
            "\tcom.amazonaws#aws-java-sdk-core;1.12.500 from central in [default]\n",
            "\tcom.amazonaws#aws-java-sdk-kms;1.12.500 from central in [default]\n",
            "\tcom.amazonaws#aws-java-sdk-s3;1.12.500 from central in [default]\n",
            "\tcom.amazonaws#jmespath-java;1.12.500 from central in [default]\n",
            "\tcom.github.universal-automata#liblevenshtein;3.0.0 from central in [default]\n",
            "\tcom.google.android#annotations;4.1.1.4 from central in [default]\n",
            "\tcom.google.api#api-common;2.6.2 from central in [default]\n",
            "\tcom.google.api#gax;2.23.2 from central in [default]\n",
            "\tcom.google.api#gax-grpc;2.23.2 from central in [default]\n",
            "\tcom.google.api#gax-httpjson;0.108.2 from central in [default]\n",
            "\tcom.google.api-client#google-api-client;2.2.0 from central in [default]\n",
            "\tcom.google.api.grpc#gapic-google-cloud-storage-v2;2.20.1-alpha from central in [default]\n",
            "\tcom.google.api.grpc#grpc-google-cloud-storage-v2;2.20.1-alpha from central in [default]\n",
            "\tcom.google.api.grpc#proto-google-cloud-storage-v2;2.20.1-alpha from central in [default]\n",
            "\tcom.google.api.grpc#proto-google-common-protos;2.14.2 from central in [default]\n",
            "\tcom.google.api.grpc#proto-google-iam-v1;1.9.2 from central in [default]\n",
            "\tcom.google.apis#google-api-services-storage;v1-rev20220705-2.0.0 from central in [default]\n",
            "\tcom.google.auth#google-auth-library-credentials;1.16.0 from central in [default]\n",
            "\tcom.google.auth#google-auth-library-oauth2-http;1.16.0 from central in [default]\n",
            "\tcom.google.auto.value#auto-value;1.10.1 from central in [default]\n",
            "\tcom.google.auto.value#auto-value-annotations;1.10.1 from central in [default]\n",
            "\tcom.google.cloud#google-cloud-core;2.12.0 from central in [default]\n",
            "\tcom.google.cloud#google-cloud-core-grpc;2.12.0 from central in [default]\n",
            "\tcom.google.cloud#google-cloud-core-http;2.12.0 from central in [default]\n",
            "\tcom.google.cloud#google-cloud-storage;2.20.1 from central in [default]\n",
            "\tcom.google.code.findbugs#jsr305;3.0.2 from central in [default]\n",
            "\tcom.google.code.gson#gson;2.10.1 from central in [default]\n",
            "\tcom.google.errorprone#error_prone_annotations;2.18.0 from central in [default]\n",
            "\tcom.google.guava#failureaccess;1.0.1 from central in [default]\n",
            "\tcom.google.guava#guava;31.1-jre from central in [default]\n",
            "\tcom.google.guava#listenablefuture;9999.0-empty-to-avoid-conflict-with-guava from central in [default]\n",
            "\tcom.google.http-client#google-http-client;1.43.0 from central in [default]\n",
            "\tcom.google.http-client#google-http-client-apache-v2;1.43.0 from central in [default]\n",
            "\tcom.google.http-client#google-http-client-appengine;1.43.0 from central in [default]\n",
            "\tcom.google.http-client#google-http-client-gson;1.43.0 from central in [default]\n",
            "\tcom.google.http-client#google-http-client-jackson2;1.43.0 from central in [default]\n",
            "\tcom.google.j2objc#j2objc-annotations;1.3 from central in [default]\n",
            "\tcom.google.oauth-client#google-oauth-client;1.34.1 from central in [default]\n",
            "\tcom.google.protobuf#protobuf-java;3.21.12 from central in [default]\n",
            "\tcom.google.protobuf#protobuf-java-util;3.21.12 from central in [default]\n",
            "\tcom.google.re2j#re2j;1.6 from central in [default]\n",
            "\tcom.johnsnowlabs.nlp#spark-nlp_2.12;5.3.3 from central in [default]\n",
            "\tcom.johnsnowlabs.nlp#tensorflow-cpu_2.12;0.4.4 from central in [default]\n",
            "\tcom.microsoft.onnxruntime#onnxruntime;1.17.0 from central in [default]\n",
            "\tcom.navigamez#greex;1.0 from central in [default]\n",
            "\tcom.typesafe#config;1.4.2 from central in [default]\n",
            "\tcommons-codec#commons-codec;1.15 from central in [default]\n",
            "\tcommons-logging#commons-logging;1.1.3 from central in [default]\n",
            "\tdk.brics.automaton#automaton;1.11-8 from central in [default]\n",
            "\tio.grpc#grpc-alts;1.53.0 from central in [default]\n",
            "\tio.grpc#grpc-api;1.53.0 from central in [default]\n",
            "\tio.grpc#grpc-auth;1.53.0 from central in [default]\n",
            "\tio.grpc#grpc-context;1.53.0 from central in [default]\n",
            "\tio.grpc#grpc-core;1.53.0 from central in [default]\n",
            "\tio.grpc#grpc-googleapis;1.53.0 from central in [default]\n",
            "\tio.grpc#grpc-grpclb;1.53.0 from central in [default]\n",
            "\tio.grpc#grpc-netty-shaded;1.53.0 from central in [default]\n",
            "\tio.grpc#grpc-protobuf;1.53.0 from central in [default]\n",
            "\tio.grpc#grpc-protobuf-lite;1.53.0 from central in [default]\n",
            "\tio.grpc#grpc-services;1.53.0 from central in [default]\n",
            "\tio.grpc#grpc-stub;1.53.0 from central in [default]\n",
            "\tio.grpc#grpc-xds;1.53.0 from central in [default]\n",
            "\tio.opencensus#opencensus-api;0.31.1 from central in [default]\n",
            "\tio.opencensus#opencensus-contrib-http-util;0.31.1 from central in [default]\n",
            "\tio.opencensus#opencensus-proto;0.2.0 from central in [default]\n",
            "\tio.perfmark#perfmark-api;0.26.0 from central in [default]\n",
            "\tit.unimi.dsi#fastutil;7.0.12 from central in [default]\n",
            "\tjavax.annotation#javax.annotation-api;1.3.2 from central in [default]\n",
            "\tjoda-time#joda-time;2.8.1 from central in [default]\n",
            "\torg.apache.httpcomponents#httpclient;4.5.13 from central in [default]\n",
            "\torg.apache.httpcomponents#httpcore;4.4.13 from central in [default]\n",
            "\torg.checkerframework#checker-qual;3.31.0 from central in [default]\n",
            "\torg.codehaus.mojo#animal-sniffer-annotations;1.22 from central in [default]\n",
            "\torg.conscrypt#conscrypt-openjdk-uber;2.5.2 from central in [default]\n",
            "\torg.projectlombok#lombok;1.16.8 from central in [default]\n",
            "\torg.rocksdb#rocksdbjni;6.29.5 from central in [default]\n",
            "\torg.threeten#threetenbp;1.6.5 from central in [default]\n",
            "\tsoftware.amazon.ion#ion-java;1.0.2 from central in [default]\n",
            "\t:: evicted modules:\n",
            "\tcommons-logging#commons-logging;1.2 by [commons-logging#commons-logging;1.1.3] in [default]\n",
            "\tcommons-codec#commons-codec;1.11 by [commons-codec#commons-codec;1.15] in [default]\n",
            "\tcom.google.protobuf#protobuf-java-util;3.0.0-beta-3 by [com.google.protobuf#protobuf-java-util;3.21.12] in [default]\n",
            "\tcom.google.protobuf#protobuf-java;3.0.0-beta-3 by [com.google.protobuf#protobuf-java;3.21.12] in [default]\n",
            "\tcom.google.code.gson#gson;2.3 by [com.google.code.gson#gson;2.10.1] in [default]\n",
            "\t---------------------------------------------------------------------\n",
            "\t|                  |            modules            ||   artifacts   |\n",
            "\t|       conf       | number| search|dwnlded|evicted|| number|dwnlded|\n",
            "\t---------------------------------------------------------------------\n",
            "\t|      default     |   83  |   0   |   0   |   5   ||   78  |   0   |\n",
            "\t---------------------------------------------------------------------\n",
            ":: retrieving :: org.apache.spark#spark-submit-parent-9880cb02-82c9-474a-8f38-96695ac04bf4\n",
            "\tconfs: [default]\n",
            "\t0 artifacts copied, 78 already retrieved (0kB/18ms)\n",
            "24/05/22 13:11:33 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable\n",
            "Setting default log level to \"WARN\".\n",
            "To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).\n",
            "24/05/22 13:11:34 WARN Utils: Service 'SparkUI' could not bind on port 4040. Attempting port 4041.\n"
          ]
        }
      ],
      "source": [
        "import sparknlp\n",
        "\n",
        "# let's start Spark with Spark NLP\n",
        "spark = sparknlp.start()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NlTYNUNzWVb-"
      },
      "source": [
        "- Let's use `loadSavedModel` functon in `T5Transformer` which allows us to load the ONNX model\n",
        "- Most params will be set automatically. They can also be set later after loading the model in `T5Transformer` during runtime, so don't worry about setting them now\n",
        "- `loadSavedModel` accepts two params, first is the path to the exported model. The second is the SparkSession that is `spark` variable we previously started via `sparknlp.start()`\n",
        "- NOTE: `loadSavedModel` accepts local paths in addition to distributed file systems such as `HDFS`, `S3`, `DBFS`, etc. This feature was introduced in Spark NLP 4.2.2 release. Keep in mind the best and recommended way to move/share/reuse Spark NLP models is to use `write.save` so you can use `.load()` from any file systems natively.st and recommended way to move/share/reuse Spark NLP models is to use `write.save` so you can use `.load()` from any file systems natively."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "ziFTyKiqWVb-"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Using CPUs\n",
            "Using CPUs\n"
          ]
        }
      ],
      "source": [
        "from sparknlp.annotator import *\n",
        "\n",
        "# T5 = T5Transformer.loadSavedModel(EXPORT_PATH, spark)\\\n",
        "#   .setUseCache(True) \\\n",
        "#   .setTask(\"summarize:\") \\\n",
        "#   .setMaxOutputLength(200)\n",
        "T5 = T5Transformer.loadSavedModel(EXPORT_PATH, spark)\\\n",
        "    .setUseCache(True) \\\n",
        "    .setTask(\"question:\") \\\n",
        "    .setMaxOutputLength(200) \\\n",
        "    .setInputCols([\"documents\"]) \\\n",
        "    .setOutputCol(\"answers\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "82HtR456WVb-"
      },
      "source": [
        "Let's save it on disk so it is easier to be moved around and also be used later via `.load` function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "PxAkFbdbWVb-"
      },
      "outputs": [],
      "source": [
        "T5.write().overwrite().save(f\"{MODEL_NAME}_spark_nlp\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lmp6655gWVb_"
      },
      "source": [
        "Let's clean up stuff we don't need anymore"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "9Cy7lDrHWVb_"
      },
      "outputs": [],
      "source": [
        "!rm -rf {EXPORT_PATH}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QOXCjyznWVb_"
      },
      "source": [
        "Awesome  😎 !\n",
        "\n",
        "This is your ONNX T5 model from HuggingFace 🤗  loaded and saved by Spark NLP 🚀"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "id": "hx0vKU-iWVb_",
        "outputId": "e160068d-716c-4133-81bf-d63db1482562"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "total 366308\n",
            "-rw-r--r-- 1 root root 232819989 May 22 13:12 decoder.onxx\n",
            "-rw-r--r-- 1 root root 141478076 May 22 13:12 encoder.onxx\n",
            "drwxr-xr-x 2 root root      4096 May 22 13:12 metadata\n",
            "-rw-r--r-- 1 root root    791656 May 22 13:12 t5_spp\n"
          ]
        }
      ],
      "source": [
        "! ls -l {MODEL_NAME}_spark_nlp"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xDNStfQ2WVb_"
      },
      "source": [
        "Now let's see how we can use it on other machines, clusters, or any place you wish to use your new and shiny T5 model 😊"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2GSHLsL3WVb_"
      },
      "source": [
        "That's it! You can now go wild and use hundreds of T5 models from HuggingFace 🤗 in Spark NLP 🚀\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'/data/HW/proj2/best_model'"
            ]
          },
          "execution_count": 15,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "MODEL_NAME"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "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.8.19"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}