examples/notebooks/roberta-python-pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Copyright 2025 Modular, Inc: Licensed under the Apache License v2.0 with LLVM Exceptions.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MAX Engine and PyTorch model example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The [Python API for MAX Engine](https://docs.modular.com/engine/reference/python/engine) enables you\n",
    "to upgrade your runtime performance for PyTorch, TensorFlow, and ONNX models,\n",
    "on a wide range of hardware, with just three lines of code (not counting the\n",
    "import):\n",
    "\n",
    "```python\n",
    "from max import engine\n",
    "\n",
    "# Load your model:\n",
    "session = engine.InferenceSession()\n",
    "model = session.load(model_path)\n",
    "\n",
    "# Prepare the inputs, then run an inference:\n",
    "outputs = model.execute(**inputs)\n",
    "\n",
    "# Process the output here.\n",
    "```\n",
    "\n",
    "That's all you need! Everything else is the usual code to prepare your\n",
    "inputs and process the outputs.\n",
    "\n",
    "But, it's always nice to see a fully working example. So the\n",
    "rest of this page shows how to run an inference using a version of\n",
    "[RoBERTa from Cardiff\n",
    "NLP](https://huggingface.co/cardiffnlp/twitter-roberta-base-emotion-multilabel-latest),\n",
    "which is a language model trained on tweets to perform sentiment analysis.\n",
    "\n",
    "This example uses is a PyTorch model (which must be converted to TorchScript\n",
    "format), and it's just as easy to load a model from ONNX or TensorFlow (in\n",
    "SavedModel format)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install the MAX Engine Python package"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Naturally, you first need to install the `max` Python package.\n",
    "This package is not hosted in a package repository (PyPI), and can only be\n",
    "installed with the `modular` CLI tool.\n",
    "\n",
    "For instructions, see\n",
    "[Get started with MAX Engine](https://docs.modular.com/max/get-started)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install the MAX Engine Python package\n",
    "!python3 -m pip install -q --find-links \"$(modular config max.path)/wheels\" max-engine\n",
    "# Install other packages\n",
    "!python3 -m pip install -q transformers\n",
    "!python3 -m pip install -q torch --index-url https://download.pytorch.org/whl/cpu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import Python modules"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To start coding, we need some libraries that help us get the model and process\n",
    "the input/output data.\n",
    "\n",
    "NOTE: Make sure you have these packages installed:\n",
    "\n",
    "```sh\n",
    "python3 -m pip install torch transformers\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# suppress extraneous logging\n",
    "import os\n",
    "os.environ[\"TRANSFORMERS_VERBOSITY\"] = \"critical\"\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
    "\n",
    "from max import engine\n",
    "from max.dtype import DType"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we download the [RoBERTa model](https://huggingface.co/cardiffnlp/twitter-roberta-base-emotion-multilabel-latest)\n",
    "from HuggingFace and save it in the PyTorch TorchScript format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HF_MODEL_NAME = \"cardiffnlp/twitter-roberta-base-emotion-multilabel-latest\"\n",
    "hf_model = AutoModelForSequenceClassification.from_pretrained(HF_MODEL_NAME)\n",
    "hf_model.config.return_dict = False\n",
    "\n",
    "# Converting model to TorchScript\n",
    "model_path = Path(\"roberta.torchscript\")\n",
    "batch = 1\n",
    "seqlen = 128\n",
    "inputs = {\n",
    "    \"input_ids\": torch.zeros((batch, seqlen), dtype=torch.int64),\n",
    "    \"attention_mask\": torch.zeros((batch, seqlen), dtype=torch.int64),\n",
    "}\n",
    "with torch.no_grad():\n",
    "    traced_model = torch.jit.trace(\n",
    "        hf_model, example_kwarg_inputs=inputs, strict=False\n",
    "    )\n",
    "\n",
    "torch.jit.save(traced_model, model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load and compile the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we load and compile the model in MAX Engine using an\n",
    "[`InferenceSession`](https://docs.modular.com/engine/reference/python/engine.html#max.engine.InferenceSession).\n",
    "\n",
    "When loading a TorchScript model, you need to specify the shape and data type\n",
    "for all input tensors. This is required because TorchScript models do not\n",
    "include input spec annotations (unlike TensorFlow and ONNX models), which MAX\n",
    "Engine needs to compile the model.\n",
    "\n",
    "To define the input specs, you need to create a list of\n",
    "[`TorchInputSpec`](https://docs.modular.com/engine/reference/python/engine#max.engine.TorchInputSpec)\n",
    "objects (one item for each input), and pass the list to\n",
    "[`InferenceSession.load`](https://docs.modular.com/engine/reference/python/engine.html#max.engine.InferenceSession.load).\n",
    "For example, here's how to declare the input specs for the RoBERTa TorchScript\n",
    "model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We use the same `inputs` that we used above to trace the model\n",
    "input_spec_list = [\n",
    "    engine.TorchInputSpec(shape=tensor.size(), dtype=DType.int64)\n",
    "    for tensor in inputs.values()\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can load the model: (If you're loading a TensorFlow SavedModel or ONNX\n",
    "model, then you don't need the `input_specs` argument.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session = engine.InferenceSession()\n",
    "model = session.load(model_path, input_specs=input_spec_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's two lines down, just one to go.\n",
    "\n",
    "NOTE: The first time you load a model, it might take a few minutes to compile it,\n",
    "but this up-front cost will pay dividends in latency savings provided by\n",
    "our next-generation graph compiler."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare the input"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This part is your usual pre-processing. \n",
    "For the RoBERTa model, we need to process the text input into a sequence of tokens, so we'll do that with [`transformers.AutoTokenizer`](https://huggingface.co/docs/transformers/main/en/model_doc/auto#transformers.AutoTokenizer).\n",
    "\n",
    "First, let's take a look at the model's inputs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for tensor in model.input_metadata:\n",
    "    print(f'name: {tensor.name}, shape: {tensor.shape}, dtype: {tensor.dtype}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tells us the model needs 2 inputs. (If your model shows a dimension size\n",
    "is `None`, that means it's dynamic.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "INPUT=\"There are many exciting developments in the field of AI Infrastructure!\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_NAME)\n",
    "inputs = tokenizer(INPUT, return_tensors=\"pt\", padding='max_length', truncation=True, max_length=seqlen)\n",
    "print(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run an inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now for that third line of code, we pass the inputs to\n",
    "[`execute()`](https://docs.modular.com/engine/reference/python/engine#max.engine.Model.execute). This\n",
    "function requires all inputs as keyword arguments, so we'll\n",
    "unpack the `inputs` dictionary as we pass it through:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs = model.execute(**inputs)\n",
    "print(outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it!\n",
    "\n",
    "The output from [`execute()`](https://docs.modular.com/engine/reference/python/engine.html#max.engine.Model.execute) is a dictionary of output tensors, each in an `ndarray`. Let's now figure out what they say."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Process the outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, we'll use some help from the [transformers library](https://huggingface.co/docs/transformers/main/en/model_doc/roberta#transformers.TFRobertaForSequenceClassification) to convert the output ids to labels:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract class prediction from output\n",
    "predicted_class_id = outputs[\"result0\"].argmax(axis=-1)[0]\n",
    "classification = hf_model.config.id2label[predicted_class_id]\n",
    "\n",
    "print(f\"The sentiment is: {classification}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ta-da! 🎉\n",
    "\n",
    "If you're running this notebook yourself, beware that this notebook does not\n",
    "illustrate MAX Engine's runtime performance. For actual benchmark results,\n",
    "check out our [performance dashboard](https://performance.modular.com).\n",
    "\n",
    "For more details about the inferencing API, see the [Python API reference](https://docs.modular.com/max/api/python/engine)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}