examples/18_Export_and_run_models_with_ONNX.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4Pjmz-RORV8E"
      },
      "source": [
        "# Export and run models with ONNX\n",
        "\n",
        "The [ONNX runtime](https://onnx.ai/) provides a common serialization format for machine learning models. ONNX supports a number of [different platforms/languages](https://onnxruntime.ai/docs/how-to/install.html#requirements) and has features built in to help reduce inference time. \n",
        "\n",
        "PyTorch has robust support for exporting Torch models to ONNX. This enables exporting Hugging Face Transformer and/or other downstream models directly to ONNX. \n",
        "\n",
        "ONNX opens an avenue for direct inference using a number of languages and platforms. For example, a model could be run directly on Android to limit data sent to a third party service. ONNX is an exciting development with a lot of promise. Microsoft has also released [Hummingbird](https://github.com/microsoft/hummingbird) which enables exporting traditional models (sklearn, decision trees, logistical regression..) to ONNX. \n",
        "\n",
        "This notebook will cover how to export models to ONNX using txtai. These models will then be directly run in Python, JavaScript, Java and Rust. Currently, txtai supports all these languages through it's API and that is still the recommended approach. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dk31rbYjSTYm"
      },
      "source": [
        "# Install dependencies\n",
        "\n",
        "Install `txtai` and all dependencies. Since this notebook uses ONNX quantization, we need to install the pipeline extras package."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XMQuuun2R06J"
      },
      "source": [
        "%%capture\n",
        "!pip install datasets git+https://github.com/neuml/txtai#egg=txtai[pipeline]"
      ],
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PNPJ95cdTKSS"
      },
      "source": [
        "# Run a model with ONNX\n",
        "\n",
        "Let's get right to it! The following example exports a sentiment analysis model to ONNX and runs an inference session.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "USb4JXZHxqTA",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "28d3e70e-efa9-4b07-a602-6ffd89d1279f"
      },
      "source": [
        "import numpy as np\n",
        "\n",
        "from onnxruntime import InferenceSession, SessionOptions\n",
        "from transformers import AutoTokenizer\n",
        "from txtai.pipeline import HFOnnx\n",
        "\n",
        "# Normalize logits using sigmoid function\n",
        "sigmoid = lambda x: 1.0 / (1.0 + np.exp(-x))\n",
        "\n",
        "# Export to ONNX\n",
        "onnx = HFOnnx()\n",
        "model = onnx(\"distilbert-base-uncased-finetuned-sst-2-english\", \"text-classification\")\n",
        "\n",
        "# Start inference session\n",
        "options = SessionOptions()\n",
        "session = InferenceSession(model, options)\n",
        "\n",
        "# Tokenize\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"distilbert-base-uncased-finetuned-sst-2-english\")\n",
        "tokens = tokenizer([\"I am happy\", \"I am mad\"], return_tensors=\"np\")\n",
        "\n",
        "# Print results\n",
        "outputs = session.run(None, dict(tokens))\n",
        "print(sigmoid(outputs[0]))"
      ],
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[[0.01295124 0.9909526 ]\n",
            " [0.9874723  0.0297817 ]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jkmQoQvlmHfQ"
      },
      "source": [
        "And just like that, there are results! The text classification model is judging sentiment using two labels, 0 for negative to 1 for positive. The results above shows the probability of each label per text snippet.\n",
        "\n",
        "The ONNX pipeline loads the model, converts the graph to ONNX and returns. Note that no output file was provided, in this case the ONNX model is returned as a byte array. If an output file is provided, this method returns the output path."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yFAOHVmXml8o"
      },
      "source": [
        "# Train and Export a model for Text Classification\n",
        "\n",
        "Next we'll combine the ONNX pipeline with a Trainer pipeline to create a \"train and export to ONNX\" workflow."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 579
        },
        "id": "Wh8TkszumlIe",
        "outputId": "864f2074-ae50-40d6-bc34-2b1d86a71488"
      },
      "source": [
        "from datasets import load_dataset\n",
        "from txtai.pipeline import HFTrainer\n",
        "\n",
        "trainer = HFTrainer()\n",
        "\n",
        "# Hugging Face dataset\n",
        "ds = load_dataset(\"glue\", \"sst2\")\n",
        "data = ds[\"train\"].select(range(10000)).flatten_indices()\n",
        "\n",
        "# Train new model using 10,000 SST2 records (in-memory)\n",
        "model, tokenizer = trainer(\"google/electra-base-discriminator\", data, columns=(\"sentence\", \"label\"))\n",
        "\n",
        "# Export model trained in-memory to ONNX (still in-memory)\n",
        "output = onnx((model, tokenizer), \"text-classification\", quantize=True)\n",
        "\n",
        "# Start inference session\n",
        "options = SessionOptions()\n",
        "session = InferenceSession(output, options)\n",
        "\n",
        "# Tokenize\n",
        "tokens = tokenizer([\"I am happy\", \"I am mad\"], return_tensors=\"np\")\n",
        "\n",
        "# Print results\n",
        "outputs = session.run(None, dict(tokens))\n",
        "print(sigmoid(outputs[0]))"
      ],
      "execution_count": 27,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:datasets.arrow_dataset:Loading cached processed dataset at /root/.cache/huggingface/datasets/glue/sst2/1.0.0/dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad/cache-e28d0e20a676bad0.arrow\n",
            "WARNING:datasets.arrow_dataset:Loading cached processed dataset at /root/.cache/huggingface/datasets/glue/sst2/1.0.0/dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad/cache-d7b5d80ca22204f9.arrow\n",
            "Some weights of the model checkpoint at google/electra-base-discriminator were not used when initializing ElectraForSequenceClassification: ['discriminator_predictions.dense_prediction.weight', 'discriminator_predictions.dense_prediction.bias', 'discriminator_predictions.dense.weight', 'discriminator_predictions.dense.bias']\n",
            "- This IS expected if you are initializing ElectraForSequenceClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
            "- This IS NOT expected if you are initializing ElectraForSequenceClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
            "Some weights of ElectraForSequenceClassification were not initialized from the model checkpoint at google/electra-base-discriminator and are newly initialized: ['classifier.out_proj.bias', 'classifier.dense.weight', 'classifier.out_proj.weight', 'classifier.dense.bias']\n",
            "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
            "/usr/local/lib/python3.7/dist-packages/transformers/optimization.py:310: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
            "  FutureWarning,\n",
            "You're using a ElectraTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "    <div>\n",
              "      \n",
              "      <progress value='3750' max='3750' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
              "      [3750/3750 07:56, Epoch 3/3]\n",
              "    </div>\n",
              "    <table border=\"1\" class=\"dataframe\">\n",
              "  <thead>\n",
              " <tr style=\"text-align: left;\">\n",
              "      <th>Step</th>\n",
              "      <th>Training Loss</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <td>500</td>\n",
              "      <td>0.396800</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>1000</td>\n",
              "      <td>0.330900</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>1500</td>\n",
              "      <td>0.232400</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>2000</td>\n",
              "      <td>0.188200</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>2500</td>\n",
              "      <td>0.173600</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>3000</td>\n",
              "      <td>0.068600</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <td>3500</td>\n",
              "      <td>0.069800</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table><p>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[[0.01525715 0.975399  ]\n",
            " [0.97395283 0.04432926]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lE7dPj3tsn5S"
      },
      "source": [
        "The results are similar to the previous step, although this model is only trained on a fraction of the sst2 dataset. Lets save this model for later."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Q_kAFYd_s_Bi"
      },
      "source": [
        "onnx = HFOnnx()\n",
        "text = onnx((model, tokenizer), \"text-classification\", \"text-classify.onnx\", quantize=True)"
      ],
      "execution_count": 29,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ugNZO4c-uAS-"
      },
      "source": [
        "# Export a Sentence Embeddings model\n",
        "\n",
        "The ONNX pipeline also supports exporting sentence embeddings models trained with the [sentence-transformers](https://github.com/UKPLab/sentence-transformers) package. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x9B7qOk_uQRN"
      },
      "source": [
        "embeddings = onnx(\"sentence-transformers/paraphrase-MiniLM-L6-v2\", \"pooling\", \"embeddings.onnx\", quantize=True)"
      ],
      "execution_count": 30,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rirMSM2kvgJF"
      },
      "source": [
        "Now let's run the model with ONNX."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6MBraENcu8Oz",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ba4528d2-6d6a-4181-e9c4-3d2a98d6663a"
      },
      "source": [
        "from sklearn.metrics.pairwise import cosine_similarity\n",
        "\n",
        "options = SessionOptions()\n",
        "session = InferenceSession(embeddings, options)\n",
        "\n",
        "tokens = tokenizer([\"I am happy\", \"I am glad\"], return_tensors=\"np\")\n",
        "\n",
        "outputs = session.run(None, dict(tokens))[0]\n",
        "\n",
        "print(cosine_similarity(outputs))"
      ],
      "execution_count": 31,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[[0.99999994 0.8474637 ]\n",
            " [0.8474637  0.9999997 ]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pwgU4vu8vk0T"
      },
      "source": [
        "The code above tokenizes two separate text snippets (\"I am happy\" and \"I am glad\") and runs it through the ONNX model. \n",
        "\n",
        "This outputs two embeddings arrays and those arrays are compared using cosine similarity. As we can see, the two text snippets have close semantic meaning."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t_OQaQeIb7UB"
      },
      "source": [
        "# Load an ONNX model with txtai\n",
        "\n",
        "txtai has built-in support for ONNX models. Loading an ONNX model is seamless and Embeddings and Pipelines support it. The following section shows how to load a classification pipeline and embeddings model backed by ONNX."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vhsFzCRBby-h",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "745d69b0-035b-4e44-a881-57f9ece171ab"
      },
      "source": [
        "from txtai.embeddings import Embeddings\n",
        "from txtai.pipeline import Labels\n",
        "\n",
        "labels = Labels((\"text-classify.onnx\", \"google/electra-base-discriminator\"), dynamic=False)\n",
        "print(labels([\"I am happy\", \"I am mad\"]))\n",
        "\n",
        "embeddings = Embeddings({\"path\": \"embeddings.onnx\", \"tokenizer\": \"sentence-transformers/paraphrase-MiniLM-L6-v2\"})\n",
        "print(embeddings.similarity(\"I am happy\", [\"I am glad\"]))"
      ],
      "execution_count": 32,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[[(1, 0.999687910079956), (0, 0.0003121310146525502)], [(0, 0.9991233944892883), (1, 0.0008765518432483077)]]\n",
            "[(0, 0.8298245072364807)]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xx8G29hkwdNY"
      },
      "source": [
        "# JavaScript\n",
        "\n",
        "So far, we've exported models to ONNX and run them through Python. This already has a lot of advantages, which include fast inference times, quantization and less software dependencies. But ONNX really shines when we run a model trained in Python in other languages/platforms.\n",
        "\n",
        "Let's try running the models trained above in JavaScript. First step is getting the Node.js environment and dependencies setup.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_RK79O9c4Z_y"
      },
      "source": [
        "%%capture\n",
        "import os\n",
        "\n",
        "os.chdir(\"/content\")\n",
        "!mkdir js\n",
        "os.chdir(\"/content/js\")\n",
        "\n",
        "# Copy ONNX models\n",
        "!cp ../text-classify.onnx .\n",
        "!cp ../embeddings.onnx .\n",
        "\n",
        "# Get tokenizers project\n",
        "!git clone https://github.com/huggingface/tokenizers.git\n",
        "\n",
        "os.chdir(\"/content/js/tokenizers/bindings/node\")\n",
        "\n",
        "# Install Rust to compile tokenizer bindings\n",
        "!apt-get install rustc cargo\n",
        "\n",
        "# Build tokenizers package locally as binary version on npm doesn't work for latest version of Node.js\n",
        "!npm install --also=dev\n",
        "!npm run dev\n",
        "\n",
        "os.chdir(\"/content/js\")"
      ],
      "execution_count": 33,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0HtVEl74xrZ7",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "7ffea868-6dd7-4603-e04c-ce8d4c557ff6"
      },
      "source": [
        "%%writefile package.json\n",
        "{\n",
        "  \"name\": \"onnx-test\",\n",
        "  \"private\": true,\n",
        "  \"version\": \"1.0.0\",\n",
        "  \"description\": \"ONNX Runtime Node.js test\",\n",
        "  \"main\": \"index.js\",\n",
        "  \"dependencies\": {\n",
        "    \"onnxruntime-node\": \">=1.12.1\",\n",
        "    \"tokenizers\": \"file:tokenizers/bindings/node\"\n",
        "  }\n",
        "}"
      ],
      "execution_count": 34,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing package.json\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%capture\n",
        "\n",
        "# Install all dependencies\n",
        "!npm install"
      ],
      "metadata": {
        "id": "4naPtk-iBI-g"
      },
      "execution_count": 35,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "At85iA8U63iV"
      },
      "source": [
        "Next we'll write the inference code in JavaScript to an index.js file."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RImohEnFyFg0",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "094ef937-0650-4d5b-f4e3-79f114bd9807",
        "cellView": "form"
      },
      "source": [
        "#@title\n",
        "%%writefile index.js\n",
        "const ort = require('onnxruntime-node');\n",
        "const { promisify } = require('util');\n",
        "const { Tokenizer } = require(\"tokenizers/dist/bindings/tokenizer\");\n",
        "\n",
        "function sigmoid(data) {\n",
        "    return data.map(x => 1 / (1 + Math.exp(-x)))\n",
        "}\n",
        "\n",
        "function softmax(data) { \n",
        "    return data.map(x => Math.exp(x) / (data.map(y => Math.exp(y))).reduce((a,b) => a+b)) \n",
        "}\n",
        "\n",
        "function similarity(v1, v2) {\n",
        "    let dot = 0.0;\n",
        "    let norm1 = 0.0;\n",
        "    let norm2 = 0.0;\n",
        "\n",
        "    for (let x = 0; x < v1.length; x++) {\n",
        "        dot += v1[x] * v2[x];\n",
        "        norm1 += Math.pow(v1[x], 2);\n",
        "        norm2 += Math.pow(v2[x], 2);\n",
        "    }\n",
        "\n",
        "    return dot / (Math.sqrt(norm1) * Math.sqrt(norm2));\n",
        "}\n",
        "\n",
        "function tokenizer() {\n",
        "    let tokenizer = Tokenizer.fromPretrained(\"bert-base-uncased\");\n",
        "    return promisify(tokenizer.encode.bind(tokenizer));\n",
        "}\n",
        "\n",
        "async function predict(session, text) {\n",
        "    try {\n",
        "        // Tokenize input\n",
        "        let encode = tokenizer();\n",
        "        let output = await encode(text);\n",
        "\n",
        "        let ids = output.getIds().map(x => BigInt(x))\n",
        "        let mask = output.getAttentionMask().map(x => BigInt(x))\n",
        "        let tids = output.getTypeIds().map(x => BigInt(x))\n",
        "\n",
        "        // Convert inputs to tensors    \n",
        "        let tensorIds = new ort.Tensor('int64', BigInt64Array.from(ids), [1, ids.length]);\n",
        "        let tensorMask = new ort.Tensor('int64', BigInt64Array.from(mask), [1, mask.length]);\n",
        "        let tensorTids = new ort.Tensor('int64', BigInt64Array.from(tids), [1, tids.length]);\n",
        "\n",
        "        let inputs = null;\n",
        "        if (session.inputNames.length > 2) {\n",
        "            inputs = { input_ids: tensorIds, attention_mask: tensorMask, token_type_ids: tensorTids};\n",
        "        }\n",
        "        else {\n",
        "            inputs = { input_ids: tensorIds, attention_mask: tensorMask};\n",
        "        }\n",
        "\n",
        "        return await session.run(inputs);\n",
        "    } catch (e) {\n",
        "        console.error(`failed to inference ONNX model: ${e}.`);\n",
        "    }\n",
        "}\n",
        "\n",
        "async function main() {\n",
        "    let args = process.argv.slice(2);\n",
        "    if (args.length > 1) {\n",
        "        // Run sentence embeddings\n",
        "        const session = await ort.InferenceSession.create('./embeddings.onnx');\n",
        "\n",
        "        let v1 = await predict(session, args[0]);\n",
        "        let v2 = await predict(session, args[1]);\n",
        "\n",
        "        // Unpack results\n",
        "        v1 = v1.embeddings.data;\n",
        "        v2 = v2.embeddings.data;\n",
        "\n",
        "        // Print similarity\n",
        "        console.log(similarity(Array.from(v1), Array.from(v2)));\n",
        "    }\n",
        "    else {\n",
        "        // Run text classifier\n",
        "        const session = await ort.InferenceSession.create('./text-classify.onnx');\n",
        "        let results = await predict(session, args[0]);\n",
        "\n",
        "        // Normalize results using softmax and print\n",
        "        console.log(softmax(results.logits.data));\n",
        "    }\n",
        "}\n",
        "\n",
        "main();"
      ],
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing index.js\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rZI9PJzi6_bO"
      },
      "source": [
        "## Run Text Classification in JavaScript with ONNX"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bdz68KZT1Jfm",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "48c4a427-3108-436c-eac9-564835ea061c"
      },
      "source": [
        "!node . \"I am happy\"\n",
        "!node . \"I am mad\""
      ],
      "execution_count": 37,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Float32Array(2) [ \u001b[33m0.0003121308400295675\u001b[39m, \u001b[33m0.9996878504753113\u001b[39m ]\n",
            "Float32Array(2) [ \u001b[33m0.9991234540939331\u001b[39m, \u001b[33m0.0008765519596636295\u001b[39m ]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "swSEmqto33VP"
      },
      "source": [
        "First off, have to say this is 🔥🔥🔥! Just amazing that this model can be fully run in JavaScript. It's a great time to be in NLP!\n",
        "\n",
        "The steps above installed a JavaScript environment with dependencies to run ONNX and tokenize data in JavaScript. The text classification model previously created is loaded into the JavaScript ONNX runtime and inference is run.\n",
        "\n",
        "As a reminder, the text classification model is judging sentiment using two labels, 0 for negative to 1 for positive. The results above shows the probability of each label per text snippet."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Az9YaDc6u9P"
      },
      "source": [
        "## Build sentence embeddings and compare similarity in JavaScript with ONNX"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "10jcUbUx6MAI",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ac751cee-5f44-4dad-c164-5d124be75ec3"
      },
      "source": [
        "!node . \"I am happy\", \"I am glad\""
      ],
      "execution_count": 38,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[33m0.8285076844387538\u001b[39m\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8Jyk-9Ko78Ma"
      },
      "source": [
        "Once again....wow!! The sentence embeddings model produces vectors that can be used to compare semantic similarity, -1 being most dissimilar and 1 being most similar.\n",
        "\n",
        "While the results don't match the exported model exactly, it's very close. Worth mentioning again that this is 100% JavaScript, no API or remote calls, all within node."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BQeMBNWO9Hpr"
      },
      "source": [
        "# Java\n",
        "\n",
        "Let's try the same thing with Java. The following sections initialize a Java build environment and writes out the code necessary to run the ONNX inference."
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%capture\n",
        "import os\n",
        "\n",
        "os.chdir(\"/content\")\n",
        "!mkdir java\n",
        "os.chdir(\"/content/java\")\n",
        "\n",
        "# Copy ONNX models\n",
        "!cp ../text-classify.onnx .\n",
        "!cp ../embeddings.onnx .\n",
        "\n",
        "# Save copy of Bert Tokenizer\n",
        "tokenizer.save_pretrained(\"bert\")\n",
        "\n",
        "!mkdir -p src/main/java\n",
        "\n",
        "# Install gradle\n",
        "!wget https://services.gradle.org/distributions/gradle-7.5.1-bin.zip\n",
        "!unzip -o gradle-7.5.1-bin.zip"
      ],
      "metadata": {
        "id": "L1YMoO7WwkEk"
      },
      "execution_count": 39,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gjZ2p7Jf9mOV",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0cc3f4ef-bac6-4c13-dc6b-8482353bb741"
      },
      "source": [
        "%%writefile build.gradle\n",
        "apply plugin: \"java\"\n",
        "\n",
        "repositories {\n",
        "    mavenCentral()\n",
        "}\n",
        "\n",
        "dependencies {\n",
        "    implementation \"com.robrua.nlp:easy-bert:1.0.3\"\n",
        "    implementation \"com.microsoft.onnxruntime:onnxruntime:1.12.1\"\n",
        "}\n",
        "\n",
        "java {\n",
        "    toolchain {\n",
        "        languageVersion = JavaLanguageVersion.of(8)\n",
        "    }\n",
        "}\n",
        "\n",
        "jar {\n",
        "    archiveBaseName = \"onnxjava\"\n",
        "}\n",
        "\n",
        "task onnx(type: JavaExec) {\n",
        "    description = \"Runs ONNX demo\"\n",
        "    classpath = sourceSets.main.runtimeClasspath\n",
        "    main = \"OnnxDemo\"\n",
        "}"
      ],
      "execution_count": 40,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing build.gradle\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9wlVWVky9NZ3"
      },
      "source": [
        "%%capture\n",
        "\n",
        "# Create environment\n",
        "!gradle-7.5.1/bin/gradle wrapper"
      ],
      "execution_count": 41,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vnxKGSuz_fnj",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c1b73bf8-c5f4-46fa-df90-29bd60507287",
        "cellView": "form"
      },
      "source": [
        "#@title\n",
        "%%writefile src/main/java/OnnxDemo.java\n",
        "import java.io.File;\n",
        "\n",
        "import java.nio.LongBuffer;\n",
        "\n",
        "import java.util.Arrays;\n",
        "import java.util.ArrayList;\n",
        "import java.util.HashMap;\n",
        "import java.util.List;\n",
        "import java.util.Map;\n",
        "\n",
        "import ai.onnxruntime.OnnxTensor;\n",
        "import ai.onnxruntime.OrtEnvironment;\n",
        "import ai.onnxruntime.OrtSession;\n",
        "import ai.onnxruntime.OrtSession.Result;\n",
        "\n",
        "import com.robrua.nlp.bert.FullTokenizer;\n",
        "\n",
        "class Tokens {\n",
        "    public long[] ids;\n",
        "    public long[] mask;\n",
        "    public long[] types;\n",
        "}\n",
        "\n",
        "class Tokenizer {\n",
        "    private FullTokenizer tokenizer;\n",
        "\n",
        "    public Tokenizer(String path) {\n",
        "        File vocab = new File(path);\n",
        "        this.tokenizer = new FullTokenizer(vocab, true);\n",
        "    }\n",
        "\n",
        "    public Tokens tokenize(String text) {\n",
        "        // Build list of tokens\n",
        "        List<String> tokensList = new ArrayList();\n",
        "        tokensList.add(\"[CLS]\"); \n",
        "        tokensList.addAll(Arrays.asList(tokenizer.tokenize(text)));\n",
        "        tokensList.add(\"[SEP]\");\n",
        "\n",
        "        int[] ids = tokenizer.convert(tokensList.toArray(new String[0]));\n",
        "\n",
        "        Tokens tokens = new Tokens();\n",
        "\n",
        "        // input ids    \n",
        "        tokens.ids = Arrays.stream(ids).mapToLong(i -> i).toArray();\n",
        "\n",
        "        // attention mask\n",
        "        tokens.mask = new long[ids.length];\n",
        "        Arrays.fill(tokens.mask, 1);\n",
        "\n",
        "        // token type ids\n",
        "        tokens.types = new long[ids.length];\n",
        "        Arrays.fill(tokens.types, 0);\n",
        "\n",
        "        return tokens;\n",
        "    }\n",
        "}\n",
        "\n",
        "class Inference {\n",
        "    private Tokenizer tokenizer;\n",
        "    private OrtEnvironment env;\n",
        "    private OrtSession session;\n",
        "\n",
        "    public Inference(String model) throws Exception {\n",
        "        this.tokenizer = new Tokenizer(\"bert/vocab.txt\");\n",
        "        this.env = OrtEnvironment.getEnvironment();\n",
        "        this.session = env.createSession(model, new OrtSession.SessionOptions());\n",
        "    }\n",
        "\n",
        "    public float[][] predict(String text) throws Exception {\n",
        "        Tokens tokens = this.tokenizer.tokenize(text);\n",
        "\n",
        "        Map<String, OnnxTensor> inputs = new HashMap<String, OnnxTensor>();\n",
        "        inputs.put(\"input_ids\", OnnxTensor.createTensor(env, LongBuffer.wrap(tokens.ids),  new long[]{1, tokens.ids.length}));\n",
        "        inputs.put(\"attention_mask\", OnnxTensor.createTensor(env, LongBuffer.wrap(tokens.mask),  new long[]{1, tokens.mask.length}));\n",
        "        inputs.put(\"token_type_ids\", OnnxTensor.createTensor(env, LongBuffer.wrap(tokens.types),  new long[]{1, tokens.types.length}));\n",
        "\n",
        "        return (float[][])session.run(inputs).get(0).getValue();\n",
        "    }\n",
        "}\n",
        "\n",
        "class Vectors {\n",
        "    public static double similarity(float[] v1, float[] v2) {\n",
        "        double dot = 0.0;\n",
        "        double norm1 = 0.0;\n",
        "        double norm2 = 0.0;\n",
        "\n",
        "        for (int x = 0; x < v1.length; x++) {\n",
        "            dot += v1[x] * v2[x];\n",
        "            norm1 += Math.pow(v1[x], 2);\n",
        "            norm2 += Math.pow(v2[x], 2);\n",
        "        }\n",
        "\n",
        "        return dot / (Math.sqrt(norm1) * Math.sqrt(norm2));\n",
        "    }\n",
        "\n",
        "    public static float[] softmax(float[] input) {\n",
        "        double[] t = new double[input.length];\n",
        "        double sum = 0.0;\n",
        "\n",
        "        for (int x = 0; x < input.length; x++) {\n",
        "            double val = Math.exp(input[x]);\n",
        "            sum += val;\n",
        "            t[x] = val;\n",
        "        }\n",
        "\n",
        "        float[] output = new float[input.length];\n",
        "        for (int x = 0; x < output.length; x++) {\n",
        "            output[x] = (float) (t[x] / sum);\n",
        "        }\n",
        "\n",
        "        return output;\n",
        "    }\n",
        "}\n",
        "\n",
        "public class OnnxDemo {\n",
        "    public static void main(String[] args) {\n",
        "        try {\n",
        "            if (args.length < 2) {\n",
        "              Inference inference = new Inference(\"text-classify.onnx\");\n",
        "\n",
        "              float[][] v1 = inference.predict(args[0]);\n",
        "\n",
        "              System.out.println(Arrays.toString(Vectors.softmax(v1[0])));\n",
        "            }\n",
        "            else {\n",
        "              Inference inference = new Inference(\"embeddings.onnx\");\n",
        "              float[][] v1 = inference.predict(args[0]);\n",
        "              float[][] v2 = inference.predict(args[1]);\n",
        "\n",
        "              System.out.println(Vectors.similarity(v1[0], v2[0]));\n",
        "            }\n",
        "        }\n",
        "        catch (Exception ex) {\n",
        "            ex.printStackTrace();\n",
        "        }\n",
        "    }\n",
        "}"
      ],
      "execution_count": 42,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing src/main/java/OnnxDemo.java\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qQuuXw97Z_I7"
      },
      "source": [
        "## Run Text Classification in Java with ONNX"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hFXyH96gAZpu",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "efd5b783-b23a-407c-8577-c18e3a6cb984"
      },
      "source": [
        "!./gradlew -q --console=plain onnx --args='\"I am happy\"' 2> /dev/null\n",
        "!./gradlew -q --console=plain onnx --args='\"I am mad\"' 2> /dev/null"
      ],
      "execution_count": 43,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[3.1213084E-4, 0.99968785]\n",
            "\u001b[m[0.99912345, 8.7655196E-4]\n",
            "\u001b[m"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pE3FSsAAaJHe"
      },
      "source": [
        "The command above tokenizes the input and runs inference with a text classification model previously created using a Java ONNX inference session. \n",
        "\n",
        "As a reminder, the text classification model is judging sentiment using two labels, 0 for negative to 1 for positive. The results above shows the probability of each label per text snippet."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bux8v0C4aDyP"
      },
      "source": [
        "## Build sentence embeddings and compare similarity in Java with ONNX"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "f6zE9VrwCcUa",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "988e59d0-943f-45b6-d37e-5fc1ebbbcefe"
      },
      "source": [
        "!./gradlew -q --console=plain onnx --args='\"I am happy\" \"I am glad\"' 2> /dev/null"
      ],
      "execution_count": 44,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.8298244656285757\n",
            "\u001b[m"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0uepOZvJDOCB"
      },
      "source": [
        "The sentence embeddings model produces vectors that can be used to compare semantic similarity, -1 being most dissimilar and 1 being most similar. \n",
        "\n",
        "This is 100% Java, no API or remote calls, all within the JVM. Still think it's amazing!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "faRu9EAJDUXw"
      },
      "source": [
        "# Rust\n",
        "\n",
        "Last but not least, let's try Rust. The following sections initialize a Rust build environment and writes out the code necessary to run the ONNX inference."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X3Xp1KLhelqw"
      },
      "source": [
        "%%capture\n",
        "import os\n",
        "\n",
        "os.chdir(\"/content\")\n",
        "!mkdir rust\n",
        "os.chdir(\"/content/rust\")\n",
        "\n",
        "# Copy ONNX models\n",
        "!cp ../text-classify.onnx .\n",
        "!cp ../embeddings.onnx .\n",
        "\n",
        "# Install Rust\n",
        "!apt-get install rustc cargo\n",
        "\n",
        "!mkdir -p src"
      ],
      "execution_count": 45,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c7hz--Gne6Oa",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d98ad709-5675-4193-e598-e6cbe12edda3"
      },
      "source": [
        "%%writefile Cargo.toml\n",
        "[package]\n",
        "name = \"onnx-test\"\n",
        "version = \"1.0.0\"\n",
        "description = \"\"\"\n",
        "ONNX Runtime Rust test\n",
        "\"\"\"\n",
        "edition = \"2018\"\n",
        "\n",
        "[dependencies]\n",
        "onnxruntime = { version = \"0.0.14\"}\n",
        "tokenizers = { version = \"0.13.1\"}"
      ],
      "execution_count": 46,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing Cargo.toml\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_8fdRvO1fFBm",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "53168684-12c2-46fc-e0e6-54c4c6d03cb1",
        "cellView": "form"
      },
      "source": [
        "#@title\n",
        "%%writefile src/main.rs\n",
        "use onnxruntime::environment::Environment;\n",
        "use onnxruntime::GraphOptimizationLevel;\n",
        "use onnxruntime::ndarray::{Array2, Axis};\n",
        "use onnxruntime::tensor::OrtOwnedTensor;\n",
        "\n",
        "use std::env;\n",
        "\n",
        "use tokenizers::tokenizer::{Result, Tokenizer};\n",
        "\n",
        "fn tokenize(text: String, inputs: usize) -> Vec<Array2<i64>> {\n",
        "    // Load tokenizer from HF Hub\n",
        "    let tokenizer = Tokenizer::from_pretrained(\"bert-base-uncased\", None).unwrap();\n",
        "\n",
        "    // Encode input text\n",
        "    let encoding = tokenizer.encode(text, true).unwrap();\n",
        "\n",
        "    let v1: Vec<i64> = encoding.get_ids().to_vec().into_iter().map(|x| x as i64).collect();\n",
        "    let v2: Vec<i64> = encoding.get_attention_mask().to_vec().into_iter().map(|x| x as i64).collect();\n",
        "    let v3: Vec<i64> = encoding.get_type_ids().to_vec().into_iter().map(|x| x as i64).collect();\n",
        "\n",
        "    let ids = Array2::from_shape_vec((1, v1.len()), v1).unwrap();\n",
        "    let mask = Array2::from_shape_vec((1, v2.len()), v2).unwrap();\n",
        "    let tids = Array2::from_shape_vec((1, v3.len()), v3).unwrap();\n",
        "\n",
        "    return if inputs > 2 { vec![ids, mask, tids] } else { vec![ids, mask] };\n",
        "}\n",
        "\n",
        "fn predict(text: String, softmax: bool) -> Vec<f32> {\n",
        "    // Start onnx session\n",
        "    let environment = Environment::builder()\n",
        "        .with_name(\"test\")\n",
        "        .build().unwrap();\n",
        "\n",
        "    // Derive model path\n",
        "    let model = if softmax { \"text-classify.onnx\" } else { \"embeddings.onnx\" };\n",
        "\n",
        "    let mut session = environment\n",
        "        .new_session_builder().unwrap()\n",
        "        .with_optimization_level(GraphOptimizationLevel::Basic).unwrap()\n",
        "        .with_number_threads(1).unwrap()\n",
        "        .with_model_from_file(model).unwrap();\n",
        "\n",
        "    let inputs = tokenize(text, session.inputs.len());\n",
        "\n",
        "    // Run inference and print result\n",
        "    let outputs: Vec<OrtOwnedTensor<f32, _>> = session.run(inputs).unwrap();\n",
        "    let output: &OrtOwnedTensor<f32, _> = &outputs[0];\n",
        "\n",
        "    let probabilities: Vec<f32>;\n",
        "    if softmax {\n",
        "        probabilities = output\n",
        "            .softmax(Axis(1))\n",
        "            .iter()\n",
        "            .copied()\n",
        "            .collect::<Vec<_>>();\n",
        "    }\n",
        "    else {\n",
        "        probabilities= output\n",
        "            .iter()\n",
        "            .copied()\n",
        "            .collect::<Vec<_>>();\n",
        "    }\n",
        "\n",
        "    return probabilities;\n",
        "}\n",
        "\n",
        "fn similarity(v1: &Vec<f32>, v2: &Vec<f32>) -> f64 {\n",
        "    let mut dot = 0.0;\n",
        "    let mut norm1 = 0.0;\n",
        "    let mut norm2 = 0.0;\n",
        "\n",
        "    for x in 0..v1.len() {\n",
        "        dot += v1[x] * v2[x];\n",
        "        norm1 += v1[x].powf(2.0);\n",
        "        norm2 += v2[x].powf(2.0);\n",
        "    }\n",
        "\n",
        "    return dot as f64 / (norm1.sqrt() * norm2.sqrt()) as f64\n",
        "}\n",
        "\n",
        "fn main() -> Result<()> {\n",
        "    // Tokenize input string\n",
        "    let args: Vec<String> = env::args().collect();\n",
        "\n",
        "    if args.len() <= 2 {\n",
        "      let v1 = predict(args[1].to_string(), true);\n",
        "      println!(\"{:?}\", v1);\n",
        "    }\n",
        "    else {\n",
        "      let v1 = predict(args[1].to_string(), false);\n",
        "      let v2 = predict(args[2].to_string(), false);\n",
        "      println!(\"{:?}\", similarity(&v1, &v2));\n",
        "    }\n",
        "\n",
        "    Ok(())\n",
        "}"
      ],
      "execution_count": 47,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Writing src/main.rs\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OdfQFY-MiA-n"
      },
      "source": [
        "## Run Text Classification in Rust with ONNX"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b0ymX4ftgWcT",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "84b42d3c-82d4-46fc-fb84-94967bd5330f"
      },
      "source": [
        "!cargo run \"I am happy\" 2> /dev/null\n",
        "!cargo run \"I am mad\" 2> /dev/null"
      ],
      "execution_count": 48,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[0.00030939875, 0.9996906]\n",
            "[0.99912345, 0.0008765513]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NKccz6bBiIgW"
      },
      "source": [
        "The command above tokenizes the input and runs inference with a text classification model previously created using a Rust ONNX inference session. \n",
        "\n",
        "As a reminder, the text classification model is judging sentiment using two labels, 0 for negative to 1 for positive. The results above shows the probability of each label per text snippet."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1D1kN0yNiEg7"
      },
      "source": [
        "## Build sentence embeddings and compare similarity in Rust with ONNX"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A9p6F_ODhenH",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "b43ad47e-e1f3-4748-d854-a0dc2024b780"
      },
      "source": [
        "!cargo run \"I am happy\" \"I am glad\" 2> /dev/null"
      ],
      "execution_count": 49,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.8298246060854143\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TQ7Wvn0OiRr4"
      },
      "source": [
        "The sentence embeddings model produces vectors that can be used to compare semantic similarity, -1 being most dissimilar and 1 being most similar. \n",
        "\n",
        "Once again, this is 100% Rust, no API or remote calls. And yes, still think it's amazing!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-_FNKUWtjLsO"
      },
      "source": [
        "# Wrapping up\n",
        "\n",
        "This notebook covered how to export models to ONNX using txtai. These models were then run in Python, JavaScript, Java and Rust. Golang was also evaluated but there doesn't currently appear to be a stable enough ONNX runtime available. \n",
        "\n",
        "This method provides a way to train and run machine learning models using a number of programming languages on a number of platforms.\n",
        "\n",
        "The following is a non-exhaustive list of use cases. \n",
        "\n",
        "*   Build locally executed models for mobile/edge devices\n",
        "*   Run models with Java/JavaScript/Rust development stacks when teams prefer not to add Python to the mix\n",
        "*   Export models to ONNX for Python inference to improve CPU performance and/or reduce number of software dependencies"
      ]
    }
  ]
}