first_5_steps/5_serving_tt_models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A Practical Guide to Serving AI Models on Tenstorrent Hardware: Deploying BERT with FastAPI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook serves as a practical guide to demonstrate how you can deploy an AI model on Tenstorrent hardware for an inference service using FastAPI.\n",
    "\n",
    "The tutorial will walk through an example of running the [BERT](https://en.wikipedia.org/wiki/BERT_(language_model)) model on Tenstorrent AI accelerator hardware. The model weights will be directly downloaded from the [HuggingFace library](https://huggingface.co/docs/transformers/model_doc/bert) and executed through the PyBUDA SDK. We will use FastAPI to build a RESTful API."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Import libraries\n",
    "\n",
    "Make sure that you have an activate Python environment with the latest version of PyBUDA installed.\n",
    "\n",
    "We will start by first pip installing a few libraries required to build a RESTful API: `fastapi`, `uvicorn`, and `nest-asyncio`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install a pip package in the current Jupyter kernel\n",
    "import sys\n",
    "!{sys.executable} -m pip install fastapi==0.85.1 uvicorn==0.19.0 nest-asyncio==1.5.8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import the pybuda library and additional libraries required for this tutorial\n",
    "import os\n",
    "from threading import Thread\n",
    "from typing import Dict, Tuple, Union\n",
    "\n",
    "import nest_asyncio\n",
    "import pybuda\n",
    "import requests\n",
    "import torch\n",
    "import uvicorn\n",
    "from fastapi import FastAPI\n",
    "from transformers import BertForSequenceClassification, BertTokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Build a Handler class\n",
    "\n",
    "We're going to build a Handler class that will act as the interface to be deployed.\n",
    "\n",
    "The class will hold the following methods:\n",
    "\n",
    "* `initailize` -- initialize / compile the model\n",
    "* `preprocess` -- preprocess the user input for the model\n",
    "* `inference` -- run inference on the model\n",
    "* `postprocess` -- postprocess the model outputs (logits)\n",
    "* `handle` -- pull all of the steps together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BERTHandler:\n",
    "    \"\"\"\n",
    "    A class to represent a BERT model RESTful API handler.\n",
    "\n",
    "    ...\n",
    "\n",
    "    Attributes\n",
    "    ----------\n",
    "    initialized : bool\n",
    "        Flag to mark if model as been compiled or not\n",
    "    device0 : pybuda.TTDevice\n",
    "        Tenstorrent device object which represents the hardware target to deploy model on\n",
    "    seqlen : int\n",
    "        Input sequence length\n",
    "\n",
    "    Methods\n",
    "    -------\n",
    "    initialize():\n",
    "        Initializes the model by downloading the weights, selecting the hardware target, and compiling the model\n",
    "    preprocess(input_text):\n",
    "        Preprocess the input (apply tokenization)\n",
    "    inference(processed_inputs):\n",
    "        Run inference on device\n",
    "    postprocess(logits):\n",
    "        Run post-processing on logits from model\n",
    "    handle(input_text):\n",
    "        Run all of the steps on user inputs\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, seqlen: int = 128):\n",
    "        \"\"\"\n",
    "        Constructs all the necessary attributes for the BERTHandler object.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        seqlen : int, optional\n",
    "            Input sequence length, by default 128\n",
    "        batch_size : int, optional\n",
    "            Input batch size, by default 1\n",
    "        \"\"\"\n",
    "        self.initialized = False\n",
    "        self.device0 = None\n",
    "        self.seqlen = seqlen\n",
    "\n",
    "    def initialize(self):\n",
    "        \"\"\"\n",
    "        Initialize and compile model pipeline.\n",
    "        \"\"\"\n",
    "\n",
    "        # Set logging levels\n",
    "        os.environ[\"LOGURU_LEVEL\"] = \"ERROR\"\n",
    "        os.environ[\"LOGGER_LEVEL\"] = \"ERROR\"\n",
    "\n",
    "        # Load BERT tokenizer and model from HuggingFace for text classification task\n",
    "        model_ckpt = \"assemblyai/bert-large-uncased-sst2\"\n",
    "        model = BertForSequenceClassification.from_pretrained(model_ckpt)\n",
    "        self.tokenizer = BertTokenizer.from_pretrained(model_ckpt)\n",
    "\n",
    "        # Initialize TTDevice object\n",
    "        tt0 = pybuda.TTDevice(\n",
    "            name=\"tt_device_0\",  # here we can give our device any name we wish, for tracking purposes\n",
    "            arch=pybuda.BackendDevice.Grayskull  # we set the target device architecture to compile for\n",
    "        )\n",
    "\n",
    "        # Create PyBUDA module\n",
    "        pybuda_module = pybuda.PyTorchModule(\n",
    "            name = \"pt_bert_text_classification\",  # give the module a name, this will be used for tracking purposes\n",
    "            module=model  # specify the model that is being targeted for compilation\n",
    "        )\n",
    "\n",
    "        # Place module on device\n",
    "        tt0.place_module(module=pybuda_module)\n",
    "        self.device0 = tt0\n",
    "\n",
    "        # Load data sample to compile model\n",
    "        sample_input = self.preprocess(\"sample input text\")\n",
    "\n",
    "        # Push input to model\n",
    "        self.device0.push_to_inputs(*sample_input)\n",
    "\n",
    "        # Compile & initialize the pipeline for inference, with given shapes\n",
    "        output_q = pybuda.run_inference()\n",
    "        _ = output_q.get()\n",
    "\n",
    "        # Configure initialization flag\n",
    "        self.initialized = True\n",
    "        print(\"BERTHandler initialized.\")\n",
    "\n",
    "    def preprocess(self, input_text: str) -> Tuple[torch.Tensor, torch.Tensor]:\n",
    "        \"\"\"\n",
    "        Preprocess the user inputs.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        input_text : str\n",
    "            User input\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        Tuple[torch.Tensor, torch.Tensor]\n",
    "            Processed outputs: `input_ids` and `attention_mask`\n",
    "        \"\"\"\n",
    "\n",
    "        input_tokens = self.tokenizer(\n",
    "            input_text,\n",
    "            max_length=self.seqlen,  # set the maximum input context length\n",
    "            padding=\"max_length\",  # pad to max length for fixed input size\n",
    "            truncation=True,  # truncate to max length\n",
    "            return_tensors=\"pt\",  # return PyTorch tensor\n",
    "        )\n",
    "\n",
    "        return (input_tokens[\"input_ids\"], input_tokens[\"attention_mask\"])\n",
    "\n",
    "    def inference(self, processed_inputs: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Run inference on Tenstorrent hardware.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        processed_inputs : Tuple[torch.Tensor, torch.Tensor]\n",
    "            Processed inputs: `input_ids` and `attention_mask`\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        torch.Tensor\n",
    "            Output logits from model\n",
    "        \"\"\"\n",
    "\n",
    "        self.device0.push_to_inputs(*processed_inputs)\n",
    "        output_q = pybuda.run_inference()\n",
    "        output = output_q.get()\n",
    "        logits = output[0].value().detach()\n",
    "        return logits\n",
    "\n",
    "    def postprocessing(self, logits: torch.Tensor) -> Dict[str, Union[str, float]]:\n",
    "        \"\"\"\n",
    "        Post-process logits and return dictionary with prediction and confidence score.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        logits : torch.Tensor\n",
    "            Predicted logits from model\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        Dict[str, Union[str, float]]\n",
    "            Output dictionary with predicted class and confidence score\n",
    "        \"\"\"\n",
    "\n",
    "        probabilities = torch.softmax(logits, dim=1)\n",
    "        confidences, predicted_classes = torch.max(probabilities, dim=1)\n",
    "        confidences = confidences.cpu().tolist()[0]\n",
    "        predicted_classes = predicted_classes.cpu()\n",
    "        output = {\n",
    "            \"predicted sentiment\": \"positive\" if predicted_classes else \"negative\",\n",
    "            \"confidence\": confidences\n",
    "        }\n",
    "\n",
    "        return output\n",
    "\n",
    "    def handle(self, text_input: str) -> Dict[str, Union[str, float]]:\n",
    "        \"\"\"\n",
    "        Handler function which runs end-to-end model pipeline\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        text_input : str\n",
    "            User input\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        Dict[str, Union[str, float]]\n",
    "            Output dictionary with predicted class and confidence score\n",
    "        \"\"\"\n",
    "\n",
    "        # Data preprocessing\n",
    "        processed_text = self.preprocess(text_input)\n",
    "\n",
    "        # Run inference\n",
    "        logits = self.inference(processed_text)\n",
    "\n",
    "        # Data postprocessing\n",
    "        output = self.postprocessing(logits)\n",
    "\n",
    "        return output\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Create FastAPI App\n",
    "\n",
    "We're going to use FastAPI to develop a simple RESTful API. You can experiment with alterative frameworks such as Flask and TorchServe to build your own application!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create FastAPI app\n",
    "app = FastAPI(\n",
    "    title=\"BERT Sentiment Analysis\",\n",
    "    description=\"Inference engine to classify texts.\",\n",
    "    version=\"1.0\",\n",
    ")\n",
    "\n",
    "# Initialize model on startup\n",
    "@app.on_event(\"startup\")\n",
    "async def startup():\n",
    "    global model\n",
    "    model = BERTHandler()\n",
    "    model.initialize()\n",
    "\n",
    "# Safely shutdown on exit\n",
    "@app.on_event(\"shutdown\")\n",
    "async def shutdown():\n",
    "    pybuda.shutdown()\n",
    "    pybuda.pybuda_reset()\n",
    "\n",
    "# Call handler on post request\n",
    "@app.post(\"/sentiment_v1/\")\n",
    "async def sentiment_v1(input_text: str) -> Dict[str, Union[str, float]]:\n",
    "    response = model.handle(input_text)\n",
    "    return response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Launch App on LocalHost\n",
    "\n",
    "Launch the app on your LocalHost. The model will first need to initialize and compile which can take 1-2 minutes.\n",
    "\n",
    "You can query the model once you see the following message:\n",
    "\n",
    "```\n",
    "INFO:     Application startup complete.\n",
    "INFO:     Uvicorn running on http://localhost:8000 (Press CTRL+C to quit)\n",
    "BERTHandler initialized.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run asyncio in Jupyter\n",
    "nest_asyncio.apply()\n",
    "\n",
    "# Define `uvicorn` command to launch app on LocalHost\n",
    "def run():\n",
    "    uvicorn.run(app, port=8000, host=\"localhost\")\n",
    "\n",
    "# Start app on thread\n",
    "thread = Thread(target=run)\n",
    "thread.start()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: Query the model\n",
    "\n",
    "Send post requets to your deployed model with the following API.\n",
    "\n",
    "Try changing the `INPUT_TEXT` to different sentiment sentences and observe the outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ↓↓↓↓↓↓↓↓ CONFIGURE INPUT ↓↓↓↓↓↓↓↓\n",
    "INPUT_TEXT = \"TT-BUDA is awesome!\"\n",
    "# ↑↑↑↑↑↑↑↑ CONFIGURE INPUT ↑↑↑↑↑↑↑↑\n",
    "\n",
    "# Set localhost url for app\n",
    "url = \"http://localhost:8000/sentiment_v1/\"\n",
    "\n",
    "# Issue post request\n",
    "input_text = {\"input_text\": INPUT_TEXT}\n",
    "response = requests.post(url, params=input_text).json()\n",
    "\n",
    "# Display outputs\n",
    "print(f\"Statement: {INPUT_TEXT}\\nPredicted sentiment: {response['predicted sentiment']}\\nConfidence: {response['confidence']*100:.0f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Congratulations on deploying your first RESTful API on Tenstorrent AI hardware!\n",
    "\n",
    "With this framework, you can now build your own AI applications on Tenstorrent AI hardware and deploy them in real life."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "v71_env_md",
   "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": 2
}