examples/1-getting-started.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AIVM Tutorial: Encrypted Neural Network Inference\n",
    "\n",
    "This tutorial demonstrates how to use the AIVM (AI Virtual Machine) client to perform encrypted inference on neural networks. We'll explore two use cases:\n",
    "1. Digit recognition using LeNet5 on MNIST dataset\n",
    "2. SMS spam detection using BERT Tiny\n",
    "\n",
    "## Setup and Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cell 1: Import necessary libraries\n",
    "import time\n",
    "import torch\n",
    "import torchvision.datasets as dset\n",
    "import torchvision.transforms as transforms\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import aivm_client as aic # Import the Nillion-AIVM client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This cell imports the required Python libraries for handling neural networks, data processing, and visualization. The `aivm_client` is the main interface for interacting with the AIVM system.\n",
    "\n",
    "## MNIST Dataset Loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_mnist():\n",
    "    trans = transforms.Compose(\n",
    "        [\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Resize((28, 28)),\n",
    "            transforms.Normalize((0.5,), (1.0,)),\n",
    "        ]\n",
    "    )\n",
    "    train_set = dset.MNIST(\n",
    "        root=\"/tmp/mnist\", train=True, transform=trans, download=True\n",
    "    )\n",
    "    return train_set\n",
    "\n",
    "dataset = load_mnist()\n",
    "\n",
    "inputs, labels =  dataset[20]\n",
    "inputs = inputs.reshape(1, 1, 28, 28)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This cell defines a function to load the MNIST dataset of handwritten digits. The data is:\n",
    "- Transformed to tensors\n",
    "- Resized to $28 \\times 28$ pixels\n",
    "- Normalized to the range $[-0.5, 0.5]$\n",
    "We then load a single example (index 20) and reshape it for model input.\n",
    "\n",
    "## Visualization\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the inputs as a grayscale image\n",
    "plt.imshow(inputs.detach().numpy().squeeze(), cmap='gray')\n",
    "plt.title(f'Grayscale Image of a {labels}')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "This cell visualizes the selected MNIST digit as a grayscale image, helping us verify the input data.\n",
    "\n",
    "## Encrypted Inference with LeNet5\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "encrypted_inputs = aic.LeNet5Cryptensor(inputs)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "This cell encrypts the input tensor using AIVM's encryption scheme specifically designed for `LeNet5` architecture.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "result = aic.get_prediction(encrypted_inputs, \"LeNet5MNIST\")\n",
    "results = torch.argmax(result, dim=1)\n",
    "print(\"Predicted Label:\", results.item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "This cell performs encrypted inference on the `LeNet5` model:\n",
    "1. Sends encrypted data to the AIVM\n",
    "2. Receives encrypted predictions\n",
    "3. Decrypts and processes the results to get the predicted digit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "This cell demonstrates an alternative method for inference using the `forward` method directly on the encrypted tensor.\n",
    "\n",
    "## Batch Processing\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "for i in range(100):\n",
    "    inputs, labels =  dataset[i]\n",
    "    inputs = inputs.reshape(1, 1, 28, 28)\n",
    "    encrypted_inputs = aic.LeNet5Cryptensor(inputs)\n",
    "    result = aic.get_prediction(encrypted_inputs, \"LeNet5MNIST\")\n",
    "    results = torch.argmax(result, dim=1)\n",
    "    print(\"Predicted Label:\", results.item(), \"True Label:\", labels, \"Correct:\", results.item() == labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "This cell demonstrates batch processing by:\n",
    "1. Loading 100 different MNIST images\n",
    "2. Performing encrypted inference on each\n",
    "3. Comparing predictions with true labels to assess accuracy\n",
    "\n",
    "## SMS Spam Detection\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_inputs = aic.tokenize(\"Your free ringtone is waiting to be collected. Simply text the password 'MIX' to 85069 to verify. Get Usher and Britney. FML, PO Box 5249, MK17 92H. 450Ppw 16\")\n",
    "\n",
    "encrypted_inputs = aic.BertTinyCryptensor(*tokenized_inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "This cell shows how to:\n",
    "1. Tokenize text input for BERT processing\n",
    "2. Encrypt the tokenized input for spam detection\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "result = aic.get_prediction(encrypted_inputs, \"BertTinySMS\")\n",
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "This cell performs encrypted inference using a BERT Tiny model for spam detection.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"SPAM\" if torch.argmax(result) else \"HAM\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "This final cell interprets the model's prediction as either \"HAM\" (legitimate message) or \"SPAM\".\n",
    "\n",
    "## Key Concepts\n",
    "- The notebook demonstrates privacy-preserving inference using AIVM\n",
    "- Input data is encrypted before being processed by the models\n",
    "- Two different use cases showcase the versatility of the system:\n",
    "  - Computer vision (MNIST digit recognition)\n",
    "  - Natural language processing (SMS spam detection)\n",
    "- Performance timing is measured to evaluate the overhead of encrypted computation"
   ]
  }
 ],
 "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.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}