torch-neuron/inference/bertlargeuncased/BertLargeUncased.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5d416726",
   "metadata": {},
   "source": [
    "# BertLargeUncased - Pytorch\n",
    "This notebook shows how to compile a pre-trainded BertLarge/Pytorch to AWS Inferentia (inf1 instances) using NeuronSDK. The original implementation is provided by HuggingFace.\n",
    "\n",
    "**Reference:** https://huggingface.co/bert-large-uncased"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66ea195b",
   "metadata": {},
   "source": [
    "## 1) Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57bf6e99",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set Pip repository  to point to the Neuron repository\n",
    "%pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com\n",
    "# now restart the kernel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc531ad2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Install Neuron PyTorch\n",
    "%pip install -U torch-neuron==1.10.1.2.2.0.0 neuron-cc[tensorflow] \"protobuf<4\" \"transformers==4.6.0\"\n",
    "# use --force-reinstall if you're facing some issues while loading the modules\n",
    "# now restart the kernel again"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98e94754",
   "metadata": {},
   "source": [
    "## 2) Initialize libraries and prepare input samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7d590bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.neuron\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoConfig\n",
    "\n",
    "# Build tokenizer and model\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"bert-large-uncased\")\n",
    "\n",
    "# Setup some example inputs\n",
    "sequence_0 = \"The company HuggingFace is based in New York City\"\n",
    "sequence_1 = \"Apples are especially bad for your health\"\n",
    "sequence_2 = \"HuggingFace's headquarters are situated in Manhattan\"\n",
    "\n",
    "max_length=128\n",
    "paraphrase = tokenizer.encode_plus(sequence_0, sequence_2, max_length=max_length, padding='max_length', truncation=True, return_tensors=\"pt\")\n",
    "not_paraphrase = tokenizer.encode_plus(sequence_0, sequence_1, max_length=max_length, padding='max_length', truncation=True, return_tensors=\"pt\")\n",
    "\n",
    "# Convert example inputs to a format that is compatible with TorchScript tracing\n",
    "example_inputs_paraphrase = paraphrase['input_ids'], paraphrase['attention_mask'], paraphrase['token_type_ids']\n",
    "example_inputs_not_paraphrase = not_paraphrase['input_ids'], not_paraphrase['attention_mask'], not_paraphrase['token_type_ids']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78c268ce",
   "metadata": {},
   "source": [
    "## 3) Load a pre-trained model and check if it is .jit traceable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27e80b5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'Loading a pre-trained model')\n",
    "model = AutoModelForSequenceClassification.from_pretrained(\"bert-large-uncased\", return_dict=False)\n",
    "model.eval()\n",
    "    \n",
    "y = model(**paraphrase) # warmup the model\n",
    "try:\n",
    "    traced_model = torch.jit.trace(model, example_inputs_paraphrase)\n",
    "    print(\"Cool! Model is jit traceable\")\n",
    "except Exception as e:\n",
    "    print(\"Ops. Something went wrong. Model is not traceable\")\n",
    "## ok the model is .jit traceable. now let's compile it with NeuronSDK"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b46d0b80",
   "metadata": {},
   "source": [
    "## 4) Analyze & compile the model for Inferentia with NeuronSDK"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cc67351",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.neuron.analyze_model(model, example_inputs_paraphrase)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2943b83",
   "metadata": {},
   "outputs": [],
   "source": [
    "neuron_model = torch.neuron.trace(model, example_inputs_paraphrase)\n",
    "neuron_model.save(f\"neuron_bert_large_uncased.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1db4378",
   "metadata": {},
   "source": [
    "### 4.1) Verify the optimized model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b42f845a",
   "metadata": {},
   "outputs": [],
   "source": [
    "y = neuron_model(*example_inputs_paraphrase) # warmup\n",
    "%timeit neuron_model(*example_inputs_paraphrase)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a32d3a0",
   "metadata": {},
   "source": [
    "## 5) A simple test to check the predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c499d396",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BERT says that \"The company HuggingFace is based in New York City\" and \"HuggingFace's headquarters are situated in Manhattan\" are paraphrase\n",
      "BERT says that \"The company HuggingFace is based in New York City\" and \"Apples are especially bad for your health\" are not paraphrase\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.neuron\n",
    "# Load TorchScript back\n",
    "model_neuron = torch.jit.load('neuron_bert_large_uncased.pt')\n",
    "\n",
    "# Verify the TorchScript works on both example inputs\n",
    "paraphrase_classification_logits_neuron = model_neuron(*example_inputs_paraphrase)\n",
    "not_paraphrase_classification_logits_neuron = model_neuron(*example_inputs_not_paraphrase)\n",
    "\n",
    "classes = ['not paraphrase', 'paraphrase']\n",
    "paraphrase_prediction = paraphrase_classification_logits_neuron[0][0].argmax().item()\n",
    "not_paraphrase_prediction = not_paraphrase_classification_logits_neuron[0][0].argmax().item()\n",
    "print('BERT says that \"{}\" and \"{}\" are {}'.format(sequence_0, sequence_2, classes[paraphrase_prediction]))\n",
    "print('BERT says that \"{}\" and \"{}\" are {}'.format(sequence_0, sequence_1, classes[not_paraphrase_prediction]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d766fd6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Environment (conda_aws_neuron_pytorch_p36)",
   "language": "python",
   "name": "conda_aws_neuron_pytorch_p36"
  },
  "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.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}