NER_Transformer_Pipeline.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "NER Transformer Pipeline.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "gfrlsDrEx-Vq",
        "outputId": "58b7c2b9-ceb2-4827-82c3-ec5765167c15"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Ignoring colorama: markers 'platform_system == \"Windows\" and python_full_version >= \"3.6.0\" and python_version >= \"3.6\"' don't match your environment\n",
            "Ignoring numpy: markers 'python_version >= \"3.8\" and python_full_version >= \"3.6.0\"' don't match your environment\n",
            "Ignoring pyreadline3: markers 'sys_platform == \"win32\" and python_version >= \"3.8\" and (python_version >= \"2.7\" and python_full_version < \"3.0.0\" or python_full_version >= \"3.5.0\")' don't match your environment\n",
            "Requirement already satisfied: certifi==2021.10.8 in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 1)) (2021.10.8)\n",
            "Requirement already satisfied: charset-normalizer==2.0.12 in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 4)) (2.0.12)\n",
            "Collecting click==8.0.4\n",
            "  Downloading click-8.0.4-py3-none-any.whl (97 kB)\n",
            "\u001b[K     |████████████████████████████████| 97 kB 4.3 MB/s \n",
            "\u001b[?25hCollecting coloredlogs==15.0.1\n",
            "  Downloading coloredlogs-15.0.1-py2.py3-none-any.whl (46 kB)\n",
            "\u001b[K     |████████████████████████████████| 46 kB 3.4 MB/s \n",
            "\u001b[?25hRequirement already satisfied: filelock==3.6.0 in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 16)) (3.6.0)\n",
            "Requirement already satisfied: flatbuffers==2.0 in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 19)) (2.0)\n",
            "Collecting huggingface-hub==0.4.0\n",
            "  Downloading huggingface_hub-0.4.0-py3-none-any.whl (67 kB)\n",
            "\u001b[K     |████████████████████████████████| 67 kB 5.3 MB/s \n",
            "\u001b[?25hCollecting humanfriendly==10.0\n",
            "  Downloading humanfriendly-10.0-py2.py3-none-any.whl (86 kB)\n",
            "\u001b[K     |████████████████████████████████| 86 kB 6.1 MB/s \n",
            "\u001b[?25hCollecting idna==3.3\n",
            "  Downloading idna-3.3-py3-none-any.whl (61 kB)\n",
            "\u001b[K     |████████████████████████████████| 61 kB 8.6 MB/s \n",
            "\u001b[?25hRequirement already satisfied: joblib==1.1.0 in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 31)) (1.1.0)\n",
            "Collecting onnx==1.11.0\n",
            "  Downloading onnx-1.11.0-cp37-cp37m-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (12.8 MB)\n",
            "\u001b[K     |████████████████████████████████| 12.8 MB 52.2 MB/s \n",
            "\u001b[?25hCollecting onnxconverter-common==1.9.0\n",
            "  Downloading onnxconverter_common-1.9.0-py2.py3-none-any.whl (78 kB)\n",
            "\u001b[K     |████████████████████████████████| 78 kB 7.5 MB/s \n",
            "\u001b[?25hCollecting onnxruntime-tools==1.7.0\n",
            "  Downloading onnxruntime_tools-1.7.0-py3-none-any.whl (212 kB)\n",
            "\u001b[K     |████████████████████████████████| 212 kB 53.9 MB/s \n",
            "\u001b[?25hCollecting onnxruntime==1.10.0\n",
            "  Downloading onnxruntime-1.10.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (4.9 MB)\n",
            "\u001b[K     |████████████████████████████████| 4.9 MB 40.7 MB/s \n",
            "\u001b[?25hRequirement already satisfied: packaging==21.3 in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 106)) (21.3)\n",
            "Collecting protobuf==3.19.4\n",
            "  Downloading protobuf-3.19.4-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.1 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.1 MB 40.3 MB/s \n",
            "\u001b[?25hCollecting psutil==5.9.0\n",
            "  Downloading psutil-5.9.0-cp37-cp37m-manylinux_2_12_x86_64.manylinux2010_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (280 kB)\n",
            "\u001b[K     |████████████████████████████████| 280 kB 51.2 MB/s \n",
            "\u001b[?25hCollecting py-cpuinfo==8.0.0\n",
            "  Downloading py-cpuinfo-8.0.0.tar.gz (99 kB)\n",
            "\u001b[K     |████████████████████████████████| 99 kB 10.2 MB/s \n",
            "\u001b[?25hCollecting py3nvml==0.2.7\n",
            "  Downloading py3nvml-0.2.7-py3-none-any.whl (55 kB)\n",
            "\u001b[K     |████████████████████████████████| 55 kB 3.8 MB/s \n",
            "\u001b[?25hRequirement already satisfied: pyparsing==3.0.7 in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 174)) (3.0.7)\n",
            "Collecting pyyaml==6.0\n",
            "  Downloading PyYAML-6.0-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl (596 kB)\n",
            "\u001b[K     |████████████████████████████████| 596 kB 49.4 MB/s \n",
            "\u001b[?25hCollecting regex==2022.3.2\n",
            "  Downloading regex-2022.3.2-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (749 kB)\n",
            "\u001b[K     |████████████████████████████████| 749 kB 50.1 MB/s \n",
            "\u001b[?25hCollecting requests==2.27.1\n",
            "  Downloading requests-2.27.1-py2.py3-none-any.whl (63 kB)\n",
            "\u001b[K     |████████████████████████████████| 63 kB 1.7 MB/s \n",
            "\u001b[?25hCollecting sacremoses==0.0.47\n",
            "  Downloading sacremoses-0.0.47-py2.py3-none-any.whl (895 kB)\n",
            "\u001b[K     |████████████████████████████████| 895 kB 50.7 MB/s \n",
            "\u001b[?25hCollecting sentencepiece==0.1.96\n",
            "  Downloading sentencepiece-0.1.96-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.2 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.2 MB 42.8 MB/s \n",
            "\u001b[?25hCollecting six==1.16.0\n",
            "  Downloading six-1.16.0-py2.py3-none-any.whl (11 kB)\n",
            "Requirement already satisfied: tabulate==0.8.9 in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 340)) (0.8.9)\n",
            "Collecting tf2onnx==1.8.4\n",
            "  Downloading tf2onnx-1.8.4-py3-none-any.whl (345 kB)\n",
            "\u001b[K     |████████████████████████████████| 345 kB 51.7 MB/s \n",
            "\u001b[?25hCollecting tokenizers==0.11.6\n",
            "  Downloading tokenizers-0.11.6-cp37-cp37m-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (6.5 MB)\n",
            "\u001b[K     |████████████████████████████████| 6.5 MB 39.8 MB/s \n",
            "\u001b[?25hCollecting torch==1.10.2\n",
            "  Downloading torch-1.10.2-cp37-cp37m-manylinux1_x86_64.whl (881.9 MB)\n",
            "\u001b[K     |██████████████████████████████▎ | 834.1 MB 1.3 MB/s eta 0:00:37tcmalloc: large alloc 1147494400 bytes == 0x559735774000 @  0x7fbb3be19615 0x559732ea917c 0x559732f8947a 0x559732eabf9d 0x559732f9dd4d 0x559732f1fec8 0x559732f1aa2e 0x559732ead88a 0x559732f1fd30 0x559732f1aa2e 0x559732ead88a 0x559732f1c719 0x559732f9eb76 0x559732f1bd95 0x559732f9eb76 0x559732f1bd95 0x559732f9eb76 0x559732f1bd95 0x559732eadce9 0x559732ef1579 0x559732eac902 0x559732f1fc4d 0x559732f1aa2e 0x559732ead88a 0x559732f1c719 0x559732f1aa2e 0x559732ead88a 0x559732f1b8f6 0x559732ead7aa 0x559732f1bb4f 0x559732f1aa2e\n",
            "\u001b[K     |████████████████████████████████| 881.9 MB 1.8 kB/s \n",
            "\u001b[?25hRequirement already satisfied: tqdm==4.63.0 in /usr/local/lib/python3.7/dist-packages (from -r requirements.txt (line 404)) (4.63.0)\n",
            "Collecting transformers==4.17.0\n",
            "  Downloading transformers-4.17.0-py3-none-any.whl (3.8 MB)\n",
            "\u001b[K     |████████████████████████████████| 3.8 MB 40.7 MB/s \n",
            "\u001b[?25hCollecting typing-extensions==4.1.1\n",
            "  Downloading typing_extensions-4.1.1-py3-none-any.whl (26 kB)\n",
            "Collecting urllib3==1.26.8\n",
            "  Downloading urllib3-1.26.8-py2.py3-none-any.whl (138 kB)\n",
            "\u001b[K     |████████████████████████████████| 138 kB 51.2 MB/s \n",
            "\u001b[?25hCollecting xmltodict==0.12.0\n",
            "  Downloading xmltodict-0.12.0-py2.py3-none-any.whl (9.2 kB)\n",
            "Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.7/dist-packages (from click==8.0.4->-r requirements.txt (line 7)) (4.11.3)\n",
            "Requirement already satisfied: numpy>=1.16.6 in /usr/local/lib/python3.7/dist-packages (from onnx==1.11.0->-r requirements.txt (line 54)) (1.21.5)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata->click==8.0.4->-r requirements.txt (line 7)) (3.7.0)\n",
            "Building wheels for collected packages: py-cpuinfo\n",
            "  Building wheel for py-cpuinfo (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for py-cpuinfo: filename=py_cpuinfo-8.0.0-py3-none-any.whl size=22257 sha256=2add07fc332ed6fca8852fa31e25aa08b5316886bae8ea0609529b14f65a57dd\n",
            "  Stored in directory: /root/.cache/pip/wheels/d2/f1/1f/041add21dc9c4220157f1bd2bd6afe1f1a49524c3396b94401\n",
            "Successfully built py-cpuinfo\n",
            "Installing collected packages: typing-extensions, urllib3, idna, xmltodict, six, requests, regex, pyyaml, protobuf, humanfriendly, click, tokenizers, sacremoses, py3nvml, py-cpuinfo, psutil, onnx, huggingface-hub, coloredlogs, transformers, torch, tf2onnx, sentencepiece, onnxruntime-tools, onnxruntime, onnxconverter-common\n",
            "  Attempting uninstall: typing-extensions\n",
            "    Found existing installation: typing-extensions 3.10.0.2\n",
            "    Uninstalling typing-extensions-3.10.0.2:\n",
            "      Successfully uninstalled typing-extensions-3.10.0.2\n",
            "  Attempting uninstall: urllib3\n",
            "    Found existing installation: urllib3 1.24.3\n",
            "    Uninstalling urllib3-1.24.3:\n",
            "      Successfully uninstalled urllib3-1.24.3\n",
            "  Attempting uninstall: idna\n",
            "    Found existing installation: idna 2.10\n",
            "    Uninstalling idna-2.10:\n",
            "      Successfully uninstalled idna-2.10\n",
            "  Attempting uninstall: six\n",
            "    Found existing installation: six 1.15.0\n",
            "    Uninstalling six-1.15.0:\n",
            "      Successfully uninstalled six-1.15.0\n",
            "  Attempting uninstall: requests\n",
            "    Found existing installation: requests 2.23.0\n",
            "    Uninstalling requests-2.23.0:\n",
            "      Successfully uninstalled requests-2.23.0\n",
            "  Attempting uninstall: regex\n",
            "    Found existing installation: regex 2019.12.20\n",
            "    Uninstalling regex-2019.12.20:\n",
            "      Successfully uninstalled regex-2019.12.20\n",
            "  Attempting uninstall: pyyaml\n",
            "    Found existing installation: PyYAML 3.13\n",
            "    Uninstalling PyYAML-3.13:\n",
            "      Successfully uninstalled PyYAML-3.13\n",
            "  Attempting uninstall: protobuf\n",
            "    Found existing installation: protobuf 3.17.3\n",
            "    Uninstalling protobuf-3.17.3:\n",
            "      Successfully uninstalled protobuf-3.17.3\n",
            "  Attempting uninstall: click\n",
            "    Found existing installation: click 7.1.2\n",
            "    Uninstalling click-7.1.2:\n",
            "      Successfully uninstalled click-7.1.2\n",
            "  Attempting uninstall: psutil\n",
            "    Found existing installation: psutil 5.4.8\n",
            "    Uninstalling psutil-5.4.8:\n",
            "      Successfully uninstalled psutil-5.4.8\n",
            "  Attempting uninstall: torch\n",
            "    Found existing installation: torch 1.10.0+cu111\n",
            "    Uninstalling torch-1.10.0+cu111:\n",
            "      Successfully uninstalled torch-1.10.0+cu111\n",
            "\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.8.0 requires tf-estimator-nightly==2.8.0.dev2021122109, which is not installed.\n",
            "torchvision 0.11.1+cu111 requires torch==1.10.0, but you have torch 1.10.2 which is incompatible.\n",
            "torchtext 0.11.0 requires torch==1.10.0, but you have torch 1.10.2 which is incompatible.\n",
            "torchaudio 0.10.0+cu111 requires torch==1.10.0, but you have torch 1.10.2 which is incompatible.\n",
            "google-colab 1.0.0 requires requests~=2.23.0, but you have requests 2.27.1 which is incompatible.\n",
            "google-colab 1.0.0 requires six~=1.15.0, but you have six 1.16.0 which is incompatible.\n",
            "flask 1.1.4 requires click<8.0,>=5.1, but you have click 8.0.4 which is incompatible.\n",
            "datascience 0.10.6 requires folium==0.2.1, but you have folium 0.8.3 which is incompatible.\n",
            "arviz 0.11.4 requires typing-extensions<4,>=3.7.4.3, but you have typing-extensions 4.1.1 which is incompatible.\n",
            "albumentations 0.1.12 requires imgaug<0.2.7,>=0.2.5, but you have imgaug 0.2.9 which is incompatible.\u001b[0m\n",
            "Successfully installed click-8.0.4 coloredlogs-15.0.1 huggingface-hub-0.4.0 humanfriendly-10.0 idna-3.3 onnx-1.11.0 onnxconverter-common-1.9.0 onnxruntime-1.10.0 onnxruntime-tools-1.7.0 protobuf-3.19.4 psutil-5.9.0 py-cpuinfo-8.0.0 py3nvml-0.2.7 pyyaml-6.0 regex-2022.3.2 requests-2.27.1 sacremoses-0.0.47 sentencepiece-0.1.96 six-1.16.0 tf2onnx-1.8.4 tokenizers-0.11.6 torch-1.10.2 transformers-4.17.0 typing-extensions-4.1.1 urllib3-1.26.8 xmltodict-0.12.0\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "google",
                  "psutil",
                  "six"
                ]
              }
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "!pip install -r requirements.txt"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers.onnx.features import FeaturesManager\n",
        "distilbert_features = list(FeaturesManager.get_supported_features_for_model_type(\"bert\").keys())\n",
        "print(distilbert_features)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dVAGpmFkysBO",
        "outputId": "1635fdf4-75f4-47e9-d541-978424153e6e"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['default', 'masked-lm', 'causal-lm', 'sequence-classification', 'token-classification', 'question-answering']\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!python -m transformers.onnx --model=dslim/bert-base-NER --feature=token-classification onnx/"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gcDGA4jzC5ve",
        "outputId": "5cad565c-2cb5-4832-c41d-a02e07b32be9"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading: 100% 59.0/59.0 [00:00<00:00, 40.1kB/s]\n",
            "Downloading: 100% 829/829 [00:00<00:00, 669kB/s]\n",
            "Downloading: 100% 208k/208k [00:00<00:00, 661kB/s]\n",
            "Downloading: 100% 2.00/2.00 [00:00<00:00, 1.56kB/s]\n",
            "Downloading: 100% 112/112 [00:00<00:00, 85.1kB/s]\n",
            "Downloading: 100% 413M/413M [00:10<00:00, 41.9MB/s]\n",
            "Using framework PyTorch: 1.10.2+cu102\n",
            "Overriding 1 configuration item(s)\n",
            "\t- use_cache -> False\n",
            "Validating ONNX model...\n",
            "\t-[✓] ONNX model output names match reference model ({'logits'})\n",
            "\t- Validating ONNX Model output \"logits\":\n",
            "\t\t-[✓] (2, 8, 9) matches (2, 8, 9)\n",
            "\t\t-[✓] all values close (atol: 1e-05)\n",
            "All good, model saved at: onnx/model.onnx\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "\n",
        "from onnxruntime import (\n",
        "    InferenceSession, SessionOptions, GraphOptimizationLevel\n",
        ")\n",
        "from transformers import (\n",
        "    TokenClassificationPipeline, AutoTokenizer, AutoModelForTokenClassification\n",
        ")"
      ],
      "metadata": {
        "id": "mgN5xItaE1fJ"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "options = SessionOptions() # initialize session options\n",
        "options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL\n",
        "\n",
        "session = InferenceSession(\n",
        "    \"onnx/model.onnx\", sess_options=options, providers=[\"CPUExecutionProvider\"]\n",
        ")\n",
        "\n",
        "# disable session.run() fallback mechanism, it prevents for a reset of the execution provider\n",
        "session.disable_fallback()"
      ],
      "metadata": {
        "id": "nZaBik05FcOj"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from onnxruntime import get_all_providers\n",
        "\n",
        "get_all_providers()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "n3jneuKiF7OY",
        "outputId": "905710d9-4a92-469f-a7de-a30793688c3d"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['TensorrtExecutionProvider',\n",
              " 'CUDAExecutionProvider',\n",
              " 'MIGraphXExecutionProvider',\n",
              " 'ROCMExecutionProvider',\n",
              " 'OpenVINOExecutionProvider',\n",
              " 'DnnlExecutionProvider',\n",
              " 'NupharExecutionProvider',\n",
              " 'VitisAIExecutionProvider',\n",
              " 'NnapiExecutionProvider',\n",
              " 'CoreMLExecutionProvider',\n",
              " 'ArmNNExecutionProvider',\n",
              " 'ACLExecutionProvider',\n",
              " 'DmlExecutionProvider',\n",
              " 'RknpuExecutionProvider',\n",
              " 'CPUExecutionProvider']"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class OnnxTokenClassificationPipeline(TokenClassificationPipeline):\n",
        "\n",
        "    def __init__(self, *args, **kwargs):\n",
        "        super().__init__(*args, **kwargs)\n",
        "        \n",
        "    \n",
        "    def _forward(self, model_inputs):\n",
        "        \"\"\"\n",
        "        Forward pass through the model. This method is not to be called by the user directly and is only used\n",
        "        by the pipeline to perform the actual predictions.\n",
        "\n",
        "        This is where we will define the actual process to do inference with the ONNX model and the session created\n",
        "        before.\n",
        "        \"\"\"\n",
        "\n",
        "        # This comes from the original implementation of the pipeline\n",
        "        special_tokens_mask = model_inputs.pop(\"special_tokens_mask\")\n",
        "        offset_mapping = model_inputs.pop(\"offset_mapping\", None)\n",
        "        sentence = model_inputs.pop(\"sentence\")\n",
        "\n",
        "        inputs = {k: v.cpu().detach().numpy() for k, v in model_inputs.items()} # dict of numpy arrays\n",
        "        outputs_name = session.get_outputs()[0].name # get the name of the output tensor\n",
        "\n",
        "        logits = session.run(output_names=[outputs_name], input_feed=inputs)[0] # run the session\n",
        "        logits = torch.tensor(logits) # convert to torch tensor to be compatible with the original implementation\n",
        "\n",
        "        return {\n",
        "            \"logits\": logits,\n",
        "            \"special_tokens_mask\": special_tokens_mask,\n",
        "            \"offset_mapping\": offset_mapping,\n",
        "            \"sentence\": sentence,\n",
        "            **model_inputs,\n",
        "        }\n",
        "\n",
        "    # We need to override the preprocess method because the onnx model is waiting for the attention masks as inputs\n",
        "    # along with the embeddings.\n",
        "    def preprocess(self, sentence, offset_mapping=None):\n",
        "        truncation = True if self.tokenizer.model_max_length and self.tokenizer.model_max_length > 0 else False\n",
        "        model_inputs = self.tokenizer(\n",
        "            sentence,\n",
        "            return_attention_mask=True, # This is the only difference from the original implementation\n",
        "            return_tensors=self.framework,\n",
        "            truncation=truncation,\n",
        "            return_special_tokens_mask=True,\n",
        "            return_offsets_mapping=self.tokenizer.is_fast,\n",
        "        )\n",
        "        if offset_mapping:\n",
        "            model_inputs[\"offset_mapping\"] = offset_mapping\n",
        "\n",
        "        model_inputs[\"sentence\"] = sentence\n",
        "\n",
        "        return model_inputs"
      ],
      "metadata": {
        "id": "HI-vOZd4GOx1"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model_name_from_hub = \"dslim/bert-base-NER\"\n",
        "\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_name_from_hub)\n",
        "model = AutoModelForTokenClassification.from_pretrained(model_name_from_hub)\n",
        "\n",
        "onnx_pipeline = OnnxTokenClassificationPipeline(\n",
        "    task=\"ner\", \n",
        "    model=model,\n",
        "    tokenizer=tokenizer,\n",
        "    framework=\"pt\",\n",
        "    aggregation_strategy=\"simple\",\n",
        ")"
      ],
      "metadata": {
        "id": "DLyG9SShGdQn"
      },
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "sequence = \"Apple was founded in 1976 by Steve Jobs, Steve Wozniak and Ronald Wayne to develop and sell Wozniak's Apple I personal computer\"\n",
        "\n",
        "onnx_pipeline(sequence)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cUM2waEKGkfq",
        "outputId": "13d61e49-1501-4df4-b5dd-d930ae193873"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[{'end': 5,\n",
              "  'entity_group': 'ORG',\n",
              "  'score': 0.9978969,\n",
              "  'start': 0,\n",
              "  'word': 'Apple'},\n",
              " {'end': 39,\n",
              "  'entity_group': 'PER',\n",
              "  'score': 0.9981243,\n",
              "  'start': 29,\n",
              "  'word': 'Steve Jobs'},\n",
              " {'end': 54,\n",
              "  'entity_group': 'PER',\n",
              "  'score': 0.9741297,\n",
              "  'start': 41,\n",
              "  'word': 'Steve Wozniak'},\n",
              " {'end': 71,\n",
              "  'entity_group': 'PER',\n",
              "  'score': 0.99970996,\n",
              "  'start': 59,\n",
              "  'word': 'Ronald Wayne'},\n",
              " {'end': 99,\n",
              "  'entity_group': 'PER',\n",
              "  'score': 0.86664414,\n",
              "  'start': 92,\n",
              "  'word': 'Wozniak'},\n",
              " {'end': 109,\n",
              "  'entity_group': 'MISC',\n",
              "  'score': 0.99852806,\n",
              "  'start': 102,\n",
              "  'word': 'Apple I'}]"
            ]
          },
          "metadata": {},
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "pytorch_pipeline = TokenClassificationPipeline(\n",
        "    task=\"ner\", \n",
        "    model=model,\n",
        "    tokenizer=tokenizer,\n",
        "    framework=\"pt\",\n",
        "    aggregation_strategy=\"simple\",\n",
        ")"
      ],
      "metadata": {
        "id": "PUpZqZpzGsAx"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "sequences = {\n",
        "    \"short_sequence\": \"Hello my name is Thomas and I love HuggingFace.\",\n",
        "    \"medium_sequence\": \"Winston Churchill was born in 1874 in Stoke-on-Trent, England, to a German father, William and Elizabeth Churchill.\",\n",
        "    \"long_sequence\": \"\"\"The first person to reach the summit of Everest was the South Nepalese Everest Gurun, \n",
        "                who was a member of the Royal Nepal Expedition, led by the Nepalese Mountaineer, Sir Edmund Hillary. \n",
        "                Hilary lived in the Himalayas for a time. He sadly died in 1953 at the age of 88.\"\"\"\n",
        "}"
      ],
      "metadata": {
        "id": "6sroCNTFHV0F"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import timeit\n",
        "\n",
        "results = [[\"Sequence Length\", \"PyTorch\", \"ONNX\"]]\n",
        "for k, v in sequences.items():\n",
        "    results.append(\n",
        "        [k, timeit.timeit(lambda: pytorch_pipeline(v), number=300), timeit.timeit(lambda: onnx_pipeline(v), number=300)]\n",
        "    )"
      ],
      "metadata": {
        "id": "jW6Q4LfWHdCS"
      },
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from tabulate import tabulate\n",
        "\n",
        "print(tabulate(results, headers=\"firstrow\"))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WSHHl3s2Ja2i",
        "outputId": "c4505448-d5aa-4694-f7d2-e37378dddcef"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Sequence Length      PyTorch     ONNX\n",
            "-----------------  ---------  -------\n",
            "short_sequence       31.3778  21.4062\n",
            "medium_sequence      44.9581  30.7068\n",
            "long_sequence        83.5834  58.7696\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(f\"For a short sequence: ONNX is {results[1][1]/results[1][2]:.2f}x faster than PyTorch\")\n",
        "print(f\"For a medium sequence: ONNX is {results[2][1]/results[2][2]:.2f}x faster than PyTorch\")\n",
        "print(f\"For a long sequence: ONNX is {results[3][1]/results[3][2]:.2f}x faster than PyTorch\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IE6S3TxQK_fH",
        "outputId": "38c9ad26-3d1e-4297-9917-da83ecc6e58c"
      },
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "For a short sequence: ONNX is 1.47x faster than PyTorch\n",
            "For a medium sequence: ONNX is 1.46x faster than PyTorch\n",
            "For a long sequence: ONNX is 1.42x faster than PyTorch\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "Zi9UJHNKLATh"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}