examples/python/transformers/onnx/HuggingFace_ONNX_in_Spark_NLP_RoBERTa.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "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_RoBERTa.ipynb)\n",
    "\n",
    "# Import ONNX RoBERTa 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. Please make sure you have upgraded to the latest Spark NLP release.\n",
    "- You can import models for RoBERTa from HuggingFace and they have to be in `Fill Mask` category. Meaning, you cannot use RoBERTa models trained/fine-tuned on a specific task such as token/sequence classification."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export and Save HuggingFace model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 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.34.1`. This doesn't mean it won't work with the future releases, but we wanted you to know which versions have been tested successfully."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.1/7.1 MB\u001b[0m \u001b[31m86.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m252.3/252.3 kB\u001b[0m \u001b[31m16.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
      "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
      "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m268.8/268.8 kB\u001b[0m \u001b[31m18.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.8/7.8 MB\u001b[0m \u001b[31m104.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m83.8/83.8 kB\u001b[0m \u001b[31m7.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m451.2/451.2 kB\u001b[0m \u001b[31m37.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.9/5.9 MB\u001b[0m \u001b[31m104.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m212.7/212.7 kB\u001b[0m \u001b[31m15.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m46.0/46.0 kB\u001b[0m \u001b[31m2.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m486.2/486.2 kB\u001b[0m \u001b[31m35.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m14.6/14.6 MB\u001b[0m \u001b[31m101.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m55.5/55.5 kB\u001b[0m \u001b[31m5.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m74.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.1/1.1 MB\u001b[0m \u001b[31m77.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m86.8/86.8 kB\u001b[0m \u001b[31m8.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m110.5/110.5 kB\u001b[0m \u001b[31m12.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m212.5/212.5 kB\u001b[0m \u001b[31m20.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m134.3/134.3 kB\u001b[0m \u001b[31m15.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Building wheel for optimum (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\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.12.0 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-metadata 1.13.1 requires protobuf<5,>=3.20.3, but you have protobuf 3.20.2 which is incompatible.\u001b[0m\u001b[31m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!pip install -q --upgrade transformers[onnx]==4.34.1 optimum"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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 [roberta-base](https://huggingface.co/roberta-base) model from HuggingFace as an example and load it as a `ORTModelForFeatureExtraction`, representing an ONNX model.\n",
    "- In addition to the RoBERTa model, we also need to save the tokenizer. This is the same for every model, these are assets (saved in `/assets`) needed for tokenization inside Spark NLP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Framework not specified. Using pt to export to ONNX.\n",
      "Some weights of the model checkpoint at roberta-base were not used when initializing RobertaModel: ['lm_head.layer_norm.weight', 'lm_head.dense.weight', 'lm_head.dense.bias', 'lm_head.layer_norm.bias', 'lm_head.decoder.weight', 'lm_head.bias']\n",
      "- This IS expected if you are initializing RobertaModel 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 RobertaModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Using framework PyTorch: 2.0.1+cu118\n",
      "Overriding 1 configuration item(s)\n",
      "\t- use_cache -> False\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "============= Diagnostic Run torch.onnx.export version 2.0.1+cu118 =============\n",
      "verbose: False, log level: Level.ERROR\n",
      "======================= 0 NONE 0 NOTE 0 WARNING 0 ERROR ========================\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from optimum.onnxruntime import ORTModelForFeatureExtraction\n",
    "\n",
    "MODEL_NAME = \"roberta-base\"\n",
    "EXPORT_PATH = f\"onnx_models/{MODEL_NAME}\"\n",
    "\n",
    "ort_model = ORTModelForFeatureExtraction.from_pretrained(MODEL_NAME, export=True)\n",
    "\n",
    "# Save the ONNX model\n",
    "ort_model.save_pretrained(EXPORT_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create directory for assets\n",
    "!mkdir {EXPORT_PATH}/assets\n",
    "\n",
    "# let's make sure we sort the vocabs based on their ids first\n",
    "vocabs = ort_model.preprocessors[0].get_vocab()\n",
    "vocabs = sorted(vocabs, key=vocabs.get)\n",
    "\n",
    "# let's save the vocab as txt file\n",
    "with open(f'{EXPORT_PATH}/vocab.txt', 'w') as f:\n",
    "    for item in vocabs:\n",
    "        f.write(\"%s\\n\" % item)\n",
    "\n",
    "# let's copy both vocab.txt and merges.txt files to saved_model/1/assets\n",
    "!cp {EXPORT_PATH}/vocab.txt {EXPORT_PATH}/assets\n",
    "!cp {EXPORT_PATH}/merges.txt {EXPORT_PATH}/assets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's have a look inside these two directories and see what we are dealing with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 488536\n",
      "drwxr-xr-x 2 root root      4096 Jul 11 10:38 assets\n",
      "-rw-r--r-- 1 root root       644 Jul 11 10:37 config.json\n",
      "-rw-r--r-- 1 root root    456318 Jul 11 10:37 merges.txt\n",
      "-rw-r--r-- 1 root root 496467018 Jul 11 10:37 model.onnx\n",
      "-rw-r--r-- 1 root root       280 Jul 11 10:37 special_tokens_map.json\n",
      "-rw-r--r-- 1 root root       351 Jul 11 10:37 tokenizer_config.json\n",
      "-rw-r--r-- 1 root root   2108619 Jul 11 10:37 tokenizer.json\n",
      "-rw-r--r-- 1 root root    798293 Jul 11 10:37 vocab.json\n",
      "-rw-r--r-- 1 root root    407065 Jul 11 10:38 vocab.txt\n"
     ]
    }
   ],
   "source": [
    "!ls -l {EXPORT_PATH}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 848\n",
      "-rw-r--r-- 1 root root 456318 Jul 11 10:38 merges.txt\n",
      "-rw-r--r-- 1 root root 407065 Jul 11 10:38 vocab.txt\n"
     ]
    }
   ],
   "source": [
    "!ls -l {EXPORT_PATH}/assets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import and Save RoBERTa 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": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Installing PySpark 3.2.3 and Spark NLP 5.0.0\n",
      "setup Colab for PySpark 3.2.3 and Spark NLP 5.0.0\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m281.5/281.5 MB\u001b[0m \u001b[31m2.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m498.7/498.7 kB\u001b[0m \u001b[31m30.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m199.7/199.7 kB\u001b[0m \u001b[31m17.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Building wheel for pyspark (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
     ]
    }
   ],
   "source": [
    "! wget -q http://setup.johnsnowlabs.com/colab.sh -O - | bash"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start Spark with Spark NLP included via our simple `start()` function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sparknlp\n",
    "# let's start Spark with Spark NLP\n",
    "spark = sparknlp.start()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Let's use `loadSavedModel` functon in `BertEmbeddings` 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 `BertEmbeddings` 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",
    "- `setStorageRef` is very important. When you are training a task like NER or any Text Classification, we use this reference to bound the trained model to this specific embeddings so you won't load a different embeddings by mistake and see terrible results 😊\n",
    "- It's up to you what you put in `setStorageRef` but it cannot be changed later on. We usually use the name of the model to be clear, but you can get creative if you want!\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.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sparknlp.annotator import *\n",
    "\n",
    "# All these params should be identical to the original ONNX model\n",
    "roberta = RoBertaEmbeddings.loadSavedModel(f\"{EXPORT_PATH}\", spark)\\\n",
    "    .setInputCols([\"document\",'token'])\\\n",
    "    .setOutputCol(\"roberta\")\\\n",
    "    .setCaseSensitive(True)\\\n",
    "    .setStorageRef('roberta-base')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "roberta.write().overwrite().save(f\"{MODEL_NAME}_spark_nlp\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's clean up stuff we don't need anymore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf {EXPORT_PATH}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Awesome  😎 !\n",
    "\n",
    "This is your ONNX RoBERTa model from HuggingFace 🤗  loaded and saved by Spark NLP 🚀"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 484920\n",
      "drwxr-xr-x 4 root root      4096 Jul 11 10:43 fields\n",
      "drwxr-xr-x 2 root root      4096 Jul 11 10:43 metadata\n",
      "-rw-r--r-- 1 root root 496542911 Jul 11 10:43 roberta_onnx\n"
     ]
    }
   ],
   "source": [
    "! ls -l {MODEL_NAME}_spark_nlp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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 RoBERTa model 😊"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sparknlp\n",
    "\n",
    "from sparknlp.base import *\n",
    "from sparknlp.annotator import *\n",
    "\n",
    "document_assembler = DocumentAssembler()\\\n",
    "    .setInputCol(\"text\")\\\n",
    "    .setOutputCol(\"document\")\n",
    "\n",
    "tokenizer = Tokenizer()\\\n",
    "    .setInputCols([\"document\"])\\\n",
    "    .setOutputCol(\"token\")\n",
    "\n",
    "roberta_loaded = RoBertaEmbeddings.load(f\"{MODEL_NAME}_spark_nlp\")\\\n",
    "    .setInputCols([\"document\",'token'])\\\n",
    "    .setOutputCol(\"roberta\")\\\n",
    "\n",
    "pipeline = Pipeline(\n",
    "    stages = [\n",
    "        document_assembler,\n",
    "        tokenizer,\n",
    "        roberta_loaded\n",
    "  ])\n",
    "\n",
    "data = spark.createDataFrame([['William Henry Gates III (born October 28, 1955) is an American business magnate, software developer, investor,and philanthropist.']]).toDF(\"text\")\n",
    "model = pipeline.fit(data)\n",
    "result = model.transform(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------+\n",
      "|          embeddings|\n",
      "+--------------------+\n",
      "|[0.38130012, -0.5...|\n",
      "|[0.44632936, -0.2...|\n",
      "|[0.51096636, 0.08...|\n",
      "|[-0.35451564, -0....|\n",
      "|[0.39058837, 0.43...|\n",
      "|[-0.49727106, 0.9...|\n",
      "|[0.16452815, -0.4...|\n",
      "|[0.17555067, -0.2...|\n",
      "|[-0.024700366, 0....|\n",
      "|[-0.3068737, 0.37...|\n",
      "|[0.58636177, 0.59...|\n",
      "|[-0.6334504, 0.09...|\n",
      "|[-0.6704677, -0.3...|\n",
      "|[-0.046930477, -0...|\n",
      "|[0.15983628, -0.4...|\n",
      "|[0.4974299, -0.17...|\n",
      "|[1.1059247, 0.549...|\n",
      "|[-0.19082792, -0....|\n",
      "|[-0.17161375, -0....|\n",
      "|[0.71347564, 0.23...|\n",
      "+--------------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "result.selectExpr(\"explode(roberta.embeddings) as embeddings\").show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it! You can now go wild and use hundreds of RoBERTa models from HuggingFace 🤗 in Spark NLP 🚀\n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}