torch-neuronx/inference/hf_pretrained_perceiver_language_inference.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7e9a2b5b",
   "metadata": {},
   "source": [
    "# HF Pretrained Perceiver Language Inference on Trn1 / Inf2"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2f4b19d5",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "This notebook demonstrates how to compile and run a Perceiver Language model for accelerated inference on Neuron. This notebook will use the [`PerceiverForMaksedLM`](https://huggingface.co/deepmind/language-perceiver) model.\n",
    "\n",
    "This Jupyter notebook should be run on a Trn1 or Inf2 instance (`trn1.2xlarge` or `inf2.xlarge` or larger)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "23575b62",
   "metadata": {},
   "source": [
    "## Install Dependencies\n",
    "This tutorial requires the following pip packages:\n",
    "\n",
    "- `torch-neuronx`\n",
    "- `neuronx-cc`\n",
    "- `transformers`\n",
    "\n",
    "Most of these packages will be installed when configuring your environment using the Trn1 setup guide."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64134be4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%env TOKENIZERS_PARALLELISM=True #Supresses tokenizer warnings making errors easier to detect\n",
    "!pip install transformers==4.32.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dc8c2b1",
   "metadata": {},
   "source": [
    "## Compile the model into an AWS Neuron optimized TorchScript"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ca0f1b2",
   "metadata": {},
   "source": [
    "In the following section, we load the model, get s sample input, run inference on CPU, compile the model for Neuron using `torch_neuronx.trace()`, and save the optimized model as `TorchScript`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ce06089",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, PerceiverForMaskedLM\n",
    "import torch\n",
    "import torch_neuronx\n",
    "from torch import nn\n",
    "\n",
    "# We will use this simple wrapper to convert the keyword arguments for the model into positional\n",
    "# arguments, which is required for `torch.jit.trace()`, which is used by `torch_neuronx.trace()`,\n",
    "class MaskedLMWrapper(nn.Module):\n",
    "    def __init__(self, perc):\n",
    "        super().__init__()\n",
    "        self.perc = perc\n",
    "\n",
    "    def forward(self, attention_mask, input_ids):\n",
    "        return self.perc(\n",
    "            attention_mask=attention_mask,\n",
    "            input_ids=input_ids\n",
    "        )\n",
    "    \n",
    "\n",
    "# Create the model and image processor\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"deepmind/language-perceiver\")\n",
    "model = PerceiverForMaskedLM.from_pretrained(\"deepmind/language-perceiver\")\n",
    "model.eval()\n",
    "\n",
    "# Wrap the model\n",
    "model = MaskedLMWrapper(model)\n",
    "\n",
    "# Create the masked input\n",
    "text = \"This is an incomplete sentence where some words are missing.\"\n",
    "encoding = tokenizer(text, padding=\"max_length\", return_tensors=\"pt\")\n",
    "# mask bytes corresponding to \" missing.\". Note that the model performs much better if the masked span starts with a space.\n",
    "encoding[\"input_ids\"][0, 52:61] = tokenizer.mask_token_id\n",
    "\n",
    "inputs = (\n",
    "    encoding[\"attention_mask\"],\n",
    "    encoding[\"input_ids\"],\n",
    ")\n",
    "\n",
    "# Run inference on CPU\n",
    "output_cpu = model(*inputs)['logits']\n",
    "\n",
    "# Compile the model\n",
    "model_neuron = torch_neuronx.trace(model, inputs)\n",
    "\n",
    "# Save the TorchScript for inference deployment\n",
    "filename = 'model.pt'\n",
    "torch.jit.save(model_neuron, filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "509f20ef",
   "metadata": {},
   "source": [
    "## Run inference and compare results\n",
    "\n",
    "In this section we load the compiled model, run inference on Neuron, and compare the CPU and Neuron outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1832297",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the TorchScript compiled model\n",
    "model_neuron = torch.jit.load(filename)\n",
    "\n",
    "# Run inference using the Neuron model\n",
    "output_neuron = model_neuron(*inputs)['logits']\n",
    "\n",
    "# Compare the results\n",
    "print(f\"CPU tensor:    {output_cpu[0][0][0:10]}\")\n",
    "print(f\"Neuron tensor: {output_neuron[0][0][0:10]}\")\n",
    "\n",
    "# Compare the predicted output\n",
    "cpu_masked_tokens_predictions = output_cpu[0, 52:61].argmax(dim=-1).tolist()\n",
    "cpu_prediction = tokenizer.decode(cpu_masked_tokens_predictions)\n",
    "print(f\"CPU prediction:      '{cpu_prediction}'\")\n",
    "\n",
    "neuron_masked_tokens_predictions = output_neuron[0, 52:61].argmax(dim=-1).tolist()\n",
    "neuron_prediction = tokenizer.decode(neuron_masked_tokens_predictions)\n",
    "print(f\"Neuron prediction:   '{neuron_prediction}'\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "aws_neuron_venv_pytorch",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}