integrations/huggingface-transformers/tutorials/text-classification/docs-text-classification-python-custom-teacher-tweeteval.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kSNEB-3orJ9C"
   },
   "source": [
    "# **Single Sequence Text Classification: Sparse Transfer Learning with the Python API**\n",
    "\n",
    "In this example, you will fine-tune a 90% pruned BERT model onto the TweetEval Emotions dataset (a single-sequence multi-class classification problem) using SparseML's Hugging Face Integration.\n",
    "\n",
    "### **Sparse Transfer Learning Overview**\n",
    "\n",
    "Sparse Transfer Learning is very similiar to typical fine-tuning you are used to when training models. However, with Sparse Transfer Learning, we start the training process from a pre-sparsified checkpoint and maintain the sparsity structure while the fine tuning occurs.\n",
    "\n",
    "At the end, you will have a sparse model trained on your dataset, ready to be deployed with DeepSparse for GPU-class performance on CPUs!\n",
    "\n",
    "### **Pre-Sparsified BERT**\n",
    "SparseZoo, Neural Magic's open source repository of pre-sparsified models, contains a 90% pruned version of BERT, which has been sparsified on the upstream Wikipedia and BookCorpus datasets with the\n",
    "masked language modeling objective. [Check out the model card](https://sparsezoo.neuralmagic.com/models/nlp%2Fmasked_language_modeling%2Fobert-base%2Fpytorch%2Fhuggingface%2Fwikipedia_bookcorpus%2Fpruned90-none). We will use this model as the starting point for the transfer learning process.\n",
    "\n",
    "\n",
    "***Let's dive in!***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Y0WybTbssU0g"
   },
   "source": [
    "## **Installation**\n",
    "\n",
    "Install SparseML via `pip`.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "id": "NjynYRLMIONF"
   },
   "outputs": [],
   "source": [
    "!pip install sparseml[transformers]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_jY0SKdXFGO3"
   },
   "source": [
    "If you are running on Google Colab, restart the runtime after this step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "id": "XXj0S5Jdq2M-"
   },
   "outputs": [],
   "source": [
    "import sparseml\n",
    "from sparsezoo import Model\n",
    "from sparseml.transformers.utils import SparseAutoModel\n",
    "from sparseml.transformers.sparsification import Trainer, TrainingArguments\n",
    "import numpy as np\n",
    "from transformers import (\n",
    "    AutoModelForSequenceClassification,\n",
    "    AutoConfig, \n",
    "    AutoTokenizer, \n",
    "    EvalPrediction, \n",
    "    default_data_collator\n",
    ")\n",
    "from datasets import load_dataset, load_metric"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2rS2Q5kxFcW3"
   },
   "source": [
    "## **Step 1: Load a Dataset**\n",
    "\n",
    "SparseML is integrated with Hugging Face, so we can use the `datasets` class to load datasets from the Hugging Face hub or from local files. \n",
    "\n",
    "[TweetEval Emotions](https://huggingface.co/datasets/tweet_eval/viewer/emotion/train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CkvbT1i9p87z"
   },
   "outputs": [],
   "source": [
    "# load_dataset from HF hub\n",
    "dataset = load_dataset(\"tweet_eval\", \"emotion\")\n",
    "\n",
    "# alternatively, load from a csv\n",
    "dataset[\"train\"].to_csv(\"tweet_eval_emotion-train.csv\")\n",
    "dataset[\"validation\"].to_csv(\"tweet_eval_emotion-validation.csv\")\n",
    "\n",
    "data_files = {\n",
    "  \"train\": \"tweet_eval_emotion-train.csv\",\n",
    "  \"validation\": \"tweet_eval_emotion-validation.csv\"\n",
    "}\n",
    "dataset_from_json = load_dataset(\"csv\", data_files=data_files)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nGfpSK3CI15n"
   },
   "outputs": [],
   "source": [
    "!head tweet_eval_emotion-train.csv --lines=5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fnKtfIU6IzOA"
   },
   "outputs": [],
   "source": [
    "print(dataset_from_json)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7ti52fgQqdSU"
   },
   "outputs": [],
   "source": [
    "# configs for below\n",
    "INPUT_COL_1 = \"text\"\n",
    "INPUT_COL_2 = None\n",
    "LABEL_COL = \"label\"\n",
    "NUM_LABELS = len(dataset_from_json[\"train\"].unique(LABEL_COL))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3BfXUE9HHFoq"
   },
   "source": [
    "## **Step 2: Setup Evaluation Metric**\n",
    "\n",
    "Tweet Eval emotion is a simple multi-class classification problem (we are predicting one of 4 emotions). We will use the `accuracy` function as the evaluation metric.\n",
    "\n",
    "Since SparseML is integrated with Hugging Face, we can pass `compute_metrics` function for evaluation (which will be passed to the `Trainer` class below)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tPIVsdAzIgf8"
   },
   "outputs": [],
   "source": [
    "metric = load_metric(\"accuracy\")\n",
    "\n",
    "# setup metrics function\n",
    "def compute_metrics(p: EvalPrediction):\n",
    "  preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions\n",
    "  preds = np.argmax(preds, axis=1)\n",
    "  result = metric.compute(predictions=preds, references=p.label_ids)\n",
    "  if len(result) > 1:\n",
    "      result[\"combined_score\"] = np.mean(list(result.values())).item()\n",
    "  return result\n",
    "\n",
    "print(dataset[\"train\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zERYC_SL4J2A"
   },
   "source": [
    "## **Step 3: Train the Teacher**\n",
    "\n",
    "To support the sparse transfer learning process, we can optionally provide a teacher model to apply model distillation. Since SparseML is integrated with Hugging Face, we can use Hugging Face's native flows to train the teacher.\n",
    "\n",
    "In this example, we will first train a dense teacher model from scratch - which we can then distill during the sparse transfer learning process.\n",
    "\n",
    "Although we will be fine-tuning BERT, we can use a different model as the teacher. In this case, we will use `roberta-base` as the teacher, training with native Hugging Face objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XqaExx7Hp5Zl"
   },
   "outputs": [],
   "source": [
    "TEACHER = \"roberta-base\"\n",
    "teacher_config = AutoConfig.from_pretrained(TEACHER, num_labels=NUM_LABELS)\n",
    "teacher_tokenizer = AutoTokenizer.from_pretrained(TEACHER)\n",
    "teacher = AutoModelForSequenceClassification.from_pretrained(TEACHER, config=teacher_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Qcb_MA2-JuUZ"
   },
   "source": [
    "### **Teacher Training: Tokenize the Dataset**\n",
    "\n",
    "Run the RoBERTa tokenizer on the dataset. This is native Hugging Face functionality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OT0ogPOCrPaT"
   },
   "outputs": [],
   "source": [
    "MAX_LEN = 128\n",
    "def teacher_preprocess_fn(examples):\n",
    "  args = None\n",
    "  if INPUT_COL_2 is None:\n",
    "    args = (examples[INPUT_COL_1], )\n",
    "  else:\n",
    "    args = (examples[INPUT_COL_1], examples[INPUT_COL_2])\n",
    "  result = teacher_tokenizer(*args, \n",
    "                   padding=\"max_length\", \n",
    "                   max_length=min(teacher_tokenizer.model_max_length, MAX_LEN), \n",
    "                   truncation=True)\n",
    "  return result\n",
    "\n",
    "# tokenize the dataset\n",
    "teacher_tokenized_dataset = dataset_from_json.map(\n",
    "    teacher_preprocess_fn,\n",
    "    batched=True,\n",
    "    desc=\"Running teacher tokenizer on dataset\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pTb5lf-kJ_st"
   },
   "source": [
    "### **Teacher Training: Fine-Tune the Teacher**\n",
    "\n",
    "We use the native Hugging Face `Trainer` (which we import as `HFTrainer`) to train the model. Check out the [Hugging Face documentation](https://huggingface.co/docs/transformers/main_classes/trainer) for more details on the `Trainer` as needed.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "id": "oVrGtQwbrTlO"
   },
   "outputs": [],
   "source": [
    "from transformers import Trainer as HFTrainer\n",
    "from transformers import TrainingArguments as HFTrainingArguments\n",
    "\n",
    "# setup trainer arguments\n",
    "teacher_training_args = HFTrainingArguments(\n",
    "    output_dir=\"./teacher_training\",\n",
    "    do_train=True,\n",
    "    do_eval=True,\n",
    "    num_train_epochs=6.0,\n",
    "    learning_rate=5e-6,\n",
    "    lr_scheduler_type=\"linear\",\n",
    "    logging_strategy=\"epoch\",\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    save_total_limit=1,\n",
    "    per_device_train_batch_size=32,\n",
    "    per_device_eval_batch_size=32)\n",
    "\n",
    "# initialize trainer\n",
    "teacher_trainer = HFTrainer(\n",
    "    model=teacher,\n",
    "    args=teacher_training_args,\n",
    "    train_dataset=teacher_tokenized_dataset[\"train\"],\n",
    "    eval_dataset=teacher_tokenized_dataset[\"validation\"],\n",
    "    tokenizer=teacher_tokenizer,\n",
    "    data_collator=default_data_collator,\n",
    "    compute_metrics=compute_metrics)\n",
    "\n",
    "# run training\n",
    "%rm -rf ./teacher_training\n",
    "teacher_trainer.train(resume_from_checkpoint=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8bcO7Kclv5uX"
   },
   "source": [
    "## **Step 4: Sparse Transfer Learning**\n",
    "\n",
    "Now that we have the teacher trained, we can sparse transfer learn from the pre-sparsified version of BERT with distillation support. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1GEhYi53HoAH"
   },
   "source": [
    "First, we need to select a sparse checkpoint to begin the training process. In this case, we are fine-tuning a 90% pruned version of BERT onto the TweetEval Emotion dataset. This model is available in SparseZoo, identified by the following stub:\n",
    "```\n",
    "zoo:nlp/masked_language_modeling/obert-base/pytorch/huggingface/wikipedia_bookcorpus/pruned90-none\n",
    "```\n",
    "\n",
    "Next, we need to create a sparsification recipe for usage in the training process. Recipes are YAML files that encode the sparsity related algorithms and parameters to be applied by SparseML. For Sparse Transfer Learning, we need to use a recipe that instructs SparseML to maintain sparsity during the training process and to apply quantization over the final few epochs. \n",
    "\n",
    "In SparseZoo, there is a transfer recipe which was used to fine-tune BERT onto the SST2 task (which is also a single sequence classification problem). Since the TweetEval Emotion task is a similiar problem to SST2, we will use the SST2 recipe, which is identified by the following stub:\n",
    "\n",
    "```\n",
    "zoo:nlp/sentiment_analysis/obert-base/pytorch/huggingface/sst2/pruned90_quant-none\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "U_iyuuB4Wq7N"
   },
   "source": [
    "Use the `sparsezoo` python client to download the models and recipe using their SparseZoo stubs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ykg8fEN2Q5o_"
   },
   "outputs": [],
   "source": [
    "# downloads 90% pruned upstream BERT trained on MLM objective\n",
    "model_stub = \"zoo:nlp/masked_language_modeling/obert-base/pytorch/huggingface/wikipedia_bookcorpus/pruned90-none\" \n",
    "model_path = Model(model_stub, download_path=\"./model\").training.path\n",
    "\n",
    "# downloads transfer recipe for SST2 (pruned90_quant)\n",
    "transfer_stub = \"zoo:nlp/sentiment_analysis/obert-base/pytorch/huggingface/sst2/pruned90_quant-none\"\n",
    "recipe_path = Model(transfer_stub, download_path=\"./transfer_recipe\").recipes.default.path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RLe8iEWxV_zz"
   },
   "source": [
    "We can see that the upstream model (trained on Wikipedia BookCorpus) and  configuration files have been downloaded to the local directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0NTVj1kPRSCW"
   },
   "outputs": [],
   "source": [
    "%ls ./model/training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "orjvrvdCWEUi"
   },
   "source": [
    "We can see that a transfer learning recipe has been downloaded. The `ConstantPruningModifier` instructs SparseML to maintain the sparsity structure of the network as the model trains and the `QuantizationModifier` instructs SparseML to run Quantization Aware Training at the end of training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eUYg-7eBRT5f"
   },
   "outputs": [],
   "source": [
    "%cat ./transfer_recipe/recipe/recipe_original.md"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### **Inspecting the Recipe**\n",
    "\n",
    "Here is the transfer learning recipe:\n",
    "\n",
    "```yaml\n",
    "version: 1.1.0\n",
    "\n",
    "# General Variables\n",
    "num_epochs: &num_epochs 13\n",
    "init_lr: 1.5e-4\n",
    "final_lr: 0\n",
    "\n",
    "qat_start_epoch: &qat_start_epoch 8.0\n",
    "observer_epoch: &observer_epoch 12.0\n",
    "quantize_embeddings: &quantize_embeddings 1\n",
    "\n",
    "distill_hardness: &distill_hardness 1.0\n",
    "distill_temperature: &distill_temperature 2.0\n",
    "\n",
    "weight_decay: 0.01\n",
    "\n",
    "# Modifiers:\n",
    "\n",
    "training_modifiers:\n",
    "  - !EpochRangeModifier\n",
    "      end_epoch: eval(num_epochs)\n",
    "      start_epoch: 0.0\n",
    "      \n",
    "  - !LearningRateFunctionModifier\n",
    "      start_epoch: 0\n",
    "      end_epoch: eval(num_epochs)\n",
    "      lr_func: linear\n",
    "      init_lr: eval(init_lr)\n",
    "      final_lr: eval(final_lr)\n",
    "\n",
    "quantization_modifiers:\n",
    "  - !QuantizationModifier\n",
    "      start_epoch: eval(qat_start_epoch)\n",
    "      disable_quantization_observer_epoch: eval(observer_epoch)\n",
    "      freeze_bn_stats_epoch: eval(observer_epoch)\n",
    "      quantize_embeddings: eval(quantize_embeddings)\n",
    "      quantize_linear_activations: 0\n",
    "      exclude_module_types: ['LayerNorm', 'Tanh']\n",
    "      submodules:\n",
    "        - bert.embeddings\n",
    "        - bert.encoder\n",
    "        - bert.pooler\n",
    "        - classifier\n",
    "\n",
    "distillation_modifiers:\n",
    "  - !DistillationModifier\n",
    "     hardness: eval(distill_hardness)\n",
    "     temperature: eval(distill_temperature)\n",
    "     distill_output_keys: [logits]\n",
    "\n",
    "constant_modifiers:\n",
    "  - !ConstantPruningModifier\n",
    "      start_epoch: 0.0\n",
    "      params: __ALL_PRUNABLE__\n",
    "\n",
    "regularization_modifiers:\n",
    "  - !SetWeightDecayModifier\n",
    "      start_epoch: 0.0\n",
    "      weight_decay: eval(weight_decay)\n",
    "```\n",
    "\n",
    "\n",
    "The `Modifiers` in the transfer learning recipe are the important items that encode how SparseML should modify the training process for Sparse Transfer Learning:\n",
    "- `ConstantPruningModifier` tells SparseML to pin weights at 0 over all epochs, maintaining the sparsity structure of the network\n",
    "- `QuantizationModifier` tells SparseML to quanitze the weights with quantization aware training over the last 5 epochs\n",
    "- `DistillationModifier` tells SparseML how to apply distillation during the trainign process, targeting the logits\n",
    "\n",
    "Below, SparseML's `Trainer` will parses the modifiers and updates the training process to implement the algorithms specified here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CCTxNPD4XxHs"
   },
   "source": [
    "Next, we will set up the Hugging Face `tokenizer, config, and model`. These are all native Hugging Face objects, so check out the Hugging Face docs for more details on `AutoModel`, `AutoConfig`, and `AutoTokenizer` as needed. \n",
    "\n",
    "We instantiate these classes by passing the local path to the directory containing the `pytorch_model.bin`, `tokenizer.json`, and `config.json` files from the SparseZoo download."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dhN1oGcTQ9RE"
   },
   "outputs": [],
   "source": [
    "# tokenizer\n",
    "config = AutoConfig.from_pretrained(model_path, num_labels=NUM_LABELS)\n",
    "\n",
    "# initialize config\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_path)\n",
    "\n",
    "# initialize model using familiar HF AutoModel\n",
    "model_kwargs = {\"config\": config}\n",
    "model_kwargs[\"state_dict\"], s_delayed = SparseAutoModel._loadable_state_dict(model_path)\n",
    "model = AutoModelForSequenceClassification.from_pretrained(model_path, **model_kwargs,)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ny-5BtSsYRKH"
   },
   "source": [
    "### **Sparse Transfer Learning: Tokenize The Dataset**\n",
    "\n",
    "Next, tokenize the dataset.\n",
    "\n",
    "Since we are using RoBERTa as the teacher model, which has a different tokenizer than BERT, we have to add results from teacher_tokenizer to results with `'distill_teacher': id`. SparseML parses the `result` and will send the tokens to the correct model during training.\n",
    "\n",
    "Note that if the teacher and student share a tokenizer, we can skip adding the teacher tokens and SparseML will pass the single set of tokens to each model during Training.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2EUuFSTzRAvp"
   },
   "outputs": [],
   "source": [
    "MAX_LEN = 128\n",
    "def preprocess_fn(examples):\n",
    "  args = None\n",
    "  if INPUT_COL_2 is None:\n",
    "    args = (examples[INPUT_COL_1], )\n",
    "  else:\n",
    "    args = (examples[INPUT_COL_1], examples[INPUT_COL_2])\n",
    "  result = tokenizer(*args, \n",
    "                   padding=\"max_length\", \n",
    "                   max_length=min(tokenizer.model_max_length, MAX_LEN), \n",
    "                   truncation=True)\n",
    "  \n",
    "  teacher_result = teacher_tokenizer(*args, \n",
    "                   padding=\"max_length\", \n",
    "                   max_length=min(teacher_tokenizer.model_max_length, MAX_LEN), \n",
    "                   truncation=True)\n",
    "  \n",
    "  teacher_result = {\n",
    "      f\"distill_teacher:{tokenizer_key}\": value\n",
    "      for tokenizer_key, value in teacher_result.items()\n",
    "  }\n",
    "  result.update(teacher_result)\n",
    "  return result\n",
    "\n",
    "# tokenize the dataset\n",
    "tokenized_dataset = dataset_from_json.map(\n",
    "    preprocess_fn,\n",
    "    batched=True,\n",
    "    desc=\"Running tokenizer on dataset\"\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "g5jB7s1bM4Qh"
   },
   "source": [
    "### **Sparse Transfer Learning: Fine-Tune the Model**\n",
    "\n",
    "SparseML has a custom `Trainer` class that inherits from the [Hugging Face `Trainer` Class](https://huggingface.co/docs/transformers/main_classes/trainer). As such, the SparseML `Trainer` has all of the existing functionality of the HF trainer. However, in addition, we can supply a `recipe` and (optionally) a `teacher`. \n",
    "\n",
    "\n",
    "As we saw above, the `recipe` encodes the sparsity related algorithms and hyperparameters of the training process in a YAML file. The SparseML `Trainer` parses the `recipe` and adjusts the training workflow to apply the algorithms in the recipe. We use the `recipe_args` function to modify the recipe slightly (training for fewer epochs than used for SST2).\n",
    "\n",
    "The `teacher` is an optional argument that allows SparseML to apply model distillation to support the training process. Here, we pass the RoBERTa model trained above as the teacher."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "34IXj1n6RCgQ"
   },
   "outputs": [],
   "source": [
    "# setup trainer arguments\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./training_output\",\n",
    "    do_train=True,\n",
    "    do_eval=True,\n",
    "    resume_from_checkpoint=False,\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    logging_strategy=\"epoch\",\n",
    "    save_total_limit=1,\n",
    "    per_device_train_batch_size=32,\n",
    "    per_device_eval_batch_size=32,\n",
    "    fp16=False)\n",
    "\n",
    "# initialize trainer\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    model_state_path=model_path,\n",
    "    recipe=recipe_path,\n",
    "    recipe_args='{\"num_epochs\": 10.0, \"qat_start_epoch\": 7.0, \"observer_epoch\": 9.0}',\n",
    "    teacher=teacher,\n",
    "    metadata_args=[\"per_device_train_batch_size\",\"per_device_eval_batch_size\",\"fp16\"],\n",
    "    args=training_args,\n",
    "    train_dataset=tokenized_dataset[\"train\"],\n",
    "    eval_dataset=tokenized_dataset[\"validation\"],\n",
    "    tokenizer=tokenizer,\n",
    "    data_collator=default_data_collator,\n",
    "    compute_metrics=compute_metrics)\n",
    "\n",
    "# step 5: run training\n",
    "%rm -rf training_output\n",
    "train_result = trainer.train()\n",
    "trainer.save_model()\n",
    "trainer.save_state()\n",
    "trainer.save_optimizer_and_scheduler(training_args.output_dir)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **Step 7: Export To ONNX**\n",
    "\n",
    "Run the following to export the model to ONNX. The script creates a `deployment` folder containing ONNX file and the necessary configuration files (e.g. `tokenizer.json`) for deployment with DeepSparse."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-rhWjiHBeR7M"
   },
   "outputs": [],
   "source": [
    "!sparseml.transformers.export_onnx \\\n",
    "  --model_path training_output \\\n",
    "  --task text_classification"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QicVADU5NW94"
   },
   "source": [
    "## **Optional: Deploy with DeepSparse**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-XubpXohO_8A"
   },
   "outputs": [],
   "source": [
    "%pip install deepsparse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "m_USM8mCPETg"
   },
   "outputs": [],
   "source": [
    "from deepsparse import Pipeline\n",
    "\n",
    "pipeline = Pipeline.create(\"text_classification\", model_path=\"./deployment\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YFzCkPotPXx6"
   },
   "outputs": [],
   "source": [
    "prediction = pipeline(\"@user Get Donovan out of your soccer booth. He's awful. He's bitter. He makes me want to mute the tv. #horrid\")\n",
    "print(prediction) # label 0 is anger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "g1oJZlFdPlvj"
   },
   "outputs": [],
   "source": [
    "prediction = pipeline(\"@user Welcome to #MPSVT! We are delighted to have you! #grateful #MPSVT #relationships\")\n",
    "print(prediction) # label 1 is joy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TS-OCQJ_PpoL"
   },
   "outputs": [],
   "source": [
    "prediction = pipeline(\"“The #optimist proclaims that we live in the best of all possible worlds; and the #pessimist fears this is true.” ~ James Branch Cabell\")\n",
    "print(prediction) # label 2 is optimism"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qowh540KPGwh"
   },
   "outputs": [],
   "source": [
    "prediction = pipeline(\"In need of a change! #restless\")\n",
    "print(prediction) # label 3 is sadness"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "authorship_tag": "ABX9TyNUQDzGJY5h/6OPOEWVMurw",
   "provenance": [
    {
     "file_id": "1Zawa0sifXr2wIl9tbF7ySJ7xYY0dtTzI",
     "timestamp": 1677345946788
    }
   ],
   "toc_visible": true
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}