{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AIVM Tutorial: Uploading Custom Bert Tiny Models to AIVM\n",
    "\n",
    "This tutorial demonstrates how to upload and use custom models with the AIVM (AI Virtual Machine) system, focusing on a BERT-tiny model for SMS spam detection.\n",
    "\n",
    "## Setup and Imports\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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:\n",
    "- Neural network operations (`torch`)\n",
    "- Dataset handling (`torchvision`)\n",
    "- Visualization (`matplotlib`)\n",
    "- AIVM client interface (`aivm_client`)\n",
    "\n",
    "## Model Upload\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = \"BertTinySentimentIMDB\" # Name of the model to be used\n",
    "aic.upload_bert_tiny_model(\"./imdb_bert_tiny.onnx\", MODEL_NAME) # Upload the model to the server"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This cell demonstrates how to upload a custom model to the AIVM server:\n",
    "1. Define a unique name for the model\n",
    "2. Upload an ONNX-format model file using the AIVM client\n",
    "3. The model will be available for encrypted inference after upload\n",
    "\n",
    "## IMDB Sentiment Classification Implementation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_positive_inputs = aic.tokenize(\"This movie was absolutely fantastic! I loved every minute of it.\",)\n",
    "encrypted_positive_inputs = aic.BertTinyCryptensor(tokenized_positive_inputs[0], tokenized_positive_inputs[1])\n",
    "\n",
    "tokenized_negative_inputs = aic.tokenize(\"Terrible waste of time. The plot made no sense and the acting was awful.\")\n",
    "encrypted_negative_inputs = aic.BertTinyCryptensor(tokenized_negative_inputs[0], tokenized_negative_inputs[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This cell prepares the input text for spam detection:\n",
    "1. Tokenizes example messages using BERT tokenization\n",
    "2. Creates an encrypted tensor from the tokenized input\n",
    "3. The second message is a typical spam SMS example\n",
    "\n",
    "## Encrypted Inference\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "result_positive = aic.get_prediction(encrypted_positive_inputs, MODEL_NAME)\n",
    "result_negative = aic.get_prediction(encrypted_negative_inputs, MODEL_NAME)\n",
    "\n",
    "print(\"Positive review prediction: \", result_positive)\n",
    "print(\"Negative review prediction: \", result_negative)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This cell:\n",
    "1. Performs encrypted inference using our uploaded custom model\n",
    "2. Measures the execution time\n",
    "3. Returns the raw model output (logits)\n",
    "\n",
    "Output shows:\n",
    "- CPU usage statistics\n",
    "- Wall time for inference\n",
    "- Tensor with prediction scores similar to `tensor([[-0.9099,  0.9210]])`\n",
    "\n",
    "\n",
    "\n",
    "## Result Interpretation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Positive review prediction: \", \"Positive\" if torch.argmax(result_positive) else \"Negative\")\n",
    "print(\"Negative review prediction: \", \"Positive\" if torch.argmax(result_negative) else \"Negative\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This cell interprets the model's prediction:\n",
    "- Converts numerical output to a binary classification\n",
    "- Tells whether a review is positive or negative\n",
    "\n",
    "## Key Concepts\n",
    "\n",
    "1. **Model Management**\n",
    "   - Custom models can be uploaded to AIVM in ONNX format\n",
    "   - Models are referenced by unique names\n",
    "   - Uploaded models persist on the server for future use\n",
    "\n",
    "2. **Privacy-Preserving Inference**\n",
    "   - All inference is performed on encrypted data\n",
    "   - Model weights and architecture remain secure\n",
    "   - Results are only decrypted at the client side\n",
    "\n",
    "3. **Performance Considerations**\n",
    "   - Encrypted computation introduces some overhead\n",
    "   - Wall time measurements help evaluate real-world performance\n",
    "   - System is optimized for practical use cases\n",
    "\n",
    "4. **Integration Workflow**\n",
    "   1. Upload model\n",
    "   2. Prepare and encrypt inputs\n",
    "   3. Perform inference\n",
    "   4. Interpret results\n"
   ]
  }
 ],
 "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
}
