examples/python/transformers/HuggingFace in Spark NLP - XlmRoBertaForTokenClassification.ipynb

{
 "cells": [
  {
   "attachments": {},
   "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/HuggingFace%20in%20Spark%20NLP%20-%20XlmRoBertaForTokenClassification.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import XlmRoBertaForTokenClassification models from HuggingFace 🤗  into Spark NLP 🚀 \n",
    "\n",
    "Let's keep in mind a few things before we start 😊 \n",
    "\n",
    "- This feature is only in `Spark NLP 3.3.x` and after. So please make sure you have upgraded to the latest Spark NLP release\n",
    "- You can import XLM-RoBERTa models trained/fine-tuned for token classification via `XLMRobertaForTokenClassification` or `TFXLMRobertaForTokenClassification`. These models are usually under `Token Classification` category and have `xlm-roberta` in their labels\n",
    "- Reference: [TFXLMRobertaForTokenClassification](https://huggingface.co/transformers/model_doc/xlmroberta.html#tfxlmrobertafortokenclassification)\n",
    "- Some [example models](https://huggingface.co/models?filter=xlm-roberta&pipeline_tag=token-classification)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export and Save HuggingFace model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Let's install `HuggingFace` and `TensorFlow`. You don't need `TensorFlow` to be installed for Spark NLP, however, we need it to load and save models from HuggingFace.\n",
    "- We lock TensorFlow on `2.11.0` version and Transformers on `4.39.3`. This doesn't mean it won't work with the future releases, but we wanted you to know which versions have been tested successfully.\n",
    "- XLMRobertaTokenizer requires the `SentencePiece` library, so we install that as well"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q transformers==4.39.3 tensorflow==2.11.0 sentencepiece"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- HuggingFace comes with a native `saved_model` feature inside `save_pretrained` function for TensorFlow based models. We will use that to save it as TF `SavedModel`.\n",
    "- We'll use [xlm-roberta-large-finetuned-conll03-english](https://huggingface.co/xlm-roberta-large-finetuned-conll03-english) model from HuggingFace as an example\n",
    "- In addition to `TFXLMRobertaForTokenClassification` we also need to save the `XLMRobertaTokenizer`. This is the same for every model, these are assets needed for tokenization inside Spark NLP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: ./xlm-roberta-large-finetuned-conll03-english/saved_model/1/assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: ./xlm-roberta-large-finetuned-conll03-english/saved_model/1/assets\n"
     ]
    }
   ],
   "source": [
    "from transformers import TFXLMRobertaForTokenClassification, XLMRobertaTokenizer \n",
    "import tensorflow as tf\n",
    "\n",
    "MODEL_NAME = 'xlm-roberta-large-finetuned-conll03-english'\n",
    "\n",
    "tokenizer = XLMRobertaTokenizer.from_pretrained(MODEL_NAME)\n",
    "tokenizer.save_pretrained('./{}_tokenizer/'.format(MODEL_NAME))\n",
    "\n",
    "# just in case if there is no TF/Keras file provided in the model\n",
    "# we can just use `from_pt` and convert PyTorch to TensorFlow\n",
    "try:\n",
    "  print('try downloading TF weights')\n",
    "  model = TFXLMRobertaForTokenClassification.from_pretrained(MODEL_NAME)\n",
    "except:\n",
    "  print('try downloading PyTorch weights')\n",
    "  model = TFXLMRobertaForTokenClassification.from_pretrained(MODEL_NAME, from_pt=True)\n",
    "\n",
    "# Define TF Signature\n",
    "@tf.function(\n",
    "  input_signature=[\n",
    "      {\n",
    "          \"input_ids\": tf.TensorSpec((None, None), tf.int32, name=\"input_ids\"),\n",
    "          \"attention_mask\": tf.TensorSpec((None, None), tf.int32, name=\"attention_mask\")\n",
    "      }\n",
    "  ]\n",
    ")\n",
    "def serving_fn(input):\n",
    "    return model(input)\n",
    "\n",
    "model.save_pretrained(\"./{}\".format(MODEL_NAME), saved_model=True, signatures={\"serving_default\": serving_fn})"
   ]
  },
  {
   "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 4395720\n",
      "-rw-r--r--  1 maziyar  staff        1046 Dec 15 18:44 config.json\n",
      "drwxr-xr-x  3 maziyar  staff          96 Dec 15 18:44 \u001b[34msaved_model\u001b[m\u001b[m\n",
      "-rw-r--r--  1 maziyar  staff  2235947880 Dec 15 18:44 tf_model.h5\n"
     ]
    }
   ],
   "source": [
    "!ls -l {MODEL_NAME}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 37200\n",
      "drwxr-xr-x  2 maziyar  staff        64 Dec 15 18:44 \u001b[34massets\u001b[m\u001b[m\n",
      "-rw-r--r--  1 maziyar  staff        54 Dec 15 18:44 fingerprint.pb\n",
      "-rw-r--r--  1 maziyar  staff    321614 Dec 15 18:44 keras_metadata.pb\n",
      "-rw-r--r--  1 maziyar  staff  18717362 Dec 15 18:44 saved_model.pb\n",
      "drwxr-xr-x  4 maziyar  staff       128 Dec 15 18:44 \u001b[34mvariables\u001b[m\u001b[m\n"
     ]
    }
   ],
   "source": [
    "!ls -l {MODEL_NAME}/saved_model/1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 9920\n",
      "-rw-r--r--  1 maziyar  staff  5069051 Dec 15 18:42 sentencepiece.bpe.model\n",
      "-rw-r--r--  1 maziyar  staff      279 Dec 15 18:42 special_tokens_map.json\n",
      "-rw-r--r--  1 maziyar  staff      503 Dec 15 18:42 tokenizer_config.json\n"
     ]
    }
   ],
   "source": [
    "!ls -l {MODEL_NAME}_tokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- as you can see, we need the SavedModel from `saved_model/1/` path\n",
    "- we also be needing `sentencepiece.bpe.model` file from the tokenizer\n",
    "- all we need is to copy `sentencepiece.bpe.model` file into `saved_model/1/assets` which Spark NLP will look for\n",
    "- in addition to vocabs, we also need `labels` and their `ids` which is saved inside the model's config. We will save this inside `labels.txt`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "asset_path = '{}/saved_model/1/assets'.format(MODEL_NAME)\n",
    "\n",
    "# let's copy sentencepiece.bpe.model file to saved_model/1/assets\n",
    "!cp {MODEL_NAME}_tokenizer/sentencepiece.bpe.model {asset_path}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get label2id dictionary \n",
    "labels = model.config.label2id\n",
    "# sort the dictionary based on the id\n",
    "labels = sorted(labels, key=labels.get)\n",
    "\n",
    "with open(asset_path+'/labels.txt', 'w') as f:\n",
    "    f.write('\\n'.join(labels))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Voila! We have our `vocab.txt` and `labels.txt` inside assets directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 9912\n",
      "-rw-r--r--  1 maziyar  staff       45 Dec 15 18:44 labels.txt\n",
      "-rw-r--r--  1 maziyar  staff  5069051 Dec 15 18:44 sentencepiece.bpe.model\n"
     ]
    }
   ],
   "source": [
    "! ls -l {asset_path}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import and Save XlmRoBertaForTokenClassification in Spark NLP\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 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.1 and Spark NLP 4.2.5\n",
      "setup Colab for PySpark 3.2.1 and Spark NLP 4.2.5\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 `XlmRoBertaForTokenClassification` which allows us to load TensorFlow model in SavedModel format\n",
    "- Most params can be set later when you are loading this model in `XlmRoBertaForTokenClassification` in runtime like `setMaxSentenceLength`, so don't worry what you are setting them now\n",
    "- `loadSavedModel` accepts two params, first is the path to the TF SavedModel. The second is the SparkSession that is `spark` variable we previously started via `sparknlp.start()`\n",
    "- NOTE: `loadSavedModel` accepts local paths in addition to distributed file systems such as `HDFS`, `S3`, `DBFS`, etc. This feature was introduced in Spark NLP 4.2.2 release. Keep in mind the best and recommended way to move/share/reuse Spark NLP models is to use `write.save` so you can use `.load()` from any file systems natively.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sparknlp.annotator import *\n",
    "from sparknlp.base import *\n",
    "\n",
    "tokenClassifier = XlmRoBertaForTokenClassification\\\n",
    "  .loadSavedModel('{}/saved_model/1'.format(MODEL_NAME), spark)\\\n",
    "  .setInputCols([\"document\",'token'])\\\n",
    "  .setOutputCol(\"ner\")\\\n",
    "  .setCaseSensitive(True)\\\n",
    "  .setMaxSentenceLength(128)"
   ]
  },
  {
   "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": [
    "tokenClassifier.write().overwrite().save(\"./{}_spark_nlp\".format(MODEL_NAME))"
   ]
  },
  {
   "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 {MODEL_NAME}_tokenizer {MODEL_NAME}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Awesome 😎  !\n",
    "\n",
    "This is your XlmRoBertaForTokenClassification model from HuggingFace 🤗  loaded and saved by Spark NLP 🚀 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 4432048\n",
      "drwxr-xr-x  4 maziyar  staff         128 Dec 15 18:45 \u001b[34mfields\u001b[m\u001b[m\n",
      "drwxr-xr-x  6 maziyar  staff         192 Dec 15 18:45 \u001b[34mmetadata\u001b[m\u001b[m\n",
      "-rw-r--r--  1 maziyar  staff  2254172695 Dec 15 18:45 xlm_roberta_classification_tensorflow\n",
      "-rw-r--r--  1 maziyar  staff     5069051 Dec 15 18:45 xlmroberta_spp\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 XlmRoBertaForTokenClassification model 😊 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenClassifier_loaded = XlmRoBertaForTokenClassification.load(\"./{}_spark_nlp\".format(MODEL_NAME))\\\n",
    "  .setInputCols([\"document\",'token'])\\\n",
    "  .setOutputCol(\"ner\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see what labels were used to train this model via `getClasses` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['B-LOC', 'I-ORG', 'I-MISC', 'I-LOC', 'I-PER', 'B-MISC', 'B-ORG', 'O']"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# .getClasses was introduced in spark-nlp==3.4.0\n",
    "tokenClassifier_loaded.getClasses()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is how you can use your loaded classifier model in Spark NLP 🚀 pipeline:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------+--------------------+\n",
      "|                text|              result|\n",
      "+--------------------+--------------------+\n",
      "|My name is Sarah ...|[O, O, O, I-PER, ...|\n",
      "|My name is Clara ...|[O, O, O, I-PER, ...|\n",
      "+--------------------+--------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "document_assembler = DocumentAssembler() \\\n",
    "    .setInputCol('text') \\\n",
    "    .setOutputCol('document')\n",
    "\n",
    "tokenizer = Tokenizer() \\\n",
    "    .setInputCols(['document']) \\\n",
    "    .setOutputCol('token')\n",
    "\n",
    "pipeline = Pipeline(stages=[\n",
    "    document_assembler, \n",
    "    tokenizer,\n",
    "    tokenClassifier_loaded    \n",
    "])\n",
    "\n",
    "# couple of simple examples\n",
    "example = spark.createDataFrame([[\"My name is Sarah and I live in London\"], ['My name is Clara and I live in Berkeley, California.']]).toDF(\"text\")\n",
    "\n",
    "result = pipeline.fit(example).transform(example)\n",
    "\n",
    "# result is a DataFrame\n",
    "result.select(\"text\", \"ner.result\").show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it! You can now go wild and use hundreds of `XlmRoBertaForTokenClassification` models from HuggingFace 🤗 in Spark NLP 🚀 \n"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "HuggingFace in Spark NLP - XlmRoBertaForTokenClassification.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "transformers",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}